mlx-bun 0.0.7 → 0.0.9

package/src/embed.ts

@@ -0,0 +1,77 @@
+// Text-embedding entry point. Wraps Qwen3Model.embedPooled (last-token hidden →
+// L2-normalized vector) with the Qwen3-Embedding input convention so the CLI
+// (`mlx-bun embed`), the server (`/v1/embeddings`), and in-repo experiments all
+// produce the SAME vectors — the ones verified bit-exact vs mlx-lm in
+// tests/qwen3-embed-parity.test.ts.
+
+import { Qwen3Model } from "./model/qwen3";
+import type { RuntimeModel } from "./model/factory";
+import type { LoadedTokenizer } from "./tokenizer";
+import * as ops from "./mlx/ops";
+
+/** <|endoftext|> — the token Qwen3-Embedding terminates every input with and
+ *  pools the last position of. */
+export const EMBED_EOD = 151643;
+
+/** Only plain-Qwen3 (the Qwen3-Embedding backbone) exposes a pooled-embedding
+ *  path today. Narrow the served RuntimeModel before embedding. */
+export function isEmbeddingModel(model: RuntimeModel): model is Qwen3Model {
+  return model instanceof Qwen3Model;
+}
+
+/** Qwen3-Embedding query format: a task instruction steers WHICH similarity axis
+ *  the geometry reflects ("represent this for topic routing" vs "...for
+ *  sentiment"). Documents are embedded raw; only queries get an instruction. */
+export function withInstruction(text: string, instruction?: string): string {
+  return instruction ? `Instruct: ${instruction}\nQuery:${text}` : text;
+}
+
+/** Tripwire: counts embedOne invocations (embedMany accrues via embedOne) so
+ *  callers can assert the pipeline embeds exactly as many times as expected. */
+let embedCallCount = 0;
+
+/** Reset the embed tripwire counter to 0. */
+export function resetEmbedCounter(): void {
+  embedCallCount = 0;
+}
+
+/** Read the embed tripwire counter (total embedOne calls since last reset). */
+export function getEmbedCounter(): number {
+  return embedCallCount;
+}
+
+export interface EmbedResult {
+  /** L2-normalized embedding (length = model hidden size). */
+  vector: Float32Array;
+  /** Token count of the embedded input (incl. the EOD pooling token). */
+  tokens: number;
+}
+
+/** Embed one text. `instruction` (optional) applies the query format. */
+export function embedOne(
+  model: Qwen3Model,
+  tok: LoadedTokenizer,
+  text: string,
+  instruction?: string,
+): EmbedResult {
+  embedCallCount++;
+  // addSpecialTokens=false matches the bit-exact parity path (Qwen3 has no BOS;
+  // we append the EOD pooling token ourselves).
+  const ids = [...tok.encode(withInstruction(text, instruction), false), EMBED_EOD];
+  const idArr = ops.fromInt32(ids, [1, ids.length]);
+  const vec = model.embedPooled(idArr);
+  idArr.dispose();
+  const out = vec.toFloat32();
+  vec.dispose();
+  return { vector: out, tokens: ids.length };
+}
+
+/** Embed many texts (one forward each — the runtime is single-sequence). */
+export function embedMany(
+  model: Qwen3Model,
+  tok: LoadedTokenizer,
+  texts: string[],
+  instruction?: string,
+): EmbedResult[] {
+  return texts.map((t) => embedOne(model, tok, t, instruction));
+}

package/src/eval/perplexity.ts

