mlx-bun 0.0.8 → 0.0.10

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

@@ -0,0 +1,193 @@
+// Faithful DFlash decode loop — the draft attends to a GROWING multi-layer
+// context H_ctx (the target's tapped hiddens over the accepted stream), per
+// paper §3.1 / Eq 3 + §5.1 runtime. Parallel to generate.ts (v1 single-vector).
+//
+// Each verify pass is TAPPED (captures the m layers at the verified positions);
+// after accepting kAccept drafts, the anchor + accepted tokens' m-layer hiddens
+// are appended to H_ctx (rejected tips are never added — mirrors the KV trim).
+// Losslessness is unchanged: greedy = longest-prefix match; temp>0 = min(1,p/q)
+// + residual resample. H_ctx only affects WHICH tokens get drafted (τ), never
+// what is emitted.
+
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+import { Gemma4Model } from "../../model/gemma4";
+import type { DflashDrafter } from "./module-dflash";
+import { KeyStream, processLogits, probsOf, probAtToken, sampleToken, sampleResidual, type DSparkSampleConfig } from "./sample";
+
+export interface DflashGenOptions {
+  gamma?: number;
+  maxTokens?: number;
+  eosTokenIds?: number[];
+  sample?: DSparkSampleConfig;
+}
+export interface DflashStats {
+  emitted: number; drafted: number; accepted: number; targetCalls: number;
+  decodeMs: number; prefillMs: number; acceptanceRate: number; meanAcceptLen: number; meanConf: number; decodeTps: number;
+}
+export interface DflashResult { tokens: number[]; stats: DflashStats }
+
+/** Tapped forward: run the target over `ids`, returning the final post-norm
+ *  hidden [1,L,H] AND the m-layer context [1,L,m*H] (tapLayers concatenated on
+ *  the feature axis). */
+function forwardTapped(model: Gemma4Model, ids: MlxArray, caches: import("../../model/gemma4").Cache[], tapLayers: number[]): { finalH: MlxArray; ctxML: MlxArray } {
+  model.hiddenTap = { layers: new Set(tapLayers), captured: new Map() };
+  let finalH: MlxArray;
+  let ctxML: MlxArray;
+  try {
+    finalH = model.forwardHidden(ids, caches);
+    const cap = model.hiddenTap.captured;
+    const perLayer = tapLayers.map((li) => { const a = cap.get(li); if (!a) throw new Error(`layer ${li} not captured`); return a; });
+    ctxML = ops.concatAxis(perLayer, 2); // [1,L,m*H]
+    for (const [, a] of cap) a.dispose();
+  } finally {
+    model.hiddenTap = null;
+  }
+  return { finalH, ctxML };
+}
+
+function sampleFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number, cfg: DSparkSampleConfig | null, keys: KeyStream): number {
+  const H = hidden.shape[2]!;
+  const hSl = hidden.slice([0, pos, 0], [1, pos + 1, H]);
+  const logits = model.logitsFromHidden(hSl); hSl.dispose();
+  const V = logits.shape[2]!;
+  const flat = ops.reshape(logits, [1, V]); logits.dispose();
+  let tok: number;
+  if (cfg) { const sc = processLogits(flat, cfg); tok = sampleToken(sc, keys.next()); sc.dispose(); }
+  else { const am = ops.argmaxAxis(flat, -1); tok = ops.itemUint32(am); am.dispose(); }
+  flat.dispose();
+  return tok;
+}
+
+function picksBatched(model: Gemma4Model, hidden: MlxArray, count: number): number[] {
+  const logits = model.logitsFromHidden(hidden);
+  const V = logits.shape[2]!;
+  const out: number[] = [];
+  for (let k = 0; k < count; k++) {
+    const sl = logits.slice([0, k, 0], [1, k + 1, V]);
+    const flat = ops.reshape(sl, [1, V]); sl.dispose();
+    const am = ops.argmaxAxis(flat, -1); flat.dispose();
+    out.push(ops.itemUint32(am)); am.dispose();
+  }
+  logits.dispose();
+  return out;
+}
+
+interface Verdict { kAccept: number; emit: number }
+
+function verifyGreedy(model: Gemma4Model, vHidden: MlxArray, drafts: number[], gamma: number): Verdict {
+  const gt = picksBatched(model, vHidden, gamma + 1);
+  let k = 0; while (k < gamma && drafts[k] === gt[k]) k++;
+  return { kAccept: k, emit: gt[k]! };
+}
+function verifySampling(model: Gemma4Model, vHidden: MlxArray, draftLogits: MlxArray, drafts: number[], gamma: number, cfg: DSparkSampleConfig, keys: KeyStream): Verdict {
+  const V = draftLogits.shape[2]!;
+  const tl = model.logitsFromHidden(vHidden);
+  const us = ops.randomUniform([gamma], Dtype.float32, 0, 1, keys.next());
+  const u = us.toFloat32(); us.dispose();
+  const slice = (a: MlxArray, i: number) => { const s = a.slice([0, i, 0], [1, i + 1, V]); const f = ops.reshape(s, [1, V]); s.dispose(); return f; };
+  let verdict: Verdict | null = null;
+  for (let k = 0; k < gamma; k++) {
+    const x = drafts[k]!;
+    const ps = processLogits(slice(tl, k), cfg); const P = probsOf(ps); ps.dispose();
+    const qs = processLogits(slice(draftLogits, k), cfg); const Q = probsOf(qs); qs.dispose();
+    const pX = probAtToken(P, x), qX = probAtToken(Q, x);
+    if (u[k]! < Math.min(1, qX > 0 ? pX / qX : 1)) { P.dispose(); Q.dispose(); continue; }
+    const corr = sampleResidual(P, Q, keys.next()); P.dispose(); Q.dispose();
+    verdict = { kAccept: k, emit: corr }; break;
+  }
+  if (!verdict) { const bs = processLogits(slice(tl, gamma), cfg); verdict = { kAccept: gamma, emit: sampleToken(bs, keys.next()) }; bs.dispose(); }
+  tl.dispose();
+  return verdict;
+}
+
+export function dflashGenerate(model: Gemma4Model, drafter: DflashDrafter, promptTokens: number[], options: DflashGenOptions = {}): DflashResult {
+  const gamma = options.gamma ?? drafter.cfg.gamma;
+  const { maxTokens = 256, eosTokenIds = model.config.eosTokenIds } = options;
+  const sampleCfg = options.sample && options.sample.temperature > 0 ? options.sample : null;
+  const keys = new KeyStream(options.sample?.seed ?? 0);
+  const tapLayers = drafter.cfg.tapLayers;
+
+  const caches = model.makeCache();
+  const stats: DflashStats = { emitted: 0, drafted: 0, accepted: 0, targetCalls: 0, decodeMs: 0, prefillMs: 0, acceptanceRate: 0, meanAcceptLen: 0, meanConf: 0, decodeTps: 0 };
+  const out: number[] = [];
+  let confSum = 0, rounds = 0, acceptLenSum = 0;
+  let H_ctx: MlxArray | null = null;
+
+  try {
+    // 1. prefill (tapped): seed H_ctx with the prompt's m-layer hiddens.
+    const t0 = performance.now();
+    const ids = ops.fromInt32(promptTokens, [1, promptTokens.length]);
+    const { finalH, ctxML } = forwardTapped(model, ids, caches, tapLayers);
+    ids.dispose();
+    stats.targetCalls++;
+    H_ctx = ctxML; // [1, Lp, m*H]
+    const Lp = finalH.shape[1]!;
+    let next = sampleFromHidden(model, finalH, Lp - 1, sampleCfg, keys);
+    finalH.dispose();
+    stats.prefillMs = performance.now() - t0;
+
+    if (eosTokenIds.includes(next)) return finalize(stats, out, rounds, acceptLenSum, confSum);
+    out.push(next); stats.emitted++;
+
+    const tDecode = performance.now();
+    outer: while (stats.emitted < maxTokens) {
+      // 2a. draft against the current (growing) H_ctx.
+      const block = drafter.forwardInfer(model, H_ctx!, next, gamma, { sample: sampleCfg ?? undefined, keys });
+      const drafts = block.tokens;
+      stats.drafted += gamma;
+      for (const c of block.conf) confSum += c;
+
+      // 2b. verify (tapped) — get accepted tokens' m-layer hiddens to append.
+      const vIds = ops.fromInt32([next, ...drafts], [1, gamma + 1]);
+      const { finalH: vHidden, ctxML: vCtxML } = forwardTapped(model, vIds, caches, tapLayers);
+      vIds.dispose();
+      stats.targetCalls++;
+
+      const { kAccept, emit } = sampleCfg
+        ? verifySampling(model, vHidden, block.draftLogits, drafts, gamma, sampleCfg, keys)
+        : verifyGreedy(model, vHidden, drafts, gamma);
+      block.draftLogits.dispose();
+      stats.accepted += kAccept; rounds++; acceptLenSum += kAccept + 1;
+
+      // 2c. emit accepted drafts
+      for (let k = 0; k < kAccept; k++) {
+        const isEos = eosTokenIds.includes(drafts[k]!);
+        if (!isEos) out.push(drafts[k]!);
+        stats.emitted++;
+        if (isEos || stats.emitted >= maxTokens) { vHidden.dispose(); vCtxML.dispose(); break outer; }
+      }
+      const emitIsEos = eosTokenIds.includes(emit);
+      if (!emitIsEos) out.push(emit);
+      stats.emitted++;
+
+      // 2d. append accepted stream's m-layer hiddens (anchor + kAccept drafts =
+      // vCtxML[:, 0..kAccept]) to H_ctx; never the rejected tips.
+      const add = vCtxML.slice([0, 0, 0], [1, kAccept + 1, vCtxML.shape[2]!]);
+      const grown = ops.concatAxis([H_ctx!, add], 1);
+      H_ctx!.dispose(); add.dispose(); vCtxML.dispose();
+      H_ctx = grown;
+
+      // 2e. roll back rejected KV tips (bypass trim, past-window safe)
+      if (kAccept < gamma) { const n = gamma - kAccept; for (const c of caches) c.trim(n, true); }
+
+      vHidden.dispose();
+      if (emitIsEos || stats.emitted >= maxTokens) break;
+      next = emit;
+    }
+    stats.decodeMs = performance.now() - tDecode;
+  } finally {
+    H_ctx?.dispose();
+    for (const c of caches) c.dispose();
+  }
+  return finalize(stats, out, rounds, acceptLenSum, confSum);
+}
+
+function finalize(stats: DflashStats, out: number[], rounds: number, acceptLenSum: number, confSum: number): DflashResult {
+  stats.acceptanceRate = stats.drafted ? stats.accepted / stats.drafted : 0;
+  stats.meanAcceptLen = rounds ? acceptLenSum / rounds : 0;
+  stats.meanConf = stats.drafted ? confSum / stats.drafted : 0;
+  stats.decodeTps = stats.decodeMs ? ((stats.emitted - 1) / stats.decodeMs) * 1000 : 0;
+  return { tokens: out, stats };
+}

