mlx-bun 0.4.0 → 1.0.0

package/src/spec/drafter.ts

@@ -1,304 +0,0 @@
-// Gemma-4 "-assistant" drafter — port of optiq's
-// runtime/spec/drafters/gemma_assistant.py (itself a port of Ollama's
-// x/models/gemma4/assistant.go, MIT).
-//
-// Q-only 4-layer transformer: K/V come from the TARGET's last sliding-
-// and full-attention donor caches at call time; pre/post projections
-// bridge target hidden space (backbone) <-> drafter space.
-//
-// Two output-head variants ship across the Gemma-4 drafter family — we
-// pick by TENSOR PRESENCE, not config (the larger artifacts declare
-// num_centroids in config.json but don't ship the centroid tensors —
-// the field is vestigial):
-//   - CENTROID head (small E2B/E4B drafters, `gemma4_assistant`):
-//     2048 centroids, top-32 × 128 tokens via a precomputed ordering.
-//   - TIED-EMBEDDING head (larger 12B/26B `gemma4_unified_assistant`):
-//     no centroids, no separate lm_head — logits = embed_tokens · h over
-//     the full vocab (the dense form the centroid head approximates).
-//
-// Implementation deviation (centroid path, argmax-equivalent): the
-// reference scatters candidate scores into full-vocab logits initialised
-// to -1e30 and argmaxes; we argmax over the 4096 candidate scores and map
-// through candidate ids directly — same winner, no 262k materialisation.
-// The drafter's numeric details only influence acceptance RATE, never
-// output correctness (the target's verify decides every token).
-
-import type { ModelConfig } from "../config";
-import { Weights } from "../weights";
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import * as ops from "../mlx/ops";
-import { disposeResources } from "../engine/resources";
-import type { SharedKv } from "../model/gemma4";
-
-/** The graph borrows donor attention; storage owns validity and encoding. */
-export interface AssistantAttention {
-  attend(query: MlxArray, scale: number, window: number | null): MlxArray;
-}
-export interface AssistantDonors { sliding: AssistantAttention; full: AssistantAttention; }
-
-/** Legacy plain donor adapter. The common graph never reads cache geometry. */
-export function plainAssistantDonors(shared: { sliding: [MlxArray, MlxArray]; full: [MlxArray, MlxArray] },
-  position: number): AssistantDonors {
-  const bind = ([keys, values]: [MlxArray, MlxArray]): AssistantAttention => ({
-    attend(query, scale, window) {
-      const length = keys.shape[2]!;
-      const from = window === null ? 0 : Math.max(0, Math.max(0, position - window + 1) - (position - length + 1));
-      using floats = from ? MlxArray.fromFloat32(Float32Array.from({ length }, (_, i) => i < from ? -1e9 : 0), [1,1,1,length]) : null;
-      using mask = floats?.astype(keys.dtype);
-      return ops.sdpa(query, keys, values, scale, mask ? "array" : "", mask ?? null);
-    },
-  });
-  return { sliding: bind(shared.sliding), full: bind(shared.full) };
-}
-
-export interface DrafterRowsStep { tokens: MlxArray; nextHidden: MlxArray; }
-
-export interface DrafterStep {
-  /** Drafted token id. */
-  token: number;
-  /** Post-projected hidden (1,1,backbone) — feeds the next draft step. */
-  nextHidden: MlxArray;
-}
-
-interface Block {
-  layerType: string;
-  headDim: number;
-  ropeBase: number | null;
-  ropeFreqs: MlxArray | null;
-  qProj: MlxArray;
-  oProj: MlxArray;
-  qNorm: MlxArray;
-  inputNorm: MlxArray;
-  postAttnNorm: MlxArray;
-  preFfNorm: MlxArray;
-  postFfNorm: MlxArray;
-  gateProj: MlxArray;
-  upProj: MlxArray;
-  downProj: MlxArray;
-  layerScalar: MlxArray;
-}
-
-export class GemmaAssistantDrafter {
-  readonly hidden: number;
-  readonly nHeads: number;
-  readonly slidingWindow: number;
-  readonly eps: number;
-  readonly numCentroids: number;
-  readonly topK: number;
-  readonly vocabPerCentroid: number;
-  #w: Weights;
-  #owned: MlxArray[] = [];
-  #blocks: Block[] = [];
-  readonly useCentroids: boolean;
-  #prePT: MlxArray; // pre_projection transposed [in, out]
-  #postPT: MlxArray;
-  #centroidsT: MlxArray | null = null; // [hidden, numCentroids]
-  #ordering: MlxArray | null = null;   // [numCentroids, vocabPerCentroid] int32
-  #embed: MlxArray;      // drafter token embeddings [vocab, hidden]
-  #norm: MlxArray;
-
-  private constructor(w: Weights, cfg: Record<string, any>) {
-    const t = cfg.text_config ?? cfg;
-    this.hidden = t.hidden_size;
-    this.nHeads = t.num_attention_heads;
-    this.slidingWindow = t.sliding_window;
-    this.eps = t.rms_norm_eps;
-    this.numCentroids = cfg.num_centroids;
-    this.topK = cfg.centroid_intermediate_top_k;
-    this.vocabPerCentroid = Math.floor(t.vocab_size / this.numCentroids);
-    this.#w = w;
-
-    const own = (array: MlxArray) => { this.#owned.push(array); return array; };
-    try {
-      const T = (name: string) => w.tensor(name);
-      // Keep the oracle's strided weight.T: a contiguous transpose changes
-      // M4 matmul rounding and allocates a second copy of every linear weight.
-      const transposed = (name: string): MlxArray => {
-        const c = own(ops.transposeAxes(T(name), [1, 0]));
-        c.eval();
-        return c;
-      };
-      this.#prePT = transposed("pre_projection.weight");
-      this.#postPT = transposed("post_projection.weight");
-      this.#embed = T("model.embed_tokens.weight");
-      this.#norm = T("model.norm.weight");
-      // Head variant: only the small drafters ship the centroid tensors.
-      // The 12B/26B artifacts ship a tied embed_tokens head instead (their
-      // config still declares num_centroids — ignore it, trust the tensors).
-      this.useCentroids = w.has("masked_embedding.centroids.weight");
-      if (this.useCentroids) {
-        this.#centroidsT = transposed("masked_embedding.centroids.weight");
-        const ord64 = T("masked_embedding.token_ordering");
-        const ord32 = ord64.astype(Dtype.int32);
-        this.#ordering = own(ops.reshape(ord32, [this.numCentroids, this.vocabPerCentroid]));
-        this.#ordering.eval();
-        ord32.dispose();
-      }
-
-      const ropeP = t.rope_parameters ?? {};
-      for (let i = 0; i < t.num_hidden_layers; i++) {
-        const layerType = t.layer_types[i];
-        const isFull = layerType === "full_attention";
-        const headDim = isFull ? t.global_head_dim : t.head_dim;
-        const rp = ropeP[layerType] ?? {};
-        let ropeBase: number | null = rp.rope_theta ?? (isFull ? 1e6 : 1e4);
-        let ropeFreqs: MlxArray | null = null;
-        const partial = rp.partial_rotary_factor ?? 1.0;
-        if (isFull && partial < 1.0) {
-          // ProportionalRoPE freqs, matching the target's donor-K rope
-          const rotated = Math.floor(headDim * partial);
-          const n = headDim / 2;
-          const freqs = new Float32Array(n).fill(Infinity);
-          for (let k = 0; k < rotated / 2; k++)
-            freqs[k] = Math.pow(ropeBase!, (2 * k) / headDim);
-          ropeFreqs = own(MlxArray.fromFloat32(freqs, [n]));
-          ropeBase = null;
-        }
-        const p = `model.layers.${i}`;
-        this.#blocks.push({
-          layerType, headDim, ropeBase, ropeFreqs,
-          qProj: transposed(`${p}.self_attn.q_proj.weight`),
-          oProj: transposed(`${p}.self_attn.o_proj.weight`),
-          qNorm: T(`${p}.self_attn.q_norm.weight`),
-          inputNorm: T(`${p}.input_layernorm.weight`),
-          postAttnNorm: T(`${p}.post_attention_layernorm.weight`),
-          preFfNorm: T(`${p}.pre_feedforward_layernorm.weight`),
-          postFfNorm: T(`${p}.post_feedforward_layernorm.weight`),
-          gateProj: transposed(`${p}.mlp.gate_proj.weight`),
-          upProj: transposed(`${p}.mlp.up_proj.weight`),
-          downProj: transposed(`${p}.mlp.down_proj.weight`),
-          layerScalar: T(`${p}.layer_scalar`),
-        });
-      }
-    } catch (error) { disposeResources(this.#owned.splice(0)); throw error; }
-  }
-
-  static async load(modelDir: string): Promise<GemmaAssistantDrafter> {
-    const cfg = (await Bun.file(`${modelDir}/config.json`).json()) as Record<string, any>;
-    const weights = await Weights.open(modelDir);
-    try { return new GemmaAssistantDrafter(weights, cfg); }
-    catch (error) { weights.dispose(); throw error; }
-  }
-
-  dispose(): void { disposeResources([...this.#owned.splice(0), this.#w]); }
-
-  /** f32-internal RMSNorm (reference _RMSNorm casts through float32). */
-  #rms(x: MlxArray, weight: MlxArray): MlxArray {
-    const f = x.astype(Dtype.float32);
-    const n = ops.rmsNorm(f, weight, this.eps);
-    f.dispose();
-    const back = n.astype(x.dtype);
-    n.dispose();
-    return back;
-  }
-
-  /** The request adapter consumes the same graph as a batch of one.
-   * Donors are chronological; position is the target's last position. */
-  forward(lastTokenEmb: MlxArray, targetHidden: MlxArray,
-    sharedKv: { sliding: [MlxArray, MlxArray]; full: [MlxArray, MlxArray] }, position: number): DrafterStep {
-    const result = this.forwardRows(lastTokenEmb, targetHidden, plainAssistantDonors(sharedKv, position), position);
-    try { return { token: ops.itemUint32(result.tokens), nextHidden: result.nextHidden }; }
-    catch (error) { result.nextHidden.dispose(); throw error; }
-    finally { result.tokens.dispose(); }
-  }
-
-  #sampleRows(h: MlxArray, B: number): MlxArray {
-    if (this.useCentroids) {
-      using scores = ops.matmul(h, this.#centroidsT!);
-      using flat = ops.reshape(scores, [B, this.numCentroids]);
-      using neg = ops.neg(flat);
-      using partition = ops.argpartitionAxis(neg, this.topK - 1, -1);
-      using top = partition.slice([0, 0], [B, this.topK]);
-      using topFlat = ops.reshape(top, [B * this.topK]);
-      using candidates = ops.takeAxis(this.#ordering!, topFlat, 0);
-      const count = this.topK * this.vocabPerCentroid;
-      using ids = ops.reshape(candidates, [B, count]);
-      using flatIds = ops.reshape(ids, [B * count]);
-      using embeddings = ops.takeAxis(this.#embed, flatIds, 0);
-      using matrix = ops.reshape(embeddings, [B, count, this.hidden]);
-      using vector = ops.reshape(h, [B, 1, this.hidden]);
-      // The oracle rounds each bf16 product before reduction. Matmul's
-      // accumulation changes candidate rankings on M1, even at identical h.
-      using products = ops.mul(matrix, vector);
-      using candidateScores = ops.sumAxis(products, -1, false);
-      using maximum = ops.maxAxis(candidateScores, -1, true);
-      using winners = ops.equal(candidateScores, maximum);
-      // Sparse full-vocabulary argmax chooses the lowest token ID on a tie.
-      // Apply that rule directly to the candidates without a vocab-size scatter.
-      using sentinel = ops.fromInt32([this.#embed.shape[0]!], []);
-      using eligible = ops.where(winners, ids, sentinel);
-      using selected = ops.minAxis(eligible, -1, false);
-      return selected.astype(Dtype.uint32);
-    }
-    using rows = ops.reshape(h, [B, this.hidden]);
-    using vector = B === 1 ? ops.reshape(h, [this.hidden, 1]) : ops.transposeAxes(rows, [1, 0]);
-    using scores = ops.matmul(this.#embed, vector);
-    using flat = B === 1 ? ops.reshape(scores, [1, this.#embed.shape[0]!]) : ops.transposeAxes(scores, [1, 0]);
-    return ops.argmaxAxis(flat, -1);
-  }
-
-  /** Numerical graph only. Callers supply donor validity masks for different
-   * row positions; token selection remains on device across the draft chain. */
-  forwardRows(lastTokenEmb: MlxArray, targetHidden: MlxArray,
-    donors: AssistantDonors, position: number | MlxArray): DrafterRowsStep {
-    const B = lastTokenEmb.shape[0]!;
-    let h = ops.concatAxis([lastTokenEmb, targetHidden], 2); // (B,1,2*backbone)
-    h = disposing(h, ops.matmul(h, this.#prePT));            // (B,1,256)
-
-    for (const blk of this.#blocks) {
-      const donor = blk.layerType === "sliding_attention" ? donors.sliding : donors.full;
-      const residual = h;
-      let x = this.#rms(h, blk.inputNorm);
-      // Q-only attention
-      let q = ops.matmul(x, blk.qProj);
-      x.dispose();
-      q = disposing(q, ops.reshape(q, [B, 1, this.nHeads, blk.headDim]));
-      q = disposing(q, this.#rms(q, blk.qNorm));
-      q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
-      q = disposing(q, typeof position === "number"
-        ? ops.rope(q, blk.headDim, blk.ropeBase, position, blk.ropeFreqs)
-        : ops.ropeDynamic(q, blk.headDim, blk.ropeBase, position, blk.ropeFreqs));
-      let attn = donor.attend(q, Math.pow(blk.headDim, -0.5),
-        blk.layerType === "sliding_attention" ? this.slidingWindow : null);
-      q.dispose();
-      attn = disposing(attn, ops.transposeAxes(attn, [0, 2, 1, 3]));
-      attn = disposing(attn, ops.reshape(attn, [B, 1, this.nHeads * blk.headDim]));
-      attn = disposing(attn, ops.matmul(attn, blk.oProj));
-      attn = disposing(attn, this.#rms(attn, blk.postAttnNorm));
-      h = ops.add(residual, attn);
-      residual.dispose();
-      attn.dispose();
-
-      const res2 = h;
-      let f = this.#rms(h, blk.preFfNorm);
-      const g = ops.matmul(f, blk.gateProj);
-      const u = ops.matmul(f, blk.upProj);
-      f.dispose();
-      const act = ops.geluApprox(g);
-      g.dispose();
-      let mlp = ops.mul(act, u);
-      act.dispose();
-      u.dispose();
-      mlp = disposing(mlp, ops.matmul(mlp, blk.downProj));
-      mlp = disposing(mlp, this.#rms(mlp, blk.postFfNorm));
-      h = ops.add(res2, mlp);
-      res2.dispose();
-      mlp.dispose();
-
-      h = disposing(h, ops.mul(h, blk.layerScalar));
-    }
-
-    h = disposing(h, this.#rms(h, this.#norm));
-
-    const tokens = this.#sampleRows(h, B);
-    const nextHidden = ops.matmul(h, this.#postPT);           // (B,1,backbone)
-    h.dispose();
-    return { tokens, nextHidden };
-  }
-}
-
-function disposing(old: MlxArray, next: MlxArray): MlxArray {
-  old.dispose();
-  return next;
-}

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

@@ -1,136 +0,0 @@
-// Model-free math for the Phase-1c drafter acceptance A/B
-// (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
-// 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/speculative-decoding.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

@@ -1,112 +0,0 @@
-// 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.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 +
-// 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/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 };
-}