mlx-bun 0.3.0 → 0.5.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, {
@@ -120,7 +133,9 @@ async function specRunInner(
     sampler,
   });
   const gamma = Math.max(1, numDraftTokens);
+  const trace = runtimeConfig().value("MLX_BUN_SPEC_TRACE") === "1";
   const prefillChunk = options.prefillChunkSize ??
+    binding.prefillPolicy?.chunkSize(promptIds.length) ??
     runtimeConfig().number("MLX_BUN_RD_PREFILL_CHUNK", PREFILL_CHUNK);
 
   // Allocated INSIDE the try below (2026-07-07 review): provider.open()
@@ -129,6 +144,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 +191,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 +203,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, options.cacheSessionId);
+    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, options.cacheSessionId);
+        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 +306,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 +331,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;
@@ -416,7 +447,7 @@ async function specRunInner(
       );
       const d = drafts.length;
       if (d < 0 || d > n) throw new Error(`DraftSource returned ${d} drafts (contract: 0..${n})`);
-      if (Bun.env.MLX_BUN_SPEC_TRACE === "1") {
+      if (trace) {
         console.error(
           `[SPEC_TRACE] round=${extras.rounds + 1} ` +
           `feed=${feed.join(",")} drafts=${drafts.join(",")}`,
@@ -460,39 +491,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 +579,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,82 @@ 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;
+}
+
+/** Request-owned constraints can propose known continuations without changing
+ * their committed state. The target verifier still decides every output. */
+export interface DraftRowConstraints {
+  propose(row: number, maxTokens: number): Promise<number[]>;
+}
+
+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 {
+  /** Consume externally chosen target inputs at a committed round boundary.
+   * Context and tokens are borrowed rectangular windows; lengths retain each
+   * row's consumed prefix, excluding right padding and the pending output. */
+  consume?(tokens: MlxArray, context: MlxArray, lengths: readonly number[]): void | Promise<void>;
+  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 {
+  readonly supportsExternalTokens?: boolean;
+  /** Resolve persistence identity without allocating a draft row. */
+  checkpointNamespace?(): string;
+  /** 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; constraints?: DraftRowConstraints;
+    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,62 @@
+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 {
+    checkpointNamespace: () => namespace,
+    open: options => openStandaloneDraftRows(binding, namespace, options.sampling, options.checkpoints),
+    openPrefill: options => openStandaloneDraftRows(binding, namespace, null, options.checkpoints),
+  };
+}