@@ -0,0 +1,164 @@
+// Perplexity over a local text/JSONL dataset — mlx_lm.perplexity methodology.
+//
+// The reference (mlx_lm/perplexity.py, mlx-lm 0.31.3) does:
+//   1. tokenize dataset samples, visit them in a seeded random permutation,
+//   2. concatenate tokens into one stream and cut it into NON-OVERLAPPING
+//      rows of `sequence_length` (no sliding-window stride, no BOS re-seeding
+//      per row), keeping at most `num_samples` rows,
+//   3. per batch: logits = model(batch[:, :-1]) in f32, per-token
+//      cross-entropy vs batch[:, 1:] — EVERY position counts,
+//   4. ppl = exp(mean loss); SE via the delta method:
+//      se_ppl = ppl · std(losses, ddof=1) / √N.
+//
+// We reproduce 2-4 exactly. Step 1 differs only in the data source: the
+// reference pulls a Hugging Face dataset (network); we take a LOCAL .txt or
+// .jsonl file — never a download. The permutation RNG is a seeded
+// deterministic shuffle (not NumPy's), so row selection is reproducible per
+// seed but not bit-matched to the Python tool's sampling.
+//
+// The forward reuses the existing full-sequence machinery
+// (src/train/forward.ts trainForward → model.forwardHidden + logitsFromHidden).
+
+import { clearCache, Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { trainForward } from "../train/forward";
+import type { RuntimeModel } from "../model/factory";
+
+export interface PplDataOptions {
+  sequenceLength: number;
+  /** Max rows to keep (-1 / 0 = all available). */
+  numSamples: number;
+  /** Seed for the sample-order shuffle (reference default: 123). */
+  seed: number;
+}
+
+export interface PplResult {
+  ppl: number;
+  /** Delta-method standard error of the perplexity. */
+  standardError: number;
+  meanLoss: number;
+  tokens: number;
+  rows: number;
+}
+
+/** Deterministic PRNG (mulberry32) for the seeded sample shuffle. */
+function mulberry32(seed: number): () => number {
+  let s = seed >>> 0;
+  return () => {
+    s = (s + 0x6d2b79f5) >>> 0;
+    let t = s;
+    t = Math.imul(t ^ (t >>> 15), t | 1);
+    t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+}
+
+/** Split a raw dataset file into text samples: `.jsonl` → one sample per
+ *  line from the `text` field; anything else → the whole file is one sample. */
+export function parseSamples(raw: string, path: string): string[] {
+  if (path.endsWith(".jsonl") || path.endsWith(".ndjson")) {
+    const samples: string[] = [];
+    for (const [i, line] of raw.split("\n").entries()) {
+      if (!line.trim()) continue;
+      let row: unknown;
+      try {
+        row = JSON.parse(line);
+      } catch {
+        throw new Error(`${path}:${i + 1}: not valid JSON`);
+      }
+      const text = (row as { text?: unknown }).text;
+      if (typeof text !== "string")
+        throw new Error(`${path}:${i + 1}: expected a {"text": "..."} row`);
+      samples.push(text);
+    }
+    return samples;
+  }
+  return raw.length > 0 ? [raw] : [];
+}
+
+/** Tokenize samples (seeded shuffle), concatenate, and pack into
+ *  non-overlapping rows of `sequenceLength` tokens (reference load_data). */
+export function packRows(
+  sampleTokens: number[][],
+  opts: PplDataOptions,
+): Int32Array[] {
+  const { sequenceLength: L, numSamples, seed } = opts;
+  const order = sampleTokens.map((_, i) => i);
+  const rnd = mulberry32(seed);
+  for (let i = order.length - 1; i > 0; i--) {
+    const j = Math.floor(rnd() * (i + 1));
+    [order[i], order[j]] = [order[j]!, order[i]!];
+  }
+
+  const wanted = numSamples > 0 ? L * numSamples : Infinity;
+  const stream: number[] = [];
+  for (const idx of order) {
+    if (stream.length >= wanted) break;
+    stream.push(...sampleTokens[idx]!);
+  }
+
+  const nRows = Math.min(
+    Math.floor(stream.length / L),
+    numSamples > 0 ? numSamples : Infinity,
+  );
+  const rows: Int32Array[] = [];
+  for (let r = 0; r < nRows; r++)
+    rows.push(Int32Array.from(stream.slice(r * L, (r + 1) * L)));
+  return rows;
+}
+
+/** Evaluate perplexity over packed rows. Per batch (all rows are the same
+ *  length, so the stock unpadded forward applies): f32 logits over
+ *  batch[:, :-1], per-token CE vs batch[:, 1:], every position counted. */
+export function evalPpl(
+  model: RuntimeModel,
+  rows: Int32Array[],
+  batchSize: number,
+  onBatch?: (done: number, total: number) => void,
+): PplResult {
+  const L = rows[0]!.length;
+  const allLosses: number[] = [];
+  const numBatches = Math.ceil(rows.length / batchSize);
+
+  for (let s = 0, bi = 0; s < rows.length; s += batchSize, bi++) {
+    const batch = rows.slice(s, s + batchSize);
+    const B = batch.length;
+    const inputs: number[] = [];
+    const targets: number[] = [];
+    for (const row of batch) {
+      for (let i = 0; i < L - 1; i++) inputs.push(row[i]!);
+      for (let i = 1; i < L; i++) targets.push(row[i]!);
+    }
+    const ids = ops.fromInt32(inputs, [B, L - 1]);
+    const tgt = ops.fromInt32(targets, [B, L - 1, 1]);
+    const logits = trainForward(model, ids);
+    const f32 = logits.astype(Dtype.float32);
+    logits.dispose();
+    // CE[t] = logsumexp(logits[t]) − logits[t, target]  (f32, like nn.losses)
+    const lse = ops.logsumexpAxis(f32, 2, false); // [B, L-1]
+    const picked = ops.takeAlongAxis(f32, tgt, 2); // [B, L-1, 1]
+    f32.dispose();
+    const picked2d = ops.reshape(picked, [B, L - 1]);
+    picked.dispose();
+    const losses = ops.sub(lse, picked2d);
+    lse.dispose();
+    picked2d.dispose();
+    const vals = losses.toFloat32();
+    losses.dispose();
+    ids.dispose();
+    tgt.dispose();
+    for (const v of vals) allLosses.push(v);
+    clearCache();
+    onBatch?.(bi + 1, numBatches);
+  }
+
+  const n = allLosses.length;
+  const mean = allLosses.reduce((a, b) => a + b, 0) / n;
+  const ppl = Math.exp(mean);
+  // Sample std (ddof=1) → delta-method SE of the perplexity.
+  const varSum = allLosses.reduce((a, b) => a + (b - mean) ** 2, 0);
+  const std = n > 1 ? Math.sqrt(varSum / (n - 1)) : 0;
+  const standardError = ppl * (std / Math.sqrt(n));
+
+  return { ppl, standardError, meanLoss: mean, tokens: n, rows: rows.length };
+}

package/src/eval/runner.ts

@@ -11,6 +11,7 @@ import { ChatTemplate } from "../chat-template";
 import { loadModelConfig, type ModelConfig } from "../config";
 import { Weights } from "../weights";
 import { createModel, type RuntimeModel } from "../model/factory";
+import { DiffusionGemmaModel } from "../model/diffusion-gemma";
 import { loadTokenizer, type LoadedTokenizer } from "../tokenizer";
 import { resolveModelDir } from "./kl";
 
@@ -98,6 +99,11 @@ export interface GenOpts {
    *  (its capability numbers are non-thinking; e.g. MiniCPM5 IFEval 64.7 is
    *  non-thinking). Set MLX_BUN_EVAL_THINK=1 or pass true to evaluate thinking mode. */
   enableThinking?: boolean;
+  /** Explicit KV-quant scheme (cli `generate --l2/--l3/--kv-quant`, resolved with
+   *  serve's semantics). Overrides the MLX_BUN_EVAL_KV_QUANT env default; an empty
+   *  object forces bf16. When it names quantized KV, generation runs the product
+   *  generate() path (same tokens as `serve` with the same scheme). */
+  kvScheme?: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart">;
 }
 
 /** Bit-exact greedy decode: raw `model.forward` + argmax loop, which matches mlx-lm
@@ -105,7 +111,7 @@ export interface GenOpts {
  *  (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. */
-function greedyDecodeBitExact(tm: TaskModel, ids: number[], maxTokens: number): string {
+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 {
@@ -147,21 +153,27 @@ export async function generateText(tm: TaskModel, body: string, opts: GenOpts =
   // 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).
-  const kv = (process.env.MLX_BUN_EVAL_KV_QUANT === "1" && tm.config.kvQuant?.length)
-    ? tm.config.kvQuant : undefined;
+  // 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 =
+    (process.env.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));
 
   // 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().
-  if (!kv && !opts.sampler && !tm.samplerOverride) {
+  // 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
-    ...(kv ? { kvConfig: kv, quantizedKvStart: 0 } : {}),
+    ...(kvActive ? kvScheme : {}),
     ...(opts.sampler ?? {}), // overrides temperature when provided
     ...(tm.samplerOverride ?? {}), // arm override wins over the task's own sampler
   });

package/src/eval/tasks/bfcl.ts

@@ -15,7 +15,7 @@
 // of the task — is ported exactly. (See summary for this one note.)
 
 import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
-import { type ToolDefinition } from "../chat-template";
+import { type ToolDefinition } from "../../chat-template";
 
 // ---------------------------------------------------------------------------
 // Data shapes (the merged jsonl: each line = {query, answer}).

package/src/generate.ts

@@ -16,6 +16,8 @@ import {
 import * as ops from "./mlx/ops";
 import { CompiledDecode } from "./model/compiled-decode";
 import { Gemma4Model, KVCache, RotatingKVCache, type Cache } from "./model/gemma4";
+import { DiffusionGemmaModel } from "./model/diffusion-gemma";
+import { diffusionGenerate } from "./diffusion/diffusion-generate";
 import type { RuntimeModel } from "./model/factory";
 import type { KvQuantSpec } from "./config";
 import {
@@ -55,15 +57,47 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
    *  model's LoraState for exactly the duration of this generation —
    *  a plain field, safe because the generation queue is serialized. */
   adapters?: string[];
+  /** DiffusionGemma image-text-to-text: channel-first pixel_values [1,3,H,W].
+   *  When set, `promptTokens` are the spliced (<|image|>-expanded) ids and the
+   *  denoising engine prefills the merged vision features. Caller keeps ownership. */
+  visionPixels?: MlxArray;
+  /** Capture each emitted token's log-probability (mlx_lm.server `logprobs`).
+   *  The distribution matches mlx-lm generate_step exactly: full-vocab
+   *  log-softmax of the logits AFTER logits processors, BEFORE the sampler's
+   *  temperature/top-p/top-k/min-p/XTC (generate.py L409-422). Off by default —
+   *  the hot path pays nothing when unset. */
+  logprobs?: boolean;
+  /** Capture the top-k (token id, logprob) pairs per emitted token from the
+   *  same distribution (mlx_lm.server `top_logprobs`). 0/unset = off. */
+  topLogprobs?: number;
 }
 
-/** Port of mlx-lm maybe_quantize_kv_cache + optiq serve's per-layer
- *  patched variant (incl. patch_rotating_to_quantized: rotating caches
- *  convert too — Phase 9). kvConfig overrides kvBits, matching optiq's
- *  --kv-config precedence; shipped kv_config.json files list
- *  full-attention layers only, so rotating quantization engages through
- *  uniform kvBits (like optiq --kv-bits) or a config that names sliding
- *  layers. */
+/** Per-token logprob info (only present when GenerateOptions requested it). */
+export interface TokenLogprobs {
+  /** logprob of the emitted token (requires options.logprobs). */
+  logprob?: number;
+  /** top-k (id, logprob) pairs, sorted by logprob descending (requires
+   *  options.topLogprobs > 0). mlx-lm leaves argpartition order unspecified;
+   *  sorting is the deterministic reading of the same set. */
+  top?: { id: number; logprob: number }[];
+}
+
+/** Port of mlx-lm maybe_quantize_kv_cache + BOTH halves of optiq serve's
+ *  per-layer patched variant (incl. patch_rotating_to_quantized: rotating
+ *  caches convert too — Phase 9):
+ *  - per-layer bits/group_size selection (kvConfig overrides kvBits,
+ *    matching optiq's --kv-config precedence; shipped kv_config.json
+ *    files list full-attention layers only, so rotating quantization
+ *    engages through uniform kvBits — like optiq --kv-bits — or a config
+ *    that names sliding layers), and
+ *  - STREAMING conversion (optiq streaming_kv_quant / serve.py
+ *    patched_maybe_quantize): eval each layer's quantized triples and
+ *    clear the buffer pool before building the next layer's conversion.
+ *    Lazily batching every layer's toQuantized into one eval pins ALL
+ *    layers' bf16 K/V as graph inputs alongside ALL quantized outputs —
+ *    the exact transient optiq's fix kills (16.35 → 7.60 GB at 32k on a
+ *    24 GB Mac). Numerics untouched: same quantize math, only the eval
+ *    ordering is forced (tests/kv-quant.test.ts, tests/rotating-kvq.test.ts). */
 export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void {
   const { kvBits, kvConfig } = options;
   if (!kvBits && !kvConfig?.length) return;
@@ -81,10 +115,16 @@ export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void
     if (c.offset === 0) continue;
     if (byLayer) {
       const e = byLayer.get(i);
-      if (e) cache[i] = c.toQuantized(e.groupSize, e.bits);
+      if (!e) continue;
+      cache[i] = c.toQuantized(e.groupSize, e.bits);
     } else {
       cache[i] = c.toQuantized(options.kvGroupSize ?? 64, kvBits!);
     }
+    // Streaming half: materialize THIS layer's conversion now so its bf16
+    // source (already released by toQuantized) frees before the next layer
+    // converts — the transient stays ~one layer, not the whole cache.
+    ops.evalAll(cache[i]!.state());
+    clearCache();
   }
 }
 
@@ -106,6 +146,8 @@ export interface GenerateStats {
 export interface GeneratedToken {
   token: number;
   index: number;
+  /** Present only when GenerateOptions.logprobs / topLogprobs requested it. */
+  logprobs?: TokenLogprobs;
 }
 
 export class Generation implements AsyncIterable<GeneratedToken> {
@@ -170,7 +212,13 @@ export function generate(
   promptTokens: number[],
   options: GenerateOptions = {},
 ): Generation {
-  let inner = generateInner(model, promptTokens, options);
+  // DiffusionGemma is non-autoregressive: route to the denoising engine instead
+  // of the AR decode loop. Same Generation/GenerateStats contract so the CLI and
+  // server stream it through the existing token machinery.
+  let inner =
+    model instanceof DiffusionGemmaModel
+      ? generateDiffusionInner(model, promptTokens, options)
+      : generateInner(model, promptTokens, options);
   if (options.adapters?.length) inner = adapterScoped(model, options.adapters, inner);
   const wire =
     process.env.MLX_BUN_FORCE_WIRE === "1" ||
@@ -178,6 +226,48 @@ export function generate(
   return new Generation(wire ? wiredScoped(inner) : inner);
 }
 
+/** Non-autoregressive diffusion generation, adapted to the AR Generation
+ *  contract. Runs the denoising engine (its own prefill + canvas loop) and
+ *  streams the emitted tokens. v1: greedy (temperature 0, confidence-threshold
+ *  sampler — the OptiQ default); per-block intra-stream + temperature>0
+ *  (categorical) are follow-ups. */
+async function* generateDiffusionInner(
+  model: DiffusionGemmaModel,
+  promptTokens: number[],
+  options: GenerateOptions,
+): AsyncGenerator<GeneratedToken, GenerateStats> {
+  const t0 = performance.now();
+  const maxTokens = options.maxTokens ?? 256;
+  // A fresh random canvas seed per request unless the caller pins one.
+  const seed =
+    options.seed !== undefined
+      ? BigInt(options.seed)
+      : BigInt(Math.floor(Math.random() * 0x7fffffff));
+  // The shipped checkpoint's tokenizer stops on {1, 106}; union any caller eos.
+  const eos = [...new Set([1, 106, ...(options.eosTokenIds ?? [])])];
+  const result = diffusionGenerate(model, promptTokens, {
+    maxTokens,
+    sampler: "confidence-threshold",
+    temperature: 0,
+    eosTokenIds: eos,
+    seed,
+    visionPixels: options.visionPixels,
+  });
+  const decodeMs = performance.now() - t0;
+  let index = 0;
+  for (const token of result.tokens) yield { token, index: index++ };
+  return {
+    promptTokens: promptTokens.length,
+    cachedTokens: 0,
+    generatedTokens: result.tokens.length,
+    prefillTps: 0,
+    decodeTps: result.tokens.length / Math.max(decodeMs / 1000, 1e-9),
+    prefillMs: 0,
+    decodeMs,
+    cacheTokens: [],
+  };
+}
+
 /** Hold the model's active-adapter list for exactly this generation. */
 async function* adapterScoped(
   model: RuntimeModel,
@@ -219,6 +309,14 @@ async function* generateInner(
   const processors = makeLogitsProcessors(options);
   const needsTokenHistory = processors.length > 0;
 
+  // logprobs capture (mlx_lm.server semantics — see GenerateOptions.logprobs).
+  // Everything below is gated: when neither flag is set, no extra ops, evals,
+  // or readbacks happen on the hot path.
+  const captureSelLp = options.logprobs === true;
+  const captureTopK =
+    options.topLogprobs && options.topLogprobs > 0 ? options.topLogprobs : 0;
+  const capture = captureSelLp || captureTopK > 0;
+
   const ownsCache = !options.cache;
   const cache = options.cache ?? model.makeCache();
   const cachedTokens = cache[0]!.offset;
@@ -229,16 +327,73 @@ async function* generateInner(
   /** device-side token history (only maintained when processors need it) */
   let history: MlxArray | null = null;
 
-  // logits [1,1,V] → sampled token array [1] (all on-device)
-  const sampleStep = (logits3d: MlxArray, step: number): MlxArray => {
+  /** Device-side logprob capture for one step, computed from the SAME lp the
+   *  sampler saw (post-processors, pre-truncation) — read back lazily with the
+   *  token so decode pipelining is preserved. */
+  interface StepExtras {
+    sel: MlxArray | null; // [1,1] emitted token's logprob
+    topIdx: MlxArray | null; // [1,k] top-k token ids
+    topVals: MlxArray | null; // [1,k] their logprobs
+  }
+  const extrasArrays = (e: StepExtras | null): MlxArray[] =>
+    e ? [e.sel, e.topIdx, e.topVals].filter((a): a is MlxArray => a !== null) : [];
+  const disposeExtras = (e: StepExtras | null): void => {
+    if (e) for (const a of [e.sel, e.topIdx, e.topVals]) a?.dispose();
+  };
+  /** Read extras back to JS (forces eval — they were async-dispatched with the
+   *  token, so this normally just copies) and dispose the device arrays. */
+  const readExtras = (e: StepExtras | null): TokenLogprobs | undefined => {
+    if (!e) return undefined;
+    const out: TokenLogprobs = {};
+    if (e.sel) out.logprob = e.sel.toFloat32()[0]!;
+    if (e.topIdx && e.topVals) {
+      // uint32 ids read via the f32 cast: exact for ids < 2^24 (vocab ≪ 16M)
+      const ids = e.topIdx.toFloat32();
+      const vals = e.topVals.toFloat32();
+      out.top = Array.from(ids, (id, i) => ({ id, logprob: vals[i]! })).sort(
+        (a, b) => b.logprob - a.logprob,
+      );
+    }
+    disposeExtras(e);
+    return out;
+  };
+
+  // logits [1,1,V] → sampled token array [1] (+ optional logprob capture,
+  // all on-device)
+  const sampleStep = (
+    logits3d: MlxArray,
+    step: number,
+  ): { tok: MlxArray; extras: StepExtras | null } => {
     const V = logits3d.shape[2]!;
     let logits = ops.reshape(logits3d, [1, V]);
     for (const p of processors) logits = disposing(logits, p(history, logits));
     const lp = toLogprobs(logits);
     logits.dispose();
     const tok = sampler(lp, step);
+    let extras: StepExtras | null = null;
+    if (capture) {
+      let sel: MlxArray | null = null;
+      let topIdx: MlxArray | null = null;
+      let topVals: MlxArray | null = null;
+      if (captureSelLp) {
+        const idx = ops.reshape(tok, [1, 1]);
+        sel = ops.takeAlongAxis(lp, idx, -1);
+        idx.dispose();
+      }
+      if (captureTopK > 0) {
+        // mlx-lm _format_top_logprobs: argpartition(-logprobs, kth=k-1)[:k]
+        const k = Math.min(captureTopK, V);
+        const negLp = ops.neg(lp);
+        const part = ops.argpartitionAxis(negLp, k - 1, -1);
+        const idxView = part.slice([0, 0], [1, k]);
+        topIdx = ops.contiguous(idxView); // slices are strided views; raw readback needs contiguous
+        topVals = ops.takeAlongAxis(lp, topIdx, -1);
+        for (const a of [negLp, part, idxView]) a.dispose();
+      }
+      extras = { sel, topIdx, topVals };
+    }
     lp.dispose();
-    return tok;
+    return { tok, extras };
   };
 
   const pushHistory = (tok: MlxArray) => {
@@ -263,6 +418,8 @@ async function* generateInner(
   const forwarded: number[] = [];
   let pending: MlxArray | null = null;
   let nextPending: MlxArray | null = null;
+  let pendingExtras: StepExtras | null = null;
+  let nextExtras: StepExtras | null = null;
   let finished = false;
   let threw = false;
   const makeStats = (): GenerateStats => ({
@@ -325,7 +482,9 @@ async function* generateInner(
     h0.dispose();
     const logits0 = model.logitsFromHidden(hLast);
     hLast.dispose();
-    pending = sampleStep(logits0, 0); // token array [1]
+    const s0 = sampleStep(logits0, 0); // token array [1] (+ optional extras)
+    pending = s0.tok;
+    pendingExtras = s0.extras;
     logits0.dispose();
     // mirror mlx-lm generate_step: async-dispatch the first token's
     // compute; the prefill clock keeps running until the token ARRIVES
@@ -335,7 +494,7 @@ async function* generateInner(
     // to prompt_time, not decode — replicated so cross-stack decode
     // tok/s measure the same quantity. The boundary cost is real and
     // scales with prompt length; it belongs to "having prefilled".
-    ops.asyncEvalAll([pending]);
+    ops.asyncEvalAll([pending, ...extrasArrays(pendingExtras)]);
 
     // ---- decode (pipelined) ----
     // Compiled decode (docs/design/optimization_plan.md Phase A): replay the per-step
@@ -359,8 +518,11 @@ async function* generateInner(
     let stop = false;
     while (!stop) {
       const cur = pending!;
+      const curExtras = pendingExtras;
+      pendingExtras = null;
       // build step n+1's graph from the *unread* pending token
       nextPending = null;
+      nextExtras = null;
       if (generated + 1 < maxTokens) {
         maybeQuantizeKv(cache, options);
         pushHistory(cur);
@@ -372,8 +534,11 @@ async function* generateInner(
             logits = r.logits;
             evalWith = r.evalWith;
           } catch (e) {
-            // growth done by a partial prepare is benign for the
-            // uncompiled path (updateAndFetch re-checks capacity)
+            // Safe to re-forward the SAME token below: a failed step is
+            // transactional — segmented mode stages ring adopts and rolls
+            // back committed js-layer writes on throw (see #stepSegmented),
+            // and buffer growth done by a partial prepare is benign for the
+            // uncompiled path (updateAndFetch re-checks capacity).
             compiled = null;
             console.warn(`compiled decode disabled for this generation: ${e}`);
           }
@@ -385,9 +550,11 @@ async function* generateInner(
           logits = model.logitsFromHidden(h);
           h.dispose();
         }
-        nextPending = sampleStep(logits, generated + 1);
+        const sn = sampleStep(logits, generated + 1);
+        nextPending = sn.tok;
+        nextExtras = sn.extras;
         logits.dispose();
-        ops.asyncEvalAll([nextPending, ...evalWith]);
+        ops.asyncEvalAll([nextPending, ...extrasArrays(nextExtras), ...evalWith]);
       }
 
       // sync-read step n's token while n+1 computes
@@ -406,11 +573,17 @@ async function* generateInner(
       if (nextPending !== null) forwarded.push(token);
 
       if (eosTokenIds.includes(token)) {
+        disposeExtras(curExtras);
         nextPending?.dispose();
         nextPending = null;
+        disposeExtras(nextExtras);
+        nextExtras = null;
         stop = true;
       } else {
-        yield { token, index: generated - 1 };
+        // readExtras before the yield: if the consumer breaks at this yield,
+        // the extras are already read and disposed.
+        const logprobs = readExtras(curExtras);
+        yield { token, index: generated - 1, ...(logprobs ? { logprobs } : {}) };
         // mlx-lm generate_step: clear_cache after token 0 (drops the
         // remaining prefill transients) and every 256 tokens after
         if ((generated - 1) % 256 === 0) clearCache();
@@ -419,6 +592,8 @@ async function* generateInner(
         } else {
           pending = nextPending;
           nextPending = null;
+          pendingExtras = nextExtras;
+          nextExtras = null;
         }
       }
     }
@@ -432,6 +607,8 @@ async function* generateInner(
     if (!finished) {
       pending?.dispose();
       nextPending?.dispose();
+      disposeExtras(pendingExtras);
+      disposeExtras(nextExtras);
     }
     if (ownsCache) for (const c of cache) c.dispose();
     history?.dispose();