mlx-bun 0.2.0 → 0.4.0

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

@@ -21,187 +21,49 @@
 // (held over from the verify tap). Computationally identical to mlx-vlm's
 // commit-time "seed" append, just performed at the top of the next round.
 //
-// Weights are plain bf16 (DenseLinear); the published drafter snapshots are
-// already sanitized (norms in runtime layout — no +1.0 shift here).
+// Companion projections may be dense or quantized according to their metadata.
+// Published drafter norms are already sanitized to runtime layout; no +1.0 shift.
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
+import { materializeCopy } from "../mlx/materialize";
 import { toLogprobs } from "../sampler";
 import { loadModelConfig, type ModelConfig } from "../config";
 import { Weights } from "../weights";
-import { disposing, KVCache, RMSNorm, type Mask } from "../model/gemma4-base";
-import { DenseLinear } from "../model/universal/modules";
-import { Qwen35Model } from "../model/qwen3_5";
-import type { DraftProvider, DraftSource, TargetView } 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 { 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;
 
-/** Dense-weight clone of Qwen3Attention.forward (src/model/qwen3_5.ts —
- *  ops copied verbatim; only the projection flavor differs: the drafter is
- *  bf16, the target's class hardwires QuantizedLinear). */
-class MtpAttention {
-  readonly qProj: DenseLinear;
-  readonly kProj: DenseLinear;
-  readonly vProj: DenseLinear;
-  readonly oProj: DenseLinear;
-  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) {
-    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;
-    const dense = (path: string): DenseLinear =>
-      new DenseLinear(weights.tensor(`${path}.weight`), null);
-    this.qProj = dense(`${prefix}.q_proj`);
-    this.kProj = dense(`${prefix}.k_proj`);
-    this.vProj = dense(`${prefix}.v_proj`);
-    this.oProj = dense(`${prefix}.o_proj`);
-    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: DenseLinear;
-  readonly preFcNormEmbedding: RMSNorm;
-  readonly preFcNormHidden: RMSNorm;
-  readonly attn: MtpAttention;
-  readonly mlpGate: DenseLinear;
-  readonly mlpUp: DenseLinear;
-  readonly mlpDown: DenseLinear;
-  readonly inputNorm: RMSNorm;
-  readonly postAttnNorm: RMSNorm;
-  readonly finalNorm: RMSNorm;
-
-  constructor(weights: Weights, config: ModelConfig) {
-    const eps = config.text.rmsNormEps;
-    const dense = (path: string): DenseLinear =>
-      new DenseLinear(weights.tensor(`${path}.weight`), null);
-    this.fc = dense("fc");
-    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");
-    this.mlpGate = dense("layers.0.mlp.gate_proj");
-    this.mlpUp = dense("layers.0.mlp.up_proj");
-    this.mlpDown = dense("layers.0.mlp.down_proj");
-    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, dense weights).
-    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));
-  }
-}
-
 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 #checkpointNamespace: string;
 
