mlx-bun 0.0.10 → 0.0.11

package/src/spec/dspark/quantize-drafter.ts

@@ -0,0 +1,91 @@
+// Quantize a DeepSpec-format drafter checkpoint (Phase 1a of the DSpark
+// serving program, docs/design/dspark-serving-program.md).
+//
+// Thin policy layer over the house quantizer (src/quantize): the generic
+// shape-eligibility rule already captures every 2-D matmul weight — the
+// layer projections, fc, lm_head, markov_w2, and the two gather tables
+// (embed_tokens, markov_w1: quantized-embedding gather-dequant at load,
+// mlx QuantizedEmbedding semantics). The drafter-specific overlay:
+//
+//   KEEP bf16: confidence_head.* — threshold-sensitive (the truncation
+//   comparison runs on its sigmoid; borderline flips change proposal
+//   lengths) and tiny (~4k params), so quantizing it buys nothing.
+//   Norms / layer_scalar / biases are 1-D and never shape-eligible.
+//
+// Drafter numerics only move ACCEPTANCE, never correctness (the target
+// verifies every draft) — the quality gate for any scheme swept through
+// here is the Phase-1c acceptance A/B, not a KL battery.
+
+import { quantizeModelDir, type QuantizeResult } from "../../quantize";
+import { DeepspecDrafter } from "./deepspec-module";
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+
+/** Module bases kept full-precision regardless of shape eligibility. */
+const KEEP_BF16 = [/^confidence_head\./];
+
+/** The drafter per-tensor policy: quantize everything shape-eligible
+ *  except the KEEP_BF16 set. */
+export function drafterQuantizePredicate(base: string): boolean {
+  return !KEEP_BF16.some((re) => re.test(base));
+}
+
+export interface QuantizeDrafterOptions {
+  bits: 4 | 8;
+  groupSize: 32 | 64;
+  /** Skip the post-write load smoke (tests exercise loading separately). */
+  skipLoadSmoke?: boolean;
+  onProgress?: (stage: string, message: string, progress: number) => void;
+}
+
+/**
+ * Quantize the DeepSpec drafter at `srcDir` into `outDir` and smoke the
+ * result: reload through DeepspecDrafter (quantized detection + graph
+ * build) and eval one projectContext row through the quantized fc path.
+ */
+export async function quantizeDrafterDir(
+  srcDir: string,
+  outDir: string,
+  opts: QuantizeDrafterOptions,
+): Promise<QuantizeResult> {
+  const raw = (await Bun.file(`${srcDir}/config.json`).json()) as Record<string, any>;
+  if (raw.architectures?.[0] !== "Gemma4DSparkModel")
+    throw new Error(
+      `quantizeDrafterDir: expected a DeepSpec checkpoint (architectures[0]==="Gemma4DSparkModel"), got ${JSON.stringify(raw.architectures)} — for regular models use \`mlx-bun quantize\``,
+    );
+  if (raw.quantization ?? raw.quantization_config)
+    throw new Error("quantizeDrafterDir: source is already quantized — start from the bf16 checkpoint");
+
+  const result = await quantizeModelDir(
+    srcDir,
+    outDir,
+    {
+      bits: opts.bits,
+      groupSize: opts.groupSize,
+      quantizePredicate: drafterQuantizePredicate,
+    },
+    opts.onProgress ? (e) => opts.onProgress!(e.stage, e.message, e.progress) : undefined,
+  );
+
+  if (!opts.skipLoadSmoke) {
+    // Load smoke: quantized detection + one real quantized matmul (fc) so a
+    // wrong spec/layout fails HERE, not at serve time.
+    const d = await DeepspecDrafter.load(outDir);
+    try {
+      const width = d.tapLayers.length * d.hidden;
+      const zeros = MlxArray.fromFloat32(new Float32Array(width), [1, 1, width]);
+      const x = zeros.astype(Dtype.bfloat16);
+      zeros.dispose();
+      const proj = d.projectContext(x);
+      x.dispose();
+      const got = proj.shape;
+      proj.dispose();
+      if (got.length !== 3 || got[2] !== d.hidden)
+        throw new Error(`load smoke: projectContext shape ${JSON.stringify(got)}, expected [1,1,${d.hidden}]`);
+    } finally {
+      d.dispose();
+    }
+  }
+
+  return result;
+}

