mlx-bun 0.2.0 → 0.3.0

package/src/spec/serve-loop.ts

@@ -1,6 +1,6 @@
 // Serve-time speculative decoding — the ONE verify/accept executor behind
-// `serve --draft-model` (docs/design/mlx-lm-tool-parity-plan.md §7; sequenced
-// by docs/design/grammar-spec-batching-integration.md Phase B).
+// `serve --draft-model` (docs/reference/cli.md §7; sequenced
+// by docs/design/batching.md Phase B).
 //
 // Faithful to mlx-lm's speculative_generate_step (generate.py:473-654, read
 // from the oracle venv): per round, draft n tokens from the DraftSource, run
@@ -31,7 +31,7 @@
 // so rejected drafts never touch grammar state and no matcher rollback is
 // needed. Grammar termination mid-burst truncates the round (the all--inf
 // guarantee). Gates (no oracle): greedy grammar+spec ≡ greedy grammar-only
-// long-prefix + 100% schema validity (tests/spec-serve.test.ts).
+// long-prefix + 100% schema validity (tests/parity/spec-serve.test.ts).
 //
 // Emission discipline: tokens flow through the caller's onToken ONE AT A
 // TIME, in order (bursts of ≤ n+1 per round) — the server's stop-sequence
@@ -42,68 +42,19 @@
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { clearCache } from "../mlx/ffi";
-import { flagOn } from "../flags";
+import { cleanupFailure, disposeResources } from "../engine/resources";
+import type { DisposableResource } from "../contracts/resources";
+import { assertMlxSpeculativeBinding, bindLegacySpeculativeModel, type MlxSpeculativeBinding } from "../backends/mlx/speculative";
 import type { RuntimeModel } from "../model/factory";
 import type { Cache } from "../model/gemma4";
-import type { GenerateOptions, GenerateStats } from "../generate";
+import { maybeQuantizeKv, withModelUsageFlush, withModelWiredLimit, type GenerateOptions, type GenerateStats } from "../generate";
+import { runtimeConfig, withRuntimeConfig } from "../runtime-config";
 import { makeSampler, makeStepSampler } from "../sampler";
 import type { OnToken } from "../serve/generation-gateway";
 import type { DraftProvider } from "./source";
 
 const PREFILL_CHUNK = 2048;
 