-  private constructor(id: string, config: ModelConfig, weights: Weights, module: MtpModule) {
+  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 = [...weights.shards.files.values()]
-      .reduce((a, f) => a + f.mmap.size, 0);
+    this.weightsBytes = weightsBytes;
+    this.#resources = resources;
+    this.#checkpointNamespace = checkpointNamespace;
   }
 
   static async load(dir: string): Promise<QwenMtpProvider> {
@@ -209,28 +71,81 @@ export class QwenMtpProvider implements DraftProvider {
     if (config.modelType !== "qwen3_5_mtp")
       throw new Error(`${dir}: not a qwen3_5_mtp drafter (model_type ${config.modelType})`);
     const weights = await Weights.open(dir);
+    using resources = new DisposableStack();
+    resources.defer(() => weights.dispose());
+    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, weights, new MtpModule(weights, config),
+      config, weightsBytes, module, resources.move(), `qwen-mtp-v1:${identity}`,
     );
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
-    const target = opts.target.model;
-    if (!(target instanceof Qwen35Model))
+    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 = view.qwenMtp;
+    if (!target)
       throw new Error("qwen MTP drafting requires a qwen3_5-family target");
-    if (target.config.text.hiddenSize !== this.#config.text.hiddenSize) {
+    if (target.hiddenSize !== this.#config.text.hiddenSize) {
       throw new Error(
         `qwen MTP drafter hidden ${this.#config.text.hiddenSize} != target ` +
-        `${target.config.text.hiddenSize} — split from a different checkpoint?`,
+        `${target.hiddenSize} — split from a different checkpoint?`,
       );
     }
-    return new QwenMtpSource(target, this.#module, opts.sampler);
+    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 {
-    // Weights are mmap-backed and pinned for the process (house rule: no JS
-    // dtors into mlx); the provider owns no other native state.
+    // Release cached transpose views before their native weight maps. MLX
+    // retains buffers needed by outstanding GPU commands until completion.
+    this.#resources.dispose();
   }
 }
 
@@ -243,10 +158,13 @@ export class QwenMtpSource implements DraftSource {
    *  The seam's anchorHidden is post-final-norm and is deliberately unused. */
   readonly tapLayers: number[];
 
-  readonly #target: Qwen35Model;
+  readonly #target: QwenMtpTarget;
   readonly #module: MtpModule;
   readonly #sampler: Sampler;
-  readonly #cache = new KVCache();
+  #cache = new KVCache();
+  readonly checkpoint: NonNullable<DraftSource["checkpoint"]>;
+  #prefilledTokens = 0;
+  #hasDrafted = false;
   /** Target pre-norm hidden at the position preceding the next pending
    *  token: prefill's last tapped row, then each commit's vCtx row at the
    *  emitted position. draft() consumes it to build the pending row. [1,1,H] */
@@ -254,11 +172,32 @@ export class QwenMtpSource implements DraftSource {
   #roundAppended = 0;
   #closed = false;
 
-  constructor(target: Qwen35Model, module: MtpModule, sampler: Sampler) {
+  constructor(target: QwenMtpTarget, module: MtpModule, sampler: Sampler,
+    namespace = "qwen-mtp-v1") {
     this.#target = target;
     this.#module = module;
     this.#sampler = sampler;
-    this.tapLayers = [target.config.text.numHiddenLayers - 1];
+    this.tapLayers = [target.layerCount - 1];
+    this.checkpoint = {
+      namespace,
+      restore: (tokens, attachment) => {
+        this.#checkOpen();
+        if (this.#prefilledTokens !== 0 || this.#hasDrafted)
+          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) => {
+        this.#checkOpen();
+        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;
+      },
+    };
   }
 
   /** Drafter prefill: rows for positions 0..L-2, keyed (token_{p+1}, h_p) —
@@ -269,27 +208,43 @@ export class QwenMtpSource implements DraftSource {
     if (!ctxML)
       throw new Error("qwen MTP prefill requires the tapped pre-final-norm context");
     try {
-      if (this.#cache.offset !== 0)
-        throw new Error("qwen MTP source cannot be prefilled twice");
+      if (this.#hasDrafted)
+        throw new Error("qwen MTP prefill cannot follow drafting");
       const L = promptIds.length;
+      const start = this.#prefilledTokens;
       const H = ctxML.shape[2]!;
-      if (ctxML.shape[1]! !== L)
-        throw new Error(`qwen MTP tap covered ${ctxML.shape[1]} of ${L} prompt positions`);
-      this.#pendingTrueHidden = ctxML.slice([0, L - 1, 0], [1, L, H]);
-      for (let pos = 0; pos + 1 < L; pos += DRAFT_PREFILL_CHUNK) {
+      if (L <= start || ctxML.shape[1]! !== L - start)
+        throw new Error(`qwen MTP tap covered ${ctxML.shape[1]} of ${L - start} new prompt positions`);
+      if (start > 0) {
+        if (!this.#pendingTrueHidden || this.#cache.offset !== start - 1)
+          throw new Error("qwen MTP prefix state is not aligned");
+        this.#stepOne(promptIds[start]!, this.#pendingTrueHidden).dispose();
+      }
+      this.#pendingTrueHidden?.dispose();
+      // A slice would retain the complete prompt buffer in the saved prefix.
+      // Materialize this one row before handing off the prefill state.
+      using tail = ctxML.slice([0, L - start - 1, 0], [1, L - start, H]);
+      this.#pendingTrueHidden = materializeCopy(tail);
+      for (let pos = start; pos + 1 < L; pos += DRAFT_PREFILL_CHUNK) {
         const n = Math.min(DRAFT_PREFILL_CHUNK, L - 1 - pos);
         const shifted = promptIds.slice(pos + 1, pos + 1 + n);
         const ids = ops.fromInt32(shifted, [1, n]);
-        const embeds = this.#target.embed.encode(ids);
+        const embeds = this.#target.embed(ids);
         ids.dispose();
-        const hiddens = ctxML.slice([0, pos, 0], [1, pos + n, H]);
+        const hiddens = ctxML.slice([0, pos - start, 0], [1, pos + n - start, H]);
         const out = this.#module.forward(embeds, hiddens, this.#cache);
         embeds.dispose();
         hiddens.dispose();
         out.dispose(); // prefill outputs are not seeds; only KV matters here
+        // Bound the lazy KV graph at each chunk. Evaluating the unused output
+        // would also execute attention and the MLP, which prefill does not need.
+        ops.evalAll(this.#cache.state());
       }
       if (this.#cache.offset !== L - 1)
         throw new Error(`qwen MTP prefill offset ${this.#cache.offset}, expected ${L - 1}`);
+      // Also covers a restored-prefix bridge with no subsequent full chunk.
+      ops.evalAll([this.#pendingTrueHidden, ...this.#cache.state()]);
+      this.#prefilledTokens = L;
     } finally {
       ctxML?.dispose(); // ownership per the seam contract
     }
@@ -297,6 +252,7 @@ export class QwenMtpSource implements DraftSource {
 
   async draft(feed: number[], n: number, stepBase: number): Promise<number[]> {
     this.#checkOpen();
+    this.#hasDrafted = true;
     if (this.#roundAppended !== 0)
       throw new Error("qwen MTP draft called before the prior round committed");
     if (n <= 0) return [];
@@ -316,8 +272,8 @@ export class QwenMtpSource implements DraftSource {
       this.#pendingTrueHidden.dispose();
       this.#pendingTrueHidden = null;
       this.#roundAppended++;
-      drafts.push(this.#sample(out, stepBase));
       chained = out;
+      drafts.push(this.#sample(out, stepBase));
       while (drafts.length < n) {
         const out = this.#stepOne(drafts.at(-1)!, chained!);
         chained!.dispose();
@@ -369,7 +325,7 @@ export class QwenMtpSource implements DraftSource {
       const H = vCtxML.shape[2]!;
       if (kAccept === d && d > 0) {
         const ids = ops.fromInt32([acceptedTokens[kAccept - 1]!], [1, 1]);
-        const embeds = this.#target.embed.encode(ids);
+        const embeds = this.#target.embed(ids);
         ids.dispose();
         const hidden = vCtxML.slice([0, d - 1, 0], [1, d, H]);
         const out = this.#module.forward(embeds, hidden, this.#cache);
@@ -396,33 +352,29 @@ export class QwenMtpSource implements DraftSource {
 
   /** One module forward for (token, hidden) — appends one KV row. */
   #stepOne(token: number, hidden: MlxArray): MlxArray {
-    const ids = ops.fromInt32([token], [1, 1]);
-    const embed = this.#target.embed.encode(ids);
+    using ids = ops.fromInt32([token], [1, 1]);
+    using embed = this.#target.embed(ids);
     ids.dispose();
-    const out = this.#module.forward(embed, hidden, this.#cache);
-    embed.dispose();
-    return out;
+    return this.#module.forward(embed, hidden, this.#cache);
   }
 
   /** Sample a draft token from the module output via the TARGET's lm head
    *  and the request sampler (per-step RNG stream discipline). */
   #sample(moduleOut: MlxArray, step: number): number {
-    const logits = this.#target.logitsFromHidden(moduleOut);
+    using logits = this.#target.logitsFromHidden(moduleOut);
     // Sampler contract is [1, V] (the main decode loop's shape). moduleOut
     // is [1, 1, H] → logits [1, 1, V]; without this reshape any sampler
     // that slices 2-D (top-k) throws "[slice] Invalid number of indices…
     // dimension 3" — the serve-lane MTP 500 (chat defaults carry the
     // model's top_k=20; the greedy bench harness never hit it).
     const V = logits.shape[logits.shape.length - 1]!;
-    const flat = ops.reshape(logits, [1, V]);
+    using flat = ops.reshape(logits, [1, V]);
     logits.dispose();
-    const logprobs = toLogprobs(flat);
+    using logprobs = toLogprobs(flat);
     flat.dispose();
-    const tok = this.#sampler(logprobs, step);
+    using tok = this.#sampler(logprobs, step);
     logprobs.dispose();
-    const id = ops.itemUint32(tok);
-    tok.dispose();
-    return id;
+    return ops.itemUint32(tok);
   }
 
   #checkOpen(): void {

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); }
+}