mlx-bun 0.0.10 → 0.0.11

package/src/spec/assistant-source.ts

@@ -0,0 +1,150 @@
+// AssistantSource — the optiq KV-borrowing Gemma "-assistant" drafter
+// (src/spec/drafter.ts) behind the serve-time DraftSource seam. L2 oracle:
+// optiq spec_generate. This is the SAME drafter the standalone specGenerate
+// loop (src/spec/generate.ts) drives; here it plugs into the shared
+// verify/accept executor (src/spec/serve-loop.ts) so it composes with grammar,
+// logits processors, admission accounting, and stats like any other source.
+//
+// KV-borrowing (why the seam carries a TargetView): the drafter has NO cache
+// of its own — each draft step reads the target's donor K/V (last sliding +
+// last full layer, in chronological order) and, for the block's first step,
+// the target's hidden at the anchor position. Both arrive through the seam:
+// donor views from `target.caches`, the anchor hidden from draft()'s
+// `anchorHidden` argument (borrowed — never disposed here). See
+// [[dspark-seam-kv-borrowing]].
+//
+// Drafting is GREEDY (the drafter's own argmax head), independent of the
+// request sampler — drafts are only proposals; the target's verify (which does
+// honor the sampler) decides every emitted token, so correctness holds at any
+// temperature and only the acceptance RATE moves. Matches specGenerate.
+
+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";
+
+/** Sum of the drafter's on-disk safetensors — admission accounting only. */
+function safetensorsBytes(dir: string): number {
+  let total = 0;
+  for (const f of readdirSync(dir))
+    if (f.endsWith(".safetensors")) total += statSync(join(dir, f)).size;
+  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;
+
+  private constructor(
+    private readonly drafter: GemmaAssistantDrafter,
+    id: string,
+    weightsBytes: number,
+  ) {
+    this.id = id;
+    this.weightsBytes = weightsBytes;
+  }
+
+  static async load(modelDir: string): Promise<AssistantProvider> {
+    const drafter = await GemmaAssistantDrafter.load(modelDir);
+    const id = modelDir.split("/").filter(Boolean).at(-1)!;
+    return new AssistantProvider(drafter, id, safetensorsBytes(modelDir));
+  }
+
+  open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
+    return new AssistantSource(this.drafter, opts.target);
+  }
+
+  dispose(): void {
+    // The drafter's Weights are process-lifetime (pinned mmaps); nothing to free.
+  }
+}
+
+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 };
+
+  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);
+  }
+
+  /** 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 drafts: number[] = [];
+    const ownedHiddens: MlxArray[] = [];
+    let dTok = feed[feed.length - 1]!; // the pending token (anchorHidden's token)
+    let dHid = anchorHidden; // borrowed for k=0
+    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);
+        const step = this.drafter.forward(emb, dHid, shared, position + k);
+        emb.dispose();
+        emb = null;
+        drafts.push(step.token);
+        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();
+      }
+    }
+    return drafts;
+  }
+
+  /** The serve loop already trims the target caches on rejection; the drafter
+   *  borrows them, so it has nothing of its own to roll back. */
+  commit(_n: number, _kAccept: number, _vCtxML?: MlxArray): void {}
+
+  dispose(): void {}
+}

package/src/spec/deepspec-source.ts