package/src/spec/dspark/generate.ts

@@ -0,0 +1,290 @@
+// DSpark speculative decode loop — a fork of src/spec/generate.ts
+// (specGenerate). Same γ-draft / verify-in-one-pass / accept-longest-prefix /
+// rollback spine, but the drafter is the trainable DSpark module: one PARALLEL
+// backbone pass + a cheap SEQUENTIAL Markov head produce the γ-block, instead
+// of γ sequential KV-borrowing drafter calls. DSpark needs only the target's
+// final hidden as context (the tap), so the donor-K/V plumbing is gone.
+//
+// Two lossless verify modes — both leave the OUTPUT DISTRIBUTION exactly equal
+// to vanilla e4b decode; the drafter only moves τ (speed), never WHAT is
+// emitted:
+//   - GREEDY (temperature 0): accept a draft iff it equals the target's argmax
+//     (longest-prefix). Bit-identical to vanilla greedy e4b.
+//   - SAMPLING (temperature > 0): the speculative-sampling rule — accept x_k
+//     w.p. min(1, p_k(x_k)/q_k(x_k)); on rejection resample from the residual
+//     norm(relu(p−q)) and stop; if the whole block is accepted, sample a bonus
+//     from p_{γ+1}. p and q use the SAME top-p/top-k/temperature processing
+//     (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
+// output must match model.generate; the sampling rule's correctness is the
+// statistical "emit distribution == target" check in scripts/dspark-smoke.ts.
+// If a gate fails, the accept/reject or KV rollback is wrong — fix that first.
+
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+import { Gemma4Model } from "../../model/gemma4";
+import type { DSparkDrafter } from "./module";
+import { KeyStream, processLogits, probsOf, probAtToken, sampleToken, sampleResidual, type DSparkSampleConfig } from "./sample";
+
+export interface DSparkOptions {
+  gamma?: number;
+  maxTokens?: number;
+  eosTokenIds?: number[];
+  /** Sampling config; omit (or temperature 0) ⇒ greedy. */
+  sample?: DSparkSampleConfig;
+}
+
+export interface DSparkStats {
+  emitted: number;
+  drafted: number;
+  accepted: number;
+  targetCalls: number;
+  decodeMs: number;
+  prefillMs: number;
+  acceptanceRate: number;
+  /** Mean accepted length per verify round (τ), incl. the bonus/correction token. */
+  meanAcceptLen: number;
+  /** Mean predicted confidence over drafted tokens (diagnostic). */
+  meanConf: number;
+  decodeTps: number;
+}
+
+export interface DSparkResult {
+  tokens: number[];
+  stats: DSparkStats;
+}
+
+/** Greedy pick from one hidden position [1,1,H]. */
+function pickFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number): number {
+  const H = hidden.shape[2]!;
+  const hSl = hidden.slice([0, pos, 0], [1, pos + 1, H]);
+  const logits = model.logitsFromHidden(hSl);
+  hSl.dispose();
+  const V = logits.shape[2]!;
+  const flat = ops.reshape(logits, [1, V]);
+  logits.dispose();
+  const am = ops.argmaxAxis(flat, -1);
+  flat.dispose();
+  const t = ops.itemUint32(am);
+  am.dispose();
+  return t;
+}
+
+/** Batched verify: one lm-head over the γ+1 window, argmax per position. */
+function picksBatched(model: Gemma4Model, hidden: MlxArray, count: number): number[] {
+  const logits = model.logitsFromHidden(hidden); // [1, L, V]
+  const V = logits.shape[2]!;
+  const out: number[] = [];
+  for (let k = 0; k < count; k++) {
+    const sl = logits.slice([0, k, 0], [1, k + 1, V]);
+    const flat = ops.reshape(sl, [1, V]);
+    sl.dispose();
+    const am = ops.argmaxAxis(flat, -1);
+    flat.dispose();
+    out.push(ops.itemUint32(am));
+    am.dispose();
+  }
+  logits.dispose();
+  return out;
+}
+
+/** Result of one verify round: how many leading drafts were accepted, and the
+ *  correction/bonus token to emit after them. */
+interface RoundVerdict { kAccept: number; emit: number }
+
+/** GREEDY verify: longest prefix where draft == target argmax. */
+function verifyGreedy(model: Gemma4Model, vHidden: MlxArray, drafts: number[], gamma: number): RoundVerdict {
+  const gt = picksBatched(model, vHidden, gamma + 1);
+  let kAccept = 0;
+  while (kAccept < gamma && drafts[kAccept] === gt[kAccept]) kAccept++;
+  return { kAccept, emit: gt[kAccept]! };
+}
+
+/** SAMPLING verify (speculative sampling): accept x_k w.p. min(1, p/q); on the
+ *  first rejection resample from norm(relu(p−q)); if all accepted sample a bonus
+ *  from p_{γ+1}. p,q processed identically (top-p/top-k/temperature). */
+function verifySampling(
+  model: Gemma4Model, vHidden: MlxArray, draftLogits: MlxArray, drafts: number[],
+  gamma: number, cfg: DSparkSampleConfig, keys: KeyStream,
+): RoundVerdict {
+  const V = draftLogits.shape[2]!;
+  const targetLogits = model.logitsFromHidden(vHidden); // [1,γ+1,V]
+  // γ uniforms for the accept tests, drawn together.
+  const uArr = ops.randomUniform([gamma], Dtype.float32, 0, 1, keys.next());
+  const us = uArr.toFloat32();
+  uArr.dispose();
+
+  let verdict: RoundVerdict | null = null;
+  for (let k = 0; k < gamma; k++) {
+    const x = drafts[k]!;
+    const tScaled = processLogits(sliceTmp(targetLogits, k, V), cfg);
+    const P = probsOf(tScaled);
+    const dScaled = processLogits(sliceTmp(draftLogits, k, V), cfg);
+    const Q = probsOf(dScaled);
+    tScaled.dispose(); dScaled.dispose();
+    const pX = probAtToken(P, x);
+    const qX = probAtToken(Q, x);
+    const accept = us[k]! < Math.min(1, qX > 0 ? pX / qX : 1);
+    if (accept) {
+      P.dispose(); Q.dispose();
+      continue;
+    }
+    // reject at k → resample correction from the residual, stop.
+    const corrected = sampleResidual(P, Q, keys.next());
+    P.dispose(); Q.dispose();
+    verdict = { kAccept: k, emit: corrected };
+    break;
+  }
+  if (!verdict) {
+    // all γ accepted → bonus from p_{γ+1}
+    const bScaled = processLogits(sliceTmp(targetLogits, gamma, V), cfg);
+    const bonus = sampleToken(bScaled, keys.next());
+    bScaled.dispose();
+    verdict = { kAccept: gamma, emit: bonus };
+  }
+  targetLogits.dispose();
+  return verdict;
+}
+
+/** [1,L,V] → [1,V] at position k (a fresh owned array). */
+function sliceTmp(logits: MlxArray, k: number, V: number): MlxArray {
+  const sl = logits.slice([0, k, 0], [1, k + 1, V]);
+  const flat = ops.reshape(sl, [1, V]);
+  sl.dispose();
+  return flat;
+}
+
+export function dsparkGenerate(
+  model: Gemma4Model,
+  drafter: DSparkDrafter,
+  promptTokens: number[],
+  options: DSparkOptions = {},
+): DSparkResult {
+  const gamma = options.gamma ?? drafter.cfg.gamma;
+  const { maxTokens = 256, eosTokenIds = model.config.eosTokenIds } = options;
+  if (gamma < 1) throw new Error("gamma must be >= 1");
+  if (gamma > drafter.cfg.gamma)
+    throw new Error(`gamma ${gamma} exceeds drafter max ${drafter.cfg.gamma}`);
+  const sampleCfg = options.sample && options.sample.temperature > 0 ? options.sample : null;
+  const keys = new KeyStream(options.sample?.seed ?? 0);
+
+  const caches = model.makeCache();
+  const stats: DSparkStats = {
+    emitted: 0, drafted: 0, accepted: 0, targetCalls: 0,
+    decodeMs: 0, prefillMs: 0, acceptanceRate: 0,
+    meanAcceptLen: 0, meanConf: 0, decodeTps: 0,
+  };
+  const out: number[] = [];
+  let confSum = 0;
+  let rounds = 0;
+  let acceptLenSum = 0;
+
+  try {
+    // 1. prefill → first token. For sampling, the first token must also be
+    // sampled from the target dist (not argmax) to stay lossless.
+    const t0 = performance.now();
+    const ids = ops.fromInt32(promptTokens, [1, promptTokens.length]);
+    const hidden = model.forwardHidden(ids, caches);
+    ids.dispose();
+    stats.targetCalls++;
+    const H = hidden.shape[2]!;
+    const Lp = hidden.shape[1]!;
+    let next: number;
+    if (sampleCfg) {
+      const hSl = hidden.slice([0, Lp - 1, 0], [1, Lp, H]);
+      const logits = model.logitsFromHidden(hSl);
+      hSl.dispose();
+      const flat = ops.reshape(logits, [1, logits.shape[2]!]);
+      logits.dispose();
+      const scaled = processLogits(flat, sampleCfg);
+      flat.dispose();
+      next = sampleToken(scaled, keys.next());
+      scaled.dispose();
+    } else {
+      next = pickFromHidden(model, hidden, Lp - 1);
+    }
+    let lastHidden = hidden.slice([0, Lp - 1, 0], [1, Lp, H]); // [1,1,H]
+    hidden.dispose();
+    stats.prefillMs = performance.now() - t0;
+
+    const nextIsEos = eosTokenIds.includes(next);
+    if (!nextIsEos) out.push(next);
+    stats.emitted++;
+    if (nextIsEos) { lastHidden.dispose(); return finalize(stats, out, rounds, acceptLenSum, confSum); }
+
+    // 2. outer loop
+    const tDecode = performance.now();
+    outer: while (stats.emitted < maxTokens) {
+      // 2a. draft the γ-block: one backbone pass + sequential Markov head
+      const hCtx = ops.reshape(lastHidden, [1, H]); // [1,H]
+      const block = drafter.forwardInfer(model, hCtx, next, gamma, { sample: sampleCfg ?? undefined, keys });
+      hCtx.dispose();
+      const drafts = block.tokens;
+      stats.drafted += gamma;
+      for (const c of block.conf) confSum += c;
+
+      // 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);
+      vIds.dispose();
+      stats.targetCalls++;
+
+      // 2c. accept (greedy longest-prefix OR speculative sampling)
+      const { kAccept, emit } = sampleCfg
+        ? verifySampling(model, vHidden, block.draftLogits, drafts, gamma, sampleCfg, keys)
+        : verifyGreedy(model, vHidden, drafts, gamma);
+      block.draftLogits.dispose();
+      stats.accepted += kAccept;
+      rounds++;
+      acceptLenSum += kAccept + 1; // accepted drafts + bonus/correction token
+
+      // 2d. emit accepted drafts
+      for (let k = 0; k < kAccept; k++) {
+        const isEos = eosTokenIds.includes(drafts[k]!);
+        if (!isEos) out.push(drafts[k]!);
+        stats.emitted++;
+        if (isEos || stats.emitted >= maxTokens) { vHidden.dispose(); break outer; }
+      }
+
+      // 2e. emit the correction (or bonus when all accepted)
+      const emitIsEos = eosTokenIds.includes(emit);
+      if (!emitIsEos) out.push(emit);
+      stats.emitted++;
+
+      // 2f. roll back rejected entries. bypass=true: we always trim the last
+      // n = gamma - kAccept (≤ γ) tips immediately after the verify concat
+      // write, which is the one case where trimming past the sliding window is
+      // sound (the newest tokens are intact at the buffer tail). See
+      // RotatingKVCache.trim.
+      if (kAccept < gamma) {
+        const n = gamma - kAccept;
+        for (const c of caches) c.trim(n, true);
+      }
+
+      if (emitIsEos || stats.emitted >= maxTokens) { vHidden.dispose(); break; }
+
+      // 2g. chain state
+      next = emit;
+      lastHidden.dispose();
+      lastHidden = vHidden.slice([0, kAccept, 0], [1, kAccept + 1, H]);
+      vHidden.dispose();
+    }
+    stats.decodeMs = performance.now() - tDecode;
+    lastHidden.dispose();
+  } finally {
+    for (const c of caches) c.dispose();
+  }
+
+  return finalize(stats, out, rounds, acceptLenSum, confSum);
+}
+
+function finalize(stats: DSparkStats, out: number[], rounds: number, acceptLenSum: number, confSum: number): DSparkResult {
+  stats.acceptanceRate = stats.drafted ? stats.accepted / stats.drafted : 0;
+  stats.meanAcceptLen = rounds ? acceptLenSum / rounds : 0;
+  stats.meanConf = stats.drafted ? confSum / stats.drafted : 0;
+  stats.decodeTps = stats.decodeMs ? ((stats.emitted - 1) / stats.decodeMs) * 1000 : 0;
+  return { tokens: out, stats };
+}

