mlx-bun 0.3.0 → 0.4.0

package/src/spec/serve-loop.ts

@@ -17,11 +17,9 @@
 //    trimmable mid-generation; a serve endpoint must not 500 mid-stream, so
 //    we STOP SPECULATING and finish the generation with plain single-token
 //    decode (bit-equivalent continuation — the target's own samples).
-//  - Prompt-cache reuse is BYPASSED in v1 (fresh caches per spec request,
-//    cachedTokens=0). mlx-lm composes spec with its LRU prompt cache
-//    (target+draft caches per entry); wiring that through our PromptCache +
-//    SSD tier is a tracked follow-up in the integration plan, not silently
-//    absent.
+//  - Sources with checkpoint support use the injected prefix-cache service.
+//    Target and companion state share the existing RAM/SSD lifecycle. Other
+//    sources begin with fresh state until they supply a checkpoint codec.
 //
 // Grammar × spec (Phase C, the constrained verify walk — novel: NO runtime
 // serves both; mlx-lm has no grammar, oMLX no spec): the drafter runs FREE
@@ -40,14 +38,18 @@
 // false halts the generation mid-burst.
 
 import { MlxArray } from "../mlx/array";
+import { sampleSpeculativeRows } from "../backends/mlx/speculative-round";
 import * as ops from "../mlx/ops";
 import { clearCache } from "../mlx/ffi";
 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 { cloneKvCaches } from "../kv-store";
+import { disposeAttachments, type CheckpointAttachment } from "../backends/mlx/checkpoint-state";
 import type { Cache } from "../model/gemma4";
-import { maybeQuantizeKv, withModelUsageFlush, withModelWiredLimit, type GenerateOptions, type GenerateStats } from "../generate";
+import { withModelUsageFlush, withModelWiredLimit, type GenerateOptions, type GenerateStats } from "../generate";
+import { createKvMaintenance } from "../backends/mlx/kv-maintenance";
 import { runtimeConfig, withRuntimeConfig } from "../runtime-config";
 import { makeSampler, makeStepSampler } from "../sampler";
 import type { OnToken } from "../serve/generation-gateway";
@@ -72,6 +74,13 @@ export interface SpecServeExtras {
   acceptedByPos: number[];
 }
 