@@ -0,0 +1,143 @@
+// DeepspecSource — DeepSeek's released DSpark drafter checkpoints
+// (Gemma4DSparkModel, e.g. deepseek-ai/dspark_gemma4_12b_block7) behind the
+// serve-time DraftSource seam. The reference implementation is DeepSpec
+// (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.
+//
+// Seam mapping (why this fits without serve-loop changes): the reference
+// recycles target hidden states from each verify pass into the drafter's
+// growing context — exactly the tapped ctxML/vCtxML flow the seam carries:
+//   prefill(ids, ctxML)      → project the prompt's tapped hiddens into the
+//                              per-layer context K/V cache (their DynamicCache
+//                              of PROJECTED context rows, positions 0..Lp-1);
+//   draft(feed, n)           → one block forward vs the cache (anchor = feed
+//                              tail at position ctxLen); confidence truncation
+//                              may return 0..n tokens (ℓ=0 = plain step);
+//   commit(d, kAccept, vCtxML)→ project + append ONLY the accepted rows
+//                              [anchor + kAccept drafts] (their crop()
+//                              semantics: rejected rows never enter the
+//                              cache), positions ctxLen..ctxLen+kAccept.
+// Context rows are cached post-norm-post-rope — bit-equivalent to their
+// in-round concat (per-row RMSNorm/RoPE; argued in projectContextKV's doc).
+
+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";
+
+function safetensorsBytes(dir: string): number {
+  let total = 0;
+  for (const f of readdirSync(dir))
+    if (f.endsWith(".safetensors")) total += statSync(join(dir, f)).size;
+  return total;
+}
+
+export class DeepspecProvider implements DraftProvider {
+  readonly id: string;
+  readonly weightsBytes: number;
+  /** Trained block width (config block_size, e.g. 7) — the server pins
+   *  numDraftTokens to this. */
+  readonly gamma: number;
+
+  private constructor(private readonly drafter: DeepspecDrafter, id: string, weightsBytes: number) {
+    this.id = id;
+    this.weightsBytes = weightsBytes;
+    this.gamma = drafter.gamma;
+  }
+
+  static async load(modelDir: string): Promise<DeepspecProvider> {
+    const drafter = await DeepspecDrafter.load(modelDir);
+    const id = modelDir.split("/").filter(Boolean).at(-1)!;
+    return new DeepspecProvider(drafter, id, safetensorsBytes(modelDir));
+  }
+
+  open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
+    return new DeepspecSource(this.drafter, opts.target);
+  }
+
+  dispose(): void {
+    this.drafter.dispose(); // owned MlxArrays (transpose views + mmap'd weights)
+  }
+}
+
+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))
+      throw new Error("DeepSpec drafter requires a Gemma4 target");
+    const nLayers = target.model.layers.length;
+    if (nLayers !== drafter.cfg.num_target_layers)
+      throw new Error(
+        `DeepSpec drafter was trained for a ${drafter.cfg.num_target_layers}-layer target; ` +
+          `this model has ${nLayers} layers — wrong (target, drafter) pairing`,
+      );
+    this.tapLayers = drafter.tapLayers;
+  }
+
+  /** Project the prompt's tapped hiddens into the context cache (ownership of
+   *  ctxML transfers here). */
+  prefill(promptIds: number[], ctxML?: MlxArray): void {
+    if (!ctxML) throw new Error("DeepSpec drafter: prefill missing tapped context");
+    const Lp = ctxML.shape[1]!;
+    const projected = this.drafter.projectContext(ctxML);
+    ctxML.dispose();
+    const positions = Array.from({ length: Lp }, (_, i) => i);
+    this.ctxKV = this.drafter.projectContextKV(projected, positions);
+    projected.dispose();
+    this.ctxLen = Lp;
+  }
+
+  /** One block forward against the cached context. Confidence truncation may
+   *  return 0..n tokens (0 = skip drafting; the serve loop runs a plain
+   *  tapped step and commit still grows the context by the anchor row). */
+  draft(feed: number[], n: number, _stepBase: number, _anchorHidden?: MlxArray): number[] {
+    if (!this.ctxKV) throw new Error("DeepSpec drafter: draft before prefill");
+    const anchorTok = feed[feed.length - 1]!;
+    const block = this.drafter.draftBlock(this.ctxKV, anchorTok, this.ctxLen);
+    block.baseLogits.dispose(); // serve loop runs its own verify lm-head
+    // n < gamma near maxTokens: the block always drafts full γ; hand back
+    // only what the round may verify.
+    return block.tokens.length > n ? block.tokens.slice(0, n) : block.tokens;
+  }
+
+  /** Grow the context cache by the ACCEPTED rows of the verify window —
+   *  [anchor + kAccept drafts] from vCtxML [1,d+1,m*H] (ownership transfers
+   *  here); rejected rows never enter (reference crop() semantics). */
+  commit(_d: number, kAccept: number, vCtxML?: MlxArray): void {
+    if (!vCtxML) throw new Error("DeepSpec drafter: commit missing tapped verify context");
+    if (!this.ctxKV) throw new Error("DeepSpec drafter: commit before prefill");
+    const mH = vCtxML.shape[2]!;
+    const keep = kAccept + 1; // anchor + accepted drafts
+    const rows = vCtxML.slice([0, 0, 0], [1, keep, mH]);
+    vCtxML.dispose();
+    const projected = this.drafter.projectContext(rows);
+    rows.dispose();
+    const positions = Array.from({ length: keep }, (_, i) => this.ctxLen + i);
+    const newKV = this.drafter.projectContextKV(projected, positions);
+    projected.dispose();
+    for (let i = 0; i < this.ctxKV.length; i++) {
+      const old = this.ctxKV[i]!;
+      const add = newKV[i]!;
+      const k = ops.concatAxis([old.k, add.k], 2);
+      const v = ops.concatAxis([old.v, add.v], 2);
+      old.k.dispose(); old.v.dispose();
+      add.k.dispose(); add.v.dispose();
+      this.ctxKV[i] = { k, v };
+    }
+    this.ctxLen += keep;
+  }
+
+  dispose(): void {
+    if (this.ctxKV) for (const { k, v } of this.ctxKV) { k.dispose(); v.dispose(); }
+    this.ctxKV = null;
+  }
+}