package/src/spec/dspark/loss.ts

@@ -0,0 +1,213 @@
+// DSpark training objective (paper §3.3, Eq 9–12).
+//
+//   L = α_ce·L_ce + α_tv·L_tv + α_conf·L_conf      (defaults 0.1 / 0.9 / 1.0)
+//
+// all three position-weighted by w_k = exp(−(k−1)/γ) (emphasizes earlier
+// block positions, which contribute most to prefix-based accepted length):
+//
+//   L_ce   = −Σ_k w_k · log p^d_k(x*_k)                       (Eq 9)
+//   L_tv   =  Σ_k w_k · ‖p^d_k − p^t_k‖₁                       (Eq 10)
+//   L_conf = −Σ_k w_k · [ c*_k·log c_k + (1−c*_k)·log(1−c_k) ] (Eq 11)
+//
+// where the analytic acceptance label c*_k = 1 − ½‖p^d_k − p^t_k‖₁ (Eq 8) is
+// STOP-GRADIENT'd — it supervises the confidence head without leaking BCE
+// gradient back into the draft distribution (L_tv already drives p^d→p^t).
+//
+// p^t is the FROZEN target's true distribution at each block position; the
+// caller computes targetLogits = model.logitsFromHidden(targetHidden) (same
+// head DSpark drafts with) and passes it as a constant. M = A·γ is small
+// (γ≈5–7), so plain MLX materialization is fine; a fused TV kernel (the
+// flash-cce vocab-blocked pattern) is only worth porting if measurement
+// shows this term bottlenecks — it does not at these sizes.
+
+import { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { DSparkTrainOut } from "./module";
+
+export interface LossWeights {
+  ce: number;
+  tv: number;
+  conf: number;
+}
+
+export const DEFAULT_LOSS_WEIGHTS: LossWeights = { ce: 0.1, tv: 0.9, conf: 1.0 };
+
+/** Position weights w_k = exp(−(k−1)/γ) as a [1, γ] float32 row. */
+export function positionWeights(gamma: number): MlxArray {
+  const w = new Float32Array(gamma);
+  for (let j = 0; j < gamma; j++) w[j] = Math.exp(-j / gamma);
+  return MlxArray.fromFloat32(w, [1, gamma]).eval();
+}
+
+/** mean_A( Σ_k w_k · perPos[A,k] ) → scalar. */
+function weightedReduce(perPos: MlxArray, w: MlxArray): MlxArray {
+  const A = perPos.shape[0]!;
+  const weighted = ops.mul(perPos, w); // [A,γ] · [1,γ] broadcast
+  const overK = ops.sumAxis(weighted, 1, false); // [A]
+  weighted.dispose();
+  const summed = ops.sumAxis(overK, 0, false); // scalar
+  overK.dispose();
+  const out = ops.mulScalar(summed, 1 / A);
+  summed.dispose();
+  return out;
+}
+
+/** Per-position L1 distance ‖p^d − p^t‖₁ → [A, γ]. Caller owns the result. */
+function l1PerPos(pd: MlxArray, pt: MlxArray): MlxArray {
+  const diff = ops.sub(pd, pt);
+  const ad = ops.abs(diff);
+  diff.dispose();
+  const l1 = ops.sumAxis(ad, 2, false); // [A,γ]
+  ad.dispose();
+  return l1;
+}
+
+/**
+ * Full DSpark loss. Returns the scalar to differentiate plus the three term
+ * scalars (for logging — they share the graph; read them with .toFloat32()
+ * AFTER the backward, or eval them alongside the value).
+ *
+ *  - out.draftLogits [A,γ,V] (U+B, pre-softmax, float32)
+ *  - out.conf        [A,γ]   (sigmoid output ∈ (0,1))
+ *  - targetLogits    [A,γ,V] (frozen target logits; constant)
+ *  - xStar           [A,γ]   int32 ground-truth next tokens
+ */
+export function dsparkLoss(
+  out: DSparkTrainOut,
+  targetLogits: MlxArray,
+  xStar: MlxArray,
+  gamma: number,
+  w: MlxArray,
+  weights: LossWeights = DEFAULT_LOSS_WEIGHTS,
+): { loss: MlxArray; ce: MlxArray; tv: MlxArray; conf: MlxArray } {
+  const dl = out.draftLogits; // [A,γ,V]
+  const A = dl.shape[0]!;
+
+  const pd = ops.softmaxAxis(dl, 2, true); // [A,γ,V]
+  const ptRaw = ops.softmaxAxis(targetLogits, 2, true);
+  const pt = ops.stopGradient(ptRaw); // target is constant
+  ptRaw.dispose();
+
+  // --- L_ce: −log p^d(x*) ---
+  const lse = ops.logsumexpAxis(dl, 2, true); // [A,γ,1]
+  const logp = ops.sub(dl, lse); // [A,γ,V]
+  lse.dispose();
+  const xsIdx = ops.expandDims(xStar, 2); // [A,γ,1] int32
+  const logpStar = ops.takeAlongAxis(logp, xsIdx, 2); // [A,γ,1]
+  logp.dispose(); xsIdx.dispose();
+  const cePos0 = ops.reshape(logpStar, [A, gamma]); // [A,γ]
+  logpStar.dispose();
+  const cePos = ops.neg(cePos0); // −log p^d(x*)
+  cePos0.dispose();
+  const ceTerm = weightedReduce(cePos, w);
+  cePos.dispose();
+
+  // --- L1 / TV / c* (share the L1) ---
+  const l1 = l1PerPos(pd, pt); // [A,γ]
+  pd.dispose(); pt.dispose();
+  const tvTerm = weightedReduce(l1, w);
+
+  // c*_k = 1 − ½·l1, stop-gradient (supervision label)
+  const half = ops.mulScalar(l1, 0.5);
+  const cStarRaw = rsubScalar(half, 1.0); // 1 − ½l1
+  half.dispose();
+  const cStar = ops.stopGradient(cStarRaw);
+  cStarRaw.dispose();
+  l1.dispose();
+
+  // --- L_conf: BCE(c, c*) ---
+  const eps = 1e-6;
+  const c = clampUnit(out.conf, eps);
+  const omc = rsubScalar(c, 1.0); // 1 − c
+  const logc = ops.log(c);
+  const log1mc = ops.log(omc);
+  c.dispose(); omc.dispose();
+  const t1 = ops.mul(cStar, logc); // c*·log c
+  const omcStar = rsubScalar(cStar, 1.0); // 1 − c*
+  cStar.dispose(); logc.dispose();
+  const t2 = ops.mul(omcStar, log1mc); // (1−c*)·log(1−c)
+  omcStar.dispose(); log1mc.dispose();
+  const bcePos0 = ops.add(t1, t2);
+  t1.dispose(); t2.dispose();
+  const bcePos = ops.neg(bcePos0);
+  bcePos0.dispose();
+  const confTerm = weightedReduce(bcePos, w);
+  bcePos.dispose();
+
+  // --- weighted sum ---
+  const a = ops.mulScalar(ceTerm, weights.ce);
+  const b = ops.mulScalar(tvTerm, weights.tv);
+  const cc = ops.mulScalar(confTerm, weights.conf);
+  const ab = ops.add(a, b); a.dispose(); b.dispose();
+  const loss = ops.add(ab, cc); ab.dispose(); cc.dispose();
+  return { loss, ce: ceTerm, tv: tvTerm, conf: confTerm };
+}
+
+/**
+ * Held-out monitoring (NO gradient): analytic per-position acceptance
+ * a_k = 1 − ½‖p^d_k − p^t_k‖₁ and the expected accepted length
+ * τ = 1 + Σ_{j=1}^{γ} Π_{i≤j} a_i (the +1 is the target bonus token). This is
+ * the cheap proxy to early-stop on — τ is the number that matters, not loss.
+ * Returns means over the batch.
+ */
+export function analyticAcceptance(
+  out: DSparkTrainOut,
+  targetLogits: MlxArray,
+): { perPos: number[]; tau: number } {
+  const dl = out.draftLogits;
+  const A = dl.shape[0]!;
+  const gamma = dl.shape[1]!;
+  const pd = ops.softmaxAxis(dl, 2, true);
+  const pt = ops.softmaxAxis(targetLogits, 2, true);
+  const l1 = l1PerPos(pd, pt); // [A,γ]
+  pd.dispose(); pt.dispose();
+  const half = ops.mulScalar(l1, 0.5);
+  l1.dispose();
+  const aMat = rsubScalar(half, 1.0); // [A,γ] = a_k per anchor
+  half.dispose();
+  const flat = aMat.toFloat32(); // [A*γ]
+  aMat.dispose();
+
+  const perPos: number[] = new Array(gamma).fill(0);
+  let tauSum = 0;
+  for (let row = 0; row < A; row++) {
+    let cum = 1;
+    let tau = 1; // bonus token
+    for (let j = 0; j < gamma; j++) {
+      const a = flat[row * gamma + j]!;
+      perPos[j]! += a;
+      cum *= a;
+      tau += cum;
+    }
+    tauSum += tau;
+  }
+  for (let j = 0; j < gamma; j++) perPos[j]! /= A;
+  return { perPos, tau: tauSum / A };
+}
+
+// --- small scalar helpers (no direct rsub/clamp binding) ---------------
+
+/** value − a, elementwise (value + (−a)). */
+function rsubScalar(a: MlxArray, value: number): MlxArray {
+  const neg = ops.neg(a);
+  const res = addScalarF32(neg, value);
+  neg.dispose();
+  return res;
+}
+
+/** a + value (float32). */
+function addScalarF32(a: MlxArray, value: number): MlxArray {
+  const s = MlxArray.fromFloat32(new Float32Array([value]), [1]);
+  const out = ops.add(a, s);
+  s.dispose();
+  return out;
+}
+
+/** clamp x into [eps, 1−eps]. */
+function clampUnit(x: MlxArray, eps: number): MlxArray {
+  const lo = MlxArray.fromFloat32(new Float32Array([eps]), [1]);
+  const hi = MlxArray.fromFloat32(new Float32Array([1 - eps]), [1]);
+  const out = ops.clip(x, lo, hi);
+  lo.dispose(); hi.dispose();
+  return out;
+}