mlx-bun 0.2.0 → 0.3.0

package/src/eval/kl.ts

@@ -22,14 +22,15 @@
 
 import { existsSync } from "node:fs";
 import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
+import { clearCache, Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { Registry } from "../registry";
 import { Weights } from "../weights";
 import { loadModelConfig, type ModelConfig, type TurboQuantScheme } from "../config";
 import { createModel, type RuntimeModel } from "../model/factory";
 import { loadTokenizer, type LoadedTokenizer } from "../tokenizer";
-import { maybeQuantizeKv } from "../generate";
+import { evalCacheState, maybeQuantizeKv } from "../generate";
+import { ReferenceLogitsDump, type ReferenceRecords } from "./reference-logits";
 import {
   configureRuntime,
   runtimeKey,
@@ -310,7 +311,7 @@ export async function evaluateKlServingDecode(opts: {
 /** KV-scheme A/B on the teacher-forced serving-decode path: bf16 baseline
  *  (no kvOverride) vs an explicit KV scheme (e.g. TurboQuant), same model /
  *  same prompts / same weight load — the quality-vs-bpw curve gate
- *  (scripts/eval-turboquant-curve.ts). Unlike evaluateKlServingDecode (env-flag
+ *  (scripts/turboquant/eval-turboquant-curve.ts). Unlike evaluateKlServingDecode (env-flag
  *  self-flip), the two arms differ by an actual GenerateOptions kv scheme
  *  object, so this is the seam for schemes that aren't env-lever-shaped. */
 export async function evaluateKlKvArm(opts: {
@@ -355,3 +356,235 @@ export async function evaluateKlKvArm(opts: {
       : `bf16 vs kv-scheme (${decodeSteps} steps)`,
   };
 }
+
+// --- saved-teacher KL (--reference-logits) --------------------------------
+// The two-model mode above needs both arms resident; a bf16 teacher never
+// fits beside a candidate on 32 GB. So the teacher is dumped ONCE from
+// Python (top-k logits + full-vocab logsumexp per position; format and
+// rationale in reference-logits.ts) and every candidate is scored against
+// the dump through OUR engine.
+//
+// Per position, with the teacher's top-k slice j ∈ K:
+//   p_j     = exp(ref_logit_j − ref_logsumexp)        (full-vocab softmax)
+//   log q_j = cand_logit_j − logsumexp(cand, FULL vocab)
+//   KL      = Σ_j p_j · (log p_j − log q_j)
+// The sum is truncated to K, so it under-counts by the tail p carries
+// outside the top-k — `meanCapturedMass` (Σ_j p_j) reports exactly how much
+// of p the number covers, and a run with low captured mass is not a
+// comparable KL. Everything is float32; the candidate's logsumexp is over
+// its FULL vocab (never the top-k slice), so log q is a real log-probability.
+//
+// The forward is the SERVING path (model.makeCache + chunked forwardHidden +
+// logitsFromHidden), NOT trainForward: trainForward's cache stub has no
+// DeltaNet SSMCache.advance, which is why `mlx-bun perplexity` cannot score
+// qwen3_5 at all (docs/design/turboquant.md, "Known engine gaps"). KV stays
+// bf16 — this is a weights-only instrument.
+
+export interface RefLogitsKLResult extends KLResult {
+  /** Positions actually scored (Σ over sequences). */
+  nPositions: number;
+  /** Fraction of positions where argmax(ref) == argmax(cand) over full vocab. */
+  top1Agreement: number;
+  /** Mean Σ_{j∈topK} p_j — how much of the teacher's mass the KL sum covers. */
+  meanCapturedMass: number;
+  topK: number;
+  /** manifest.model — the teacher these numbers are against. */
+  refModel: string;
+}
+
+/** Teacher-forced candidate logits for one sequence, emitted in
+ *  position-ordered chunks: `onChunk(startPos, logits[1, C, V])`. The callee
+ *  owns the array and disposes it after `onChunk` returns. This is the seam
+ *  the model-free test replaces (no weights, no registry). */
+export type TeacherForcedLogits = (
+  tokens: number[],
+  onChunk: (startPos: number, logits: MlxArray) => void,
+) => void;
+
+/** The real engine forward: one cache for the whole sequence, prefilled in
+ *  chunks, lm_head applied per chunk so full-vocab logits for a 4k context
+ *  never exist at once (4096 × 151k × 4 B ≈ 2.5 GB). */
+export function engineTeacherForcedLogits(
+  model: RuntimeModel,
+  chunkSize: number,
+): TeacherForcedLogits {
+  return (tokens, onChunk) => {
+    const cache = model.makeCache();
+    try {
+      for (let pos = 0; pos < tokens.length; pos += chunkSize) {
+        const end = Math.min(pos + chunkSize, tokens.length);
+        const ids = ops.fromInt32(tokens.slice(pos, end), [1, end - pos]);
+        let h: MlxArray;
+        try {
+          h = model.forwardHidden(ids, cache);
+        } finally {
+          ids.dispose();
+        }
+        let logits: MlxArray;
+        try {
+          logits = model.logitsFromHidden(h);
+        } finally {
+          h.dispose();
+        }
+        try {
+          onChunk(pos, logits);
+        } finally {
+          logits.dispose();
+        }
+        evalCacheState(cache);
+        clearCache(); // mlx-lm _prefill cadence: don't let chunk transients pile up
+      }
+    } finally {
+      for (const c of cache) c.dispose();
+    }
+  };
+}
+
+interface ChunkScore {
+  kl: Float32Array;   // [rows]
+  mass: Float32Array; // [rows]
+  top1: Uint8Array;   // [rows] 1 = candidate argmax == teacher argmax
+}
+
+/** Score `rows` candidate positions against the matching teacher records.
+ *  `logits` is [1, C, V]; rows [rowLo, rowHi) of it line up with
+ *  `records` rows [0, count). */
+function scoreChunkAgainstRecords(
+  logits: MlxArray, rowLo: number, rowHi: number, records: ReferenceRecords,
+): ChunkScore {
+  const rows = rowHi - rowLo;
+  const K = records.topK;
+  const V = logits.shape[2]!;
+  if (records.count !== rows)
+    throw new Error(`reference-logits: ${records.count} records for ${rows} candidate rows`);
+
+  // Teacher indices must exist in the candidate's vocabulary.
+  let maxIdx = -1;
+  for (const v of records.indices) if (v > maxIdx) maxIdx = v;
+  if (maxIdx >= V || records.indices.some((v) => v < 0))
+    throw new Error(
+      `reference-logits: teacher index ${maxIdx} out of range for candidate vocab ${V} ` +
+      `— the dump was made with a different tokenizer`,
+    );
+
+  const owned: MlxArray[] = [];
+  try {
+    const sliced = logits.slice([0, rowLo, 0], [1, rowHi, V]);
+    owned.push(sliced);
+    const f32 = sliced.dtype === Dtype.float32 ? sliced : sliced.astype(Dtype.float32);
+    if (f32 !== sliced) owned.push(f32);
+
+    const lse = ops.logsumexpAxis(f32, -1, false); // [1, rows], full vocab
+    owned.push(lse);
+    const argmax = ops.argmaxAxis(f32, -1);        // [1, rows], full vocab
+    owned.push(argmax);
+    const idx = MlxArray.fromInt32(records.indices, [1, rows, K]);
+    owned.push(idx);
+    const gathered = ops.takeAlongAxis(f32, idx, 2); // [1, rows, K]
+    owned.push(gathered);
+
+    const candTop = gathered.toFloat32();
+    const candLse = lse.toFloat32();
+    const candArgmax = argmax.toIntTokens();
+
+    const kl = new Float32Array(rows);
+    const mass = new Float32Array(rows);
+    const top1 = new Uint8Array(rows);
+    for (let r = 0; r < rows; r++) {
+      const refLse = records.logsumexp[r]!;
+      const qLse = candLse[r]!;
+      const base = r * K;
+      let acc = 0;
+      let m = 0;
+      for (let j = 0; j < K; j++) {
+        const logP = records.logits[base + j]! - refLse;
+        const p = Math.exp(logP);
+        if (p === 0) continue;
+        m += p;
+        acc += p * (logP - (candTop[base + j]! - qLse));
+      }
+      kl[r] = acc;
+      mass[r] = m;
+      top1[r] = candArgmax[r] === records.indices[base]! ? 1 : 0;
+    }
+    return { kl, mass, top1 };
+  } finally {
+    for (const a of owned) a.dispose();
+  }
+}
+
+/** KL(saved bf16 teacher ‖ candidate) at EVERY teacher-forced position. */
+export async function evaluateKlVsReferenceLogits(opts: {
+  candidate: string;
+  /** Dump directory (runs/kl-teacher/<tag>). */
+  referenceDir: string;
+  /** Cap on sequences scored (default: all in the manifest). */
+  nSeqs?: number;
+  /** Positions per candidate forward (default 512 — see the memory note). */
+  chunkSize?: number;
+  /** Test seam: use this forward instead of loading `candidate`. */
+  forward?: TeacherForcedLogits;
+}): Promise<RefLogitsKLResult> {
+  const t0 = Date.now();
+  const dump = ReferenceLogitsDump.open(opts.referenceDir);
+  const chunkSize = Math.max(1, opts.chunkSize ?? 512);
+  const nSeqs = Math.max(1, Math.min(opts.nSeqs ?? dump.nSeqs, dump.nSeqs));
+
+  const forward = opts.forward
+    ?? engineTeacherForcedLogits((await loadRunnable(opts.candidate)).model, chunkSize);
+
+  try {
+    const perSeq: Float32Array[] = [];
+    let nPositions = 0;
+    let top1Hits = 0;
+    let massSum = 0;
+
+    for (let s = 0; s < nSeqs; s++) {
+      const tokens = dump.tokens(s);
+      const scored = Math.min(dump.positionsPerSeq, tokens.length - 1);
+      const kls = new Float32Array(scored);
+      let filled = 0;
+
+      forward(tokens, (startPos, logits) => {
+        const chunkLen = logits.shape[1]!;
+        const lo = startPos;
+        const hi = Math.min(startPos + chunkLen, scored);
+        if (hi <= lo) return; // trailing positions have no teacher record
+        const records = dump.readRecords(s, lo, hi);
+        const sc = scoreChunkAgainstRecords(logits, lo - startPos, hi - startPos, records);
+        for (let r = 0; r < sc.kl.length; r++) {
+          kls[lo + r] = sc.kl[r]!;
+          massSum += sc.mass[r]!;
+          top1Hits += sc.top1[r]!;
+        }
+        filled += hi - lo;
+      });
+
+      if (filled !== scored)
+        throw new Error(
+          `reference-logits: candidate forward covered ${filled} of ${scored} positions ` +
+          `on sequence ${s}`,
+        );
+      perSeq.push(kls);
+      nPositions += scored;
+    }
+
+    const agg = aggregate(perSeq);
+    return {
+      nPrompts: perSeq.length,
+      seqLen: dump.ctxLen,
+      meanKl: agg.mean,
+      medianKl: agg.median,
+      p95Kl: agg.p95,
+      elapsedSec: (Date.now() - t0) / 1000,
+      refLabel: `ref-logits:${dump.manifest.model} top${dump.topK} (${opts.referenceDir})`,
+      nPositions,
+      top1Agreement: nPositions ? top1Hits / nPositions : 0,
+      meanCapturedMass: nPositions ? massSum / nPositions : 0,
+      topK: dump.topK,
+      refModel: dump.manifest.model,
+    };
+  } finally {
+    dump.close();
+  }
+}

package/src/eval/reference-logits.ts

@@ -0,0 +1,207 @@
+// Reader for a saved bf16 TEACHER-LOGITS dump (format_version 1).
+//
+// Why the format exists: a bf16 teacher can never be resident on a 32 GB box
+// alongside the candidate, so the teacher's distributions are computed ONCE
+// (Python side, one model at a time) and written to disk as a top-k slice
+// plus the full-vocab logsumexp. Any candidate loaded through OUR engine can
+// then be scored against it (src/eval/kl.ts `evaluateKlVsReferenceLogits`)
+// without the teacher ever being co-resident.
+//
+// On-disk contract (fixed — this module reads it, it does not negotiate it):
+//
+//   <dir>/manifest.json  { format_version: 1, model, corpus_sha256, ctx_len,
+//                          n_seqs, top_k, positions_per_seq, created, notes }
+//   <dir>/tokens.bin     int32 LE [n_seqs][ctx_len] — teacher-forced ids
+//   <dir>/seq-<i>.bin    one record per r = 0..ctx_len-2; record r is the
+//                        distribution over the token at index r+1, i.e. the
+//                        distribution predicted from tokens[0..r] — so record
+//                        r lines up with candidate logits at POSITION r:
+//                          int32[top_k]   indices, descending by logit
+//                          float16[top_k] reference logits at those indices
+//                          float32        logsumexp over the FULL vocab
+//
+// Records are fixed-width (6·top_k + 4 bytes), so a range is one pread — the
+// scorer streams a chunk of records at a time and never holds a whole
+// sequence file (top_k 2048 × ctx 4096 ≈ 50 MB) in the JS heap.
+
+import { closeSync, existsSync, openSync, readFileSync, readSync, statSync } from "node:fs";
+
+export interface ReferenceLogitsManifest {
+  format_version: number;
+  model: string;
+  corpus_sha256: string;
+  ctx_len: number;
+  n_seqs: number;
+  top_k: number;
+  positions_per_seq: number;
+  created?: string;
+  notes?: string;
+}
+
+/** A contiguous run of records, flattened: row i of `indices`/`logits` is
+ *  `[i*topK, (i+1)*topK)`. Logits are decoded from float16 to float32. */
+export interface ReferenceRecords {
+  /** Absolute record index of row 0. */
+  start: number;
+  count: number;
+  topK: number;
+  indices: Int32Array;    // [count · topK]
+  logits: Float32Array;   // [count · topK]
+  logsumexp: Float32Array; // [count] — over the FULL teacher vocab
+}
+
+const SUPPORTED_FORMAT_VERSION = 1;
+
+/** float16 bits → float32 (host-side; same expansion as MlxArray.toFloat32Host,
+ *  kept local so reading a dump never touches the GPU). */
+function decodeFloat16(bits: number): number {
+  const sign = bits & 0x8000;
+  const exponent = (bits >>> 10) & 0x1f;
+  const fraction = bits & 0x03ff;
+  let value: number;
+  if (exponent === 0) value = fraction === 0 ? 0 : fraction * 2 ** -24;
+  else if (exponent === 0x1f) value = fraction === 0 ? Infinity : NaN;
+  else value = (1 + fraction / 1024) * 2 ** (exponent - 15);
+  return sign ? -value : value;
+}
+
+export class ReferenceLogitsDump {
+  readonly dir: string;
+  readonly manifest: ReferenceLogitsManifest;
+  readonly topK: number;
+  readonly ctxLen: number;
+  readonly nSeqs: number;
+  /** Records per sequence file — derived from the file, cross-checked
+   *  against the manifest. */
+  readonly positionsPerSeq: number;
+  readonly recordBytes: number;
+
+  #tokens: Int32Array;
+  #fd: number | null = null;
+  #fdSeq = -1;
+
+  private constructor(dir: string, manifest: ReferenceLogitsManifest, tokens: Int32Array, positionsPerSeq: number) {
+    this.dir = dir;
+    this.manifest = manifest;
+    this.topK = manifest.top_k;
+    this.ctxLen = manifest.ctx_len;
+    this.nSeqs = manifest.n_seqs;
+    this.positionsPerSeq = positionsPerSeq;
+    this.recordBytes = manifest.top_k * 6 + 4;
+    this.#tokens = tokens;
+  }
+
+  static open(dir: string): ReferenceLogitsDump {
+    const manifestPath = `${dir}/manifest.json`;
+    if (!existsSync(manifestPath))
+      throw new Error(`reference-logits: no manifest.json in ${dir}`);
+    const manifest = JSON.parse(readFileSync(manifestPath, "utf8")) as ReferenceLogitsManifest;
+    if (manifest.format_version !== SUPPORTED_FORMAT_VERSION)
+      throw new Error(
+        `reference-logits: format_version ${manifest.format_version} unsupported ` +
+        `(this reader implements ${SUPPORTED_FORMAT_VERSION})`,
+      );
+    for (const k of ["ctx_len", "n_seqs", "top_k"] as const)
+      if (!Number.isInteger(manifest[k]) || manifest[k] <= 0)
+        throw new Error(`reference-logits: manifest.${k} must be a positive integer (got ${manifest[k]})`);
+
+    const tokensBytes = readFileSync(`${dir}/tokens.bin`);
+    const wantTokenBytes = manifest.n_seqs * manifest.ctx_len * 4;
+    if (tokensBytes.byteLength !== wantTokenBytes)
+      throw new Error(
+        `reference-logits: tokens.bin is ${tokensBytes.byteLength} B, expected ` +
+        `${wantTokenBytes} B (n_seqs ${manifest.n_seqs} × ctx_len ${manifest.ctx_len} × int32)`,
+      );
+    // Copy: the Buffer's byteOffset need not be int32-aligned for a view.
+    const tokens = new Int32Array(
+      tokensBytes.buffer.slice(tokensBytes.byteOffset, tokensBytes.byteOffset + tokensBytes.byteLength),
+    );
+
+    const recordBytes = manifest.top_k * 6 + 4;
+    const seq0 = `${dir}/seq-0.bin`;
+    if (!existsSync(seq0)) throw new Error(`reference-logits: missing ${seq0}`);
+    const size = statSync(seq0).size;
+    if (size % recordBytes !== 0)
+      throw new Error(
+        `reference-logits: seq-0.bin is ${size} B, not a multiple of the ` +
+        `${recordBytes} B record (top_k ${manifest.top_k})`,
+      );
+    const positions = size / recordBytes;
+    if (positions > manifest.ctx_len - 1)
+      throw new Error(
+        `reference-logits: seq-0.bin holds ${positions} records but ctx_len ` +
+        `${manifest.ctx_len} allows at most ${manifest.ctx_len - 1}`,
+      );
+    if (Number.isInteger(manifest.positions_per_seq) && manifest.positions_per_seq !== positions)
+      console.warn(
+        `[reference-logits] manifest.positions_per_seq=${manifest.positions_per_seq} but ` +
+        `seq-0.bin holds ${positions} records — using the file`,
+      );
+    return new ReferenceLogitsDump(dir, manifest, tokens, positions);
+  }
+
+  /** Teacher-forced token ids of sequence `seq` (length ctx_len). */
+  tokens(seq: number): number[] {
+    this.#assertSeq(seq);
+    const off = seq * this.ctxLen;
+    return Array.from(this.#tokens.subarray(off, off + this.ctxLen));
+  }
+
+  /** Records [lo, hi) of sequence `seq` — one pread, decoded to float32. */
+  readRecords(seq: number, lo: number, hi: number): ReferenceRecords {
+    this.#assertSeq(seq);
+    if (lo < 0 || hi > this.positionsPerSeq || hi <= lo)
+      throw new Error(
+        `reference-logits: record range [${lo},${hi}) out of bounds for ` +
+        `${this.positionsPerSeq} records`,
+      );
+    const count = hi - lo;
+    const K = this.topK;
+    const buf = Buffer.allocUnsafe(count * this.recordBytes);
+    const fd = this.#open(seq);
+    const read = readSync(fd, buf, 0, buf.byteLength, lo * this.recordBytes);
+    if (read !== buf.byteLength)
+      throw new Error(`reference-logits: short read on seq-${seq}.bin (${read}/${buf.byteLength} B)`);
+
+    const dv = new DataView(buf.buffer, buf.byteOffset, buf.byteLength);
+    const indices = new Int32Array(count * K);
+    const logits = new Float32Array(count * K);
+    const logsumexp = new Float32Array(count);
+    for (let r = 0; r < count; r++) {
+      const base = r * this.recordBytes;
+      const row = r * K;
+      for (let j = 0; j < K; j++) indices[row + j] = dv.getInt32(base + j * 4, true);
+      const lbase = base + K * 4;
+      for (let j = 0; j < K; j++) logits[row + j] = decodeFloat16(dv.getUint16(lbase + j * 2, true));
+      logsumexp[r] = dv.getFloat32(lbase + K * 2, true);
+    }
+    return { start: lo, count, topK: K, indices, logits, logsumexp };
+  }
+
+  close(): void {
+    if (this.#fd !== null) closeSync(this.#fd);
+    this.#fd = null;
+    this.#fdSeq = -1;
+  }
+
+  #open(seq: number): number {
+    if (this.#fdSeq === seq && this.#fd !== null) return this.#fd;
+    if (this.#fd !== null) closeSync(this.#fd);
+    const path = `${this.dir}/seq-${seq}.bin`;
+    if (!existsSync(path)) throw new Error(`reference-logits: missing ${path}`);
+    const size = statSync(path).size;
+    if (size < this.positionsPerSeq * this.recordBytes)
+      throw new Error(
+        `reference-logits: seq-${seq}.bin is ${size} B, short of the ` +
+        `${this.positionsPerSeq * this.recordBytes} B seq-0.bin defines`,
+      );
+    this.#fd = openSync(path, "r");
+    this.#fdSeq = seq;
+    return this.#fd;
+  }
+
+  #assertSeq(seq: number): void {
+    if (!Number.isInteger(seq) || seq < 0 || seq >= this.nSeqs)
+      throw new Error(`reference-logits: sequence ${seq} out of range (n_seqs ${this.nSeqs})`);
+  }
+}

package/src/eval/runner.ts

@@ -1,3 +1,4 @@
+import type { GenerateOptions } from "../generate";
 // Task-eval runner: load a model and greedily complete prompts through the
 // REAL serving path. Unlike the M0 KL (compat forward() + plain caches),
 // this drives generate() with the model's kv_config, so generation runs the
@@ -5,9 +6,6 @@
 
 import { homedir } from "node:os";
 import { readFileSync } from "node:fs";
-import { generate, type GenerateOptions } from "../generate";
-import { clearCache } from "../mlx/ffi";
-import * as ops from "../mlx/ops";
 import { ChatTemplate } from "../chat-template";
 import { loadModelConfig, type ModelConfig } from "../config";
 import { runtimeValue } from "../runtime-config";
@@ -17,13 +15,12 @@ import {
   type RuntimeModel,
 } from "../model/factory";
 import { isGlm52Config } from "../model/support";
-import { DiffusionGemmaModel } from "../model/diffusion-gemma";
 import { loadTokenizer, type LoadedTokenizer } from "../tokenizer";
 import { resolveModelDir } from "./kl";
 
 export const EVAL_DATA_DIR = runtimeValue("MLX_BUN_EVAL_DATA") ?? `${homedir()}/.cache/mlx-bun/eval-data`;
 
-/** Read a jsonl dataset exported by scripts/eval/export-datasets.py. */
+/** Read a jsonl dataset exported by scripts/oracle/export-datasets.py. */
 export function loadJsonl<T = Record<string, unknown>>(name: string): T[] {
   const path = `${EVAL_DATA_DIR}/${name}.jsonl`;
   const out: T[] = [];
@@ -55,6 +52,7 @@ export function sampleIndices(total: number, n: number, seed = 42): number[] {
 type SamplerArm = Partial<Pick<GenerateOptions, "temperature" | "topP" | "topK" | "seed" | "hlg" | "curve">>;
 
 export interface TaskModel {
+  completionClient?: import("../contracts/completion").CompletionClient<{ body: string; options?: GenOpts }, string>;
   model: RuntimeModel;
   tokenizer: LoadedTokenizer;
   template: ChatTemplate | null;
@@ -117,96 +115,10 @@ export interface GenOpts {
   kvScheme?: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart" | "turboQuant">;
 }
 
-/** Bit-exact greedy decode: raw `model.forward` + argmax loop, which matches mlx-lm
- *  token-for-token (gen8k: 8000/8000). The product `generate()` decode wrapper
- *  (forwardHidden/logitsFromHidden + pipelined sampler) diverges from the raw forward
- *  on near-ties past ~32 tokens, so the EVAL decodes bit-exactly itself rather than
- *  alter the serving path. Greedy + full-precision KV only. */
-export function greedyDecodeBitExact(tm: TaskModel, ids: number[], maxTokens: number): string {
-  if (tm.activeAdapters) tm.model.loraState.active = tm.activeAdapters; // apply the mounted adapter
-  const cache = tm.model.makeCache();
-  try {
-    // Chunked, LM-head-free prefill. The old single-shot `forward(prompt)`
-    // materialized [1, L, vocab] logits just to discard them — at a 248k
-    // vocab that's ~1 GB per 1k prompt tokens, and with no allocator-cache
-    // clears the Metal buffer cache retained every giant freed tensor
-    // (the 27B capability-suite 73 GB swap-thrash). Prefill only needs the
-    // cache state, which forwardHidden advances identically (the LM head
-    // never touches the cache) — bit-exact by construction. Chunk size and
-    // clear cadence mirror mlx-lm generate_step / our generate().
-    const PREFILL_CHUNK = 2048;
-    const prompt = ids.slice(0, -1);
-    for (let s = 0; s < prompt.length; s += PREFILL_CHUNK) {
-      const chunk = prompt.slice(s, s + PREFILL_CHUNK);
-      const cIds = ops.fromInt32(chunk, [1, chunk.length]);
-      tm.model.forwardHidden(cIds, cache).dispose();
-      cIds.dispose();
-      clearCache();
-    }
-    let last = ids[ids.length - 1]!;
-    const eos = new Set(tm.config.eosTokenIds);
-    const out: number[] = [];
-    for (let i = 0; i < maxTokens; i++) {
-      const lg = tm.model.forward([last], cache);
-      const am = ops.argmaxAxis(lg, -1); // GPU argmax (mlx-lm greedy: lowest-index tie-break)
-      lg.dispose();
-      const best = ops.itemUint32(am);
-      am.dispose();
-      if (i % 256 === 0) clearCache(); // mlx-lm decode cadence
-      if (eos.has(best)) break; // mlx-lm halts on EOS and excludes it from the output
-      out.push(best);
-      last = best;
-    }
-    return tm.tokenizer.decode(out, true);
-  } finally {
-    for (const c of cache) c.dispose();
-  }
-}
-
-/** Complete `body` via the model's real quantized-KV path. Greedy by default;
- *  pass `opts.sampler` to drive temperature/HLG/etc. */
-export async function generateText(tm: TaskModel, body: string, opts: GenOpts = {}): Promise<string> {
-  const maxTokens = opts.maxTokens ?? 256;
-  const enableThinking = opts.enableThinking ?? runtimeValue("MLX_BUN_EVAL_THINK") === "1";
-  const templated = opts.useChat !== false && tm.template !== null;
-  const text = templated
-    ? tm.template!.render([{ role: "user", content: body }], { addGenerationPrompt: true, enableThinking })
-    : body;
-  // The chat template ALREADY emits the BOS; encoding it with add_special_tokens on
-  // would prepend a SECOND BOS (<s><s>…), corrupting generation. Raw bodies (no
-  // template) get the BOS added as normal.
-  const ids = tm.tokenizer.encode(text, /* addSpecialTokens */ !templated);
-
-  // optiq's PUBLISHED eval generates with a FULL-PRECISION KV cache — its eval calls
-  // plain mlx_lm.generate (kv-quant lives in the serving runtime, not the eval). So
-  // for like-for-like parity we default the eval to UNQUANTIZED KV; MMLU's argmax
-  // path was already unquantized, which is why it matched while generation didn't.
-  // MLX_BUN_EVAL_KV_QUANT=1 generates through the model's quantized KV (serving
-  // config); an explicit opts.kvScheme (cli generate's resolved tier route) wins.
-  const envKv: GenOpts["kvScheme"] | undefined =
-    (runtimeValue("MLX_BUN_EVAL_KV_QUANT") === "1" && tm.config.kvQuant?.length)
-      ? { kvConfig: tm.config.kvQuant, quantizedKvStart: 0 } : undefined;
-  const kvScheme = opts.kvScheme ?? envKv;
-  const kvActive = !!(kvScheme && (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant));
-
-  // Parity default — greedy + full-precision KV + no sampler arm: decode bit-exactly
-  // via the raw forward (matches mlx-lm token-for-token). Sampler arms / kv-quant fall
-  // through to the product generate(). DiffusionGemma is non-autoregressive (no AR
-  // forward) — always go through generate(), which routes it to the denoising engine.
-  const isDiffusion = tm.model instanceof DiffusionGemmaModel;
-  if (!isDiffusion && !kvActive && !opts.sampler && !tm.samplerOverride) {
-    return greedyDecodeBitExact(tm, ids, maxTokens);
-  }
-  const gen = generate(tm.model, ids, {
-    maxTokens,
-    temperature: 0, // greedy default — deterministic head-to-head arms
-    adapters: tm.activeAdapters ?? [], // per-request activation of the mounted eval-adapter
-    ...(kvActive ? kvScheme : {}),
-    ...(opts.sampler ?? {}), // overrides temperature when provided
-    ...(tm.samplerOverride ?? {}), // arm override wins over the task's own sampler
-  });
+export { greedyDecodeBitExact } from "../backends/mlx/evaluation";
+import { createEvalCompletionClient } from "../backends/mlx/evaluation";
 
-  const out: number[] = [];
-  for await (const { token } of gen) out.push(token);
-  return tm.tokenizer.decode(out, true);
+/** Application entry point; selection and model math belong to its client. */
+export function generateText(tm: TaskModel, body: string, opts: GenOpts = {}): Promise<string> {
+  return (tm.completionClient ?? createEvalCompletionClient(tm)).complete({ body, options: opts });
 }