mlx-bun 0.4.0 → 0.5.0

package/src/backends/mlx/request-state-policy.ts

@@ -1,3 +1,5 @@
+import { namespacedCache } from "../../engine/namespaced-cache";
+import { runtimeConfig, runtimeFlag, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
 import type { GenerateOptions } from "../../generate";
 import type { RuntimeModel } from "../../model/factory";
 import { KVCache, type Cache } from "../../model/gemma4-base";
@@ -24,18 +26,29 @@ export function maybePageKv(
   for (let i = 0; i < cache.length; i++) {
     if (cache[i] instanceof KVCache) {
       cache[i]!.dispose(); // fresh (offset 0) — nothing stored yet
-      cache[i] = new PagedKVCache(capacityTokens, blockSize);
+      cache[i] = new PagedKVCache(capacityTokens, blockSize, undefined,
+        options.kvBits === 4 || options.kvBits === 8 ? { bits: options.kvBits, groupSize: options.kvGroupSize ?? 64 } : undefined);
     }
   }
 }
 
+export function pagedPrefixNamespace(options: GenerateOptions, base: string,
+  direct = runtimeFlag("MLX_BUN_PAGED_ATTN", false)): string {
+  return JSON.stringify(["paged-v1", options.pagedKv?.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE,
+    options.kvBits ?? 0, options.kvGroupSize ?? 64, direct, base]);
+}
+
 export function bindPagedRequestState(model: RuntimeModel, options: GenerateOptions,
-  capacityTokens: number): MlxRequestStatePolicy | undefined {
+  capacityTokens: number, cache?: RowPromptCache, runtime: RuntimeConfig = runtimeConfig()): MlxRequestStatePolicy | undefined {
   if (!options.pagedKv) return undefined;
   const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
-  return { key: `paged:${blockSize}`, create() {
+  const stateOptions = { pagedKv: { blockSize }, kvBits: options.kvBits, kvGroupSize: options.kvGroupSize };
+  const direct = runtime.flag("MLX_BUN_PAGED_ATTN", false);
+  const namespace = (base: string) => pagedPrefixNamespace(stateOptions, base, direct);
+  return { key: namespace(""),
+    promptCache: cache ? namespacedCache(cache, namespace) : undefined, create: () => withRuntimeConfig(runtime, () => {
     const caches = model.makeCache();
-    try { maybePageKv(caches, { pagedKv: { blockSize } }, capacityTokens); return caches; }
+    try { maybePageKv(caches, stateOptions, capacityTokens); return caches; }
     catch (error) { disposeResources(caches); throw error; }
-  } };
+  }) };
 }

package/src/backends/mlx/rollback.ts

@@ -3,7 +3,7 @@ import type { Cache } from "../../model/gemma4-base";
 
 interface TransactionLayer<Acceptance> {
   fits(drafts: number): boolean;
-  begin(): void;
+  begin(drafts: number): void;
   resolve(accepted: Acceptance, drafts: number): void;
 }
 
@@ -20,8 +20,11 @@ export function bindRowCacheRollback(
 ): SpeculativeTransaction<readonly number[]> {
   return bindTransaction(caches.map(cache => ({
     fits: () => true,
-    begin: () => cache.specRoundBegin(),
+    // A zero-candidate round consumes only the committed pending token.
+    // It cannot reject any input, so retaining rollback state is unnecessary.
+    begin: drafts => { if (drafts > 0) cache.specRoundBegin(); },
     resolve(accepted: readonly number[], drafts: number) {
+      if (drafts === 0) return;
       if (accepted.every(count => count === drafts)) cache.specRoundCommit();
       else cache.specRoundRollback(accepted.every(count => count === accepted[0])
         ? accepted[0]! + 1 : accepted.map(count => count + 1));
@@ -75,7 +78,7 @@ function bindTransaction<Acceptance>(
       if (active !== undefined || !layers.every((layer) => layer.fits(drafts)))
         throw new Error("speculative state cannot begin this round");
       active = drafts;
-      try { for (const layer of layers) layer.begin(); }
+      try { for (const layer of layers) layer.begin(drafts); }
       catch (error) { invalid = true; throw error; }
     },
     resolve(accepted) {

package/src/backends/mlx/serial-executor.ts

@@ -1,3 +1,4 @@
+import { pagedPrefixNamespace } from "./request-state-policy";
 import type { SpecRunServices } from "../../spec/serve-loop";
 import { cleanupFailure, disposeResources, ownResource } from "../../engine/resources";
 import { bindGeneration } from "../../generate";
@@ -66,6 +67,14 @@ export interface MlxSerialServices {
 export function createMlxSerialExecutor(binding: MlxSerialBinding, services: MlxSerialServices): SerialRun {
   const { promptCache, checkpoints: ssdStore } = services;
   return (promptIds, options, onToken, vision, trace, execution) => withRuntimeConfig(binding.runtime, async () => {
+    if (promptCache.reclaim) {
+      const publish = onToken;
+      let steps = 0;
+      onToken = (token, logprobs) => {
+        if (steps++ % 256 === 0) promptCache.reclaim!();
+        return publish(token, logprobs);
+      };
+    }
     let releaseContinuation: (() => void) | undefined;
     let caches: Cache[] = [];
     let retain: (() => void) | undefined;
@@ -94,16 +103,16 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
       }
       // Cache entries are adapter-specific: KV computed under one adapter
       // must never seed another's (or the base's) prefill.
-      const cacheNs = options.adapters?.length ? services.adapterNamespace(options.adapters) : "";
+      const baseNs = options.adapters?.length ? services.adapterNamespace(options.adapters) : "";
+      const cacheNs = execution.pagedKv ? pagedPrefixNamespace(options, baseNs) : baseNs;
       // Paged-KV request scope (docs/design/kv-cache.md): media
       // prompts (bidir overlay) and LoRA-adapter requests are v1 non-goals —
       // they run the PLAIN cache path even under --paged-kv (scope the flag
       // per request, never 400). Effective value computed ONCE so the
       // prompt-cache bypass below and the generate() options can't disagree.
       const pagedKv = execution.pagedKv ? options.pagedKv : undefined;
-      // Paged requests bypass the prompt cache entirely (v1 non-goal:
-      // PagedKVCache has no cloneKvCaches/restore path — the vision
-      // precedent). Fresh caches per request, disposed on completion.
+      // The resolved policy controls reuse; paged layouts have their own
+      // numerical namespace and use the same RAM/SSD cache interface.
       const skipPromptCache = !execution.promptCache;
       const checkpointEvery = services.checkpointEveryTokens;
       const checkpointEligible = execution.checkpoint;
@@ -124,7 +133,7 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
       const generationPromptIds = checkpoint?.cacheTokens ?? promptIds;
       const entry = skipPromptCache || resuming
         ? null
-        : promptCache.take(promptIds, cacheNs);
+        : promptCache.take(promptIds, cacheNs, options.cacheSessionId);
       if (entry) { caches = entry.caches; retain = entry.retain; }
       if (!checkpoint && !entry) caches = binding.makeCache();
       closeCacheLookup?.();
@@ -204,7 +213,7 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
                 try {
                   const snapshot = ownResource(services.cloneState(caches), disposeResources);
                   try {
-                    promptCache.put(promptIds.slice(0, boundary), snapshot.borrow(), cacheNs);
+                    promptCache.put(promptIds.slice(0, boundary), snapshot.borrow(), cacheNs, undefined, undefined, options.cacheSessionId);
                     snapshot.transfer();
                   } finally { snapshot.close(); }
                 } catch (e) {
@@ -229,7 +238,7 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
       if (!skipPromptCache) {
         // put() fires onPut → the debounced write-behind SSD snapshot
         // (wired below), covering the batch lane's puts too.
-        promptCache.put(s.cacheTokens, caches, cacheNs, retain);
+        promptCache.put(s.cacheTokens, caches, cacheNs, retain, undefined, options.cacheSessionId);
         caches = []; retain = undefined; // ownership returned to the prefix store
       }
       return s;

package/src/backends/mlx/speculative-group.ts

@@ -1,6 +1,7 @@
 import {captureSpeculativeOptions, speculativePrefixNamespace} from "../../spec/cache-identity";
 import type { RuntimeModel } from "../../model/factory";
 import type { GenerateOptions } from "../../generate";
+import type { FillSession, Proposal } from "../../fill/fill-session";
 import type { DraftProvider, DraftRowGroup, DraftPrefillGroup, DraftRowCheckpoint } from "../../spec/source";
 import { makeSampler, makeStepSampler, readStepExtras, type NumberStepSampler, type Sampler } from "../../sampler";
 import { MlxArray } from "../../mlx/array";
@@ -11,16 +12,19 @@ import { cloneKvCaches } from "../../kv-store";
 import type { Cache } from "../../model/gemma4-base";
 import { targetCacheLayout } from "./cache-layout";
 import { MlxStateRows } from "./state-rows";
-import { bindLegacySpeculativeModel, type MlxSpeculativeBinding } from "./speculative";
+import { bindSpeculativeTargetModel, type MlxSpeculativeTargetBinding } from "./speculative";
 import { bindLegacyDraftTarget } from "./draft-target";
 import { bindRowCacheRollback } from "./rollback";
 import { createKvMaintenance } from "./kv-maintenance";
 import { disposeAttachments, type CheckpointAttachment } from "./checkpoint-state";
 import { advanceSpeculativeOutputs } from "./speculative-round";
 import { MlxPrefillRows, type MlxPrefillState } from "./prefill-rows";
+import type { MlxForwardWork } from "./mixed-iteration";
 import type { MlxGroupMethodHost, MlxGroupMethodRequest, MlxGroupPreparation, MlxGroupedMethod, Row } from "./batch-group";
 
 interface RequestState {
+  fill?: FillSession;
+  proposal?: { value: Proposal; emitted: number };
   sampling: NumberStepSampler;
   draftSampling: Sampler;
   retain?: () => void;
@@ -33,8 +37,8 @@ interface RequestState {
 
 /** Binding owns graph/layout selection. The executor receives only the method
  * key and lifecycle; sampling, checkpoints and numerical state remain ports. */
-export function bindSpeculativeGroupRequests(model: RuntimeModel, provider: DraftProvider, depth: number) {
-  const binding = bindLegacySpeculativeModel(model, provider);
+export function bindSpeculativeGroupRequests(model: RuntimeModel, provider: Pick<DraftProvider, "id" | "grouped">, depth: number) {
+  const binding = bindSpeculativeTargetModel(model);
   return (input: GenerateOptions): MlxGroupMethodRequest => {
     const options = captureSpeculativeOptions(input);
     return {
@@ -54,7 +58,18 @@ class SpeculativeGroup implements MlxGroupedMethod {
   #steps = 0;
 
   constructor(readonly host: MlxGroupMethodHost, readonly model: RuntimeModel,
-    readonly provider: DraftProvider, readonly binding: MlxSpeculativeBinding, readonly depth: number) {}
+    readonly provider: Pick<DraftProvider, "id" | "grouped">, readonly binding: MlxSpeculativeTargetBinding, readonly depth: number) {}
+
+  get runningTokens(): number {
+    const rows = this.host.rows;
+    if (!rows.length) return 0;
+    // Target verification consumes the pending token and the draft candidates.
+    const echo = Math.max(...rows.map(row => {
+      const proposal = this.#requests.get(row)?.proposal;
+      return proposal?.emitted ? proposal.value.ids.length - proposal.emitted : 0;
+    }));
+    return rows.length * (1 + (echo || Math.min(this.depth, Math.max(...rows.map(row => row.req.maxTokens - row.generated)))));
+  }
 
   prepare(first: Row): MlxGroupPreparation {
     const method = this;
@@ -91,7 +106,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
       try {
         caches = snapshot(); attachments.push(prefix!.capture(index).attachment);
         method.host.promptCache!.put(state.row.req.promptIds.slice(0, state.pos), caches,
-          state.request.namespace, undefined, attachments);
+          state.request.namespace, undefined, attachments, state.row.req.cacheSessionId);
         caches = []; attachments = [];
       } finally { disposeResources([...caches, { dispose: () => disposeAttachments(attachments) }]); }
     };
@@ -103,11 +118,13 @@ class SpeculativeGroup implements MlxGroupedMethod {
           acceptanceLengths: [], tokensPerForward: 0, forwardsSaved: 0, draftedByPos: [], acceptedByPos: [] };
         const sampler = makeSampler(options);
         const request: RequestState = { draftSampling: sampler, namespace: "", prefixLength: 0,
+          fill: method.provider.grouped?.supportsExternalTokens ? options.fill : undefined,
           processed: method.host.promptCache ? [] : undefined,
           sampling: makeStepSampler(options, { tokenRepresentation: "number", grammarWait: "before-sample",
             historyUpdate: "after-sample", initialHistory: row.req.promptIds, acceptGrammar: true,
             eosTokenIds: row.req.eosTokenIds, captureSelectedLogprob: options.logprobs === true,
             captureTopLogprobs: options.topLogprobs }) };
+        if (request.fill) row.fill = request.fill.stats;
         let caches: Cache[] = [], retained: (() => void) | undefined;
         let appended = false;
         const previous = prefix?.rowCount ?? 0;
@@ -118,7 +135,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
           const pendingPrompt = prefix.prefillMode !== "full" && method.binding.prefillTailSplit && prompt.length > 1;
           request.prefixLength = prompt.length - Number(pendingPrompt);
           request.namespace = speculativePrefixNamespace(prefix.namespace, row.cacheNamespace ?? "", options);
-          const hit = method.host.promptCache?.take(prompt, request.namespace);
+          const hit = method.host.promptCache?.take(prompt, request.namespace, row.req.cacheSessionId);
           if (hit) {
             try {
               disposeResources(caches); caches = hit.caches.splice(0);
@@ -134,7 +151,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
           } else { prefix.append([null]); appended = true; }
           return { row, request, solo: caches, retain: retained, pos: row.cachedTokens,
             end: request.prefixLength, pendingPrompt, transferred: false, closed: false,
-            chunkSize: options.prefillChunkSize ?? method.host.prefillChunkSize,
+            chunkSize: options.prefillChunkSize ?? row.req.prefillChunkSize ?? method.host.prefillChunkSize,
             boundary: method.host.promptCache ? Math.min(row.req.snapshotAt ?? prompt.length, prompt.length - 1) : -1 };
         } catch (error) {
           return cleanupFailure(error, () => disposeResources([
@@ -148,9 +165,9 @@ class SpeculativeGroup implements MlxGroupedMethod {
         return { start: state.pos, end, kind: end === state.end ? "final" : "drain",
           snapshot: end === state.boundary && end > state.row.cachedTokens, batchYield: true, maintain: true };
       },
-      async forward(ids, caches) {
+      async forward(ids, caches, _states, work) {
         const taps = prefix!.tapLayers;
-        const output = await method.binding.forward(ids, caches, taps.length ? [...taps] : undefined);
+        const output = await method.binding.forward(ids, caches, taps.length ? [...taps] : undefined, work);
         context = output.ctxML;
         return output.hidden;
       },
@@ -184,12 +201,14 @@ class SpeculativeGroup implements MlxGroupedMethod {
             method.host.finish(row, "stop"); return;
           }
           const more = await method.host.publish(row, row.current, readStepExtras(sampled.extras));
+          state.request.fill?.observe(row.current);
           if (more === false || row.req.grammar?.isTerminated || row.generated >= row.req.maxTokens) {
             if (method.host.promptCache && !row.req.signal?.aborted) capture(state, () => cloneKvCaches(state.solo), index);
             disposeResources(state.solo.splice(0)); state.retain?.(); state.retain = undefined;
             method.host.finish(row, more === false || row.req.grammar?.isTerminated ? "stop" : "length"); return;
           }
         }
+        method.#proposeEcho(row, state.request);
         ready.push({ row, request: state.request, caches: state.solo, retain: state.retain,
           draft: prefix!.capture(index), transferred: false });
         state.transferred = true;
@@ -211,7 +230,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
       get canAdmit() { return target.canAdmit; },
       get tokenWeight() { return target.tokenWeight + ready.reduce((sum, state) => sum + state.row.promptTokens - state.row.cachedTokens, 0); },
       admit: row => target.admit(row),
-      advance: async () => {
+      advance: async work => {
         // Publication returns to scheduling before prepared state joins active
         // decode. This preserves the first-output flush boundary at every B.
         if (ready.length) {
@@ -225,7 +244,8 @@ class SpeculativeGroup implements MlxGroupedMethod {
               maintain.prepareBatch?.(state.caches);
               this.#target ??= new MlxStateRows(state.caches.map(targetCacheLayout));
               this.#draft ??= this.provider.grouped!.open({ target: bindLegacyDraftTarget(this.model, this.#target.caches),
-                checkpoints: [], sampling: { sample: (lp, steps) => this.#sampleDraftRows(lp, steps) } });
+                checkpoints: [], sampling: { sample: (lp, steps) => this.#sampleDraftRows(lp, steps) },
+                constraints: { propose: async (row, maxTokens) => this.host.rows[row]!.req.grammar?.proposeTokens(maxTokens) ?? [] } });
               applyStateChanges([() => this.#target!.prepareAppend(state.caches),
                 () => this.#draft!.prepareAppend([state.draft]), () => ({ commit: () => {
                   state.request.retain = state.retain; state.retain = undefined;
@@ -236,7 +256,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
           }
           return target.rows.length === 0;
         }
-        const done = await target.advance();
+        const done = await target.advance(work);
         return done && ready.length === 0;
       },
       dispose() {
@@ -259,7 +279,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
     } finally { disposeResources(tokens); }
   }
 
-  async advance(): Promise<void> {
+  async advance(work?: MlxForwardWork): Promise<void> {
     const active = [...this.host.rows];
     const live = active.flatMap((row, index) => row.req.signal?.aborted ? [] : [index]);
     if (live.length !== active.length) {
@@ -269,6 +289,19 @@ class SpeculativeGroup implements MlxGroupedMethod {
     const rows = [...this.host.rows];
     if (!rows.length) return;
     const depth = Math.min(this.depth, Math.max(...rows.map(row => row.req.maxTokens - row.generated)));
+    // The ordinary MTP round already supplies samples that can establish a
+    // copied prefix. Do not verify a wide span before any sample matches it.
+    const proposals = rows.map(row => {
+      const proposal = this.#requests.get(row)!.proposal;
+      return proposal?.emitted ? proposal.value.ids.slice(proposal.emitted) : [];
+    });
+    const echoDepth = Math.max(...proposals.map(ids => ids.length));
+    let externalTokens: MlxArray | undefined;
+    const rollback = bindRowCacheRollback(this.#target!.caches, rows.length);
+    const fills = echoDepth ? rows.flatMap((row, index) => {
+      const request = this.#requests.get(row)!;
+      return proposals[index]!.length ? [request.fill!] : [];
+    }) : [];
     const halted = new Set<Row>();
     const pin = this.binding.pinVerify?.();
     try {
@@ -280,11 +313,27 @@ class SpeculativeGroup implements MlxGroupedMethod {
           const more = await this.host.publish(row, ids[0]!, metadata?.[0]);
           if (more === false) { halted.add(row); return false; }
         } },
-      })), depth, this.#draft!, {
-        transaction: bindRowCacheRollback(this.#target!.caches, rows.length),
+      })), echoDepth || depth, echoDepth ? {
+        draft: () => proposals,
+        commit: (accepted, context) => this.#draft!.consume!(externalTokens!, context, accepted.map(n => n + 1)),
+      } : this.#draft!, {
+        transaction: !echoDepth ? rollback : { ...rollback,
+          begin(depth) {
+            const start = performance.now(); rollback.begin(depth);
+            const elapsed = performance.now() - start;
+            for (const fill of fills) fill.noteVerifyEvent(elapsed);
+          },
+          resolve(accepted) {
+            const start = performance.now(); rollback.resolve(accepted);
+            const elapsed = performance.now() - start;
+            for (const fill of fills) fill.stats.checkpointMs += elapsed;
+          },
+        },
         forward: async ids => {
+          if (echoDepth) externalTokens = ops.copyOf(ids);
           const taps = this.#draft!.tapLayers;
-          const result = await this.binding.forward(ids, this.#target!.caches, taps.length ? [...taps] : undefined);
+          const result = await this.binding.forward(ids, this.#target!.caches, taps.length ? [...taps] : undefined,
+            work ? (tokens, caches, options) => work(tokens, caches, { ...options, preserveTokenGeometry: true }) : undefined);
           const context = result.ctxML ?? result.hidden;
           try { return { logits: this.binding.projectLogits(result.hidden), context }; }
           catch (error) { context.dispose(); throw error; }
@@ -295,23 +344,39 @@ class SpeculativeGroup implements MlxGroupedMethod {
       for (let index = 0; index < rows.length; index++) {
         const row = rows[index]!, output = completed.outputs[index]!, round = completed.rounds[index]!;
         const request = this.#requests.get(row)!;
+        if (request.fill) {
+          for (const token of round.acceptance.emitted.slice(0, output.generated)) {
+            const proposal = request.proposal;
+            if (proposal && token === proposal.value.ids[proposal.emitted]) {
+              if (++proposal.emitted === proposal.value.ids.length) this.#settleEcho(request);
+            } else {
+              this.#settleEcho(request); request.fill.observe(token);
+            }
+          }
+          if (output.kind === "continue") this.#proposeEcho(row, request);
+          else this.#settleEcho(request);
+        }
         if (output.kind !== "failed" && output.kind !== "cancelled")
           request.processed?.push(row.current, ...round.drafts.slice(0, output.accepted));
         const stats = row.spec!;
-        stats.drafted += round.drafts.length;
-        stats.accepted += round.acceptance.accepted;
+        stats.drafted += echoDepth ? 0 : round.drafts.length;
+        stats.accepted += echoDepth ? 0 : round.acceptance.accepted;
         stats.rejected = stats.drafted - stats.accepted;
-        stats.rounds!++; stats.targetCalls++;
-        stats.acceptanceLengths!.push(round.acceptance.accepted);
-        for (let p = 0; p < round.drafts.length; p++) stats.draftedByPos![p] = (stats.draftedByPos![p] ?? 0) + 1;
-        for (let p = 0; p < round.acceptance.accepted; p++) stats.acceptedByPos![p] = (stats.acceptedByPos![p] ?? 0) + 1;
+        stats.targetCalls++;
+        if (!echoDepth) {
+          stats.rounds!++;
+          stats.acceptanceLengths!.push(round.acceptance.accepted);
+          for (let p = 0; p < round.drafts.length; p++) stats.draftedByPos![p] = (stats.draftedByPos![p] ?? 0) + 1;
+          for (let p = 0; p < round.acceptance.accepted; p++) stats.acceptedByPos![p] = (stats.acceptedByPos![p] ?? 0) + 1;
+        }
         row.sampled += round.acceptance.emitted.length + Number(round.acceptance.sawEos);
         // A consumer may stop inside a verified burst before its terminal EOS.
         if (output.kind === "stop" && !halted.has(row) && round.acceptance.sawEos &&
           output.generated === round.acceptance.emitted.length && row.generated < row.req.maxTokens)
           row.generated++;
-        stats.tokensPerForward = row.generated / stats.rounds!;
-        stats.forwardsSaved = Math.max(0, row.generated - 1 - stats.rounds!);
+        const forwards = stats.targetCalls - 1;
+        stats.tokensPerForward = row.generated / forwards;
+        stats.forwardsSaved = Math.max(0, row.generated - 1 - forwards);
         if (output.kind === "continue") { row.current = output.pending; keep.push(index); }
         else if (output.kind === "failed") row.reject(output.error);
         else if (output.kind === "cancelled") row.reject(row.req.signal?.reason ?? new Error(output.reason));
@@ -322,7 +387,19 @@ class SpeculativeGroup implements MlxGroupedMethod {
       }
       if (keep.length !== rows.length) this.host.filterRows(keep);
       if (++this.#steps % 256 === 0) clearCache();
-    } finally { pin?.close(); }
+    } finally { externalTokens?.dispose(); pin?.close(); }
+  }
+
+  #proposeEcho(row: Row, request: RequestState): void {
+    if (request.proposal || !request.fill) return;
+    const value = request.fill.propose(row.req.maxTokens - row.generated, "verify");
+    if (value) request.proposal = { value, emitted: 0 };
+  }
+
+  #settleEcho(request: RequestState): void {
+    if (!request.proposal) return;
+    request.fill!.commit(request.proposal.value, request.proposal.emitted);
+    request.proposal = undefined;
   }
 
   /** The method aligns token coverage and captures its state. The cache owns
@@ -334,7 +411,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
       target = this.#target!.extractRow(index);
       attachments.push(this.#draft!.capture(index).attachment);
       this.host.promptCache.put([...row.req.promptIds.slice(0, request.prefixLength), ...request.processed], target,
-        request.namespace, undefined, attachments);
+        request.namespace, undefined, attachments, row.req.cacheSessionId);
       target = []; attachments = [];
     } finally { disposeResources([...target, { dispose: () => disposeAttachments(attachments) }]); }
   }
@@ -344,6 +421,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
     const retained = new Set(keep.map(index => this.host.rows[index]!));
     const releases: { dispose(): void }[] = [];
     for (const [row, request] of this.#requests) if (!retained.has(row)) {
+      this.#settleEcho(request);
       this.#requests.delete(row);
       releases.push(request.sampling, { dispose: () => request.retain?.() });
     }
@@ -351,6 +429,7 @@ class SpeculativeGroup implements MlxGroupedMethod {
   }
 
   dispose(): void {
+    for (const request of this.#requests.values()) this.#settleEcho(request);
     const resources = [...(this.#target ? [this.#target] : []), ...(this.#draft ? [this.#draft] : []),
       ...[...this.#requests.values()].flatMap(request => [request.sampling, { dispose: () => request.retain?.() }])];
     this.#target = null; this.#draft = null; this.#requests.clear();

package/src/backends/mlx/speculative.ts

@@ -10,12 +10,16 @@ import type { SpeculativeTransaction } from "../../inference/rollback";
 import { bindCacheRollback } from "./rollback";
 import type { GraphDescriptor } from "../../inference/graph";
 import { bindLegacyDraftTarget } from "./draft-target";
+import type { MlxForwardWork } from "./mixed-iteration";
+import type { PrefillPolicy } from "../../inference/prefill";
+import { resolveMlxPrefillPolicy } from "./prefill-policy";
 
 /** Bound target operations and draft construction for one speculative run.
  * A replacement graph supplies this entire port, including any hidden taps.
  * The verifier never needs a RuntimeModel or an artifact-family check. */
 export interface MlxSpeculativeBinding {
   readonly runtime?: RuntimeConfig;
+  readonly prefillPolicy?: PrefillPolicy;
   readonly memory?: MlxModelMemory;
   readonly descriptor: GraphDescriptor;
   readonly eosTokenIds: readonly number[];
@@ -23,27 +27,29 @@ export interface MlxSpeculativeBinding {
   makeCache(): Cache[];
   openDraft(sampler: Sampler, caches: Cache[]): DraftSource;
   bindRollback(caches: Cache[]): SpeculativeTransaction;
-  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[]):
+  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[], work?: MlxForwardWork):
     Promise<{ hidden: MlxArray; ctxML: MlxArray | null }>;
   projectLogits(hidden: MlxArray): MlxArray;
   /** Establish the implementation's verify kernel context; restore on close. */
   pinVerify?(): { close(): void };
 }
 
+export type MlxSpeculativeTargetBinding = Omit<MlxSpeculativeBinding, "openDraft">;
+
 /** Legacy mutable tap/kernel fields require the gateway's exclusive lease. */
-export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
+export function bindSpeculativeTargetModel(model: RuntimeModel): MlxSpeculativeTargetBinding {
   const runtime = runtimeConfig();
   return {
     runtime,
+    prefillPolicy: resolveMlxPrefillPolicy(model.config, runtime),
     memory: model,
     descriptor: Object.freeze({ id: `legacy-spec:${model.config.modelType}`, backend: "mlx",
       graphAbi: "mlx-hidden-bsh-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
     eosTokenIds: model.config.eosTokenIds,
     prefillTailSplit: runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true),
     makeCache: model.makeCache.bind(model),
-    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }),
     bindRollback: bindCacheRollback,
-    forward: (ids, caches, tapLayers) => legacyForwardWithTaps(model, ids, caches, tapLayers),
+    forward: (ids, caches, tapLayers, work) => legacyForwardWithTaps(model, ids, caches, tapLayers, work),
     projectLogits: model.logitsFromHidden.bind(model),
     ...("setSpecKernelPinned" in model ? {
       pinVerify() {
@@ -54,6 +60,12 @@ export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftP
   };
 }
 
+/** Serial compatibility binds draft construction separately from target work. */
+export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
+  return { ...bindSpeculativeTargetModel(model),
+    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }) };
+}
+
 export function assertMlxSpeculativeBinding(binding: MlxSpeculativeBinding): void {
   const descriptor = binding.descriptor;
   if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
@@ -71,8 +83,10 @@ async function legacyForwardWithTaps(
   ids: MlxArray,
   caches: Cache[],
   tapLayers: number[] | undefined,
+  work?: MlxForwardWork,
 ): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
   if (!tapLayers) {
+    if (work) return { hidden: await work(ids, caches), ctxML: null };
     const asyncModel = model as RuntimeModel & {
       forwardHiddenAsync?: (
         ids: MlxArray,
@@ -88,10 +102,13 @@ async function legacyForwardWithTaps(
   const m = model;
   const previousTap = m.hiddenTap;
   const cap = new Map<number, MlxArray>();
-  m.hiddenTap = { layers: new Set(tapLayers), captured: cap };
+  const layers = new Set(tapLayers);
+  if (!work) m.hiddenTap = { layers, captured: cap };
   let hidden: MlxArray | null = null;
   try {
-    hidden = model.forwardHidden(ids, caches);
+    hidden = work ? await work(ids, caches, { captureLayer(layer, h) {
+      if (layers.has(layer)) cap.set(layer, ops.contiguous(h));
+    } }) : 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`);
@@ -109,6 +126,6 @@ async function legacyForwardWithTaps(
     // are already gone (cap cleared, hidden nulled).
     hidden?.dispose();
     for (const [, a] of cap) a.dispose();
-    m.hiddenTap = previousTap;
+    if (!work) m.hiddenTap = previousTap;
   }
 }

package/src/backends/mlx/target-layout-capability.ts

@@ -27,7 +27,7 @@ export function targetRowLayoutFactory(cache: Cache): (() => TargetRowLayout) |
   if (cache instanceof DelayedTurboQuantKVCache || cache instanceof DelayedQuantizedKVCache ||
       cache instanceof BatchedTurboQuantKVCache || cache instanceof BatchedKVCache ||
       cache instanceof BatchedQuantizedKVCache || cache instanceof BatchedSSMCache) return () => cache.makeEmptyBatch();
-  if (cache instanceof TurboQuantKVCache) return () => new BatchedTurboQuantKVCache(cache.kBits, cache.vBits);
+  if (cache instanceof TurboQuantKVCache) return () => new BatchedTurboQuantKVCache(cache.kBits, cache.vBits, cache.fusedDecode);
   if (cache instanceof SSMCache) return () => new BatchedSSMCache();
   if (cache instanceof QuantizedKVCache) return () => new BatchedQuantizedKVCache(cache.groupSize, cache.bits);
   if (cache instanceof KVCache) return () => new BatchedKVCache();

package/src/backends/mlx/transitioning-kv-rows.ts

@@ -17,6 +17,7 @@ export interface SpeculativeTransitionedKvLayout extends TransitionedKvLayout {
 }
 export interface KvRowTransition<Layout extends TransitionedKvLayout> {
   readonly signature: string;
+  readonly conversionOffset?: number;
   maintain(rows: Cache[]): void;
   converted(row: Cache): boolean;
   makeLayout(): Layout;
@@ -77,6 +78,13 @@ export abstract class TransitioningKvRows<Layout extends TransitionedKvLayout> i
   get rowOffsets(): readonly number[] { return this.packed?.rowOffsets ?? this.positions.rowOffsets; }
   get leftPad(): readonly number[] { return this.packed?.leftPad ?? this.positions.leftPad; }
   get offset(): number { return this.packed?.offset ?? this.positions.offset; }
+  maxAppendTokens(): number {
+    if (this.packed) return this.packed.maxAppendTokens?.() ?? Infinity;
+    const end = this.transition.conversionOffset;
+    if (end === undefined) return Infinity;
+    return Math.min(Infinity, ...this.rowOffsets.map((offset, row) =>
+      offset < end && !this.transition.converted(this.rows[row]!) ? end - offset : Infinity));
+  }
   get batchSize(): number | null { return this.packed?.batchSize ?? this.positions.batchSize; }
   get ropeOffsetArr(): MlxArray | undefined { return this.packed ? this.packed.ropeOffsetArr : this.positions.ropeOffsetArr; }
   bytesPerToken(): number { return this.packed?.bytesPerToken?.() ?? Math.max(0, ...this.rows.map(row => row.bytesPerToken?.() ?? 0)); }