mlx-bun 0.2.0 → 0.3.0

package/src/spec/assistant-source.ts

@@ -21,10 +21,8 @@
 import { readdirSync, statSync } from "node:fs";
 import { join } from "node:path";
 import type { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { Gemma4Model, type Cache } from "../model/gemma4";
 import { GemmaAssistantDrafter } from "./drafter";
-import type { DraftProvider, DraftSource, TargetView } from "./source";
+import type { DraftProvider, DraftSource, TargetView, AssistantTarget } from "./source";
 
 /** Sum of the drafter's on-disk safetensors — admission accounting only. */
 function safetensorsBytes(dir: string): number {
@@ -34,19 +32,6 @@ function safetensorsBytes(dir: string): number {
   return total;
 }
 
-/** The target's LAST sliding and LAST full cache owners (port of optiq
- *  kv_view.find_donor_layers; mirrors donorIndices in generate.ts). */
-function donorIndices(model: Gemma4Model): { sliding: number; full: number } {
-  let sliding = -1;
-  let full = -1;
-  for (let i = 0; i < model.numDonors; i++) {
-    if (model.layers[i]!.layerType === "sliding_attention") sliding = i;
-    else full = i;
-  }
-  if (sliding < 0 || full < 0) throw new Error("assistant drafter: missing donor layer type");
-  return { sliding, full };
-}
-
 export class AssistantProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
@@ -75,46 +60,27 @@ export class AssistantProvider implements DraftProvider {
   }
 }
 
-class AssistantSource implements DraftSource {
+export class AssistantSource implements DraftSource {
   readonly weightsBytes = 0; // provider-owned weights; per-request adds nothing
-  private readonly model: Gemma4Model;
-  private readonly caches: Cache[];
-  private readonly donors: { sliding: number; full: number };
+  private readonly target: AssistantTarget;
 
-  constructor(private readonly drafter: GemmaAssistantDrafter, target: TargetView) {
-    if (!(target.model instanceof Gemma4Model))
-      throw new Error("assistant drafter requires a Gemma4 target");
-    this.model = target.model;
-    this.caches = target.caches;
-    this.donors = donorIndices(this.model);
+  constructor(private readonly drafter: Pick<GemmaAssistantDrafter, "forward">, target: TargetView) {
+    if (!target.assistant)
+      throw new Error("assistant drafter requires a Gemma4 target with donor views");
+    this.target = target.assistant;
   }
 
   /** No own cache to prime — the drafter reads the target's live state. */
   prefill(_promptIds: number[], _ctxML?: MlxArray): void {}
 
-  private embedScaled(token: number): MlxArray {
-    const ids = ops.fromInt32([token], [1, 1]);
-    const e = this.model.embed.encode(ids);
-    ids.dispose();
-    const s = ops.mulScalar(e, this.model.embedScale);
-    e.dispose();
-    return s;
-  }
-
-  private readDonors(): { sliding: [MlxArray, MlxArray]; full: [MlxArray, MlxArray] } {
-    const view = (i: number) =>
-      (this.caches[i] as unknown as { temporalView(): [MlxArray, MlxArray] }).temporalView();
-    return { sliding: view(this.donors.sliding), full: view(this.donors.full) };
-  }
-
   /** Draft n tokens against the target's donor K/V, conditioning the first
    *  step on the anchor (the last emitted token + its target hidden) and
    *  chaining the drafter's own post-projected hiddens after that — exactly
    *  specGenerate's 2a block. anchorHidden is borrowed; never disposed here. */
   draft(feed: number[], n: number, _stepBase: number, anchorHidden?: MlxArray): number[] {
     if (!anchorHidden) throw new Error("assistant drafter needs the target anchor hidden");
-    const position = this.caches[0]!.offset - 1;
-    const shared = this.readDonors();
+    const position = this.target.position();
+    const shared = this.target.readDonors();
     const drafts: number[] = [];
     const ownedHiddens: MlxArray[] = [];
     let dTok = feed[feed.length - 1]!; // the pending token (anchorHidden's token)
@@ -122,7 +88,7 @@ class AssistantSource implements DraftSource {
     let emb: MlxArray | null = null; // hoisted so a mid-loop throw still frees it
     try {
       for (let k = 0; k < n; k++) {
-        emb = this.embedScaled(dTok);
+        emb = this.target.embedScaled(dTok);
         const step = this.drafter.forward(emb, dHid, shared, position + k);
         emb.dispose();
         emb = null;

package/src/spec/deepspec-source.ts

@@ -4,7 +4,7 @@
 // (github.com/deepseek-ai/DeepSpec, MIT); the module port with transcription
 // notes lives in src/spec/dspark/deepspec-module.ts. Oracle: DeepSpec's eval
 // at temperature 0 (RNG-free — verify degenerates to argmax token-match),
-// scripts/oracle-dspark-deepspec.py → scripts/dspark-deepspec-compare.ts.
+// scripts/oracle/oracle-dspark-deepspec.py → dspark-deepspec-compare.ts (deleted 2026-08-23; git history).
 //
 // Seam mapping (why this fits without serve-loop changes): the reference
 // recycles target hidden states from each verify pass into the drafter's
@@ -26,7 +26,6 @@ import { readdirSync, statSync } from "node:fs";
 import { join } from "node:path";
 import type { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
-import { Gemma4Model } from "../model/gemma4";
 import { DeepspecDrafter, type ContextKV } from "./dspark/deepspec-module";
 import type { DraftProvider, DraftSource, TargetView } from "./source";
 
@@ -65,16 +64,17 @@ export class DeepspecProvider implements DraftProvider {
   }
 }
 
-class DeepspecSource implements DraftSource {
+export class DeepspecSource implements DraftSource {
   readonly weightsBytes = 0; // provider-owned weights
   readonly tapLayers: number[];
   private ctxKV: ContextKV[] | null = null; // per-layer projected context rows
   private ctxLen = 0; // context rows cached == the anchor's absolute position
 
-  constructor(private readonly drafter: DeepspecDrafter, target: TargetView) {
-    if (!(target.model instanceof Gemma4Model))
+  constructor(private readonly drafter: Pick<DeepspecDrafter,
+    "cfg" | "tapLayers" | "projectContext" | "projectContextKV" | "draftBlock">, target: TargetView) {
+    if (!target.gemmaTaps)
       throw new Error("DeepSpec drafter requires a Gemma4 target");
-    const nLayers = target.model.layers.length;
+    const nLayers = target.gemmaTaps.layerCount;
     if (nLayers !== drafter.cfg.num_target_layers)
       throw new Error(
         `DeepSpec drafter was trained for a ${drafter.cfg.num_target_layers}-layer target; ` +

package/src/spec/dflash-source.ts

@@ -22,10 +22,9 @@ import { readdirSync, statSync } from "node:fs";
 import { join } from "node:path";
 import type { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
-import { Gemma4Model } from "../model/gemma4";
 import type { DflashDrafter } from "./dspark/module-dflash";
 import { loadDsparkDrafter } from "./dspark/loader";
-import type { DraftProvider, DraftSource, TargetView } from "./source";
+import type { DraftProvider, DraftSource, TargetView, DraftProjection } from "./source";
 import { runtimeValue } from "../runtime-config";
 
 function safetensorsBytes(dir: string): number {
@@ -65,16 +64,17 @@ export class DflashProvider implements DraftProvider {
   }
 }
 
-class DflashSource implements DraftSource {
+export class DflashSource implements DraftSource {
   readonly weightsBytes = 0; // provider-owned weights
   readonly tapLayers: number[];
-  private readonly model: Gemma4Model;
+  private readonly model: DraftProjection;
+  private readonly minConf = runtimeValue("MLX_BUN_DSPARK_MINCONF");
   private hCtx: MlxArray | null = null; // [1, L, m*H] — grows with the accepted stream
 
-  constructor(private readonly drafter: DflashDrafter, target: TargetView) {
-    if (!(target.model instanceof Gemma4Model))
+  constructor(private readonly drafter: Pick<DflashDrafter, "cfg" | "forwardInfer">, target: TargetView) {
+    if (!target.gemmaTaps)
       throw new Error("DSpark drafter requires a Gemma4 target");
-    this.model = target.model;
+    this.model = target.gemmaTaps.projection;
     this.tapLayers = drafter.cfg.tapLayers;
   }
 
@@ -94,7 +94,7 @@ class DflashSource implements DraftSource {
   draft(feed: number[], n: number, _stepBase: number, _anchorHidden?: MlxArray): number[] {
     if (!this.hCtx) throw new Error("DSpark drafter: draft before prefill");
     const anchorTok = feed[feed.length - 1]!;
-    const envMin = runtimeValue("MLX_BUN_DSPARK_MINCONF");
+    const envMin = this.minConf;
     const block = this.drafter.forwardInfer(this.model, this.hCtx, anchorTok, n, {
       thresholds: this.drafter.cfg.sts?.thresholds,
       collectLogits: false, // serve loop runs its own verify lm-head, never reads this

package/src/spec/dspark/ab-stats.ts

@@ -1,6 +1,6 @@
 // Model-free math for the Phase-1c drafter acceptance A/B
-// (docs/design/dspark-serving-program.md; runner:
-// scripts/dspark-drafter-ab.ts). THE gate for every drafter-quantization
+// (docs/design/speculative-decoding.md; runner:
+// scripts/dspark.ts ab). THE gate for every drafter-quantization
 // experiment: same target, same prompts, temp 0 — drafter A vs drafter B.
 //
 // Paired design: prompts are the pairing unit, so per-prompt deltas cancel
@@ -76,7 +76,7 @@ export function summarizeArm(results: AbPromptResult[]): AbArmSummary {
 }
 
 /** Phase-1d exit: acceptance drop ≤ maxDropPts absolute AND wall-clock
- *  strictly improves (docs/design/dspark-serving-program.md 1d). */
+ *  strictly improves (docs/design/speculative-decoding.md 1d). */
 export function pairedVerdict(
   aResults: AbPromptResult[],
   bResults: AbPromptResult[],

package/src/spec/dspark/calibration.ts

@@ -2,7 +2,7 @@
 // confidence thresholds from real (confidence, accepted) verify outcomes, so
 // the scheduler in module-dflash.ts#forwardInfer prunes positions the
 // confidence head can't actually call. Pure math, no MLX/GPU here: the
-// caller (scripts/dspark-calibrate.ts) collects samples off a live model via
+// caller (scripts/dspark.ts calibrate) collects samples off a live model via
 // dflashGenerate's onRound hook, this module just fits.
 //
 // VERIFIED vs the paper + reference (2026-07-06 audit, arXiv:2607.05147 +
@@ -14,7 +14,7 @@
 // survival products, feeding Alg-1's throughput-maximizing scheduler with a
 // profiled SPS(B) cost table) — that machinery lives in DeepSeek's UNRELEASED
 // production serving layer, not in DeepSpec. Paper-faithful STS/Alg-1 is a
-// future Lab item; details in docs/design/dspark-speculative-decoding.md
+// future Lab item; details in docs/design/speculative-decoding.md
 // "DeepSpec ground truth". The isotonic-style "smallest τ meeting a target
 // precision, Laplace-smoothed" fit below is ours (the reference ships no
 // calibrator at all — raw sigmoid vs a hand-set threshold).

package/src/spec/dspark/deepspec-module.ts

@@ -38,7 +38,7 @@ function disposing(old: MlxArray, next: MlxArray): MlxArray {
 // ---------------------------------------------------------------------------
 // Weight representation: every matmul weight and gather table loads as either
 // bf16 (the released checkpoint) or affine-quantized (a
-// scripts/dspark-quantize-drafter.ts sibling). Detection is the house
+// scripts/dspark.ts quantize sibling). Detection is the house
 // pattern — a `.scales` sibling key marks a quantized module; bits/group_size
 // come from the config's `quantization` block (per-module overrides + the
 // mlx `false` = kept-full-precision convention). Both formats flow through
@@ -366,7 +366,7 @@ export class DeepspecDrafter {
   static async load(dir: string): Promise<DeepspecDrafter> {
     const raw = (await Bun.file(`${dir}/config.json`).json()) as Record<string, any>;
     const cfg = readConfig(raw);
-    // Quantized siblings (scripts/dspark-quantize-drafter.ts) carry the house
+    // Quantized siblings (scripts/dspark.ts quantize) carry the house
     // quantization block; the released bf16 checkpoint has none → null.
     const quant = parseQuantization(raw.quantization ?? raw.quantization_config);
     const w = await Weights.open(dir);

package/src/spec/dspark/generate-dflash.ts

@@ -27,7 +27,7 @@ export interface DflashGenOptions {
    *  minConf for a uniform manual override; pass thresholds:[] to disable. */
   thresholds?: number[];
   minConf?: number;
-  /** Calibration hook (scripts/dspark-calibrate.ts): called once per draft
+  /** Calibration hook (scripts/dspark.ts calibrate): called once per draft
    *  round with that round's per-position (confidence, accepted) samples,
    *  AFTER kAccept is known. Not used by ordinary decode. */
   onRound?: (samples: ConfSample[]) => void;
@@ -192,7 +192,7 @@ export function dflashGenerate(model: Gemma4Model, drafter: DflashDrafter, promp
       blockLogits = null;
       stats.accepted += kAccept; rounds++; acceptLenSum += kAccept + 1;
 
-      // Calibration hook (scripts/dspark-calibrate.ts): report this round's
+      // Calibration hook (scripts/dspark.ts calibrate): report this round's
       // per-position (confidence, accepted) outcomes now that kAccept is known.
       if (options.onRound) {
         const roundSamples: ConfSample[] = [];

package/src/spec/dspark/generate.ts

@@ -17,9 +17,9 @@
 //     (sample.ts), so the emitted stream is distributed exactly as if e4b
 //     sampled token-by-token at that temperature — independent of draft quality.
 //
-// Verify the corresponding gate before chasing τ (scripts/dspark-ab.ts): greedy
+// Verify the corresponding gate before chasing τ (dspark-ab.ts (deleted 2026-08-23; git history)): greedy
 // output must match model.generate; the sampling rule's correctness is the
-// statistical "emit distribution == target" check in scripts/dspark-smoke.ts.
+// statistical "emit distribution == target" check in dspark-smoke.ts (deleted 2026-08-23; git history).
 // If a gate fails, the accept/reject or KV rollback is wrong — fix that first.
 
 import { MlxArray } from "../../mlx/array";

package/src/spec/dspark/loader.ts

@@ -18,6 +18,6 @@ export function loadDsparkDrafter(dir: string): DflashDrafter {
   if (meta.variant === "dspark" || meta.variant === "dflash") return DflashDrafter.load(dir);
   throw new Error(
     `${dir} is a v1 single-vector DSpark checkpoint (variant=${meta.variant ?? "none"}) — ` +
-      `superseded, not serveable; retrain with scripts/dspark-train-dflash.ts`,
+      `superseded, not serveable; retrain with scripts/dspark.ts train`,
   );
 }

package/src/spec/dspark/module-dflash.ts

@@ -17,6 +17,7 @@ import { MlxArray } from "../../mlx/array";
 import { Dtype } from "../../mlx/ffi";
 import * as ops from "../../mlx/ops";
 import type { Gemma4Model } from "../../model/gemma4";
+import type { DraftProjection } from "../source";
 import { loadAdapterTensors } from "../../lora";
 import { writeShardedSafetensors, type NamedTensor } from "../../quantize/safetensors-writer";
 import { processLogits, sampleToken, KeyStream, type DSparkSampleConfig } from "./sample";
@@ -25,7 +26,7 @@ import { join } from "node:path";
 
 /** STS calibration (§3.2.1) — per-position confidence thresholds fit on real
  *  (confidence, accepted) verify outcomes (src/spec/dspark/calibration.ts,
- *  scripts/dspark-calibrate.ts). Consumed by the confidence-scheduled
+ *  scripts/dspark.ts calibrate). Consumed by the confidence-scheduled
  *  draft-length pruning (Alg 1, single-user form) in forwardInfer. Stored
  *  inside DflashConfig so it round-trips through dspark.json for free.
  *  ABSENT ⟹ no pruning — an uncalibrated checkpoint drafts exactly as before. */
@@ -300,7 +301,7 @@ export class DflashDrafter {
     return H_d;
   }
 
-  #baseLogits(model: Gemma4Model, block: MlxArray): MlxArray {
+  #baseLogits(model: DraftProjection, block: MlxArray): MlxArray {
     const hOut = ops.matmul(block, this.get("out_proj")); // [A,γ,H]
     const hBf = hOut.astype(model.embed.scales.dtype); hOut.dispose();
     const logits = model.logitsFromHidden(hBf); hBf.dispose();
@@ -416,7 +417,7 @@ export class DflashDrafter {
    *  one toFloat32 after the loop). The sampling path already forces a host
    *  read every position (sampleToken draws from a device RNG key) — that
    *  path is unchanged. */
-  forwardInfer(model: Gemma4Model, hCtx: MlxArray, anchorTok: number, gamma: number, opts: DflashDraftOpts = {}): DflashDraftBlock {
+  forwardInfer(model: DraftProjection, hCtx: MlxArray, anchorTok: number, gamma: number, opts: DflashDraftOpts = {}): DflashDraftBlock {
     const { dDraft, markovRank: r } = this.cfg;
     const V = this.dims.vocabSize;
     const sample = opts.sample && opts.sample.temperature > 0 ? opts.sample : null;

package/src/spec/dspark/module.ts

@@ -31,6 +31,7 @@ import { MlxArray } from "../../mlx/array";
 import { Dtype } from "../../mlx/ffi";
 import * as ops from "../../mlx/ops";
 import type { Gemma4Model } from "../../model/gemma4";
+import type { DraftProjection } from "../source";
 import { loadAdapterTensors } from "../../lora";
 import { writeShardedSafetensors, type NamedTensor } from "../../quantize/safetensors-writer";
 import { processLogits, sampleToken, KeyStream, type DSparkSampleConfig } from "./sample";
@@ -343,7 +344,7 @@ export class DSparkDrafter {
 
   /** Base logits U from the SHARED frozen LM head over the backbone block
    *  hidden. [A,γ,V] (model dtype). */
-  #baseLogits(model: Gemma4Model, block: MlxArray): MlxArray {
+  #baseLogits(model: DraftProjection, block: MlxArray): MlxArray {
     const hOut = ops.matmul(block, this.get("out_proj")); // [A,γ,H] (f32)
     const hBf = hOut.astype(model.embed.scales.dtype);    // head's activation dtype
     hOut.dispose();
@@ -404,7 +405,7 @@ export class DSparkDrafter {
    * reconstructs (so q is consistent). hCtx [1,H], anchorTok the bonus token.
    * Returns the raw per-position logits so the verifier can recover q.
    */
-  forwardInfer(model: Gemma4Model, hCtx: MlxArray, anchorTok: number, gamma: number, opts: DSparkDraftOpts = {}): DSparkDraftBlock {
+  forwardInfer(model: DraftProjection, hCtx: MlxArray, anchorTok: number, gamma: number, opts: DSparkDraftOpts = {}): DSparkDraftBlock {
     const { dDraft, markovRank: r } = this.cfg;
     const V = this.dims.vocabSize;
     const A = 1;
@@ -537,4 +538,3 @@ function buildNames(cfg: DSparkConfig): string[] {
     "markov.w1", "markov.w2", "conf.w", "conf.b");
   return names;
 }
-

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

@@ -1,5 +1,5 @@
 // Quantize a DeepSpec-format drafter checkpoint (Phase 1a of the DSpark
-// serving program, docs/design/dspark-serving-program.md).
+// serving program, docs/design/speculative-decoding.md).
 //
 // Thin policy layer over the house quantizer (src/quantize): the generic
 // shape-eligibility rule already captures every 2-D matmul weight — the

package/src/spec/generate.ts

@@ -11,7 +11,7 @@
 // with stock only on tie-free prompts (where batched == per-position).
 // An earlier version verified per-position to stay bit-exact to STOCK —
 // that matched no real oracle and cost an extra γ× read of the lm-head
-// weight per step; superseded. See docs/design/spec-decode-larger-targets.md.
+// weight per step; superseded. See docs/design/speculative-decoding.md.
 //
 // Limitation (same as the reference): partial-accept rollback requires
 // trimmable caches; RotatingKVCache loses trimability once its ring

package/src/spec/glm52-mtp-source.ts

@@ -57,7 +57,7 @@ export class Glm52NativeMtpProvider implements DraftProvider {
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
-    if (opts.target.model !== this.model)
+    if (opts.target.identity !== this.model)
       throw new Error("native MTP provider was opened for a different target");
     return new Glm52NativeMtpSource(this.model, this.#layer, opts.sampler);
   }

package/src/spec/ngram-source.ts

@@ -38,7 +38,7 @@
 // speculation is structurally never worse than plain decode by more than the
 // JS scan (~µs against 30k-token histories).
 
-import { flagOn } from "../flags";
+import { flagOn } from "../runtime-config";
 import type { DraftProvider, DraftSource } from "./source";
 
 export interface NgramOptions {

package/src/spec/prefix-state.ts

@@ -0,0 +1,45 @@
+import type { Cache } from "../model/gemma4";
+import type { DisposableResource } from "../contracts/resources";
+import { cleanupFailure } from "../engine/resources";
+
+/** One most-recent, exact prefill boundary. Taking transfers ownership; a
+ * recurrent target cannot be trimmed to an arbitrary common prefix. The
+ * provider owns this state and releases it before its weights at shutdown. */
+export interface SpeculativePrefixState extends DisposableResource {
+  readonly tokens: number[];
+  readonly targetIdentity: object;
+  readonly namespace: string;
+  readonly bytes: number;
+  readonly target: Cache[];
+}
+
+export class SpeculativePrefixStore<T extends SpeculativePrefixState> {
+  #entry: T | null = null;
+
+  take(prompt: readonly number[], targetIdentity: object, namespace: string,
+    maxBytes: number): T | null {
+    const entry = this.#entry;
+    this.#entry = null;
+    if (!entry) return null;
+    if (entry.bytes <= maxBytes && entry.targetIdentity === targetIdentity &&
+      entry.namespace === namespace && entry.tokens.length < prompt.length &&
+      entry.tokens.every((token, index) => prompt[index] === token))
+      return entry;
+    entry.dispose();
+    return null;
+  }
+
+  /** Takes ownership even when the entry exceeds the configured RAM budget. */
+  put(entry: T, maxBytes: number): void {
+    try { this.dispose(); }
+    catch (error) { cleanupFailure(error, () => entry.dispose()); }
+    if (entry.bytes <= maxBytes && maxBytes > 0) this.#entry = entry;
+    else entry.dispose();
+  }
+
+  dispose(): void {
+    const entry = this.#entry;
+    this.#entry = null;
+    entry?.dispose();
+  }
+}