mlx-bun 0.3.0 → 0.4.0

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

@@ -0,0 +1,219 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { materializeCopy } from "../mlx/materialize";
+import { toLogprobs } from "../sampler";
+import { BatchedKVCache } from "../model/batched-kv";
+import { KVCache } from "../model/gemma4-base";
+import type { DraftRowSampling, QwenMtpTarget } from "./source";
+import type { MtpModule } from "./qwen-mtp-module";
+import type { PreparedStateChange } from "../contracts/resources";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../engine/resources";
+
+/** A committed request's draft KV and the preceding true target hidden row.
+ * Construction borrows these arrays; extraction returns owned state. */
+export interface MtpRowState {
+  readonly cache: KVCache;
+  readonly hidden: MlxArray;
+}
+
+/** Sampling owns each request's RNG policy. Scores are [B,V] log-probabilities;
+ * the returned owned [B] token array stays on device through the draft chain. */
+export type MtpRowSampling = DraftRowSampling;
+
+/** The same draft graph and state lifecycle for one or several requests.
+ * Admission, target verification, output publication and persistence belong to
+ * the caller. Rows remain fixed within a draft/commit round. */
+export class QwenMtpRows {
+  #cache = new BatchedKVCache();
+  #hidden: MlxArray | null = null;
+  #prefilled: boolean[] = [];
+  #drafts: number[][] = [];
+  #depth = 0;
+
+  constructor(
+    readonly target: QwenMtpTarget,
+    readonly module: MtpModule,
+    readonly sampling: MtpRowSampling | null,
+    states: readonly (MtpRowState | null)[],
+  ) {
+    this.append(states);
+  }
+
+  get rowCount(): number { return this.#cache.rowOffsets.length; }
+
+  /** Add committed state at a method boundary. Inputs remain caller-owned. */
+  append(states: readonly (MtpRowState | null)[]): void {
+    if (!states.length) return;
+    applyStateChanges([() => this.prepareAppend(states)]);
+  }
+
+  prepareAppend(states: readonly (MtpRowState | null)[]): PreparedStateChange {
+    if (!states.length) return { commit() {}, dispose() {} };
+    const prefilled = [...this.#prefilled, ...states.map(state => state !== null)];
+    const example = this.#hidden ?? states.find(state => state !== null)?.hidden;
+    let hidden: MlxArray | null = null;
+    let cache: BatchedKVCache | null = new BatchedKVCache();
+    const held: { dispose(): void }[] = [];
+    try {
+      if (example) {
+        const zero = (rows: number) => {
+          const array = ops.zeros([rows, 1, this.target.hiddenSize], example.dtype);
+          held.push(array); return array;
+        };
+        hidden = ops.concatAxis([
+          ...(this.rowCount ? [this.#hidden ?? zero(this.rowCount)] : []),
+          ...states.map(state => state?.hidden ?? zero(1)),
+        ], 0);
+      }
+      cache.mergeRows([...(this.rowCount ? [this.#cache] : []), ...states.map(state => {
+        if (state) return state.cache;
+        const empty = new KVCache(); held.push(empty); return empty;
+      })]);
+    } catch (error) { return cleanupFailure(error, () => disposeResources([cache!, ...(hidden ? [hidden] : [])])); }
+    finally { disposeResources(held); }
+    return {
+      commit: () => {
+        const previousCache = this.#cache, previousHidden = this.#hidden;
+        this.#cache = cache!; this.#hidden = hidden;
+        this.#prefilled = prefilled;
+        cache = previousCache; hidden = previousHidden;
+      },
+      dispose: () => {
+        const held = [cache, hidden].filter(value => value !== null);
+        cache = null; hidden = null;
+        disposeResources(held);
+      },
+    };
+  }
+
+  /** Retire/reorder committed requests without involving their sampler policy. */
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    using indices = ops.fromInt32([...keep], [keep.length]);
+    const hidden = this.#hidden ? ops.takeAxis(this.#hidden, indices, 0) : null;
+    try { this.#cache.filterRows(keep); }
+    catch (error) { hidden?.dispose(); throw error; }
+    this.#hidden?.dispose(); this.#hidden = hidden;
+    this.#prefilled = keep.map(row => this.#prefilled[row]!);
+  }
+
+  /** Build companion KV from borrowed new target tokens/taps at the same B.
+   * The first token bridges to each restored row's preceding true hidden.
+   * Cold rows discard that bridge. All later rows pair token[p+1] with h[p].
+   * Only KV and the final true hidden are evaluated; unused attention/MLP
+   * outputs remain outside the evaluated prefill graph. */
+  prefill(tokens: MlxArray, context: MlxArray): void {
+    const [B, length] = tokens.shape as [number, number];
+    const H = context.shape[2]!;
+    if (this.#prefilled.some(Boolean)) {
+      using first = tokens.slice([0, 0], [B, 1]);
+      using embeds = this.target.embed(first);
+      const mixed = this.#prefilled.some(value => !value);
+      if (mixed) this.#cache.specRoundBegin();
+      this.module.forward(embeds, this.#hidden!, this.#cache).dispose();
+      if (mixed) this.#cache.specRoundRollback(this.#prefilled.map(Number));
+    }
+    for (let pos = 0; pos + 1 < length; pos += 2048) {
+      const end = Math.min(pos + 2048, length - 1);
+      using ids = tokens.slice([0, pos + 1], [B, end + 1]);
+      using embeds = this.target.embed(ids);
+      using hidden = context.slice([0, pos, 0], [B, end, H]);
+      this.module.forward(embeds, hidden, this.#cache).dispose();
+      ops.evalAll(this.#cache.state());
+    }
+    using tail = context.slice([0, length - 1, 0], [B, length, H]);
+    const hidden = materializeCopy(tail);
+    this.#hidden?.dispose(); this.#hidden = hidden;
+    ops.evalAll([hidden, ...this.#cache.state()]);
+    this.#prefilled.fill(true);
+  }
+
+  materialize(): void {
+    const state = [...this.#cache.state(), ...(this.#hidden ? [this.#hidden] : [])];
+    if (state.length) ops.evalAll(state);
+  }
+
+  /** Nonnegative depth, with every request's first pending token already known.
+   * Build the complete device dependency chain before reading proposal IDs. */
+  draft(pending: readonly number[], depth: number, steps: readonly number[]): number[][] {
+    const B = pending.length;
+    const tokens: MlxArray[] = [];
+    let ids: MlxArray | null = ops.fromInt32([...pending], [B, 1]);
+    let chained: MlxArray | null = null;
+    this.#cache.specRoundBegin();
+    try {
+      // A zero-proposal round still processes the pending target token. Keep
+      // its companion KV aligned so a later round can resume speculation.
+      if (depth === 0) {
+        using embeds = this.target.embed(ids!);
+        this.module.forward(embeds, this.#hidden!, this.#cache).dispose();
+        this.#drafts = Array.from({ length: B }, () => []);
+        this.#depth = 0;
+        this.#hidden!.dispose(); this.#hidden = null;
+        return this.#drafts;
+      }
+      for (let position = 0; position < depth; position++) {
+        using embeds = this.target.embed(ids!);
+        const output = this.module.forward(embeds, chained ?? this.#hidden!, this.#cache);
+        chained?.dispose(); chained = output;
+        using logits = this.target.logitsFromHidden(output);
+        using flat = ops.reshape(logits, [B, logits.shape.at(-1)!]);
+        using logprobs = toLogprobs(flat);
+        const token = this.sampling!.sample(logprobs, steps.map(step => step + position));
+        tokens.push(token);
+        ids!.dispose(); ids = ops.reshape(token, [B, 1]);
+      }
+      using packed = ops.concatAxis(tokens, 0);
+      const read = packed.toIntTokens();
+      this.#drafts = Array.from({ length: B }, (_, row) =>
+        Array.from({ length: depth }, (_, position) => read[position * B + row]!));
+      this.#depth = depth;
+      this.#hidden!.dispose(); this.#hidden = null;
+      return this.#drafts;
+    } finally {
+      ids?.dispose(); chained?.dispose();
+      for (const token of tokens) token.dispose();
+    }
+  }
+
+  /** Resolve draft storage using each request's accepted prefix. The verify
+   * context is borrowed [B,depth+1,H], before the target's final norm. */
+  commit(accepted: readonly number[], context: MlxArray): void {
+    const B = accepted.length, depth = this.#depth;
+    const appended = Math.max(1, depth);
+    const kept = accepted.map(count => Math.min(count + 1, appended));
+    if (kept.every(count => count === appended)) this.#cache.specRoundCommit();
+    else this.#cache.specRoundRollback(kept);
+
+    // Full acceptance needs the last accepted token's missing draft row,
+    // using the TRUE target hidden at its preceding position. Other requests
+    // keep their current coverage; their temporary append is never retained.
+    const extra = accepted.map(count => Number(depth > 0 && count === depth));
+    if (extra.some(Boolean)) {
+      using ids = ops.fromInt32(this.#drafts.map(row => row[depth - 1]!), [B, 1]);
+      using embeds = this.target.embed(ids);
+      using hidden = context.slice([0, depth - 1, 0], [B, depth, context.shape[2]!]);
+      this.#cache.specRoundBegin();
+      this.module.forward(embeds, hidden, this.#cache).dispose();
+      this.#cache.specRoundRollback(extra);
+    }
+    using positions = ops.fromInt32([...accepted], [B, 1, 1]);
+    this.#hidden = ops.takeAlongAxis(context, positions, 1);
+    this.#drafts = []; this.#depth = 0;
+  }
+
+  /** Only between committed rounds. Returned state is independent of the
+   * group's storage and uses the existing single-request cache representation. */
+  extractRow(row: number): MtpRowState {
+    const cache = this.#cache.extractRow(row);
+    try {
+      using view = this.#hidden!.slice([row, 0, 0], [row + 1, 1, this.#hidden!.shape[2]!]);
+      return { cache, hidden: materializeCopy(view) };
+    } catch (error) { cache.dispose(); throw error; }
+  }
+
+  dispose(): void {
+    this.#cache.dispose(); this.#hidden?.dispose(); this.#hidden = null;
+    this.#drafts = []; this.#depth = 0; this.#prefilled = [];
+  }
+}

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

@@ -30,206 +30,40 @@ import { materializeCopy } from "../mlx/materialize";
 import { toLogprobs } from "../sampler";
 import { loadModelConfig, type ModelConfig } from "../config";
 import { Weights } from "../weights";
-import { disposing, KVCache, QuantizedLinear, RMSNorm, type Mask } from "../model/gemma4-base";
-import { DenseLinear } from "../model/universal/modules";
-import type { DraftProvider, DraftSource, QwenMtpTarget } from "./source";
+import { KVCache } from "../model/gemma4-base";
+import { MtpModule } from "./qwen-mtp-module";
+import type { DraftProvider, DraftSource, DraftRowCheckpoint, DraftRowGroup, DraftPrefillGroup, DraftRowSampling, GroupedDraftProvider, QwenMtpTarget, TargetView } from "./source";
+import { QwenMtpRows, type MtpRowState } from "./qwen-mtp-rows";
+import { captureQwenMtpState, restoreQwenMtpState } from "./qwen-mtp-state";
 import type { Cache } from "../model/gemma4";
-import { cloneKvCaches } from "../kv-store";
-import { cacheBytes } from "../prompt-cache";
-import { disposeResources } from "../engine/resources";
-import { SpeculativePrefixStore, type SpeculativePrefixState } from "./prefix-state";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../engine/resources";
+import { artifactIdentity } from "../model/artifact-identity";
+import { configFingerprint } from "../model/fingerprint";
+import type { PreparedStateChange } from "../contracts/resources";
 
 type Sampler = (logprobs: MlxArray, step: number) => MlxArray;
 
 const DRAFT_PREFILL_CHUNK = 2048;
 
-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: KVCache): 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]));
-
-    q = disposing(q, ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
-    k = disposing(k, 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. */
-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 [1,S,·]: token embeddings ([1,S,H], target
-   *  embed_tokens output) paired with hiddens ([1,S,H], target pre-final-norm
-   *  or the module's own chained output). Appends S rows to `cache`; returns
-   *  the module output [1,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: KVCache): 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 silu = ops.silu(g);
-    g.dispose();
-    const act = ops.mul(silu, u);
-    silu.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));
-  }
-}
-
-class QwenMtpPrefixState implements SpeculativePrefixState {
-  constructor(readonly tokens: number[], readonly targetIdentity: object,
-    readonly namespace: string, readonly bytes: number, readonly target: Cache[],
-    public draft: KVCache | null, public hidden: MlxArray | null) {}
-
-  dispose(): void {
-    const resources = [...this.target.splice(0), this.draft, this.hidden];
-    this.draft = null;
-    this.hidden = null;
-    disposeResources(resources.filter((r) => r != null));
-  }
-}
-
 export class QwenMtpProvider 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: string;
   readonly weightsBytes: number;
   readonly #module: MtpModule;
   readonly #config: ModelConfig;
   readonly #resources: DisposableStack;
-  readonly #prefixStore = new SpeculativePrefixStore<QwenMtpPrefixState>();
+  readonly #checkpointNamespace: string;
 
-  private constructor(id: string, config: ModelConfig, weightsBytes: number, module: MtpModule, resources: DisposableStack) {
+  private constructor(id: string, config: ModelConfig, weightsBytes: number, module: MtpModule, resources: DisposableStack, checkpointNamespace: string) {
     this.id = id;
     this.#config = config;
     this.#module = module;
     this.weightsBytes = weightsBytes;
     this.#resources = resources;
+    this.#checkpointNamespace = checkpointNamespace;
   }
 
   static async load(dir: string): Promise<QwenMtpProvider> {
@@ -242,16 +76,24 @@ export class QwenMtpProvider implements DraftProvider {
     const weightsBytes = [...weights.shards.files.values()]
       .reduce((a, f) => a + f.mmap.size, 0);
     const module = new MtpModule(weights, config, resources);
+    // Stable across restarts and distinct for differently folded/quantized
+    // companions. Hash once at provider load, outside inference execution.
+    const identity = await artifactIdentity(configFingerprint(config),
+      [...weights.shards.files].map(([name, shard]) => ({ name, path: shard.path })));
     return new QwenMtpProvider(
       dir.split("/").filter(Boolean).at(-1) ?? "qwen-mtp",
-      config, weightsBytes, module, resources.move(),
+      config, weightsBytes, module, resources.move(), `qwen-mtp-v1:${identity}`,
     );
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
+    return new QwenMtpSource(this.#target(opts.target), this.#module, opts.sampler, this.#checkpointNamespace);
+  }
+
+  #target(view: TargetView): QwenMtpTarget {
     if (this.#resources.disposed)
       throw new Error("qwen MTP provider is disposed");
-    const target = opts.target.qwenMtp;
+    const target = view.qwenMtp;
     if (!target)
       throw new Error("qwen MTP drafting requires a qwen3_5-family target");
     if (target.hiddenSize !== this.#config.text.hiddenSize) {
@@ -260,14 +102,50 @@ export class QwenMtpProvider implements DraftProvider {
         `${target.hiddenSize} — split from a different checkpoint?`,
       );
     }
-    return new QwenMtpSource(target, this.#module, opts.sampler,
-      { store: this.#prefixStore, identity: opts.target.identity });
+    return target;
+  }
+
+  #openRows(view: TargetView, sampling: DraftRowSampling | null,
+    checkpoints: readonly (DraftRowCheckpoint | null)[]): DraftRowGroup & DraftPrefillGroup {
+    const target = this.#target(view);
+    const rows = new QwenMtpRows(target, this.#module, sampling, []);
+    const prepareAppend = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => {
+      const states: Array<MtpRowState | null> = [];
+      let change: PreparedStateChange | undefined;
+      try {
+        for (const checkpoint of checkpoints) states.push(checkpoint ? restoreQwenMtpState(checkpoint) : null);
+        change = rows.prepareAppend(states);
+        disposeResources(states.splice(0).flatMap(state => state ? [state.cache, state.hidden] : []));
+        return change;
+      } catch (error) {
+        return cleanupFailure(error, () => disposeResources([
+          ...states.flatMap(state => state ? [state.cache, state.hidden] : []), ...(change ? [change] : []),
+        ]));
+      }
+    };
+    const append = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => applyStateChanges([() => prepareAppend(checkpoints)]);
+    try { append(checkpoints); }
+    catch (error) { return cleanupFailure(error, () => rows.dispose()); }
+    return {
+      namespace: this.#checkpointNamespace, prefillMode: "full", tapLayers: [target.layerCount - 1],
+      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 captureQwenMtpState(state); }
+        finally { disposeResources([state.cache, state.hidden]); }
+      },
+      dispose: rows.dispose.bind(rows),
+    };
   }
 
+
   dispose(): void {
     // Release cached transpose views before their native weight maps. MLX
     // retains buffers needed by outstanding GPU commands until completion.
-    disposeResources([this.#prefixStore, { dispose: () => this.#resources.dispose() }]);
+    this.#resources.dispose();
   }
 }
 
@@ -284,7 +162,7 @@ export class QwenMtpSource implements DraftSource {
   readonly #module: MtpModule;
   readonly #sampler: Sampler;
   #cache = new KVCache();
-  readonly prefix: DraftSource["prefix"];
+  readonly checkpoint: NonNullable<DraftSource["checkpoint"]>;
   #prefilledTokens = 0;
   #hasDrafted = false;
   /** Target pre-norm hidden at the position preceding the next pending
@@ -295,52 +173,29 @@ export class QwenMtpSource implements DraftSource {
   #closed = false;
 
   constructor(target: QwenMtpTarget, module: MtpModule, sampler: Sampler,
-    prefix?: { store: SpeculativePrefixStore<QwenMtpPrefixState>; identity: object }) {
+    namespace = "qwen-mtp-v1") {
     this.#target = target;
     this.#module = module;
     this.#sampler = sampler;
     this.tapLayers = [target.layerCount - 1];
-    if (prefix) this.prefix = {
-      restore: (prompt, caches, namespace, maxBytes) => {
+    this.checkpoint = {
+      namespace,
+      restore: (tokens, attachment) => {
         this.#checkOpen();
         if (this.#prefilledTokens !== 0 || this.#hasDrafted)
-          throw new Error("Qwen MTP prefix restore requires a fresh source");
-        const entry = prefix.store.take(prompt, prefix.identity, namespace, maxBytes);
-        if (!entry) return 0;
-        try {
-          if (entry.target.length !== caches.length || !entry.draft || !entry.hidden ||
-            entry.draft.offset !== entry.tokens.length - 1)
-            throw new Error("invalid paired Qwen MTP prefix state");
-          disposeResources(caches);
-          caches.splice(0, caches.length, ...entry.target.splice(0));
-          this.#cache.dispose();
-          this.#cache = entry.draft;
-          entry.draft = null;
-          this.#pendingTrueHidden?.dispose();
-          this.#pendingTrueHidden = entry.hidden;
-          entry.hidden = null;
-          return this.#prefilledTokens = entry.tokens.length;
-        } finally { entry.dispose(); }
+          throw new Error("Qwen MTP checkpoint restore requires a fresh source");
+        const state = restoreQwenMtpState({ processedTokens: tokens, attachment });
+        try { disposeResources([this.#cache, ...(this.#pendingTrueHidden ? [this.#pendingTrueHidden] : [])]); }
+        catch (error) { return cleanupFailure(error, () => disposeResources([state.cache, state.hidden])); }
+        this.#cache = state.cache;
+        this.#pendingTrueHidden = state.hidden;
+        this.#prefilledTokens = tokens;
       },
-      capture: (tokens, caches, namespace, maxBytes) => {
+      capture: (tokens) => {
         this.#checkOpen();
-        if (tokens.length !== this.#prefilledTokens || !this.#pendingTrueHidden ||
-          this.#cache.offset !== tokens.length - 1 || this.#roundAppended !== 0)
-          throw new Error("Qwen MTP snapshot requires an aligned prefill boundary");
-        const bytes = cacheBytes([...caches, this.#cache]) + this.#pendingTrueHidden.nbytes;
-        if (bytes > maxBytes) { prefix.store.dispose(); return; }
-        const retained: Cache[] = [];
-        let draft: KVCache | null = null, hidden: MlxArray | null = null;
-        try {
-          retained.push(...cloneKvCaches(caches));
-          draft = cloneKvCaches([this.#cache])[0] as KVCache;
-          hidden = this.#pendingTrueHidden.slice([0, 0, 0], this.#pendingTrueHidden.shape);
-          const entry = new QwenMtpPrefixState([...tokens], prefix.identity, namespace,
-            bytes, retained.splice(0), draft, hidden);
-          draft = null;
-          hidden = null;
-          prefix.store.put(entry, maxBytes);
-        } finally { disposeResources([...retained, draft, hidden].filter((r) => r != null)); }
+        if (!this.#pendingTrueHidden || this.#cache.offset !== tokens - 1 || this.#roundAppended !== 0)
+          throw new Error("Qwen MTP snapshot requires an aligned committed boundary");
+        return captureQwenMtpState({ cache: this.#cache, hidden: this.#pendingTrueHidden }).attachment;
       },
     };
   }

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

@@ -0,0 +1,30 @@
+import type { MlxArray } from "../mlx/array";
+import { KVCache } from "../model/gemma4-base";
+import { cloneAttachments } from "../backends/mlx/checkpoint-state";
+import { disposeResources } from "../engine/resources";
+import type { MtpRowState } from "./qwen-mtp-rows";
+import type { DraftRowCheckpoint } from "./source";
+
+/** The existing Qwen companion attachment is also the state interchange
+ * between prepared requests and a group. Storage-tier policy stays outside. */
+export function restoreQwenMtpState(checkpoint: DraftRowCheckpoint): MtpRowState {
+  const { processedTokens: tokens, attachment } = checkpoint;
+  if (attachment.schema !== "qwen-mtp-v1" || attachment.metadata.draftOffset !== tokens - 1)
+    throw new Error("invalid paired Qwen MTP checkpoint alignment");
+  const held = cloneAttachments([attachment])[0]!.tensors;
+  let cache: KVCache | null = new KVCache();
+  try {
+    if (tokens > 1) { cache.restoreState(held[0]!, held[1]!, tokens - 1); held.splice(0, 2); }
+    const state = { cache, hidden: held.pop()! }; cache = null; return state;
+  } finally { disposeResources([...held, ...(cache ? [cache] : [])]); }
+}
+
+export function captureQwenMtpState(state: MtpRowState): DraftRowCheckpoint {
+  const held: MlxArray[] = [];
+  try {
+    if (state.cache.offset > 0) held.push(...state.cache.temporalView());
+    held.push(state.hidden.slice([0, 0, 0], [...state.hidden.shape]));
+    return { processedTokens: state.cache.offset + 1,
+      attachment: { schema: "qwen-mtp-v1", metadata: { draftOffset: state.cache.offset }, tensors: held.splice(0) } };
+  } finally { disposeResources(held); }
+}