package/src/spec/dflash-source.ts

@@ -0,0 +1,124 @@
+// DflashSource — DSpark (faithful DFlash + Markov + confidence,
+// src/spec/dspark/module-dflash.ts) behind the serve-time DraftSource seam.
+// L3 (KL/quality-gated). The DRAFT half of the standalone dflashGenerate loop
+// (src/spec/dspark/generate-dflash.ts) plugged into the shared verify/accept
+// executor (src/spec/serve-loop.ts).
+//
+// H_ctx (why the seam taps): DSpark drafts by attending to a GROWING
+// multi-layer context — the target's tapped hiddens over the accepted stream
+// (paper Eq 2-3). The seam declares `tapLayers`, so the serve loop taps the
+// target's prefill (→ seeds H_ctx here) and every verify forward (→ grown in
+// commit by the accepted window; rejected tips are never appended, mirroring
+// the target-cache trim). See [[dspark-seam-kv-borrowing]].
+//
+// v1 acceptance: the serve loop verifies with TOKEN-MATCH acceptance
+// (mlx-lm/optiq style — lossless at any temperature), NOT the paper's
+// distribution-level rejection-sampling verify. Greedy is identical; temp>0 is
+// still lossless but lower-acceptance. Drafting here is greedy (proposals; the
+// target verify decides). The richer rejection-sampling verify stays in the
+// standalone dflashGenerate for the measure script.
+
+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";
+
+function safetensorsBytes(dir: string): number {
+  let total = 0;
+  for (const f of readdirSync(dir))
+    if (f.endsWith(".safetensors")) total += statSync(join(dir, f)).size;
+  return total;
+}
+
+export class DflashProvider implements DraftProvider {
+  readonly id: string;
+  readonly weightsBytes: number;
+  /** The trained block width — the server pins numDraftTokens to this so the
+   *  serve loop never asks for more positions than the block was trained for. */
+  readonly gamma: number;
+
+  private constructor(private readonly drafter: DflashDrafter, id: string, weightsBytes: number) {
+    this.id = id;
+    this.weightsBytes = weightsBytes;
+    this.gamma = drafter.cfg.gamma;
+  }
+
+  static async load(modelDir: string): Promise<DflashProvider> {
+    const drafter = loadDsparkDrafter(modelDir); // variant dispatch (dspark|legacy dflash)
+    const id = modelDir.split("/").filter(Boolean).at(-1)!;
+    return new DflashProvider(drafter, id, safetensorsBytes(modelDir));
+  }
+
+  open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
+    return new DflashSource(this.drafter, opts.target);
+  }
+
+  dispose(): void {
+    // Unlike the assistant drafter (process-pinned mmaps), DflashDrafter
+    // materializes owned MlxArrays — free them on teardown.
+    this.drafter.dispose();
+  }
+}
+
+class DflashSource implements DraftSource {
+  readonly weightsBytes = 0; // provider-owned weights
+  readonly tapLayers: number[];
+  private readonly model: Gemma4Model;
+  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))
+      throw new Error("DSpark drafter requires a Gemma4 target");
+    this.model = target.model;
+    this.tapLayers = drafter.cfg.tapLayers;
+  }
+
+  /** Seed H_ctx from the tapped prompt context (ownership transfers here). */
+  prefill(_promptIds: number[], ctxML?: MlxArray): void {
+    if (!ctxML) throw new Error("DSpark drafter: prefill missing tapped context");
+    this.hCtx = ctxML;
+  }
+
+  /** Draft up to n tokens against the current H_ctx, conditioned on the anchor
+   *  token (the last emitted token = feed's tail). Greedy — the target verify
+   *  decides. anchorHidden is unused (DSpark reads H_ctx, not the target
+   *  hidden). Confidence-scheduled pruning (Alg 1) may return FEWER than n:
+   *  active when the checkpoint carries STS thresholds (cfg.sts, §3.2.1) or
+   *  via the MLX_BUN_DSPARK_MINCONF env override; uncalibrated checkpoints
+   *  draft fixed-length exactly as before. The server pins n ≤ cfg.gamma. */
+  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 = process.env.MLX_BUN_DSPARK_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
+      ...(envMin ? { minConf: Number(envMin) } : {}),
+    });
+    return block.tokens;
+  }
+
+  /** Grow H_ctx by the accepted window [anchor + kAccept drafts] from the
+   *  verified positions' tapped context [1,d+1,m*H]; drop the rejected tips
+   *  (ownership of vCtxML transfers here). */
+  commit(_d: number, kAccept: number, vCtxML?: MlxArray): void {
+    if (!vCtxML) throw new Error("DSpark drafter: commit missing tapped verify context");
+    if (!this.hCtx) throw new Error("DSpark drafter: commit before prefill");
+    const mH = vCtxML.shape[2]!;
+    const add = vCtxML.slice([0, 0, 0], [1, kAccept + 1, mH]); // anchor + accepted
+    const grown = ops.concatAxis([this.hCtx, add], 1);
+    this.hCtx.dispose();
+    add.dispose();
+    vCtxML.dispose();
+    this.hCtx = grown;
+  }
+
+  dispose(): void {
+    this.hCtx?.dispose();
+    this.hCtx = null;
+  }
+}

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

