mlx-bun 0.0.10 → 0.0.12

package/src/spec/source.ts

@@ -0,0 +1,118 @@
+// DraftSource — the seam between the serve-time speculative verify loop
+// (src/spec/serve-loop.ts) and whatever produces draft tokens. Designed in
+// docs/design/mlx-lm-tool-parity-plan.md §7 so every drafter shares ONE
+// verify/accept executor:
+//   - TwoModelSource (src/spec/two-model.ts) — mlx-lm parity (L1 oracle:
+//     mlx_lm.server --draft-model), a full second model. Ignores `target`.
+//   - AssistantSource (src/spec/assistant-source.ts) — the optiq KV-borrowing
+//     Gemma drafter (src/spec/drafter.ts; L2 oracle: optiq spec_generate).
+//     Reads the target's donor K/V + anchor hidden each step.
+//   - DflashSource (src/spec/dflash-source.ts) — DSpark (L3, KL/quality-gated).
+//     Taps the target's multi-layer hiddens (prefill + verify) into a growing
+//     H_ctx (docs/design/dspark-speculative-decoding.md).
+//   - NgramSource (src/spec/ngram-source.ts) — model-free prompt lookup.
+//   - Glm52NativeMtpSource (src/spec/glm52-mtp-source.ts) — the target
+//     artifact's native Colibri MTP row.
+//
+// The sources differ ONLY in what fills the draft; the serve loop, admission
+// accounting, and stats never change. The KV-borrowing sources need target
+// state the two-model source doesn't — carried by the `target` view at open()
+// and the optional prefill/draft/commit arguments below (all no-ops for
+// two-model). See [[dspark-seam-kv-borrowing]].
+
+import type { MlxArray } from "../mlx/array";
+import type { RuntimeModel } from "../model/factory";
+import type { Cache } from "../model/gemma4";
+
+/** Read-only handle to the target's live decode state, handed to every source
+ *  at open(). Two-model ignores it; KV-borrowing sources read the target's
+ *  model (embed / lm-head / the draft module) and its LIVE cache array (donor
+ *  K/V views + offset — the SAME caches the serve loop drives). */
+export interface TargetView {
+  readonly model: RuntimeModel;
+  readonly caches: Cache[];
+}
+
+/** 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 {
+  /** Target prefill shape required by this source's oracle. Most mlx-lm
+   *  sources leave the final prompt token pending; native Colibri MTP starts
+   *  from a full-prompt target forward. */
+  readonly prefillMode?: "tail-split" | "full";
+  /** Request one fixed target kernel family across the speculative verify
+   * batch. Native GLM MTP uses the direct-Colibri SPEC_PIN contract. */
+  readonly pinTargetKernelFamily?: boolean;
+
+  /** Multi-layer target tap the source needs captured on the target's prefill
+   *  AND every verify forward (DSpark's H_ctx; e4b {20,31,41,42}). When set,
+   *  the serve loop sets model.hiddenTap around those forwards and passes the
+   *  captured context [1,L,m*H] into prefill()/commit(). Undefined for sources
+   *  that don't tap (two-model, assistant). */
+  readonly tapLayers?: number[];
+
+  /** 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>;
+
+  /** Propose 0..n tokens (RETURN LENGTH IS AUTHORITATIVE — a source may
+   *  return fewer than n, e.g. DSpark's confidence-scheduled draft-length
+   *  pruning; ZERO means "skip drafting this round" — DeepSpec ℓ=0 semantics,
+   *  the serve loop degenerates to one plain target step, still tapped +
+   *  committed for context-growing sources). The serve loop verifies over
+   *  exactly the returned length. `feed` is the token(s) the draft has not
+   *  yet consumed: [pending]
+   *  after a rejection/first round, or [lastDraft, bonus] after an all-accept
+   *  round (mlx-lm's re-feed rule, generate.py:645-648). `stepBase` is the
+   *  emitted-token index, threaded to the sampler for per-step RNG streams.
+   *  `anchorHidden` is the target's final hidden [1,1,H] at the pending/anchor
+   *  position — the assistant source borrows it for its first draft step;
+   *  two-model and dflash ignore it. */
+  draft(
+    feed: number[],
+    n: number,
+    stepBase: number,
+    anchorHidden?: MlxArray,
+  ): number[] | Promise<number[]>;
+
+  /** Verify outcome for the last round: kAccept of d accepted, where d is the
+   *  length draft() actually RETURNED (≤ n). Two-model: trim the draft cache
+   *  by max(d - kAccept - 1, 0) — mlx-lm's rewind rule (generate.py:589-591).
+   *  DSpark: grow H_ctx by the accepted window from `vCtxML`, the verified
+   *  window's tapped context [1,d+1,m*H] (present iff tapLayers is set), and
+   *  drop the rejected tips. Native GLM MTP receives the target's verified
+   *  hidden window [1,d+1,H] in `verifiedHidden` and the kAccept accepted token
+   *  ids in `acceptedTokens`; it uses them to rebuild accepted MTP KV rows
+   *  from target state rather than recursively drafted hidden state. */
+  commit(
+    d: number,
+    kAccept: number,
+    vCtxML?: MlxArray,
+    verifiedHidden?: MlxArray,
+    acceptedTokens?: readonly number[],
+  ): void | Promise<void>;
+
+  /** Resident per-request draft weights, for admission accounting (0 when the
+   *  provider owns the weights). */
+  readonly weightsBytes: number;
+  dispose(): void;
+}
+
+/** Server-lifetime owner of the draft machinery (the loaded draft model);
+ *  open() mints a per-request DraftSource. */
+export interface DraftProvider {
+  /** Human-readable id (registry id / path tail) for logs + cache namespacing. */
+  readonly id: string;
+  readonly weightsBytes: number;
+  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
+     *  temperature>0 request is NOT parity). */
+    sampler: (logprobs: MlxArray, step: number) => MlxArray;
+    /** The target model + its live caches (donor views / offset). */
+    target: TargetView;
+  }): DraftSource;
+  dispose(): void;
+}