-/** Run one target forward, and when `tapLayers` is set (⟹ a Gemma4 target — a
- *  DSpark-style source), also return the captured multi-layer context
- *  [1,L,m*H] (tapLayers concatenated on the feature axis; index nLayers is the
- *  post-finalNorm sentinel). Non-tapping sources get ctxML=null and never
- *  touch model.hiddenTap. Mirrors generate-dflash.ts forwardTapped. */
-async function forwardMaybeTap(
-  model: RuntimeModel,
-  ids: MlxArray,
-  caches: Cache[],
-  tapLayers: number[] | undefined,
-): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
-  if (!tapLayers) {
-    const asyncModel = model as RuntimeModel & {
-      forwardHiddenAsync?: (
-        ids: MlxArray,
-        caches: Cache[],
-      ) => Promise<MlxArray>;
-    };
-    const hidden = typeof asyncModel.forwardHiddenAsync === "function"
-      ? await asyncModel.forwardHiddenAsync(ids, caches)
-      : model.forwardHidden(ids, caches);
-    return { hidden, ctxML: null };
-  }
-  const m = model as unknown as {
-    hiddenTap: { layers: Set<number>; captured: Map<number, MlxArray> } | null;
-  };
-  const cap = new Map<number, MlxArray>();
-  m.hiddenTap = { layers: new Set(tapLayers), captured: cap };
-  let hidden: MlxArray | null = null;
-  try {
-    hidden = model.forwardHidden(ids, caches);
-    const perLayer = tapLayers.map((li) => {
-      const a = cap.get(li);
-      if (!a) throw new Error(`spec tap: layer ${li} not captured`);
-      return a;
-    });
-    const ctxML = ops.concatAxis(perLayer, 2); // [1,L,m*H]
-    for (const [, a] of cap) a.dispose();
-    cap.clear(); // consumed — the finally must not double-dispose
-    const out = { hidden, ctxML };
-    hidden = null; // ownership returned to the caller
-    return out;
-  } finally {
-    // On any throw (forward mid-capture, missing layer, concat), free the
-    // partially-captured tap tensors and the orphaned hidden. On success both
-    // are already gone (cap cleared, hidden nulled).
-    hidden?.dispose();
-    for (const [, a] of cap) a.dispose();
-    m.hiddenTap = null;
-  }
-}
-
 export interface SpecServeExtras {
   drafted: number;
   accepted: number;
@@ -115,7 +66,7 @@ export interface SpecServeExtras {
   forwardsSaved: number;
   /** Per-draft-position counters (index = position within a round's block,
    *  0..γ-1): how many rounds drafted/accepted at that position. Drives the
-   *  Phase-1c per-position acceptance report (scripts/dspark-drafter-ab.ts)
+   *  Phase-1c per-position acceptance report (scripts/dspark.ts ab)
    *  — near-zero cost, always populated. */
   draftedByPos: number[];
   acceptedByPos: number[];
@@ -128,9 +79,36 @@ export async function specServeRun(
   promptIds: number[],
   options: GenerateOptions & { stopSequences?: string[] },
   onToken: OnToken,
+): Promise<GenerateStats> {
+  return specRun(bindLegacySpeculativeModel(model, provider), numDraftTokens, promptIds, options, onToken);
+}
+
+/** The verifier consumes a bound target and source factory, independent of
+ * concrete model classes. Compatibility callers use specServeRun above. */
+export async function specRun(
+  binding: MlxSpeculativeBinding,
+  numDraftTokens: number,
+  promptIds: number[],
+  options: GenerateOptions & { stopSequences?: string[] },
+  onToken: OnToken,
+): Promise<GenerateStats> {
+  return withRuntimeConfig(binding.runtime ?? runtimeConfig(), () => {
+    const run = () => specRunInner(binding, numDraftTokens, promptIds, options, onToken);
+    return binding.memory
+      ? withModelWiredLimit(binding.memory, () => withModelUsageFlush(binding.memory!, run))
+      : run();
+  });
+}
+
+async function specRunInner(
+  binding: MlxSpeculativeBinding,
+  numDraftTokens: number,
+  promptIds: number[],
+  options: GenerateOptions & { stopSequences?: string[] },
+  onToken: OnToken,
 ): Promise<GenerateStats> {
   const maxTokens = options.maxTokens ?? 512;
-  const eos = options.eosTokenIds ?? model.config.eosTokenIds;
+  const eos = options.eosTokenIds ?? [...binding.eosTokenIds];
   const sampler = makeSampler(options);
   const stepSampler = makeStepSampler(options, {
     tokenRepresentation: "number",
@@ -142,6 +120,8 @@ export async function specServeRun(
     sampler,
   });
   const gamma = Math.max(1, numDraftTokens);
+  const prefillChunk = options.prefillChunkSize ??
+    runtimeConfig().number("MLX_BUN_RD_PREFILL_CHUNK", PREFILL_CHUNK);
 
   // Allocated INSIDE the try below (2026-07-07 review): provider.open()
   // throws on a mismatched (target, drafter) pairing — as pre-try consts a
@@ -205,11 +185,48 @@ export async function specServeRun(
     return flat;
   };
 
+  const forwardTokens = async (tokens: number[], taps?: number[]) => {
+    const ids = ops.fromInt32(tokens, [1, tokens.length]);
+    try { return await binding.forward(ids, caches, taps); }
+    finally { ids.dispose(); }
+  };
+  let failure: { error: unknown } | undefined;
+
   try {
-    caches = model.makeCache();
-    const src = provider.open({ sampler, target: { model, caches } });
+    options.signal?.throwIfAborted();
+    if (!Number.isSafeInteger(prefillChunk) || prefillChunk < 1)
+      throw new Error("speculative prefill chunk must be a positive integer");
+    assertMlxSpeculativeBinding(binding);
+    if (options.turboQuant || options.kvConfig?.length ||
+        (options.kvBits && (options.kvBits !== 4 || options.quantizedKvStart !== 0)))
+      throw new Error("speculative KV requires the qualified uniform KV4 start=0 policy");
+    caches = binding.makeCache();
+    const src = binding.openDraft(sampler, caches);
     source = src;
     tapLayers = src.tapLayers;
+    const prefixBudget = Math.max(0, options.speculativeCacheBytes ?? 0);
+    const prefix = src.prefillMode === "full" ? src.prefix : undefined;
+    const prefixNamespace = JSON.stringify({ adapters: options.adapters ?? [],
+      kvBits: options.kvBits ?? null, kvGroupSize: options.kvGroupSize ?? 64,
+      quantizedKvStart: options.quantizedKvStart ?? null });
+    const cached = prefix?.restore(promptIds, caches, prefixNamespace, prefixBudget) ?? 0;
+    stats.cachedTokens = cached;
+    const snapshotAt = prefix && prefixBudget > 0
+      ? Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1)
+      : -1;
+    if (prefix && prefixBudget > 0 && cached > 0 && snapshotAt <= cached)
+      prefix.capture(promptIds.slice(0, cached), caches, prefixNamespace, prefixBudget);
+    const seedSource = async (end: number) => {
+      if (ctxParts.length === 1) prefillCtx = ctxParts[0]!;
+      else if (ctxParts.length > 1) {
+        prefillCtx = ops.concatAxis(ctxParts, 1);
+        for (const part of ctxParts) part.dispose();
+      }
+      ctxParts.length = 0;
+      const context = prefillCtx;
+      prefillCtx = null; // source.prefill owns context on both success and failure
+      await src.prefill(promptIds.slice(0, end), context ?? undefined);
+    };
     // ---- prefill (target, chunked; optionally tapped for DSpark), then seed
     // the source. Order: the source's prefill needs the tapped context that the
     // target prefill produces, so target-first (two-model's own draft prefill
@@ -240,77 +257,71 @@ export async function specServeRun(
     const t0 = performance.now();
     const tailSplit =
       src.prefillMode !== "full" &&
-      flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+      binding.prefillTailSplit;
     const oracleShape = tailSplit && promptIds.length > 1;
     if (oracleShape) {
       let pos = 0;
       const end = promptIds.length - 1; // last prompt token NEVER prefilled
       while (pos < end) {
-        const n = Math.min(PREFILL_CHUNK, end - pos);
+        options.signal?.throwIfAborted();
+        const n = Math.min(prefillChunk, end - pos);
         const chunk = promptIds.slice(pos, pos + n);
-        const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const { hidden: h, ctxML } = await forwardMaybeTap(
-          model,
-          ids,
-          caches,
-          tapLayers,
-        );
-        ids.dispose();
+        const { hidden: h, ctxML } = await forwardTokens(chunk, tapLayers);
         if (ctxML) ctxParts.push(ctxML);
-        if (pos + n >= end) {
-          const L = h.shape[1]!;
-          const H = h.shape[2]!;
-          anchorHidden = h.slice([0, L - 1, 0], [1, L, H]); // position len-2
-        }
-        h.dispose(); // logits never computed during the drain
+        try {
+          if (pos + n >= end) {
+            const L = h.shape[1]!;
+            const H = h.shape[2]!;
+            anchorHidden = h.slice([0, L - 1, 0], [1, L, H]); // position len-2
+          }
+        } finally { h.dispose(); } // no logits during the drain
+        maybeQuantizeKv(caches, options);
         clearCache();
         pos += n;
+        if (pos < end) await new Promise<void>((resolve) => setImmediate(resolve));
       }
     } else {
       // Legacy shape: full chunked prefill (final chunk included), token0
       // sampled from the last position below.
-      for (let off = 0; off < promptIds.length; off += PREFILL_CHUNK) {
-        const chunk = promptIds.slice(off, off + PREFILL_CHUNK);
-        const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const { hidden: h, ctxML } = await forwardMaybeTap(
-          model,
-          ids,
-          caches,
-          tapLayers,
-        );
-        ids.dispose();
+      for (let off = cached; off < promptIds.length;) {
+        options.signal?.throwIfAborted();
+        const end = Math.min(off + prefillChunk, promptIds.length,
+          snapshotAt > off ? snapshotAt : Infinity);
+        const chunk = promptIds.slice(off, end);
+        const { hidden: h, ctxML } = await forwardTokens(chunk, tapLayers);
         if (ctxML) ctxParts.push(ctxML);
-        if (off + PREFILL_CHUNK >= promptIds.length) {
+        try { if (end >= promptIds.length) {
           const L = h.shape[1]!;
           const H = h.shape[2]!;
           anchorHidden = h.slice([0, L - 1, 0], [1, L, H]); // [1,1,H], retained
-          const lg = model.logitsFromHidden(anchorHidden);
-          const V = lg.shape[lg.shape.length - 1]!;
-          lastLogits = ops.reshape(lg, [1, V]);
-          lg.dispose();
+          const lg = binding.projectLogits(anchorHidden);
+          try {
+            const V = lg.shape[lg.shape.length - 1]!;
+            lastLogits = ops.reshape(lg, [1, V]);
+          } finally { lg.dispose(); }
+        } } finally { h.dispose(); }
+        maybeQuantizeKv(caches, options);
+        if (end === snapshotAt && snapshotAt > cached) {
+          await seedSource(end);
+          prefix!.capture(promptIds.slice(0, end), caches, prefixNamespace, prefixBudget);
         }
-        h.dispose();
         clearCache();
+        off = end;
+        if (off < promptIds.length)
+          await new Promise<void>((resolve) => setImmediate(resolve));
       }
     }
     extras.targetCalls++;
+    options.signal?.throwIfAborted();
     // Seed the source: two-model prefills its own draft cache; DSpark seeds
     // H_ctx from the tapped prompt context (ownership transfers to prefill);
     // the assistant is a no-op.
-    if (tapLayers) {
-      if (ctxParts.length === 1) prefillCtx = ctxParts[0]!;
-      else if (ctxParts.length > 1) {
-        prefillCtx = ops.concatAxis(ctxParts, 1); // [1,Lp,m*H]
-        for (const p of ctxParts) p.dispose();
-      }
-      ctxParts.length = 0; // parts consumed (transferred as prefillCtx or disposed)
-      await src.prefill(promptIds, prefillCtx ?? undefined); // takes ownership of prefillCtx
-      prefillCtx = null; // ownership transferred
-    } else {
-      await src.prefill(promptIds);
-    }
+    await seedSource(promptIds.length);
+    // Conversion replaces plain cache objects. Bind only after prefill has
+    // converted every attention layer; rollback must own the active objects.
+    const rollback = binding.bindRollback(caches);
     stats.prefillMs = performance.now() - t0;
-    stats.prefillTps = (promptIds.length / Math.max(stats.prefillMs, 1e-6)) * 1000;
+    stats.prefillTps = ((promptIds.length - cached) / Math.max(stats.prefillMs, 1e-6)) * 1000;
 
     // ---- first pending token ----
     // Oracle shape: pending = the UNPROCESSED last prompt token — never
@@ -326,11 +337,12 @@ export async function specServeRun(
       pending = await samplePos(lastLogits!, 0);
       lastLogits!.dispose();
       lastLogits = null;
+      stats.generatedTokens++;
       if (eos.includes(pending)) {
+        stats.finishReason = "stop";
         stats.decodeMs = performance.now() - tDecode;
         return stats;
       }
-      stats.generatedTokens++;
       if (
         (await onToken(pending)) === false ||
         stats.generatedTokens >= maxTokens ||
@@ -347,28 +359,24 @@ export async function specServeRun(
 
     // ---- rounds ----
     decode: while (stats.generatedTokens < maxTokens) {
+      options.signal?.throwIfAborted();
       if (!speculating) {
         // plain single-token continuation (post-ring-wrap fallback)
-        const ids = ops.fromInt32([pending], [1, 1]);
-        const { hidden: h } = await forwardMaybeTap(
-          model,
-          ids,
-          caches,
-          undefined,
-        );
-        ids.dispose();
-        const lg = model.logitsFromHidden(h);
-        h.dispose();
-        const V = lg.shape[lg.shape.length - 1]!;
-        roundRow = ops.reshape(lg, [1, V]);
-        lg.dispose();
+        const { hidden: h } = await forwardTokens([pending], undefined);
+        try {
+          const lg = binding.projectLogits(h);
+          try {
+            const V = lg.shape[lg.shape.length - 1]!;
+            roundRow = ops.reshape(lg, [1, V]);
+          } finally { lg.dispose(); }
+        } finally { h.dispose(); }
         extras.targetCalls++;
         const tok = await samplePos(roundRow, stats.generatedTokens);
         roundRow.dispose();
         roundRow = null;
         clearCache();
-        if (eos.includes(tok)) break;
         stats.generatedTokens++;
+        if (eos.includes(tok)) { stats.finishReason = "stop"; break; }
         pending = tok;
         if ((await onToken(tok)) === false) break;
         if (grammar?.isTerminated) break; // grammar closed — finish "stop"
@@ -387,11 +395,7 @@ export async function specServeRun(
       // spec-eligible through the round snapshot/replay contract instead:
       // specRoundBegin before the verify forward, then commit or a bit-exact
       // rollback(keep) after the accept walk.
-      const roundFits = caches.every((c) =>
-        "maxSize" in c && typeof (c as { maxSize?: number }).maxSize === "number"
-          ? c.offset + n + 1 < (c as { maxSize: number }).maxSize
-          : c.isTrimmable() || typeof c.specRoundBegin === "function",
-      );
+      const roundFits = rollback.canBegin(n);
       if (!roundFits) {
         console.warn(
           "spec: sliding window nearly wrapped — finishing without speculation",
@@ -432,21 +436,20 @@ export async function specServeRun(
       // the emitted position is the next round's anchor hidden.
       // Arm recurrent caches' spec round around the verify forward (no-op for
       // trimmable caches, which roll back via trim() below).
-      for (const c of caches) c.specRoundBegin?.();
+      rollback.begin(d);
       const vIds = ops.fromInt32([pending, ...drafts], [1, d + 1]);
-      const pinTarget = src.pinTargetKernelFamily === true &&
-        "setSpecKernelPinned" in model;
-      if (pinTarget) model.setSpecKernelPinned(true);
+      let pin: { close(): void } | undefined;
       try {
-        const ft = await forwardMaybeTap(model, vIds, caches, tapLayers);
+        pin = src.pinTargetKernelFamily ? binding.pinVerify?.() : undefined;
+        const ft = await binding.forward(vIds, caches, tapLayers);
         // Assign these before the lm-head call so the outer cleanup owns them
         // even if logits construction throws.
         vHidden = ft.hidden;
         vCtxML = ft.ctxML;
-        vLogits = model.logitsFromHidden(vHidden);
+        vLogits = binding.projectLogits(vHidden);
       } finally {
         vIds.dispose();
-        if (pinTarget) model.setSpecKernelPinned(false);
+        pin?.close();
       }
       extras.targetCalls++;
       // The lm-head remains one logical batched operation. Native GLM MTP's
@@ -502,10 +505,18 @@ export async function specServeRun(
 
       // (d) emit the round's tokens through onToken, one at a time
       for (const tok of emitted) {
+        options.signal?.throwIfAborted();
         stats.generatedTokens++;
         if ((await onToken(tok)) === false) { halted = true; break; }
         if (stats.generatedTokens >= maxTokens) break;
       }
+      // Native generation and mlx-lm count the stopping EOS in usage while
+      // keeping it out of content. A callback or budget may stop earlier in
+      // this verified burst; that unused suffix must not enter usage.
+      if (sawEos && !halted && stats.generatedTokens < maxTokens) {
+        stats.generatedTokens++;
+        stats.finishReason = "stop";
+      }
       extras.tokensPerForward = extras.rounds > 0
         ? stats.generatedTokens / extras.rounds
         : 0;
@@ -526,14 +537,7 @@ export async function specServeRun(
         // drafts). Trimmable caches drop the rejected tail; recurrent caches
         // restore their pre-round snapshot and bit-exactly replay those
         // kAccept+1 kept tokens. On full accept the round just commits.
-        if (kAccept < d) {
-          for (const c of caches) {
-            if (c.specRoundRollback) c.specRoundRollback(kAccept + 1);
-            else c.trim(d - kAccept);
-          }
-        } else {
-          for (const c of caches) c.specRoundCommit?.();
-        }
+        rollback.resolve(kAccept);
         await src.commit(
           d,
           kAccept,
@@ -565,21 +569,15 @@ export async function specServeRun(
       ? (stats.generatedTokens / stats.decodeMs) * 1000
       : 0;
     return stats;
+  } catch (error) {
+    failure = { error };
+    throw error;
   } finally {
-    // Dispose whatever a mid-round/prefill throw left live (each is nulled the
-    // instant its normal disposal or ownership transfer ran, so no double-free).
-    anchorHidden?.dispose();
-    lastLogits?.dispose();
-    prefillCtx?.dispose();
-    for (const p of ctxParts) p.dispose();
-    vLogits?.dispose();
-    roundRow?.dispose();
-    vHidden?.dispose();
-    vCtxML?.dispose();
-    for (const c of caches) c.dispose();
-    source?.dispose();
-    stepSampler.dispose();
-    options.grammar?.dispose(); // Phase C will consume it; never leak either way
-    clearCache();
+    const resources = [anchorHidden, lastLogits, prefillCtx, ...ctxParts,
+      vLogits, roundRow, vHidden, vCtxML, ...caches, source, stepSampler,
+      options.grammar, { dispose: clearCache }]
+      .filter((resource): resource is DisposableResource => resource != null);
+    if (failure) cleanupFailure(failure.error, () => disposeResources(resources));
+    else disposeResources(resources);
   }
 }

package/src/spec/source.ts

@@ -1,6 +1,6 @@
 // DraftSource — the seam between the serve-time speculative verify loop
 // (src/spec/serve-loop.ts) and whatever produces draft tokens. Designed in
-// docs/design/mlx-lm-tool-parity-plan.md §7 so every drafter shares ONE
+// docs/reference/cli.md §7 so every drafter shares ONE
 // verify/accept executor:
 //   - TwoModelSource (src/spec/two-model.ts) — mlx-lm parity (L1 oracle:
 //     mlx_lm.server --draft-model), a full second model. Ignores `target`.
@@ -9,7 +9,7 @@
 //     Reads the target's donor K/V + anchor hidden each step.
 //   - DflashSource (src/spec/dflash-source.ts) — DSpark (L3, KL/quality-gated).
 //     Taps the target's multi-layer hiddens (prefill + verify) into a growing
-//     H_ctx (docs/design/dspark-speculative-decoding.md).
+//     H_ctx (docs/design/speculative-decoding.md).
 //   - NgramSource (src/spec/ngram-source.ts) — model-free prompt lookup.
 //   - Glm52NativeMtpSource (src/spec/glm52-mtp-source.ts) — the target
 //     artifact's native Colibri MTP row.
@@ -21,21 +21,53 @@
 // two-model). See [[dspark-seam-kv-borrowing]].
 
 import type { MlxArray } from "../mlx/array";
-import type { RuntimeModel } from "../model/factory";
+import type { Dtype } from "../mlx/ffi";
 import type { Cache } from "../model/gemma4";
 
-/** Read-only handle to the target's live decode state, handed to every source
- *  at open(). Two-model ignores it; KV-borrowing sources read the target's
- *  model (embed / lm-head / the draft module) and its LIVE cache array (donor
- *  K/V views + offset — the SAME caches the serve loop drives). */
+/** Numerical ports are backend-specific; returned arrays are caller-owned. */
+export interface DraftProjection {
+  readonly embed: {
+    encode(ids: MlxArray): MlxArray;
+    readonly scales: { readonly dtype: Dtype };
+  };
+  logitsFromHidden(hidden: MlxArray): MlxArray;
+}
+
+export interface AssistantTarget {
+  position(): number;
+  embedScaled(token: number): MlxArray;
+  /** Retained temporal views; the draft round disposes every returned array. */
+  readDonors(): { sliding: [MlxArray, MlxArray]; full: [MlxArray, MlxArray] };
+}
+
+export interface QwenMtpTarget {
+  readonly hiddenSize: number;
+  readonly layerCount: number;
+  embed(ids: MlxArray): MlxArray;
+  logitsFromHidden(hidden: MlxArray): MlxArray;
+}
+
+/** Graph-declared extensions over the target's live state. Sources request
+ * only the ports they consume; they never inspect a concrete model or cache.
+ * Absence refuses an unsupported pairing before draft-side allocation. */
 export interface TargetView {
-  readonly model: RuntimeModel;
-  readonly caches: Cache[];
+  /** Opaque identity for providers borrowing weights from one exact target. */
+  readonly identity: object;
+  readonly assistant?: AssistantTarget;
+  readonly gemmaTaps?: { readonly layerCount: number; readonly projection: DraftProjection };
+  readonly qwenMtp?: QwenMtpTarget;
 }
 
 /** A per-request draft-token producer. Created per generation (owns its own
  *  draft-side state), disposed by the serve loop's finally. */
 export interface DraftSource {
+  /** Optional paired target/draft prefill state. Restore transfers matching
+   * target state into the supplied array. Capture borrows evaluated state;
+   * the source owns its retained snapshot. Only full-prefill sources opt in. */
+  readonly prefix?: {
+    restore(prompt: readonly number[], caches: Cache[], namespace: string, maxBytes: number): number;
+    capture(tokens: number[], caches: Cache[], namespace: string, maxBytes: number): void;
+  };
   /** Target prefill shape required by this source's oracle. Most mlx-lm
    *  sources leave the final prompt token pending; native Colibri MTP starts
    *  from a full-prompt target forward. */
@@ -111,7 +143,7 @@ export interface DraftProvider {
      *  sampler as the target — mlx-lm parity; greedy drafting under a
      *  temperature>0 request is NOT parity). */
     sampler: (logprobs: MlxArray, step: number) => MlxArray;
-    /** The target model + its live caches (donor views / offset). */
+    /** Bound target capabilities, with explicit donor-view ownership. */
     target: TargetView;
   }): DraftSource;
   dispose(): void;

package/src/spec/two-model.ts

@@ -16,7 +16,7 @@
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { clearCache } from "../mlx/ffi";
-import { flagOn } from "../flags";
+import { flagOn } from "../runtime-config";
 import { loadModelConfig } from "../config";
 import { Weights } from "../weights";
 import { createModel, type RuntimeModel } from "../model/factory";