+export interface SpecRunServices {
+  readonly prefixCache?: import("../backends/mlx/checkpoint-state").MlxPrefixCache;
+  /** Mounted adapter revision, resolved while the execution lease is held. */
+  readonly cacheNamespace?: string;
+  readonly cloneState?: (caches: Cache[]) => Cache[];
+}
+
 export async function specServeRun(
   model: RuntimeModel,
   provider: DraftProvider,
@@ -79,8 +88,9 @@ export async function specServeRun(
   promptIds: number[],
   options: GenerateOptions & { stopSequences?: string[] },
   onToken: OnToken,
+  services: SpecRunServices = {},
 ): Promise<GenerateStats> {
-  return specRun(bindLegacySpeculativeModel(model, provider), numDraftTokens, promptIds, options, onToken);
+  return specRun(bindLegacySpeculativeModel(model, provider), numDraftTokens, promptIds, options, onToken, services);
 }
 
 /** The verifier consumes a bound target and source factory, independent of
@@ -91,9 +101,10 @@ export async function specRun(
   promptIds: number[],
   options: GenerateOptions & { stopSequences?: string[] },
   onToken: OnToken,
+  services: SpecRunServices = {},
 ): Promise<GenerateStats> {
   return withRuntimeConfig(binding.runtime ?? runtimeConfig(), () => {
-    const run = () => specRunInner(binding, numDraftTokens, promptIds, options, onToken);
+    const run = () => specRunInner(binding, numDraftTokens, promptIds, options, onToken, services);
     return binding.memory
       ? withModelWiredLimit(binding.memory, () => withModelUsageFlush(binding.memory!, run))
       : run();
@@ -106,8 +117,10 @@ async function specRunInner(
   promptIds: number[],
   options: GenerateOptions & { stopSequences?: string[] },
   onToken: OnToken,
+  services: SpecRunServices = {},
 ): Promise<GenerateStats> {
   const maxTokens = options.maxTokens ?? 512;
+  const maintainKv = createKvMaintenance(options);
   const eos = options.eosTokenIds ?? [...binding.eosTokenIds];
   const sampler = makeSampler(options);
   const stepSampler = makeStepSampler(options, {
@@ -129,6 +142,7 @@ async function specRunInner(
   // and turned a config error into a per-request 500. loadContext now
   // probe-opens the pairing at startup, so a throw here is belt+suspenders.
   let caches: Cache[] = [];
+  let prefixRetain: (() => void) | undefined;
   let source: import("./source").DraftSource | null = null;
   let tapLayers: number[] | undefined;
   // Target final hidden [1,1,H] at the anchor position — the assistant source
@@ -175,16 +189,6 @@ async function specRunInner(
   const samplePos = async (logits1V: MlxArray, step: number): Promise<number> =>
     (await stepSampler.sample(logits1V, step)).token;
 
-  /** The [1,V] logits row at position `pos` of a hidden window — batched
-   *  lm-head is applied by the caller ONCE; this slices its output. */
-  const logitsRow = (logitsWindow: MlxArray, pos: number): MlxArray => {
-    const V = logitsWindow.shape[logitsWindow.shape.length - 1]!;
-    const sl = logitsWindow.slice([0, pos, 0], [1, pos + 1, V]);
-    const flat = ops.reshape(sl, [1, V]);
-    sl.dispose();
-    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); }
@@ -197,25 +201,50 @@ async function specRunInner(
     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");
+    if (options.kvConfig?.length ||
+        (options.turboQuant && (options.quantizedKvStart ?? 0) !== 0) ||
+        (options.kvBits && ((options.kvBits !== 4 && options.kvBits !== 8) || options.quantizedKvStart !== 0)))
+      throw new Error("speculative KV requires uniform KV4/KV8 or TurboQuant start=0");
     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;
+    const checkpoint = src.prefillMode === "full" ? src.checkpoint : undefined;
+    const prefixCache = checkpoint ? services.prefixCache : undefined;
+    const prefixNamespace = JSON.stringify({ method: checkpoint?.namespace,
+      adapters: services.cacheNamespace ?? options.adapters ?? [], kvBits: options.kvBits ?? null,
+      kvGroupSize: options.kvGroupSize ?? 64, quantizedKvStart: options.quantizedKvStart ?? null,
+      turboQuant: options.turboQuant ?? null });
+    let cached = 0;
+    const hit = prefixCache?.take(promptIds, prefixNamespace);
+    if (hit) {
+      try {
+        disposeResources(caches);
+        caches.splice(0, caches.length, ...hit.caches.splice(0));
+        prefixRetain = hit.retain;
+        hit.retain = undefined;
+        checkpoint!.restore(hit.tokens.length, hit.attachments![0]!);
+        cached = hit.tokens.length;
+      } finally {
+        disposeResources([...hit.caches, { dispose: () => disposeAttachments(hit.attachments) },
+          { dispose: () => hit.retain?.() }]);
+      }
+    }
     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 capture = (tokens: number[]) => {
+      let state: Cache[] = [];
+      let attachments: CheckpointAttachment[] = [];
+      try {
+        state = (services.cloneState ?? cloneKvCaches)(caches);
+        attachments.push(checkpoint!.capture(tokens.length));
+        prefixCache!.put(tokens, state, prefixNamespace, undefined, attachments);
+        state = []; attachments = [];
+      } finally {
+        disposeResources([...state, { dispose: () => disposeAttachments(attachments) }]);
+      }
+    };
+    const snapshotAt = prefixCache
+      ? Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1) : -1;
     const seedSource = async (end: number) => {
       if (ctxParts.length === 1) prefillCtx = ctxParts[0]!;
       else if (ctxParts.length > 1) {
@@ -275,7 +304,7 @@ async function specRunInner(
             anchorHidden = h.slice([0, L - 1, 0], [1, L, H]); // position len-2
           }
         } finally { h.dispose(); } // no logits during the drain
-        maybeQuantizeKv(caches, options);
+        maintainKv(caches);
         clearCache();
         pos += n;
         if (pos < end) await new Promise<void>((resolve) => setImmediate(resolve));
@@ -300,10 +329,10 @@ async function specRunInner(
             lastLogits = ops.reshape(lg, [1, V]);
           } finally { lg.dispose(); }
         } } finally { h.dispose(); }
-        maybeQuantizeKv(caches, options);
+        maintainKv(caches);
         if (end === snapshotAt && snapshotAt > cached) {
           await seedSource(end);
-          prefix!.capture(promptIds.slice(0, end), caches, prefixNamespace, prefixBudget);
+          capture(promptIds.slice(0, end));
         }
         clearCache();
         off = end;
