mlx-bun 0.2.0 → 0.4.0

package/src/spec/glm52-mtp-source.ts

@@ -13,8 +13,11 @@ import * as ops from "../mlx/ops";
 import { toLogprobs } from "../sampler";
 import { Glm52DecoderLayer, Glm52Model } from "../model/glm52";
 import { MLACache } from "../model/glm52-cache";
-import { rmsNormF32Mlx } from "../model/glm52-mla";
-import type { DraftProvider, DraftSource } from "./source";
+import { Glm52MtpGraph } from "./glm52-mtp-graph";
+import type { DraftProvider, DraftSource, GroupedDraftProvider, TargetView, DraftRowSampling, DraftRowCheckpoint, DraftRowGroup, DraftPrefillGroup } from "./source";
+import { Glm52MtpRows, type Glm52MtpRowState } from "./glm52-mtp-rows";
+import { captureGlm52MtpState, restoreGlm52MtpState } from "./glm52-mtp-state";
+import { applyStateChanges, disposeResources } from "../engine/resources";
 
 type Sampler = (logprobs: MlxArray, step: number) => MlxArray;
 
@@ -26,6 +29,10 @@ function lastToken(feed: readonly number[]): number {
 }
 
 export class Glm52NativeMtpProvider implements DraftProvider {
+  readonly grouped: GroupedDraftProvider = {
+    open: options => this.#openRows(options.target, options.sampling, options.checkpoints),
+    openPrefill: options => this.#openRows(options.target, null, options.checkpoints),
+  };
   readonly id = "glm52-native-mtp";
   readonly weightsBytes = 0;
   readonly #layer: Glm52DecoderLayer;
@@ -57,11 +64,39 @@ export class Glm52NativeMtpProvider implements DraftProvider {
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
-    if (opts.target.model !== this.model)
+    if (opts.target.identity !== this.model)
       throw new Error("native MTP provider was opened for a different target");
     return new Glm52NativeMtpSource(this.model, this.#layer, opts.sampler);
   }
 
+  #openRows(target: TargetView, sampling: DraftRowSampling | null,
+    checkpoints: readonly (DraftRowCheckpoint | null)[]): DraftRowGroup & DraftPrefillGroup {
+    if (target.identity !== this.model) throw new Error("native MTP provider was opened for a different target");
+    const rows = new Glm52MtpRows(new Glm52MtpGraph(this.model, this.#layer), sampling);
+    const prepareAppend = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => {
+      const states: (Glm52MtpRowState | null)[] = [];
+      try {
+        for (const checkpoint of checkpoints) states.push(checkpoint ? restoreGlm52MtpState(checkpoint, () => rows.makeCache()) : null);
+        return rows.prepareAppend(states);
+      } finally { disposeResources(states.flatMap(state => state ? [state.cache, state.hidden] : [])); }
+    };
+    const append = (states: readonly (DraftRowCheckpoint | null)[]) => applyStateChanges([() => prepareAppend(states)]);
+    try { append(checkpoints); } catch (error) { rows.dispose(); throw error; }
+    return {
+      namespace: "glm52-native-mtp-v1", prefillMode: "full", tapLayers: [],
+      get rowCount() { return rows.rowCount; },
+      append, prepareAppend, prefill: (tokens, context) => rows.prefill(tokens, context!),
+      materialize: rows.materialize.bind(rows), filterRows: rows.filterRows.bind(rows),
+      draft: rows.draft.bind(rows), commit: rows.commit.bind(rows),
+      capture(row) {
+        const state = rows.extractRow(row);
+        try { return captureGlm52MtpState(state); }
+        finally { state.cache.dispose(); state.hidden.dispose(); }
+      },
+      dispose: rows.dispose.bind(rows),
+    };
+  }
+
   dispose(): void {
     // The target model owns every shared/MTP weight and the expert tier.
   }
@@ -72,6 +107,7 @@ export class Glm52NativeMtpSource implements DraftSource {
   readonly pinTargetKernelFamily = true;
   readonly weightsBytes = 0;
   #cache: MLACache;
+  readonly #graph: Glm52MtpGraph;
   #roundStart = 0;
   #lastDraftCount = 0;
   #closed = false;
@@ -81,6 +117,7 @@ export class Glm52NativeMtpSource implements DraftSource {
     readonly layer: Glm52DecoderLayer,
     readonly sampler: Sampler,
   ) {
+    this.#graph = new Glm52MtpGraph(model, layer);
     this.#cache = new MLACache({
       kvLoraRank: model.glmConfig.kvLoraRank,
       ropeHeadDim: model.glmConfig.qkRopeHeadDim,
@@ -256,142 +293,22 @@ export class Glm52NativeMtpSource implements DraftSource {
     hidden: MlxArray,
     sampleStep: number,
   ): Promise<{ hidden: MlxArray; token: number }> {
-    const config = this.model.glmConfig;
-    const ids = ops.fromInt32([token], [1, 1]);
-    let embedded: MlxArray | null = null;
-    let embeddedNorm: MlxArray | null = null;
-    let hiddenNorm: MlxArray | null = null;
-    let joined: MlxArray | null = null;
-    let projected: MlxArray | null = null;
-    let output: MlxArray | null = null;
-    let headInput: MlxArray | null = null;
-    let logits: MlxArray | null = null;
-    let logprobs: MlxArray | null = null;
-    let sampled: MlxArray | null = null;
-    try {
-      embedded = this.model.weights.embedding(
-        ids,
-        "model.embed_tokens.weight",
-        config.vocabSize,
-        config.hiddenSize,
-      );
-      embeddedNorm = rmsNormF32Mlx(
-        embedded,
-        this.model.weights.tensor(
-          `model.layers.${config.numHiddenLayers}.enorm.weight`,
-        ),
-        config.rmsNormEps,
-      );
-      hiddenNorm = rmsNormF32Mlx(
-        hidden,
-        this.model.weights.tensor(
-          `model.layers.${config.numHiddenLayers}.hnorm.weight`,
-        ),
-        config.rmsNormEps,
-      );
-      joined = ops.concatAxis([embeddedNorm, hiddenNorm], 2);
-      projected = this.model.weights.linear(
-        joined,
-        `model.layers.${config.numHiddenLayers}.eh_proj.weight`,
-        config.hiddenSize,
-        2 * config.hiddenSize,
-      );
-      output = await this.layer.forwardAsync(projected, this.#cache, null);
-      headInput = rmsNormF32Mlx(
-        output,
-        this.model.weights.tensor(
-          `model.layers.${config.numHiddenLayers}.shared_head.norm.weight`,
-        ),
-        config.rmsNormEps,
-      );
-      logits = this.model.logitsFromHidden(headInput);
-      // Sampler contract is [1, V] — same 3-D top-k hazard as
-      // qwen-mtp-source #sample (2026-08-20); see the comment there.
-      {
-        const V = logits.shape[logits.shape.length - 1]!;
-        const flat = ops.reshape(logits, [1, V]);
-        logits.dispose();
-        logits = flat;
-      }
-      logprobs = toLogprobs(logits);
-      sampled = this.sampler(logprobs, sampleStep);
-      return {
-        hidden: ops.contiguous(output),
-        token: ops.itemUint32(sampled),
-      };
-    } finally {
-      ids.dispose();
-      embedded?.dispose();
-      embeddedNorm?.dispose();
-      hiddenNorm?.dispose();
-      joined?.dispose();
-      projected?.dispose();
-      output?.dispose();
-      headInput?.dispose();
-      logits?.dispose();
-      logprobs?.dispose();
-      sampled?.dispose();
-    }
+    using ids = ops.fromInt32([token], [1, 1]);
+    using output = await this.#graph.forward(ids, hidden, this.#cache);
+    using logits = this.#graph.project(output);
+    // Sampling consumes [B,V], independently of the graph's token dimension.
+    using flat = ops.reshape(logits, [1, logits.shape.at(-1)!]);
+    using logprobs = toLogprobs(flat);
+    using sampled = this.sampler(logprobs, sampleStep);
+    return { hidden: ops.contiguous(output), token: ops.itemUint32(sampled) };
   }
 
-  async #absorb(
-    acceptedTokens: readonly number[],
-    verifiedHidden: MlxArray,
-  ): Promise<void> {
-    const config = this.model.glmConfig;
-    const k = acceptedTokens.length;
-    const ids = ops.fromInt32([...acceptedTokens], [1, k]);
-    let embedded: MlxArray | null = null;
-    let embeddedNorm: MlxArray | null = null;
-    let trueRows: MlxArray | null = null;
-    let hiddenNorm: MlxArray | null = null;
-    let joined: MlxArray | null = null;
-    let projected: MlxArray | null = null;
-    let output: MlxArray | null = null;
-    try {
-      embedded = this.model.weights.embedding(
-        ids,
-        "model.embed_tokens.weight",
-        config.vocabSize,
-        config.hiddenSize,
-      );
-      embeddedNorm = rmsNormF32Mlx(
-        embedded,
-        this.model.weights.tensor(
-          `model.layers.${config.numHiddenLayers}.enorm.weight`,
-        ),
-        config.rmsNormEps,
-      );
-      trueRows = verifiedHidden.slice(
-        [0, 0, 0],
-        [1, k, config.hiddenSize],
-      );
-      hiddenNorm = rmsNormF32Mlx(
-        trueRows,
-        this.model.weights.tensor(
-          `model.layers.${config.numHiddenLayers}.hnorm.weight`,
-        ),
-        config.rmsNormEps,
-      );
-      joined = ops.concatAxis([embeddedNorm, hiddenNorm], 2);
-      projected = this.model.weights.linear(
-        joined,
-        `model.layers.${config.numHiddenLayers}.eh_proj.weight`,
-        config.hiddenSize,
-        2 * config.hiddenSize,
-      );
-      output = await this.layer.forwardAsync(projected, this.#cache, null);
-      output.eval();
-    } finally {
-      ids.dispose();
-      embedded?.dispose();
-      embeddedNorm?.dispose();
-      trueRows?.dispose();
-      hiddenNorm?.dispose();
-      joined?.dispose();
-      projected?.dispose();
-      output?.dispose();
-    }
+  async #absorb(acceptedTokens: readonly number[], verifiedHidden: MlxArray): Promise<void> {
+    const count = acceptedTokens.length;
+    using ids = ops.fromInt32([...acceptedTokens], [1, count]);
+    using trueRows = verifiedHidden.slice([0, 0, 0], [1, count, this.model.glmConfig.hiddenSize]);
+    using output = await this.#graph.forward(ids, trueRows, this.#cache);
+    output.eval();
   }
 
   #checkOpen(): void {

package/src/spec/glm52-mtp-state.ts

@@ -0,0 +1,33 @@
+import type { MlxArray } from "../mlx/array";
+import type { MLACache } from "../model/glm52-cache";
+import { cloneAttachments } from "../backends/mlx/checkpoint-state";
+import { disposeResources } from "../engine/resources";
+import type { DraftRowCheckpoint } from "./source";
+import type { Glm52MtpRowState } from "./glm52-mtp-rows";
+
+export function captureGlm52MtpState(state: Glm52MtpRowState): DraftRowCheckpoint {
+  const held: MlxArray[] = [];
+  try {
+    if (state.cache.offset) {
+      const cache = state.cache.fetch(); held.push(cache.latent, cache.rope);
+    }
+    held.push(state.hidden.slice([0, 0, 0], [...state.hidden.shape]));
+    return { processedTokens: state.processedTokens, attachment: {
+      schema: "glm52-native-mtp-v1", metadata: { processedTokens: state.processedTokens, draftOffset: state.cache.offset },
+      tensors: held.splice(0),
+    } };
+  } finally { disposeResources(held); }
+}
+
+export function restoreGlm52MtpState(checkpoint: DraftRowCheckpoint, makeCache: () => MLACache): Glm52MtpRowState {
+  const { attachment, processedTokens } = checkpoint;
+  if (attachment.schema !== "glm52-native-mtp-v1" || attachment.metadata.processedTokens !== processedTokens)
+    throw new Error("invalid native GLM MTP checkpoint alignment");
+  const held = cloneAttachments([attachment])[0]!.tensors;
+  let cache: MLACache | null = makeCache();
+  try {
+    const offset = attachment.metadata.draftOffset as number;
+    if (offset) { cache.restoreCompressedState(held[0]!, held[1]!, null, offset); held.splice(0, 2); }
+    const state = { cache, hidden: held.pop()!, processedTokens }; cache = null; return state;
+  } finally { disposeResources([...held, ...(cache ? [cache] : [])]); }
+}

package/src/spec/ngram-source.ts

@@ -38,8 +38,80 @@
 // speculation is structurally never worse than plain decode by more than the
 // JS scan (~µs against 30k-token histories).
 
-import { flagOn } from "../flags";
-import type { DraftProvider, DraftSource } from "./source";
+import { flagOn } from "../runtime-config";
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { CheckpointAttachment } from "../backends/mlx/checkpoint-state";
+import type { DraftProvider, DraftSource, DraftRowCheckpoint, DraftRowGroup, DraftPrefillGroup, GroupedDraftProvider } from "./source";
+import { applyStateChanges } from "../engine/resources";
+
+/** The matching policy is independent of the executor's feed convention. */
+function proposeNgram(h: readonly number[], max: number, min: number, n: number): number[] {
+  const L = h.length;
+  for (let k = Math.min(max, L - 1); k >= min; k--) {
+    const tailAt = L - k;
+    search: for (let i = 0; i + k < L; i++) {
+      for (let j = 0; j < k; j++) if (h[i + j] !== h[tailAt + j]) continue search;
+      return h.slice(i + k, i + k + n);
+    }
+  }
+  return [];
+}
+
+const namespace = (max: number, min: number) => `ngram-history-v1:${max}:${min}`;
+const captureHistory = (history: readonly number[], max: number, min: number): CheckpointAttachment => ({
+  schema: "ngram-history-v1", metadata: { max, min }, tensors: [ops.fromInt32([...history], [history.length])],
+});
+function restoreHistory(state: DraftRowCheckpoint, max: number, min: number): number[] {
+  const { attachment, processedTokens } = state;
+  if (attachment.schema !== "ngram-history-v1" || attachment.metadata.max !== max || attachment.metadata.min !== min ||
+      attachment.tensors.length !== 1 || attachment.tensors[0]!.shape.length !== 1 || attachment.tensors[0]!.shape[0] !== processedTokens)
+    throw new Error("ngram checkpoint does not match its token coverage or lookup policy");
+  return attachment.tensors[0]!.toIntTokens();
+}
+
+/** Each row owns only its committed token history. Verification and sampling
+ * remain in the shared method; target storage and scheduling are not needed. */
+class NgramRows implements DraftRowGroup, DraftPrefillGroup {
+  readonly prefillMode = "tail-split" as const;
+  readonly tapLayers: readonly number[] = [];
+  readonly namespace: string;
+  #histories: number[][] = [];
+  #drafts: number[][] = [];
+  constructor(readonly max: number, readonly min: number, checkpoints: readonly (DraftRowCheckpoint | null)[]) {
+    this.namespace = namespace(max, min);
+    this.append(checkpoints);
+  }
+  get rowCount() { return this.#histories.length; }
+  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]) {
+    let next = [...this.#histories, ...checkpoints.map(state => state ? restoreHistory(state, this.max, this.min) : [])];
+    return { commit: () => { const old = this.#histories; this.#histories = next; next = old; }, dispose: () => { next = []; } };
+  }
+  append(checkpoints: readonly (DraftRowCheckpoint | null)[]) { applyStateChanges([() => this.prepareAppend(checkpoints)]); }
+  filterRows(keep: readonly number[]) { this.#histories = keep.map(row => this.#histories[row]!); this.#drafts = []; }
+  prefill(tokens: MlxArray): void {
+    using packed = ops.contiguous(tokens);
+    const ids = packed.toIntTokens(), length = tokens.shape[1]!;
+    for (let row = 0; row < this.rowCount; row++) this.#histories[row]!.push(...ids.slice(row * length, (row + 1) * length));
+  }
+  materialize(): void {}
+  draft(pending: readonly number[], depth: number): number[][] {
+    this.#drafts = this.#histories.map((history, row) => {
+      history.push(pending[row]!);
+      return proposeNgram(history, this.max, this.min, depth);
+    });
+    return this.#drafts;
+  }
+  commit(accepted: readonly number[]) {
+    for (let row = 0; row < this.rowCount; row++) this.#histories[row]!.push(...this.#drafts[row]!.slice(0, accepted[row]!));
+    this.#drafts = [];
+  }
+  capture(row: number): DraftRowCheckpoint {
+    const history = this.#histories[row]!;
+    return { processedTokens: history.length, attachment: captureHistory(history, this.max, this.min) };
+  }
+  dispose() { this.#histories = []; this.#drafts = []; }
+}
 
 export interface NgramOptions {
   /** Longest suffix k-gram tried first (Saxena max_ngram_size). */
@@ -53,6 +125,11 @@ export class NgramProvider implements DraftProvider {
   readonly weightsBytes = 0;
   readonly max: number;
   readonly min: number;
+  readonly grouped: GroupedDraftProvider = {
+    supportsTargetAdapters: true,
+    open: options => new NgramRows(this.max, this.min, options.checkpoints),
+    openPrefill: options => new NgramRows(this.max, this.min, options.checkpoints),
+  };
 
   constructor(opts: NgramOptions = {}) {
     this.max = Math.max(1, opts.max ?? 3);
@@ -73,24 +150,33 @@ class NgramSource implements DraftSource {
   readonly weightsBytes = 0;
   #hist: number[] = [];
   #lastDrafts: number[] = [];
+  readonly #tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+  readonly checkpoint: NonNullable<DraftSource["checkpoint"]>;
 
   constructor(
     private readonly max: number,
     private readonly min: number,
-  ) {}
+  ) {
+    this.checkpoint = {
+      namespace: namespace(max, min),
+      capture: () => captureHistory(this.#hist, this.max, this.min),
+      restore: (processedTokens, attachment) => {
+        this.#hist = restoreHistory({ processedTokens, attachment }, this.max, this.min);
+      },
+    };
+  }
 
   /** Test hook — the reconstructed prompt+emitted stream (see header). */
   get history(): readonly number[] {
     return this.#hist;
   }
 
-  prefill(promptIds: number[]): void {
+  prefill(promptIds: number[], _context?: Parameters<DraftSource["prefill"]>[1], processedTokens?: number): void {
     // Mirror the serve loop's prefill shape (two-model.ts does the same): under
     // the oracle tail split the last prompt token is never prefilled — it IS
     // the first feed. Legacy shape (kill switch / 1-token prompt): full prompt,
     // and the sampled token0 arrives as the first feed.
-    const tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
-    const upTo = tailSplit && promptIds.length > 1 ? promptIds.length - 1 : promptIds.length;
+    const upTo = processedTokens ?? (this.#tailSplit && promptIds.length > 1 ? promptIds.length - 1 : promptIds.length);
     this.#hist = promptIds.slice(0, upTo);
   }
 
@@ -111,21 +197,7 @@ class NgramSource implements DraftSource {
   /** Longest-k-first, first-occurrence prompt lookup (Saxena/vLLM order):
    *  find the trailing k-gram earlier in history, propose what followed it. */
   #propose(n: number): number[] {
-    const h = this.#hist;
-    const L = h.length;
-    const kTop = Math.min(this.max, L - 1);
-    for (let k = kTop; k >= this.min; k--) {
-      const tailAt = L - k;
-      // First occurrence, left to right; i + k < L both excludes the trailing
-      // self-match and guarantees at least one continuation token.
-      search: for (let i = 0; i + k < L; i++) {
-        for (let j = 0; j < k; j++) {
-          if (h[i + j] !== h[tailAt + j]) continue search;
-        }
-        return h.slice(i + k, i + k + n);
-      }
-    }
-    return [];
+    return proposeNgram(this.#hist, this.max, this.min, n);
   }
 
   dispose(): void {}

package/src/spec/projected-draft-rows.ts

@@ -0,0 +1,73 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { KVCache, type Cache } from "../model/gemma4-base";
+import { MlxProjectedContextRows } from "../backends/mlx/projected-context-rows";
+import { captureCacheAttachment, restoreCacheAttachment } from "../backends/mlx/cache-attachment";
+import { applyStateChanges, disposeResources } from "../engine/resources";
+import type { PreparedStateChange } from "../contracts/resources";
+import type { DraftPrefillGroup, DraftRowGroup, DraftRowCheckpoint, GroupedDraftProvider, TargetView } from "./source";
+
+export interface ProjectedDraftMethod {
+  namespace: string;
+  schema: string;
+  layers: number;
+  tapLayers: readonly number[];
+  project(hidden: MlxArray, positions: number | MlxArray): { k: MlxArray; v: MlxArray }[];
+  draft(context: ReturnType<MlxProjectedContextRows["readAttention"]>, pending: readonly number[], positions: number | MlxArray, depth: number): number[][];
+}
+
+/** Provider numerics compose with shared context membership and persistence. */
+class ProjectedDraftRows implements DraftPrefillGroup, DraftRowGroup {
+  readonly prefillMode = "tail-split" as const;
+  readonly tapLayers: readonly number[];
+  readonly context: MlxProjectedContextRows;
+  constructor(readonly method: ProjectedDraftMethod,
+    checkpoints: readonly (DraftRowCheckpoint | null)[]) {
+    this.tapLayers = method.tapLayers;
+    this.context = new MlxProjectedContextRows(method.layers, { project: method.project });
+    try { this.append(checkpoints); } catch (error) { this.dispose(); throw error; }
+  }
+  get namespace(): string { return this.method.namespace; }
+  get rowCount(): number { return this.context.state.rowCount; }
+  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange {
+    const restored: Cache[][] = [];
+    const empty = () => Array.from({ length: this.method.layers }, () => new KVCache());
+    try {
+      for (const checkpoint of checkpoints) restored.push(checkpoint ? restoreCacheAttachment(checkpoint.attachment, empty) : empty());
+      return this.context.prepareAppend(restored);
+    } finally { disposeResources(restored.flat()); }
+  }
+  append(checkpoints: readonly (DraftRowCheckpoint | null)[]): void {
+    if (checkpoints.length) applyStateChanges([() => this.prepareAppend(checkpoints)]);
+  }
+  prefill(tokens: MlxArray, hidden?: MlxArray): void {
+    this.context.append(hidden!, Array(this.rowCount).fill(tokens.shape[1]!)); this.materialize();
+  }
+  materialize(): void { this.context.materialize(); }
+  filterRows(keep: readonly number[]): void { this.context.state.filterRows(keep); }
+  capture(row: number): DraftRowCheckpoint {
+    const caches = this.context.state.extractRow(row);
+    try {
+      const processedTokens = caches[0]!.offset;
+      return { processedTokens, attachment: captureCacheAttachment(this.method.schema, caches, { processedTokens }) };
+    } finally { disposeResources(caches); }
+  }
+  draft(pending: readonly number[], depth: number): number[][] {
+    if (!depth) return pending.map(() => []);
+    const positions = this.context.positions, donors = this.context.readAttention();
+    using offsets = positions.every(position => position === positions[0]) ? null : ops.fromInt32([...positions], [pending.length]);
+    try { return this.method.draft(donors, pending, offsets ?? positions[0]!, depth); }
+    finally { donors.dispose(); }
+  }
+  commit(accepted: readonly number[], hidden: MlxArray): void {
+    this.context.append(hidden, accepted.map(count => count + 1));
+  }
+  dispose(): void { this.context.dispose(); }
+}
+
+export function projectedDraftGroups(bind: (target: TargetView) => ProjectedDraftMethod): GroupedDraftProvider {
+  const open = ({ target, checkpoints }: {
+    target: TargetView; checkpoints: readonly (DraftRowCheckpoint | null)[];
+  }) => new ProjectedDraftRows(bind(target), checkpoints);
+  return { open, openPrefill: open };
+}

package/src/spec/qwen-mtp-module.ts

@@ -0,0 +1,171 @@
+// Qwen MTP numerical graph. Request state, sampling and checkpoint retention
+// belong to its callers; the cache supplies positions and attention storage.
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { ModelConfig } from "../config";
+import type { Weights } from "../weights";
+import { disposing, QuantizedLinear, RMSNorm, type Cache, type Mask } from "../model/gemma4-base";
+import { DenseLinear } from "../model/universal/modules";
+import { compiledSwiglu } from "../model/qwen3_5";
+
+type MtpLinear = DenseLinear | QuantizedLinear;
+
+/** Projection tensors belong to Weights; dense transpose views belong to
+ *  the provider's resource stack. Quantized heads keep no dense copy. */
+function loadMtpLinear(
+  weights: Weights, path: string, config: ModelConfig, resources: DisposableStack,
+): MtpLinear {
+  if (weights.has(`${path}.scales`)) return QuantizedLinear.load(weights, path, config);
+  const layer = new DenseLinear(weights.tensor(`${path}.weight`), null);
+  resources.use(layer.wT);
+  return layer;
+}
+
+/** Qwen3Attention.forward with the companion's dense or quantized projections.
+ *  Attention operations follow src/model/qwen3_5.ts verbatim. */
+class MtpAttention {
+  readonly qProj: MtpLinear;
+  readonly kProj: MtpLinear;
+  readonly vProj: MtpLinear;
+  readonly oProj: MtpLinear;
+  readonly qNorm: RMSNorm;
+  readonly kNorm: RMSNorm;
+  readonly nHeads: number;
+  readonly nKvHeads: number;
+  readonly headDim: number;
+  readonly scale: number;
+  readonly ropeDims: number;
+  readonly ropeBase: number;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string, resources: DisposableStack) {
+    const t = config.text;
+    this.nHeads = t.numAttentionHeads;
+    this.nKvHeads = t.numKeyValueHeads;
+    this.headDim = t.headDim;
+    this.scale = Math.pow(this.headDim, -0.5);
+    this.ropeDims = Math.trunc(this.headDim * t.partialRotaryFactor);
+    this.ropeBase = t.ropeParameters.full_attention?.ropeTheta ?? 10000;
+    this.qProj = loadMtpLinear(weights, `${prefix}.q_proj`, config, resources);
+    this.kProj = loadMtpLinear(weights, `${prefix}.k_proj`, config, resources);
+    this.vProj = loadMtpLinear(weights, `${prefix}.v_proj`, config, resources);
+    this.oProj = loadMtpLinear(weights, `${prefix}.o_proj`, config, resources);
+    this.qNorm = new RMSNorm(weights.tensor(`${prefix}.q_norm.weight`), t.rmsNormEps);
+    this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
+  }
+
+  forward(x: MlxArray, mask: Mask, cache: Cache): MlxArray {
+    const [B, L] = x.shape as [number, number, number];
+    const qp = this.qProj.forward(x);
+    const qpr = disposing(qp, ops.reshape(qp, [B, L, this.nHeads, this.headDim * 2]));
+    const [qHeads, gateHeads] = ops.split(qpr, [this.headDim], -1) as [MlxArray, MlxArray];
+    qpr.dispose();
+    const gate = disposing(gateHeads, ops.reshape(gateHeads, [B, L, this.nHeads * this.headDim]));
+
+    let k = this.kProj.forward(x);
+    let v = this.vProj.forward(x);
+
+    let q = this.qNorm.forward(qHeads);
+    qHeads.dispose();
+    q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
+    k = disposing(k, ops.reshape(k, [B, L, this.nKvHeads, this.headDim]));
+    k = disposing(k, this.kNorm.forward(k));
+    k = disposing(k, ops.transposeAxes(k, [0, 2, 1, 3]));
+    v = disposing(v, ops.reshape(v, [B, L, this.nKvHeads, this.headDim]));
+    v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
+
+    const offsets = cache.ropeOffsetArr;
+    q = disposing(q, offsets
+      ? ops.ropeDynamic(q, this.ropeDims, this.ropeBase, offsets, null)
+      : ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
+    k = disposing(k, offsets
+      ? ops.ropeDynamic(k, this.ropeDims, this.ropeBase, offsets, null)
+      : ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
+
+    const [keys, values] = cache.updateAndFetch(k, v);
+    k.dispose();
+    v.dispose();
+    const attn = ops.sdpa(q, keys, values, this.scale, mask.mode, mask.arr);
+    keys.dispose();
+    values.dispose();
+    q.dispose();
+
+    const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+    attn.dispose();
+    const merged = ops.reshape(attnT, [B, L, -1]);
+    attnT.dispose();
+    const sig = ops.sigmoid(gate);
+    gate.dispose();
+    const gated = ops.mul(merged, sig);
+    merged.dispose();
+    sig.dispose();
+    const out = this.oProj.forward(gated);
+    gated.dispose();
+    return out;
+  }
+}
+
+/** The one MTP decoder block: fc-merge → attention → swiglu MLP → norm. */
+export class MtpModule {
+  readonly fc: MtpLinear;
+  readonly preFcNormEmbedding: RMSNorm;
+  readonly preFcNormHidden: RMSNorm;
+  readonly attn: MtpAttention;
+  readonly mlpGate: MtpLinear;
+  readonly mlpUp: MtpLinear;
+  readonly mlpDown: MtpLinear;
+  readonly inputNorm: RMSNorm;
+  readonly postAttnNorm: RMSNorm;
+  readonly finalNorm: RMSNorm;
+
+  constructor(weights: Weights, config: ModelConfig, resources: DisposableStack) {
+    const eps = config.text.rmsNormEps;
+    this.fc = loadMtpLinear(weights, "fc", config, resources);
+    this.preFcNormEmbedding = new RMSNorm(weights.tensor("pre_fc_norm_embedding.weight"), eps);
+    this.preFcNormHidden = new RMSNorm(weights.tensor("pre_fc_norm_hidden.weight"), eps);
+    this.attn = new MtpAttention(weights, config, "layers.0.self_attn", resources);
+    this.mlpGate = loadMtpLinear(weights, "layers.0.mlp.gate_proj", config, resources);
+    this.mlpUp = loadMtpLinear(weights, "layers.0.mlp.up_proj", config, resources);
+    this.mlpDown = loadMtpLinear(weights, "layers.0.mlp.down_proj", config, resources);
+    this.inputNorm = new RMSNorm(weights.tensor("layers.0.input_layernorm.weight"), eps);
+    this.postAttnNorm = new RMSNorm(weights.tensor("layers.0.post_attention_layernorm.weight"), eps);
+    this.finalNorm = new RMSNorm(weights.tensor("norm.weight"), eps);
+  }
+
+  /** One block forward over [B,S,·]: token embeddings ([B,S,H], target
+   *  embed_tokens output) paired with hiddens ([B,S,H], target pre-final-norm
+   *  or the module's own chained output). Appends S rows to `cache`; returns
+   *  the module output [B,S,H] (post final norm — what the target lm_head
+   *  consumes AND what chains into the next step's `hidden`). */
+  forward(tokenEmbeds: MlxArray, hiddens: MlxArray, cache: Cache): MlxArray {
+    const embNorm = this.preFcNormEmbedding.forward(tokenEmbeds);
+    const hidNorm = this.preFcNormHidden.forward(hiddens);
+    const joined = ops.concatAxis([embNorm, hidNorm], 2);
+    embNorm.dispose();
+    hidNorm.dispose();
+    const x = this.fc.forward(joined);
+    joined.dispose();
+
+    // Decoder layer (Qwen3Layer.forward shape).
+    const L = x.shape[1]!;
+    const mask = cache.makeMask(L, null);
+    const xn = this.inputNorm.forward(x);
+    const r = this.attn.forward(xn, mask, cache);
+    xn.dispose();
+    mask.arr?.dispose();
+    const h = ops.add(x, r);
+    x.dispose();
+    r.dispose();
+    const hn = this.postAttnNorm.forward(h);
+    const g = this.mlpGate.forward(hn);
+    const u = this.mlpUp.forward(hn);
+    hn.dispose();
+    const act = compiledSwiglu(g, u);
+    g.dispose(); u.dispose();
+    const m = this.mlpDown.forward(act);
+    act.dispose();
+    const out = ops.add(h, m);
+    h.dispose();
+    m.dispose();
+    return disposing(out, this.finalNorm.forward(out));
+  }
+}