package/src/spec/two-model.ts

@@ -0,0 +1,171 @@
+// Two-model speculative drafting — mlx-lm parity (`mlx_lm.server
+// --draft-model`). A full second model drafts autoregressively; the serve
+// loop (src/spec/serve-loop.ts) verifies with the target. L1 oracle:
+// mlx_lm.server with the same target/draft pair, greedy, token-for-token
+// (spec-vs-spec — both batch the verify lm-head, so neither is bit-exact to
+// stock decode at knife-edges; see src/spec/generate.ts header).
+//
+// Faithfulness notes (read from the oracle venv's generate.py):
+//  - drafts are sampled with the REQUEST sampler (generate.py:593-601) —
+//    greedy drafting under a temperature>0 request is not parity;
+//  - draft-cache rewind is max(n - kAccept - 1, 0) (generate.py:589-591),
+//    with the all-accept re-feed of the last draft handled by the serve
+//    loop's `feed` (generate.py:645-648);
+//  - logits processors do NOT run on draft steps (they're target-side).
+
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { clearCache } from "../mlx/ffi";
+import { flagOn } from "../flags";
+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";
+
+/** 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
+ *  causal attention. */
+const PREFILL_CHUNK = 2048;
+
+export class TwoModelProvider implements DraftProvider {
+  readonly id: string;
+  readonly weightsBytes: number;
+
+  private constructor(
+    readonly model: RuntimeModel,
+    private readonly weights: Weights,
+    id: string,
+  ) {
+    this.id = id;
+    this.weightsBytes = [...weights.shards.files.values()].reduce(
+      (a, f) => a + f.mmap.size,
+      0,
+    );
+  }
+
+  /** Load the draft model from a snapshot dir. Vocab-size mismatch with the
+   *  target is a WARNING (mlx-lm parity, server.py:363-368); the tokenizer-
+   *  family hard check lives at the server (it owns both tokenizers). */
+  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)!);
+  }
+
+  open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
+    return new TwoModelSource(this.model, opts.sampler);
+  }
+
+  dispose(): void {
+    this.weights.dispose();
+  }
+}
+
+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();
+  }
+
+  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.)
+    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) {
+      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];
+    }
+    return drafts;
+  }
+
+  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();
+  }
+
+  dispose(): void {
+    for (const c of this.caches) c.dispose();
+  }
+}

package/src/ssd-cache.ts

@@ -21,7 +21,7 @@ import { join } from "node:path";
 import { randomUUID } from "node:crypto";
 import { commonPrefixLength } from "./prompt-cache";
 import {
-  saveKvCache, loadKvCache, readKvHeader,
+  saveKvCache, saveKvCacheAsync, loadKvCache, readKvHeader,
   type KvSaveMeta, type KvLoadExpect, type LoadedKvCache,
 } from "./kv-store";
 import type { Cache } from "./model/gemma4-base";
@@ -32,8 +32,27 @@ export interface SsdIndexEntry {
   tokens: number[];
   bytes: number;
   mtimeMs: number;
+  /** Whether every serialized cache can trim (derived from kinds+state at
+   *  store/scan time). An untrimmable entry (wrapped ring, SSM) can only
+   *  seed prompts it matches WITHOUT a trim — find() must not prefer it
+   *  over a shorter usable one (the 2026-07-06 restart-0 defect: the
+   *  [prompt+gen] file always won, always got rejected after restore). */
+  trimmable: boolean;
 }
 