@@ -460,39 +489,12 @@ async function specRunInner(
       // accept while it reproduces the draft. Sampling is sequential because
       // the processor history (and, in Phase C, the grammar mask) at position
       // i depends on the tokens accepted at positions < i.
-      let kAccept = 0;
-      let correction: number | null = null; // target's token at first mismatch (or bonus)
-      let sawEos = false;
+      const [result] = await sampleSpeculativeRows([{
+        pending, step: stats.generatedTokens, remaining: maxTokens - stats.generatedTokens,
+        eosTokenIds: eos, sampling: stepSampler, grammarDone: () => grammar?.isTerminated ?? false,
+      }], [drafts], vLogits!);
+      const { accepted: kAccept, correction, sawEos, emitted, grammarDone } = result!.acceptance;
       let halted = false;
-      const emitted: number[] = [];
-      let grammarDone = false;
-      for (let i = 0; i <= d; i++) {
-        roundRow = logitsRow(vLogits!, i);
-        const tok = await samplePos(roundRow, stats.generatedTokens + emitted.length);
-        roundRow.dispose();
-        roundRow = null;
-        if (i < d && tok === drafts[i]) {
-          kAccept++;
-          // EOS is never content — even when it arrives as an ACCEPTED DRAFT
-          // (the correction/bonus branch below always excluded it; this
-          // branch pushed-then-broke, leaking the EOS through onToken —
-          // caught by the 2026-07-07 live oracle gate: streams were
-          // bit-identical, ours emitted one extra "content" token: <|eot_id|>).
-          if (eos.includes(tok)) { sawEos = true; break; }
-          emitted.push(tok);
-          // grammar termination mid-burst truncates the round — nothing may
-          // be sampled past a satisfied grammar (the all--inf guarantee).
-          if (grammar?.isTerminated) { grammarDone = true; break; }
-          if (stats.generatedTokens + emitted.length >= maxTokens) break;
-          continue;
-        }
-        // mismatch (target replaces the draft) or bonus (i === n)
-        correction = tok;
-        if (!eos.includes(tok)) emitted.push(tok);
-        else sawEos = true;
-        if (grammar?.isTerminated) grammarDone = true;
-        break;
-      }
       vLogits!.dispose();
       vLogits = null;
       extras.accepted += kAccept;
@@ -575,7 +577,7 @@ async function specRunInner(
   } finally {
     const resources = [anchorHidden, lastLogits, prefillCtx, ...ctxParts,
       vLogits, roundRow, vHidden, vCtxML, ...caches, source, stepSampler,
-      options.grammar, { dispose: clearCache }]
+      options.grammar, { dispose: () => prefixRetain?.() }, { dispose: clearCache }]
       .filter((resource): resource is DisposableResource => resource != null);
     if (failure) cleanupFailure(failure.error, () => disposeResources(resources));
     else disposeResources(resources);

package/src/spec/source.ts

@@ -23,6 +23,9 @@
 import type { MlxArray } from "../mlx/array";
 import type { Dtype } from "../mlx/ffi";
 import type { Cache } from "../model/gemma4";
+import type { CheckpointAttachment } from "../backends/mlx/checkpoint-state";
+import type { MlxDraftRows } from "../backends/mlx/speculative-round";
+import type { PreparedStateChange } from "../contracts/resources";
 
 /** Numerical ports are backend-specific; returned arrays are caller-owned. */
 export interface DraftProjection {
@@ -33,11 +36,14 @@ export interface DraftProjection {
   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] };
+/** Read-only target state and embedding, independent of draft scheduling. */
+export interface AssistantRowsTarget {
+  readonly hiddenSize: number;
+  embed(ids: MlxArray): MlxArray;
+  readDonors(): import("./drafter").AssistantDonors & {
+    readonly positions: readonly number[];
+    dispose(): void;
+  };
 }
 
 export interface QwenMtpTarget {
@@ -53,7 +59,7 @@ export interface QwenMtpTarget {
 export interface TargetView {
   /** Opaque identity for providers borrowing weights from one exact target. */
   readonly identity: object;
-  readonly assistant?: AssistantTarget;
+  readonly assistantRows?: AssistantRowsTarget;
   readonly gemmaTaps?: { readonly layerCount: number; readonly projection: DraftProjection };
   readonly qwenMtp?: QwenMtpTarget;
 }
@@ -61,12 +67,13 @@ export interface TargetView {
 /** 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;
+  /** Method-owned companion state, independent of retention and storage.
+   * Capture returns owned immutable views. Restore borrows its attachment;
+   * the source retains the views it needs before returning. */
+  readonly checkpoint?: {
+    readonly namespace: string;
+    capture(processedTokens: number): import("../backends/mlx/checkpoint-state").CheckpointAttachment;
+    restore(processedTokens: number, attachment: import("../backends/mlx/checkpoint-state").CheckpointAttachment): void;
   };
   /** Target prefill shape required by this source's oracle. Most mlx-lm
    *  sources leave the final prompt token pending; native Colibri MTP starts
@@ -86,8 +93,9 @@ export interface DraftSource {
   /** Process the prompt (two-model: prefill the draft model's cache;
    *  assistant/dflash: read the target's state — mostly a no-op, but DSpark
    *  seeds H_ctx from `ctxML`, the tapped prefill context [1,Lp,m*H], present
-   *  iff tapLayers is set). */
-  prefill(promptIds: number[], ctxML?: MlxArray): void | Promise<void>;
+   *  iff tapLayers is set). A composed prefill may supply its exact processed
+   *  token count, so a history source need not infer the target's boundary. */
+  prefill(promptIds: number[], ctxML?: MlxArray, processedTokens?: number): void | Promise<void>;
 
   /** Propose 0..n tokens (RETURN LENGTH IS AUTHORITATIVE — a source may
    *  return fewer than n, e.g. DSpark's confidence-scheduled draft-length
@@ -132,12 +140,69 @@ export interface DraftSource {
   dispose(): void;
 }
 
+export interface DraftRowSampling {
+  /** Borrow [B,V] log-probabilities; return owned [B] IDs on device. */
+  sample(logprobs: MlxArray, steps: readonly number[]): MlxArray;
+}
+
+export interface DraftRowCheckpoint {
+  readonly processedTokens: number;
+  readonly attachment: CheckpointAttachment;
+}
+
+/** Method-owned state membership, independent of scheduling and cache tiers.
+ * Open/append borrow checkpoints; capture returns owned immutable state.
+ * Membership changes and capture occur only between committed rounds.
+ * The caller retains the provider's residency lease for the group's lifetime. */
+export interface DraftRowGroup extends MlxDraftRows {
+  readonly namespace: string;
+  readonly tapLayers: readonly number[];
+  readonly rowCount: number;
+  append(checkpoints: readonly DraftRowCheckpoint[]): void;
+  prepareAppend(checkpoints: readonly DraftRowCheckpoint[]): PreparedStateChange;
+  filterRows(keep: readonly number[]): void;
+  capture(row: number): DraftRowCheckpoint;
+  dispose(): void;
+}
+
+/** Draft preparation consumes each target chunk at the same batch geometry.
+ * Inputs are borrowed and contain only new tokens/context, with equal chunk
+ * lengths across rows. A null checkpoint appends a cold request. The provider
+ * owns alignment and retained hidden state; scheduling and persistence do not.
+ * Membership and capture occur between prefill calls. Capture returns owned
+ * state after at least one target token has been consumed for that row. */
+export interface DraftPrefillGroup {
+  readonly namespace: string;
+  readonly prefillMode: NonNullable<DraftSource["prefillMode"]>;
+  readonly tapLayers: readonly number[];
+  readonly rowCount: number;
+  append(checkpoints: readonly (DraftRowCheckpoint | null)[]): void;
+  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange;
+  prefill(tokens: MlxArray, context?: MlxArray): void | Promise<void>;
+  /** Resolve borrowed backing before its external restore lease is released. */
+  materialize(): void;
+  filterRows(keep: readonly number[]): void;
+  capture(row: number): DraftRowCheckpoint;
+  dispose(): void;
+}
+
+export interface GroupedDraftProvider {
+  /** Draft state remains valid when target forwards run under a mounted
+   * adapter context. Unqualified learned draft graphs leave this absent. */
+  readonly supportsTargetAdapters?: boolean;
+  openPrefill(options: { target: TargetView;
+    checkpoints: readonly (DraftRowCheckpoint | null)[] }): DraftPrefillGroup;
+  open(options: { target: TargetView; sampling: DraftRowSampling;
+    checkpoints: readonly DraftRowCheckpoint[] }): DraftRowGroup;
+}
+
 /** Server-lifetime owner of the draft machinery (the loaded draft model);
- *  open() mints a per-request DraftSource. */
+ * open() mints a per-request DraftSource; grouped opens shared row state. */
 export interface DraftProvider {
   /** Human-readable id (registry id / path tail) for logs + cache namespacing. */
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped?: GroupedDraftProvider;
   open(opts: {
     /** The request's sampler over logprobs [1,V] → token array [1] (the SAME
      *  sampler as the target — mlx-lm parity; greedy drafting under a

package/src/spec/two-model-rows.ts

@@ -0,0 +1,61 @@
+import type { MlxAutoregressiveBinding } from "../backends/mlx/autoregressive";
+import { MlxAutoregressiveDraftRows, type MlxDraftLayout } from "../backends/mlx/autoregressive-draft-rows";
+import { prefillCacheLayout } from "../backends/mlx/cache-layout";
+import { MlxStateRows } from "../backends/mlx/state-rows";
+import { captureCacheAttachment, restoreCacheAttachment } from "../backends/mlx/cache-attachment";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../engine/resources";
+import type { Cache } from "../model/gemma4-base";
+import type { DraftPrefillGroup, DraftRowCheckpoint, DraftRowGroup, DraftRowSampling, GroupedDraftProvider } from "./source";
+import type { PreparedStateChange } from "../contracts/resources";
+import type { MlxArray } from "../mlx/array";
+
+export interface StandaloneDraftRows extends DraftPrefillGroup {
+  draft: DraftRowGroup["draft"];
+  commit(accepted: readonly number[], context?: MlxArray): Promise<void>;
+}
+
+/** The standalone model supplies graph/state capabilities once. Target
+ * verification, scheduling and cache tiers use the existing provider ports. */
+export function openStandaloneDraftRows(binding: MlxAutoregressiveBinding, namespace: string,
+  sampling: DraftRowSampling | null, checkpoints: readonly (DraftRowCheckpoint | null)[]): StandaloneDraftRows {
+    const empty = binding.makeCache(), layouts: MlxDraftLayout[] = [];
+    try { for (const cache of empty) layouts.push(prefillCacheLayout(cache) as MlxDraftLayout); }
+    catch (error) { disposeResources(layouts); throw error; }
+    finally { disposeResources(empty); }
+    const state = new MlxStateRows(layouts), rows = new MlxAutoregressiveDraftRows(binding.graph, state, sampling);
+    const prepareAppend = (checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange => {
+      const restored: Cache[][] = [], pending: Array<number | null> = [];
+      try {
+        for (const checkpoint of checkpoints) {
+          restored.push(checkpoint ? restoreCacheAttachment(checkpoint.attachment, binding.makeCache) : binding.makeCache());
+          pending.push(checkpoint && checkpoint.attachment.metadata.pendingLast !== false ? Number(checkpoint.attachment.metadata.pendingLast) : null);
+        }
+        return rows.prepareAppend(restored, pending);
+      } finally { disposeResources(restored.flat()); }
+    };
+    const append = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => applyStateChanges([() => prepareAppend(checkpoints)]);
+    try { append(checkpoints); }
+    catch (error) { return cleanupFailure(error, () => rows.dispose()); }
+    return {
+      namespace, prefillMode: "tail-split", tapLayers: [],
+      get rowCount() { return state.rowCount; },
+      append, prepareAppend, prefill: tokens => rows.prefill(tokens), materialize: () => rows.materialize(),
+      filterRows: keep => rows.filterRows(keep), draft: rows.draft.bind(rows), commit: accepted => rows.commit(accepted),
+      capture(row): DraftRowCheckpoint {
+        const caches = state.extractRow(row), pendingLast = rows.pendingLast[row] ?? null;
+        try {
+          const processedTokens = caches[0]!.offset + Number(pendingLast !== null);
+          return { processedTokens, attachment: captureCacheAttachment("standalone-draft-v1", caches,
+            { processedTokens, pendingLast: pendingLast ?? false }) };
+        } finally { disposeResources(caches); }
+      },
+      dispose: () => rows.dispose(),
+    };
+}
+
+export function standaloneDraftGroups(binding: MlxAutoregressiveBinding, namespace: string): GroupedDraftProvider {
+  return {
+    open: options => openStandaloneDraftRows(binding, namespace, options.sampling, options.checkpoints),
+    openPrefill: options => openStandaloneDraftRows(binding, namespace, null, options.checkpoints),
+  };
+}

package/src/spec/two-model.ts

@@ -20,9 +20,11 @@ import { flagOn } from "../runtime-config";
 import { loadModelConfig } from "../config";
 import { Weights } from "../weights";
 import { createModel, type RuntimeModel } from "../model/factory";
-import type { Cache } from "../model/gemma4";
-import { toLogprobs } from "../sampler";
-import type { DraftProvider, DraftSource } from "./source";
+import type { DraftProvider, DraftSource, GroupedDraftProvider } from "./source";
+import { standaloneDraftGroups, openStandaloneDraftRows } from "./two-model-rows";
+import { bindLegacyAutoregressiveModel } from "../backends/mlx/autoregressive";
+import { artifactIdentity } from "../model/artifact-identity";
+import { configFingerprint } from "../model/fingerprint";
 
 /** mlx-lm's prefill_step_size — the draft drain chunks at the same stride
  *  as the target's (serve-loop.ts); chunking is numerically exact for
@@ -32,13 +34,16 @@ const PREFILL_CHUNK = 2048;
 export class TwoModelProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped: GroupedDraftProvider;
 
   private constructor(
     readonly model: RuntimeModel,
     private readonly weights: Weights,
     id: string,
+    private readonly namespace: string,
   ) {
     this.id = id;
+    this.grouped = standaloneDraftGroups(bindLegacyAutoregressiveModel(model), namespace);
     this.weightsBytes = [...weights.shards.files.values()].reduce(
       (a, f) => a + f.mmap.size,
       0,
@@ -51,20 +56,25 @@ export class TwoModelProvider implements DraftProvider {
   static async load(modelDir: string, targetVocabSize?: number): Promise<TwoModelProvider> {
     const config = await loadModelConfig(modelDir);
     const weights = await Weights.open(modelDir);
-    const model = createModel(weights, config);
-    if (
-      targetVocabSize !== undefined &&
-      config.text.vocabSize !== targetVocabSize
-    )
-      console.warn(
-        `draft model vocab size ${config.text.vocabSize} != target ${targetVocabSize} — ` +
-          `speculation will accept nothing useful (mlx-lm warns identically)`,
-      );
-    return new TwoModelProvider(model, weights, modelDir.split("/").filter(Boolean).at(-1)!);
+    try {
+      const model = createModel(weights, config);
+      if (
+        targetVocabSize !== undefined &&
+        config.text.vocabSize !== targetVocabSize
+      )
+        console.warn(
+          `draft model vocab size ${config.text.vocabSize} != target ${targetVocabSize} — ` +
+            `speculation will accept nothing useful (mlx-lm warns identically)`,
+        );
+      const identity = await artifactIdentity(configFingerprint(config),
+        [...weights.shards.files].map(([name, shard]) => ({ name, path: shard.path })));
+      return new TwoModelProvider(model, weights, modelDir.split("/").filter(Boolean).at(-1)!,
+        `standalone-draft-v1:${identity}`);
+    } catch (error) { weights.dispose(); throw error; }
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
-    return new TwoModelSource(this.model, opts.sampler);
+    return new TwoModelSource(this.model, opts.sampler, this.namespace);
   }
 
   dispose(): void {
@@ -72,100 +82,32 @@ export class TwoModelProvider implements DraftProvider {
   }
 }
 
+/** Compatibility adapter: the existing request loop also uses the shared
+ * draft graph at B=1. Proposal generation and state recovery have one owner. */
 class TwoModelSource implements DraftSource {
-  readonly weightsBytes = 0; // provider-owned weights; per-request adds caches only
-  private caches: Cache[];
-
-  constructor(
-    private readonly model: RuntimeModel,
-    private readonly sampler: (lp: MlxArray, step: number) => MlxArray,
-  ) {
-    this.caches = this.model.makeCache();
+  readonly weightsBytes = 0;
+  private readonly rows: ReturnType<typeof openStandaloneDraftRows>;
+  private processed = 0;
+  constructor(model: RuntimeModel, sampler: (lp: MlxArray, step: number) => MlxArray, namespace: string) {
+    this.rows = openStandaloneDraftRows(bindLegacyAutoregressiveModel(model), namespace,
+      { sample: (logprobs, steps) => sampler(logprobs, steps[0]!) }, [null]);
   }
-
-  prefill(promptIds: number[]): void {
-    // Oracle convention (mlx-lm speculative_generate_step._prefill, re-anchored
-    // 2026-07-07): the DRAFT model also drains only to len-1 (`while y.size >
-    // 1`) — the last prompt token stays unprocessed and is consumed by the
-    // FIRST draft() round's feed ([pending] = that same token under the serve
-    // loop's oracle shape), mirroring _draft_generate's first _step. Draft-side
-    // full-prefill was the residual knife-edge flipper in the 2026-07-07 live
-    // oracle gate (γ=2 haiku cell). MLX_BUN_PREFILL_TAIL_SPLIT=0 reverts.
-    // (mlx-lm chunks its drain at prefill_step_size=2048; chunking is
-    // numerically exact for causal attention — same KV, same positions —
-    // so we chunk too. Pre-2026-07-07 this was a SINGLE forward of the
-    // whole head: nothing enforced the "<2048 serve regime" the old
-    // comment assumed, and a 32k prompt through the spec lane ran one
-    // 32k-position draft forward — a large activation/mask transient on a
-    // 24 GB box, and untested numerics vs the chunked oracle.)
+  async prefill(promptIds: number[]): Promise<void> {
     const tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
     const upTo = tailSplit && promptIds.length > 1 ? promptIds.length - 1 : promptIds.length;
-    for (let at = 0; at < upTo; at += PREFILL_CHUNK) {
+    for (let at = this.processed; at < upTo; at += PREFILL_CHUNK) {
       const end = Math.min(at + PREFILL_CHUNK, upTo);
-      const ids = ops.fromInt32(promptIds.slice(at, end), [1, end - at]);
-      const h = this.model.forwardHidden(ids, this.caches);
-      ids.dispose();
-      // The head's last-position logits are not needed: the serve loop's
-      // first `feed` is consumed at the top of draft().
-      h.dispose();
-      clearCache();
-    }
-  }
-
-  /** One forward over `feed`, sample the first draft from its last position,
-   *  then n-1 single-token steps — mirrors _draft_generate's autoregressive
-   *  chain (with async_eval pipelining upstream; eager here, the draft is
-   *  small and the serve loop overlaps nothing with it in v1). */
-  draft(feed: number[], n: number, stepBase: number): number[] {
-    const drafts: number[] = [];
-    let input = feed;
-    for (let k = 0; k < n; k++) {
-      const ids = ops.fromInt32(input, [1, input.length]);
-      const h = this.model.forwardHidden(ids, this.caches);
-      ids.dispose();
-      const L = h.shape[1]!;
-      const H = h.shape[2]!;
-      const hLast = h.slice([0, L - 1, 0], [1, L, H]);
-      h.dispose();
-      const logits = this.model.logitsFromHidden(hLast);
-      hLast.dispose();
-      const V = logits.shape[logits.shape.length - 1]!;
-      const flat = ops.reshape(logits, [1, V]);
-      logits.dispose();
-      const lp = toLogprobs(flat);
-      flat.dispose();
-      const tokArr = this.sampler(lp, stepBase + k);
-      lp.dispose();
-      const tok = ops.itemUint32(tokArr);
-      tokArr.dispose();
-      drafts.push(tok);
-      input = [tok];
+      using ids = ops.fromInt32(promptIds.slice(at, end), [1, end - at]);
+      await this.rows.prefill(ids); clearCache();
     }
-    return drafts;
+    this.processed = upTo;
   }
-
-  commit(n: number, kAccept: number): void {
-    const rewind = Math.max(n - kAccept - 1, 0);
-    if (rewind > 0) {
-      for (const c of this.caches) {
-        if (!c.isTrimmable()) {
-          // The serve loop stops speculating on the TARGET-side trim check
-          // before this can matter; a wrapped draft ring here just means the
-          // source is done being useful. Recreate is the safe recovery.
-          this.resetCaches();
-          return;
-        }
-        c.trim(rewind);
-      }
-    }
-  }
-
-  private resetCaches(): void {
-    for (const c of this.caches) c.dispose();
-    this.caches = this.model.makeCache();
+  async draft(feed: number[], count: number, step: number): Promise<number[]> {
+    return (await this.rows.draft([feed.at(-1)!], count, [step]))[0]!;
   }
-
-  dispose(): void {
-    for (const c of this.caches) c.dispose();
+  async commit(_drafted: number, accepted: number): Promise<void> {
+    await this.rows.commit([accepted]);
+    this.processed += accepted + 1;
   }
+  dispose(): void { this.rows.dispose(); }
 }