package/src/spec/generate.ts

@@ -128,18 +128,36 @@ export function specGenerate(
     full: (caches[donors.full] as never as { temporalView(): [MlxArray, MlxArray] }).temporalView(),
   });
 
+  // Round/prefill scratch tensors, hoisted so the finally frees them on a
+  // mid-round throw (the 2f not-trimmable throw, or any mlx op) — try-body
+  // locals are invisible to the finally. Each is nulled at its normal
+  // disposal, so the finally never double-frees. Same discipline as
+  // src/spec/serve-loop.ts / generate-dflash.ts (2026-07-06 review fix);
+  // see memory [[mlx-inline-slice-leak-pattern]].
+  let prefillH: MlxArray | null = null;
+  let lastHidden: MlxArray | null = null;
+  let vHidden: MlxArray | null = null;
+
   try {
-    // 1. prefill
+    // 1. prefill — FULL-prompt single forward, deliberately. Per-oracle
+    // convention (2026-07-07 re-anchor): mlx-lm's speculative path drains to
+    // len-1 with no step-0 (the SERVE loop, whose oracle it is, now mirrors
+    // that — src/spec/serve-loop.ts); THIS loop's oracle is optiq
+    // spec_generate, which prefills the whole prompt in one forward and
+    // samples token0 from it (runtime/spec/runtime.py:162-170, read from the
+    // installed oracle venv) — the 2026-06-14 bit-exact gate was established
+    // against exactly that shape. Do not tail-split here.
     const t0 = performance.now();
     const ids = ops.fromInt32(promptTokens, [1, promptTokens.length]);
-    const hidden = model.forwardHidden(ids, caches);
+    prefillH = model.forwardHidden(ids, caches);
     ids.dispose();
     stats.targetCalls++;
-    const H = hidden.shape[2]!;
-    const Lp = hidden.shape[1]!;
-    let next = pickFromHidden(model, hidden, Lp - 1);
-    let lastHidden = hidden.slice([0, Lp - 1, 0], [1, Lp, H]);
-    hidden.dispose();
+    const H = prefillH.shape[2]!;
+    const Lp = prefillH.shape[1]!;
+    let next = pickFromHidden(model, prefillH, Lp - 1);
+    lastHidden = prefillH.slice([0, Lp - 1, 0], [1, Lp, H]);
+    prefillH.dispose();
+    prefillH = null;
     stats.prefillMs = performance.now() - t0;
 
     // EOS convention: an EOS id stops generation and counts toward
@@ -158,30 +176,38 @@ export function specGenerate(
       const position = caches[0]!.offset - 1;
       const shared = readDonors();
 
-      // 2a. draft γ tokens against the frozen donor views
+      // 2a. draft γ tokens against the frozen donor views. try/finally so a
+      // mid-draft throw still frees the chained hiddens + donor views + the
+      // in-flight emb (assistant-source.ts draft() discipline).
       const drafts: number[] = [];
       let dTok = next;
-      let dHid = lastHidden; // borrowed for k=0
+      let dHid = lastHidden!; // borrowed for k=0
       const ownedHiddens: MlxArray[] = [];
-      for (let k = 0; k < gamma; k++) {
-        const emb = embedScaled(dTok);
-        const step = drafter.forward(emb, dHid, shared, position + k);
-        emb.dispose();
-        drafts.push(step.token);
-        stats.drafted++;
-        ownedHiddens.push(step.nextHidden);
-        dTok = step.token;
-        dHid = step.nextHidden;
-      }
-      for (const a of ownedHiddens) a.dispose();
-      for (const [k, v] of [shared.sliding, shared.full]) {
-        k.dispose();
-        v.dispose();
+      let emb: MlxArray | null = null;
+      try {
+        for (let k = 0; k < gamma; k++) {
+          emb = embedScaled(dTok);
+          const step = drafter.forward(emb, dHid, shared, position + k);
+          emb.dispose();
+          emb = null;
+          drafts.push(step.token);
+          stats.drafted++;
+          ownedHiddens.push(step.nextHidden);
+          dTok = step.token;
+          dHid = step.nextHidden;
+        }
+      } finally {
+        emb?.dispose();
+        for (const a of ownedHiddens) a.dispose();
+        for (const [k, v] of [shared.sliding, shared.full]) {
+          k.dispose();
+          v.dispose();
+        }
       }
 
       // 2b. verify all γ drafts (+ the pending token) in one forward
       const vIds = ops.fromInt32([next, ...drafts], [1, gamma + 1]);
-      const vHidden = model.forwardHidden(vIds, caches);
+      vHidden = model.forwardHidden(vIds, caches);
       vIds.dispose();
       stats.targetCalls++;
 
@@ -199,6 +225,7 @@ export function specGenerate(
         stats.emitted++;
         if (isEos || stats.emitted >= maxTokens) {
           vHidden.dispose();
+          vHidden = null;
           break outer;
         }
       }
@@ -221,18 +248,25 @@ export function specGenerate(
 
       if (emitIsEos || stats.emitted >= maxTokens) {
         vHidden.dispose();
+        vHidden = null;
         break;
       }
 
       // 2g. chain state
       next = emit;
-      lastHidden.dispose();
+      lastHidden!.dispose();
       lastHidden = vHidden.slice([0, kAccept, 0], [1, kAccept + 1, H]);
       vHidden.dispose();
+      vHidden = null;
     }
     stats.decodeMs = performance.now() - tDecode;
     lastHidden.dispose();
+    lastHidden = null;
   } finally {
+    // Survivors of a mid-round throw (each nulled at its normal disposal).
+    prefillH?.dispose();
+    lastHidden?.dispose();
+    vHidden?.dispose();
     for (const c of caches) c.dispose();
   }
 

package/src/spec/ngram-source.ts

@@ -0,0 +1,132 @@
+// NgramSource — MODEL-FREE prompt-lookup speculative drafting behind the
+// DraftSource seam (src/spec/source.ts). No weights, no caches, no tokenizer:
+// drafts are copied from the request's own token stream. When the trailing
+// k-gram of (prompt + emitted-so-far) has occurred earlier in that stream,
+// the tokens that followed the earlier occurrence are proposed as the draft;
+// the shared verify/accept executor (src/spec/serve-loop.ts) makes the result
+// LOSSLESS by construction (drafts are only proposals — the target's own
+// samples decide every emitted token, at any temperature; only the acceptance
+// rate moves). Best case: agentic/RAG/code-edit traffic that re-emits spans
+// already present in context.
+//
+// Prior art (ported, not invented): "Prompt Lookup Decoding" (Apoorv Saxena,
+// github.com/apoorvumang/prompt-lookup-decoding) and vLLM's `ngram` proposer
+// (vllm/v1 spec_decode ngram_proposer) — both match the LONGEST k-gram first
+// (k = max..min) and take the FIRST occurrence scanning left-to-right; we
+// follow that exactly. Defaults mirror Saxena (max_ngram_size=3,
+// num_pred_tokens=10 → our --num-draft-tokens default for this kind).
+//
+// Token-history reconstruction (the one subtle part): the seam hands sources
+// the prompt at prefill() and only the FEED tail each round — never the full
+// emitted stream — so this source rebuilds it from the feed/commit discipline
+// (same information two-model keeps in its draft KV):
+//   - prefill: history = prompt to len-1 under the oracle tail-split shape
+//     (the last prompt token arrives as the first round's feed = [pending]);
+//     full prompt under the legacy shape (token0 arrives as its feed).
+//   - draft(feed): history += feed. feed is [correction] after a rejected
+//     round, or [lastDraft, bonus] after an all-accept round (mlx-lm's
+//     re-feed rule, generate.py:645-648) — so the all-accept round's LAST
+//     draft is deliberately NOT pushed at commit (it arrives here instead).
+//   - commit(d, kAccept): history += drafts[0 .. min(kAccept, d-1)) — the
+//     accepted drafts except that re-fed last one. Rejected tips never enter
+//     history; the correction token arrives via the next feed.
+// Invariant at matching time (just after the feed push): history ==
+// prompt + every emitted token, ending at the pending/anchor token.
+//
+// A round with no match returns [] — the serve loop's d=0 semantics degrade
+// it to one plain target step (bit-equivalent to non-spec decode), so ngram
+// speculation is structurally never worse than plain decode by more than the
+// JS scan (~µs against 30k-token histories).
+
+import { flagOn } from "../flags";
+import type { DraftProvider, DraftSource } from "./source";
+
+export interface NgramOptions {
+  /** Longest suffix k-gram tried first (Saxena max_ngram_size). */
+  max?: number;
+  /** Shortest k-gram tried before giving up (vLLM prompt_lookup_min). */
+  min?: number;
+}
+
+export class NgramProvider implements DraftProvider {
+  readonly id = "ngram";
+  readonly weightsBytes = 0;
+  readonly max: number;
+  readonly min: number;
+
+  constructor(opts: NgramOptions = {}) {
+    this.max = Math.max(1, opts.max ?? 3);
+    this.min = Math.max(1, Math.min(opts.min ?? 1, this.max));
+  }
+
+  open(_opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
+    // Sampler and target both ignored: proposals come from lookup, not a
+    // model, so drafting is sampler-independent (like the assistant source —
+    // correctness holds at any temperature, only acceptance moves).
+    return new NgramSource(this.max, this.min);
+  }
+
+  dispose(): void {}
+}
+
+class NgramSource implements DraftSource {
+  readonly weightsBytes = 0;
+  #hist: number[] = [];
+  #lastDrafts: number[] = [];
+
+  constructor(
+    private readonly max: number,
+    private readonly min: number,
+  ) {}
+
+  /** Test hook — the reconstructed prompt+emitted stream (see header). */
+  get history(): readonly number[] {
+    return this.#hist;
+  }
+
+  prefill(promptIds: number[]): void {
+    // Mirror the serve loop's prefill shape (two-model.ts does the same): under
+    // the oracle tail split the last prompt token is never prefilled — it IS
+    // the first feed. Legacy shape (kill switch / 1-token prompt): full prompt,
+    // and the sampled token0 arrives as the first feed.
+    const tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+    const upTo = tailSplit && promptIds.length > 1 ? promptIds.length - 1 : promptIds.length;
+    this.#hist = promptIds.slice(0, upTo);
+  }
+
+  draft(feed: number[], n: number, _stepBase: number): number[] {
+    this.#hist.push(...feed);
+    this.#lastDrafts = this.#propose(n);
+    return this.#lastDrafts;
+  }
+
+  commit(d: number, kAccept: number): void {
+    // Accepted drafts join history, EXCEPT an all-accept round's last draft —
+    // the serve loop re-feeds it ([lastDraft, bonus]) and the next draft()'s
+    // feed push would double it (see header). d=0 rounds push nothing.
+    const upTo = Math.max(0, Math.min(kAccept, d - 1));
+    for (let i = 0; i < upTo; i++) this.#hist.push(this.#lastDrafts[i]!);
+  }
+
+  /** Longest-k-first, first-occurrence prompt lookup (Saxena/vLLM order):
+   *  find the trailing k-gram earlier in history, propose what followed it. */
+  #propose(n: number): number[] {
+    const h = this.#hist;
+    const L = h.length;
+    const kTop = Math.min(this.max, L - 1);
+    for (let k = kTop; k >= this.min; k--) {
+      const tailAt = L - k;
+      // First occurrence, left to right; i + k < L both excludes the trailing
+      // self-match and guarantees at least one continuation token.
+      search: for (let i = 0; i + k < L; i++) {
+        for (let j = 0; j < k; j++) {
+          if (h[i + j] !== h[tailAt + j]) continue search;
+        }
+        return h.slice(i + k, i + k + n);
+      }
+    }
+    return [];
+  }
+
+  dispose(): void {}
+}