+/** Trimmability from a file header alone (mirror of the live caches'
+ *  isTrimmable(): kv/qkv/turboquant always trim; rotating kinds only pre-wrap
+ *  (offset < maxSize); GLM compressed kinds always trim; ssm never). Keep in
+ *  sync with gemma4-base/qwen3-delta/glm52-cache. */
+const headerTrimmable = (caches: { kind: string; offset: number; maxSize?: number }[]): boolean =>
+  caches.every((c) =>
+    c.kind === "kv" || c.kind === "qkv" || c.kind === "turboquant" ||
+      c.kind === "mla" || c.kind === "mla-dsa" || c.kind === "mtp-mla"
+      ? true
+      : c.kind === "rotating" || c.kind === "rotating-qkv"
+        ? c.offset < (c.maxSize ?? 0)
+        : false);
+
 export interface SsdStoreOptions {
   dir: string;
   maxBytes: number;
@@ -94,6 +113,7 @@ export class SsdCacheStore {
         try {
           const h = readKvHeader(path); // verifies the header hash
           if (
+            h.modelId !== this.#opts.modelId ||
             h.configFingerprint !== this.#opts.configFingerprint ||
             h.tokenizerHash !== this.#opts.tokenizerHash
           )
@@ -101,6 +121,7 @@ export class SsdCacheStore {
           const st = statSync(path);
           this.#index.push({
             path, ns: h.ns ?? "", tokens: h.tokens, bytes: st.size, mtimeMs: st.mtimeMs,
+            trimmable: headerTrimmable(h.caches),
           });
         } catch {
           try { rmSync(path, { force: true }); } catch {}
@@ -110,13 +131,15 @@ export class SsdCacheStore {
     return this.#index.length;
   }
 
-  /** Best stored entry for `prompt` in `ns`: the longest common prefix
-   *  (capped at prompt.length−1 — at least one token must forward). An
-   *  entry LONGER than the matched prefix is still returned (the divergent-
-   *  tail pattern: same long context, different question) — the caller
-   *  trims after restore when every cache kind is trimmable, and falls
-   *  back to a fresh prefill when not (ring post-wrap, SSM). Pure index
-   *  lookup; no I/O. */
+  /** Best USABLE stored entry for `prompt` in `ns`: the longest common
+   *  prefix (capped at prompt.length−1 — at least one token must forward).
+   *  An entry LONGER than the matched prefix is returned only when it can
+   *  actually seed the prompt: the divergent tail must be trimmable after
+   *  restore. Untrimmable entries (ring post-wrap, SSM) that would need a
+   *  trim are SKIPPED here — before this gate, the big [prompt+gen] file
+   *  always outranked the usable boundary snapshot, got restored, and was
+   *  thrown away by take()'s backstop (the 2026-07-06 restart-0 + wasted
+   *  84 s restore-then-re-prefill path). Pure index lookup; no I/O. */
   find(prompt: number[], ns = ""): { entry: SsdIndexEntry; prefixLen: number } | null {
     let best: SsdIndexEntry | null = null;
     let bestLen = 0;
@@ -124,18 +147,22 @@ export class SsdCacheStore {
       if (e.ns !== ns) continue;
       const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
       if (p === 0 || p <= bestLen) continue;
+      if (e.tokens.length - p > 0 && !e.trimmable) continue;
       bestLen = p;
       best = e;
     }
     return best ? { entry: best, prefixLen: bestLen } : null;
   }
 
-  /** Restore an entry zero-copy (COW mmap; pages fault in lazily). Bumps
-   *  LRU mtime. On ANY failure the file is dropped and null returned. */
+  /** Restore an entry via kv-store's STREAMED COPY (bounded host
+   *  transient — live entry + one tensor; no mapping survives the call,
+   *  the caches own their bytes). Bumps LRU mtime. On ANY failure the
+   *  file is dropped and null returned. */
   restore(entry: SsdIndexEntry, model: { makeCache(): Cache[] }): LoadedKvCache | null {
     const t0 = performance.now();
     try {
       const expect: KvLoadExpect = {
+        modelId: this.#opts.modelId,
         configFingerprint: this.#opts.configFingerprint,
         tokenizerHash: this.#opts.tokenizerHash,
         ns: entry.ns,
@@ -147,7 +174,10 @@ export class SsdCacheStore {
       this.stats.restores++;
       this.stats.restoreMsLast = performance.now() - t0;
       return loaded;
-    } catch {
+    } catch (err) {
+      // Loud: a failed restore silently degraded to a full re-prefill for
+      // a month before the 2026-07-06 bench made it observable.
+      console.warn(`[ssd-cache] restore failed, dropping ${entry.path}: ${(err as Error).message}`);
       this.remove(entry.path);
       return null;
     }
@@ -165,35 +195,81 @@ export class SsdCacheStore {
     const path = join(dir, `${randomUUID()}.mlxkv`);
     try {
       mkdirSync(dir, { recursive: true });
-      const meta: KvSaveMeta = {
-        modelId: this.#opts.modelId,
-        configFingerprint: this.#opts.configFingerprint,
-        tokenizerHash: this.#opts.tokenizerHash,
-        ns,
-      };
-      saveKvCache(path, tokens, caches, meta);
-      const st = statSync(path);
-      if (st.size > this.#opts.maxBytes) {
-        rmSync(path, { force: true });
-        return false;
-      }
-      // Supersede shadowed ancestors of this prefix (same ns).
+      saveKvCache(path, tokens, caches, this.#meta(ns));
+      return this.#indexStored(path, tokens, caches, ns);
+    } catch (err) {
+      return this.#storeFailed(path, err);
+    }
+  }
+
+  /** Non-blocking store (the write-behind persistence path): same file,
+   *  index, and supersede semantics as store(), but the tensor flush
+   *  yields the event loop between tensors so serving interleaves.
+   *  Caller passes zero-copy CLONES it owns (a consistent snapshot no
+   *  matter what the live entry does meanwhile) and disposes them after.
+   *  `waitTurn` gates every per-tensor step (see saveKvCacheAsync) — the
+   *  server passes the gateway's onIdle so the flush only progresses while
+   *  the engine is idle, never mid-decode. */
+  async storeAsync(tokens: number[], caches: Cache[], ns = "", waitTurn?: () => Promise<void>): Promise<boolean> {
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      mkdirSync(dir, { recursive: true });
+      await saveKvCacheAsync(path, tokens, caches, this.#meta(ns), waitTurn);
+      return this.#indexStored(path, tokens, caches, ns);
+    } catch (err) {
+      return this.#storeFailed(path, err);
+    }
+  }
+
+  #meta(ns: string): KvSaveMeta {
+    return {
+      modelId: this.#opts.modelId,
+      configFingerprint: this.#opts.configFingerprint,
+      tokenizerHash: this.#opts.tokenizerHash,
+      ns,
+    };
+  }
+
+  #indexStored(path: string, tokens: number[], caches: Cache[], ns: string): boolean {
+    const st = statSync(path);
+    if (st.size > this.#opts.maxBytes) {
+      rmSync(path, { force: true });
+      console.warn(`[ssd-cache] entry not stored: ${st.size} bytes exceeds the ${this.#opts.maxBytes}-byte cap`);
+      return false;
+    }
+    const trimmable = caches.every((c) => c.isTrimmable());
+    // Exact duplicates (same tokens) are replaced regardless of
+    // trimmability — the new file serves exactly the old one's matches.
+    for (const e of [...this.#index]) {
+      if (e.ns !== ns || e.tokens.length !== tokens.length) continue;
+      if (commonPrefixLength(e.tokens, tokens) === tokens.length) this.remove(e.path);
+    }
+    // Supersede shadowed ancestors of this prefix (same ns) — ONLY when
+    // the new entry can serve their prefixes (all caches trimmable; the
+    // same guard as PromptCache.put()). An untrimmable [prompt+gen]
+    // entry deleting the boundary snapshot was the gemma restart-0
+    // defect (2026-07-06): the only restorable file died the moment the
+    // unrestorable one landed.
+    if (trimmable) {
       for (const e of [...this.#index]) {
         if (e.ns !== ns || e.tokens.length >= tokens.length) continue;
         if (commonPrefixLength(e.tokens, tokens) === e.tokens.length) this.remove(e.path);
       }
-      this.#index.push({ path, ns, tokens, bytes: st.size, mtimeMs: Date.now() });
-      this.stats.spills++;
-      this.evictToCap();
-      return true;
-    } catch (err) {
-      try { rmSync(path, { force: true }); rmSync(`${path}.tmp`, { force: true }); } catch {}
-      if (!this.#warnedWriteFailure) {
-        this.#warnedWriteFailure = true;
-        console.warn(`[ssd-cache] store failed (disk full or unwritable?) — cold tier disabled for this entry: ${(err as Error).message}`);
-      }
-      return false;
     }
+    this.#index.push({ path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable });
+    this.stats.spills++;
+    this.evictToCap();
+    return true;
+  }
+
+  #storeFailed(path: string, err: unknown): boolean {
+    try { rmSync(path, { force: true }); rmSync(`${path}.tmp`, { force: true }); } catch {}
+    if (!this.#warnedWriteFailure) {
+      this.#warnedWriteFailure = true;
+      console.warn(`[ssd-cache] store failed (disk full or unwritable?) — cold tier disabled for this entry: ${(err as Error).message}`);
+    }
+    return false;
   }
 
   /** Enforce the byte cap: unlink oldest-mtime entries until under it. */