@@ -0,0 +1,136 @@
+// 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
+// experiment: same target, same prompts, temp 0 — drafter A vs drafter B.
+//
+// Paired design: prompts are the pairing unit, so per-prompt deltas cancel
+// prompt difficulty AND machine-load drift between arms (the
+// dirty-machine-numbers rule: paired deltas survive memory pressure,
+// absolutes don't — final pre/post pairs still belong on a clean machine).
+
+/** One prompt × one drafter arm. */
+export interface AbPromptResult {
+  prompt: string;
+  generatedTokens: number;
+  drafted: number;
+  accepted: number;
+  targetCalls: number;
+  wallMs: number;
+  draftedByPos: number[];
+  acceptedByPos: number[];
+}
+
+/** Aggregates over one arm's prompt set. */
+export interface AbArmSummary {
+  nPrompts: number;
+  generatedTokens: number;
+  /** accepted / drafted over the whole set (per-token acceptance). */
+  acceptance: number;
+  /** committed tokens per target forward (τ). */
+  tau: number;
+  /** generated tokens per wall second over the whole set. */
+  tokPerSec: number;
+  /** acceptance at each draft position (0..γ-1). */
+  acceptanceByPos: number[];
+}
+
+export interface AbVerdict {
+  a: AbArmSummary;
+  b: AbArmSummary;
+  /** B − A, absolute points (e.g. -2.1 = B accepts 2.1 points less). */
+  acceptanceDeltaPts: number;
+  /** B / A wall-clock throughput ratio (>1 = B faster). */
+  speedRatio: number;
+  /** Prompts where B's per-prompt acceptance ≥ A's − tolerance. */
+  pairedAcceptanceHolds: number;
+  /** Phase-1d exit criteria: acceptance drop ≤ `maxDropPts` AND
+   *  wall-clock strictly improves. */
+  pass: boolean;
+  verdict: string;
+}
+
+export function summarizeArm(results: AbPromptResult[]): AbArmSummary {
+  const total = (f: (r: AbPromptResult) => number) => results.reduce((s, r) => s + f(r), 0);
+  const generatedTokens = total((r) => r.generatedTokens);
+  const drafted = total((r) => r.drafted);
+  const accepted = total((r) => r.accepted);
+  const targetCalls = total((r) => r.targetCalls);
+  const wallMs = total((r) => r.wallMs);
+
+  const gamma = Math.max(0, ...results.map((r) => r.draftedByPos.length));
+  const acceptanceByPos: number[] = [];
+  for (let i = 0; i < gamma; i++) {
+    const d = total((r) => r.draftedByPos[i] ?? 0);
+    const a = total((r) => r.acceptedByPos[i] ?? 0);
+    acceptanceByPos.push(d > 0 ? a / d : 0);
+  }
+
+  return {
+    nPrompts: results.length,
+    generatedTokens,
+    acceptance: drafted > 0 ? accepted / drafted : 0,
+    tau: targetCalls > 0 ? generatedTokens / targetCalls : 0,
+    tokPerSec: wallMs > 0 ? (generatedTokens / wallMs) * 1000 : 0,
+    acceptanceByPos,
+  };
+}
+
+/** Phase-1d exit: acceptance drop ≤ maxDropPts absolute AND wall-clock
+ *  strictly improves (docs/design/dspark-serving-program.md 1d). */
+export function pairedVerdict(
+  aResults: AbPromptResult[],
+  bResults: AbPromptResult[],
+  opts: { maxDropPts?: number } = {},
+): AbVerdict {
+  if (aResults.length !== bResults.length)
+    throw new Error(`pairedVerdict: arm sizes differ (${aResults.length} vs ${bResults.length})`);
+  const maxDropPts = opts.maxDropPts ?? 3;
+  const a = summarizeArm(aResults);
+  const b = summarizeArm(bResults);
+
+  const acceptanceDeltaPts = (b.acceptance - a.acceptance) * 100;
+  const speedRatio = a.tokPerSec > 0 ? b.tokPerSec / a.tokPerSec : 0;
+
+  let pairedAcceptanceHolds = 0;
+  for (let i = 0; i < aResults.length; i++) {
+    const ra = aResults[i]!;
+    const rb = bResults[i]!;
+    const accA = ra.drafted > 0 ? ra.accepted / ra.drafted : 0;
+    const accB = rb.drafted > 0 ? rb.accepted / rb.drafted : 0;
+    if (accB >= accA - maxDropPts / 100) pairedAcceptanceHolds++;
+  }
+
+  const acceptanceOk = acceptanceDeltaPts >= -maxDropPts;
+  const speedOk = speedRatio > 1;
+  const pass = acceptanceOk && speedOk;
+  const verdict = pass
+    ? `PASS — acceptance ${fmtDelta(acceptanceDeltaPts)} pts (within ${maxDropPts}), wall-clock ×${speedRatio.toFixed(2)}`
+    : !acceptanceOk
+      ? `FAIL — acceptance dropped ${fmtDelta(acceptanceDeltaPts)} pts (limit ${maxDropPts}); B does not preserve A's acceptance`
+      : `FAIL — acceptance held (${fmtDelta(acceptanceDeltaPts)} pts) but wall-clock did not improve (×${speedRatio.toFixed(2)})`;
+
+  return { a, b, acceptanceDeltaPts, speedRatio, pairedAcceptanceHolds, pass, verdict };
+}
+
+function fmtDelta(x: number): string {
+  return `${x >= 0 ? "+" : ""}${x.toFixed(1)}`;
+}
+
+/** Render the two-arm report (pure string — unit-testable, printed by the
+ *  runner). */
+export function renderReport(labelA: string, labelB: string, v: AbVerdict): string {
+  const pct = (x: number) => `${(x * 100).toFixed(1)}%`;
+  const row = (label: string, s: AbArmSummary) =>
+    `  ${label.padEnd(28)} acceptance ${pct(s.acceptance).padStart(6)}  τ ${s.tau.toFixed(2)}  ` +
+    `${s.tokPerSec.toFixed(1)} tok/s  (${s.generatedTokens} tok over ${s.nPrompts} prompts)`;
+  const posRow = (label: string, s: AbArmSummary) =>
+    `  ${label.padEnd(28)} by-pos [${s.acceptanceByPos.map(pct).join(" ")}]`;
+  return [
+    row(`A: ${labelA}`, v.a),
+    row(`B: ${labelB}`, v.b),
+    posRow(`A: ${labelA}`, v.a),
+    posRow(`B: ${labelB}`, v.b),
+    `  paired: B holds A's acceptance on ${v.pairedAcceptanceHolds}/${v.a.nPrompts} prompts`,
+    `  ${v.verdict}`,
+  ].join("\n");
+}

package/src/spec/dspark/calibration.ts

@@ -0,0 +1,112 @@
+// Threshold calibration for the confidence scheduler — fits per-position
+// 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
+// dflashGenerate's onRound hook, this module just fits.
+//
+// VERIFIED vs the paper + reference (2026-07-06 audit, arXiv:2607.05147 +
+// github.com/deepseek-ai/DeepSpec): the RELEASED reference schedules exactly
+// this way — per-position threshold truncation on sigmoid confidence
+// (default 0 = off) — so this calibrator fits the reference-shaped consumer.
+// The PAPER's §3.2.1 "STS" is a different mechanism (Sequential Temperature
+// Scaling: per-position temperatures by ECE grid search on cumulative
+// 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
+// "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).
+
+import type { StsCalibration } from "./module-dflash";
+
+/** One verify-time observation: the drafter's predicted confidence at draft
+ *  position `pos`, and whether the target ACTUALLY accepted that position
+ *  (k < kAccept in the round that produced it). */
+export interface ConfSample {
+  pos: number;
+  conf: number;
+  accepted: boolean;
+}
+
+/**
+ * Fit per-position STS thresholds from observed (conf, accepted) pairs.
+ *
+ * For each position k, we want the smallest τ_k such that pruning on
+ * `conf_k < τ_k` still keeps empirical acceptance precision at or above
+ * `target` among the samples that WOULD be kept (conf ≥ τ_k). Smaller τ_k
+ * keeps more of the block (better speedup); we want the smallest one that
+ * still meets the target, i.e. we scan candidate thresholds from low to
+ * high (in the order of observed conf values) and stop at the first that
+ * clears the bar.
+ *
+ * Candidate thresholds are exactly the position's observed conf values
+ * (a threshold only ever needs to sit at an observed value — moving it
+ * between two observed values changes no sample's keep/drop outcome, so a
+ * finer sweep can't do better). Sorted ascending; a candidate τ's precision
+ * is P(accepted | conf ≥ τ) over the position's samples, with Laplace
+ * (+1 accept, +2 total) smoothing so single-sample / early candidates near
+ * the tail can't spike to a false 100%.
+ *
+ * Edge rules (in priority order):
+ *  1. Position 0 → threshold 0. Alg-1 never prunes position 0 anyway (the
+ *     scheduler in forwardInfer special-cases `k > 0`), so any fit would be
+ *     dead code; 0 documents that explicitly and short-circuits the fit.
+ *  2. Fewer than `minSamples` observations at position k → threshold 0
+ *     (insufficient evidence: don't prune on a guess).
+ *  3. No candidate τ clears `target` (including τ = 1, i.e. "prune always")
+ *     → threshold 1.0 (the head can't call this position at all; prune it
+ *     unconditionally rather than risk emitting a token the head has no
+ *     signal on).
+ *  4. Otherwise → the smallest candidate τ meeting the target.
+ *
+ * Pure and deterministic: same `samples` array (any order) → same output,
+ * since scoring only depends on the multiset of (conf, accepted) pairs per
+ * position, and ties in conf sort stably by value (order-independent).
+ */
+export function fitStsThresholds(
+  samples: ConfSample[],
+  gamma: number,
+  target = 0.5,
+  minSamples = 50,
+): StsCalibration {
+  const byPos: ConfSample[][] = Array.from({ length: gamma }, () => []);
+  for (const s of samples) {
+    // Non-finite conf would poison the fit: a NaN candidate τ makes every
+    // `conf >= τ` comparison false, and the empty keep-set's Laplace
+    // precision (0+1)/(0+2)=0.5 clears the default target — serializing NaN
+    // (→ null in JSON) into the thresholds. Drop such samples at intake; a
+    // NaN-emitting confidence head surfaces via the sample counts instead.
+    if (Number.isFinite(s.conf) && s.pos >= 0 && s.pos < gamma) byPos[s.pos]!.push(s);
+  }
+
+  const thresholds: number[] = new Array(gamma).fill(0);
+
+  for (let k = 1; k < gamma; k++) {
+    const pos = byPos[k]!;
+    if (pos.length < minSamples) {
+      thresholds[k] = 0; // rule 2 — insufficient evidence, don't prune
+      continue;
+    }
+
+    // Candidates = observed conf values, ascending (smallest-τ-first scan).
+    const cands = pos.map((s) => s.conf).slice().sort((a, b) => a - b);
+
+    let chosen: number | null = null;
+    for (const tau of cands) {
+      let accept = 0, total = 0;
+      for (const s of pos) {
+        if (s.conf >= tau) { total++; if (s.accepted) accept++; }
+      }
+      // Laplace (+1/+2) smoothing: precision = (accept+1)/(total+2).
+      const precision = (accept + 1) / (total + 2);
+      if (precision >= target) { chosen = tau; break; }
+    }
+
+    thresholds[k] = chosen ?? 1.0; // rule 3 if nothing clears the bar
+  }
+  thresholds[0] = 0; // rule 1 — always, regardless of samples
+
+  return { thresholds, target, samples: samples.length };
+}