mlx-bun 0.0.12 → 0.2.0

package/src/eval/tasks/mmlu.ts

@@ -10,6 +10,7 @@
 // table. MLX_BUN_MMLU_FROZEN=0 falls back to our own sampling of the full 14k set.
 
 import { loadJsonl, sampleIndices, type TaskModel } from "../runner";
+import { runtimeValue } from "../../runtime-config";
 
 interface MmluRow {
   question: string;
@@ -85,7 +86,7 @@ export async function evaluateMmlu(
   // parity, not sampling. Set MLX_BUN_MMLU_FROZEN=0 (or frozen:false) to fall back to
   // our own stratified sampling of the full 14k test set (a DIFFERENT subset: JS RNG
   // ≠ numpy RandomState, so it does not match optiq's questions even at the same seed).
-  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_MMLU_FROZEN !== "0");
+  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_MMLU_FROZEN") !== "0");
 
   const testRows = loadJsonl<MmluRow>(useFrozen ? "mmlu_optiq_frozen" : "mmlu_test");
   const devRows = loadJsonl<MmluRow>(useFrozen ? "mmlu_optiq_dev" : "mmlu_dev");

package/src/expert-io.ts

@@ -4,6 +4,7 @@ import { dirname, join } from "node:path";
 import { nativePackDir } from "./native-pack";
 import type { MlxArray } from "./mlx/array";
 import type { MlxHandle } from "./mlx/ffi";
+import { runtimeValue } from "./runtime-config";
 
 const { cstring, i32, ptr: pointer, u32, u64 } = FFIType;
 const EBUSY = 16;
@@ -69,7 +70,8 @@ export interface ExpertIOHintSnapshot {
 
 function resolveLibrary(explicit?: string): string {
   if (explicit) return explicit;
-  if (process.env.MLX_BUN_EXPERT_IO_DYLIB) return process.env.MLX_BUN_EXPERT_IO_DYLIB;
+  const configured = runtimeValue("MLX_BUN_EXPERT_IO_DYLIB");
+  if (configured) return configured;
   const candidates = [
     join(dirname(process.execPath), "libmlx_bun_expert_io.dylib"),
     join(nativePackDir(), "libmlx_bun_expert_io.dylib"),

package/src/expert-offload.ts

@@ -16,6 +16,7 @@ import { readFileSync } from "node:fs";
 import { MmapFile } from "./mmap";
 import { MlxArray, SAFETENSORS_TO_MLX } from "./mlx/array";
 import type { SafetensorsDtype } from "./safetensors";
+import { runtimeValue } from "./runtime-config";
 
 export interface OffloadEntry {
   name: string;
@@ -71,4 +72,5 @@ export function expertOffloadArray(name: string): MlxArray | null {
   return MlxArray.fromPointer(mm.pointer(e.offset), e.shape, SAFETENSORS_TO_MLX[e.dtype]);
 }
 
-if (process.env.MLX_BUN_EXPERT_OFFLOAD) activateExpertOffload(process.env.MLX_BUN_EXPERT_OFFLOAD);
+const expertOffloadPath = runtimeValue("MLX_BUN_EXPERT_OFFLOAD");
+if (expertOffloadPath) activateExpertOffload(expertOffloadPath);

package/src/expert-trace.ts

@@ -19,6 +19,7 @@
 
 import { Dtype } from "./mlx/ffi";
 import type { MlxArray } from "./mlx/array";
+import { runtimeValue } from "./runtime-config";
 
 let sink: Bun.FileSink | null = null;
 let seq = 0;
@@ -61,7 +62,8 @@ export function recordRouting(layer: number, indices: MlxArray): void {
   if (++written % 2000 === 0) sink.flush();
 }
 
-if (process.env.MLX_BUN_EXPERT_TRACE) beginExpertTrace(process.env.MLX_BUN_EXPERT_TRACE);
+const expertTracePath = runtimeValue("MLX_BUN_EXPERT_TRACE");
+if (expertTracePath) beginExpertTrace(expertTracePath);
 
 process.on("exit", () => { try { endExpertTrace(); } catch { /* best-effort */ } });
 process.on("SIGINT", () => { endExpertTrace(); process.exit(0); });

package/src/fit.ts

@@ -10,6 +10,16 @@
 
 import { totalmem } from "node:os";
 import type { ModelConfig } from "./config";
+import type { MemoryPlan } from "./memory-plan";
+import {
+  kvBytesAt,
+  kvGeometry,
+  kvQuantBytesPerElement,
+  type KvSchemeOptions,
+} from "./kv-scheme";
+
+export { kvBytesAt, kvQuantBytesPerElement } from "./kv-scheme";
+export type FitKvScheme = KvSchemeOptions;
 
 /** Decode-efficiency vs theoretical bandwidth ceiling, measured on the
  *  reference machine (24.9 tok/s vs 30.3 ceiling @600 ctx). */
@@ -26,8 +36,6 @@ export const TRANSIENT_PER_TOKEN = 0.55e6;
  *  recommendedMaxWorkingSetSize is ~75% on consumer SKUs). */
 export const WIRED_FRACTION = 0.75;
 export const DEFAULT_CHUNK = 2048;
-/** KV cache element size (bf16; quantized KV lands in Phase 6). */
-const KV_BYTES = 2;
 
 export interface MachineSpec {
   name: string;
@@ -130,43 +138,9 @@ export function chooseAutoModel<T extends AutoPickCandidate>(
   return bySizeDesc.find(fitsCoexistBudget) ?? bySizeDesc.find(fitsFullBudget);
 }
 
-export interface FitReport {
-  fits: boolean;
-  contextTokens: number;
-  weightsBytes: number;
-  kvBytes: number;
-  transientBytes: number;
-  totalBytes: number;
-  usableBytes: number;
-  maxSafeContext: number;
-  predictedDecodeTps: number;
-}
-
-interface KvGeometry {
-  fullLayers: number;
-  slidingLayers: number;
-  fullBytesPerToken: number;
-  slidingBytesPerToken: number;
-  window: number;
-}
-
-function kvGeometry(config: ModelConfig): KvGeometry {
-  const t = config.text;
-  const slidingLayers = t.layerTypes.filter((l) => l === "sliding_attention").length;
-  const fullLayers = t.numHiddenLayers - slidingLayers;
-  const kEqVFactor = t.attentionKEqV ? 2 : 2; // k and v stored separately either way
-  return {
-    fullLayers,
-    slidingLayers,
-    fullBytesPerToken: fullLayers * kEqVFactor * t.numGlobalKeyValueHeads * t.globalHeadDim * KV_BYTES,
-    slidingBytesPerToken: slidingLayers * 2 * t.numKeyValueHeads * t.headDim * KV_BYTES,
-    window: t.slidingWindow,
-  };
-}
-
-export function kvBytesAt(config: ModelConfig, ctx: number): number {
-  const g = kvGeometry(config);
-  return g.fullBytesPerToken * ctx + g.slidingBytesPerToken * Math.min(ctx, g.window);
+export interface FitReport extends MemoryPlan {
+  readonly strategy: "generic-kv";
+  readonly predictedDecodeTps: number;
 }
 
 export function fit(
@@ -182,17 +156,20 @@ export function fit(
    *  replaces the machine-derived usable ceiling (ram × WIRED_FRACTION)
    *  outright — the budget IS the usable envelope. */
   usableBytes?: number,
+  /** Active KV-quant scheme; a quantized cache holds more context in the
+   *  same budget, so the solved ceiling (and admission) must bill it. */
+  kvScheme?: FitKvScheme,
 ): FitReport {
   const usable = usableBytes ?? machine.ramBytes * WIRED_FRACTION;
   const transient = Math.min(chunk, ctx) * TRANSIENT_PER_TOKEN;
-  const kv = kvBytesAt(config, ctx);
+  const kv = kvBytesAt(config, ctx, kvScheme);
   const total = weightsBytes + kv + transient;
 
   // solve max context: weights + kv(ctx) + transient ≤ usable.
   // Below the window both KV terms are linear in ctx; above it the
   // sliding term saturates and only full-attention layers keep growing.
-  const g = kvGeometry(config);
-  const fixed = weightsBytes + chunk * TRANSIENT_PER_TOKEN;
+  const g = kvGeometry(config, kvScheme);
+  const fixed = weightsBytes + chunk * TRANSIENT_PER_TOKEN + g.linearStateBytes;
   let maxCtx = 0;
   if (usable > fixed) {
     const budget = usable - fixed;
@@ -216,11 +193,14 @@ export function fit(
     (isMoe ? MOE_DECODE_EFFICIENCY : DECODE_EFFICIENCY);
 
   return {
+    schemaVersion: 1,
+    strategy: "generic-kv",
     fits: total <= usable,
     contextTokens: ctx,
     weightsBytes,
     kvBytes: kv,
     transientBytes: transient,
+    reserveBytes: 0,
     totalBytes: total,
     usableBytes: usable,
     maxSafeContext: maxCtx,
@@ -231,6 +211,7 @@ export function fit(
 /** The SKU matrix: which Apple Silicon configs run this model at `ctx`. */
 export function skuMatrix(
   config: ModelConfig, weightsBytes: number, ctx: number, expertsBytes = 0,
+  kvScheme?: FitKvScheme,
 ): { sku: string; ramGB: number; fits: boolean; maxContext: number; decodeTps: number }[] {
   const rows: ReturnType<typeof skuMatrix> = [];
   for (const sku of APPLE_SKUS) {
@@ -240,7 +221,7 @@ export function skuMatrix(
         ramBytes: ram * 2 ** 30,
         bandwidthGBs: sku.bandwidthGBs,
       };
-      const r = fit(config, weightsBytes, ctx, m, DEFAULT_CHUNK, expertsBytes);
+      const r = fit(config, weightsBytes, ctx, m, DEFAULT_CHUNK, expertsBytes, undefined, kvScheme);
       rows.push({
         sku: sku.chip,
         ramGB: ram,

package/src/flags.ts

@@ -1,13 +1,11 @@
-// Generic on/off resolver for the per-fork decode flags (--compiled-decode,
-// --compiled-activations, --fused-sdpa, …). An explicit "1"/"0" in the env
-// always wins; otherwise the caller's `defaultOn` applies. The tier aliases
-// (--l1/--l2, applyDecodeRoute in cli.ts) set the env vars; model code
-// reads them through this helper so a hand-set flag and a tier preset resolve
-// identically.
+import { runtimeFlag, type RuntimeKey } from "./runtime-config";
+
+// Generic on/off resolver for per-fork decode flags. Process environment is
+// captured once by runtime-config; CLI tier aliases install an explicit
+// immutable override before model modules are loaded.
 
 /** Resolve an on/off env flag. An explicit "1"/"0" always wins; otherwise the
  *  flag takes `defaultOn`. */
-export function flagOn(name: string, defaultOn: boolean): boolean {
-  const v = process.env[name];
-  return v === "1" ? true : v === "0" ? false : defaultOn;
+export function flagOn(name: RuntimeKey, defaultOn: boolean): boolean {
+  return runtimeFlag(name, defaultOn);
 }

package/src/generate.ts

@@ -8,8 +8,10 @@
 
 import { MlxArray, gpuStream } from "./mlx/array";
 import {
+  activeMemory,
   clearCache,
   maxRecommendedWorkingSetSize,
+  peakMemory,
   setWiredLimit,
   synchronize,
 } from "./mlx/ffi";
@@ -20,11 +22,17 @@ import { PagedKVCache } from "./model/paged-kv";
 import { DiffusionGemmaModel } from "./model/diffusion-gemma";
 import { diffusionGenerate } from "./diffusion/diffusion-generate";
 import type { RuntimeModel } from "./model/factory";
+import type { PromptResponseTrace } from "./serve/prompt-response-trace";
 import type { KvQuantSpec, TurboQuantScheme } from "./config";
 import { flagOn } from "./flags";
+import { runtimeValue } from "./runtime-config";
 import {
-  makeLogitsProcessors, makeSampler,
-  type LogitsProcessorOptions, type SamplerOptions, toLogprobs,
+  disposeStepExtras,
+  makeStepSampler,
+  stepExtrasArrays,
+  type LogitsProcessorOptions,
+  type SamplerOptions,
+  type StepExtras,
 } from "./sampler";
 import type { GrammarController } from "./grammar";
 
@@ -296,6 +304,11 @@ export interface GenerateStats {
   };
 }
 
+export interface GenerateDiagnostics {
+  trace?: PromptResponseTrace;
+  mechanism?: "serial" | "continuous";
+}
+
 export interface GeneratedToken {
   token: number;
   index: number;
@@ -346,7 +359,11 @@ export class Generation implements AsyncIterable<GeneratedToken> {
 // Scope semantics match the reference: set → generate → synchronize →
 // restore; nothing stays pinned between generations. Re-entrant: only
 // the outermost wiring scope touches the limit.
-const WIRE_THRESHOLD = 0.75;
+// macOS 26.6 reports a 24.96 GiB recommended set on this 24 GB machine.
+// The old 0.75 fraction therefore stopped wiring the 13-16 GiB Qwen/GLM
+// models even though they page heavily without an explicit wired limit.
+// Keep smaller 8-9 GiB models unwired while covering the large-model class.
+const WIRE_THRESHOLD = 0.5;
 let wiredScopeDepth = 0;
 let wiredOldLimit = 0;
 
@@ -370,7 +387,7 @@ export function wiredWorkingSetBytes(model: WiredModelMemory): number {
 export function modelNeedsWiredLimit(
   model: WiredModelMemory,
   recommendedBytes = maxRecommendedWorkingSetSize(),
-  force = process.env.MLX_BUN_FORCE_WIRE === "1",
+  force = runtimeValue("MLX_BUN_FORCE_WIRE") === "1",
 ): boolean {
   return force ||
     wiredWorkingSetBytes(model) > WIRE_THRESHOLD * recommendedBytes;
@@ -387,10 +404,23 @@ function exitWiredScope(): void {
   }
 }
 
+/** Hold the process wired limit while one model owns GPU execution. */
+export function acquireModelWiredLimit(model: WiredModelMemory): () => void {
+  if (!modelNeedsWiredLimit(model)) return () => {};
+  enterWiredScope();
+  let released = false;
+  return () => {
+    if (released) return;
+    released = true;
+    exitWiredScope();
+  };
+}
+
 export function generate(
   model: RuntimeModel,
   promptTokens: number[],
   options: GenerateOptions = {},
+  diagnostics: GenerateDiagnostics = {},
 ): Generation {
   // DiffusionGemma is non-autoregressive: route to the denoising engine instead
   // of the AR decode loop. Same Generation/GenerateStats contract so the CLI and
@@ -398,7 +428,7 @@ export function generate(
   let inner =
     model instanceof DiffusionGemmaModel
       ? generateDiffusionInner(model, promptTokens, options)
-      : generateInner(model, promptTokens, options);
+      : generateInner(model, promptTokens, options, diagnostics);
   if (options.adapters?.length) inner = adapterScoped(model, options.adapters, inner);
   const memoryModel: WiredModelMemory = model;
   if (memoryModel.expertRuntime?.flushUsage)
@@ -496,12 +526,11 @@ export async function withModelWiredLimit<T>(
   model: WiredModelMemory,
   run: () => Promise<T>,
 ): Promise<T> {
-  if (!modelNeedsWiredLimit(model)) return run();
-  enterWiredScope();
+  const release = acquireModelWiredLimit(model);
   try {
     return await run();
   } finally {
-    exitWiredScope();
+    release();
   }
 }
 
@@ -542,15 +571,13 @@ async function* generateInner(
   model: RuntimeModel,
   promptTokens: number[],
   options: GenerateOptions,
+  diagnostics: GenerateDiagnostics,
 ): AsyncGenerator<GeneratedToken, GenerateStats> {
   const {
     maxTokens = 512,
     eosTokenIds = model.config.eosTokenIds,
     prefillChunkSize = 2048,
   } = options;
-  const sampler = makeSampler(options);
-  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,
@@ -558,8 +585,23 @@ async function* generateInner(
   const captureSelLp = options.logprobs === true;
   const captureTopK =
     options.topLogprobs && options.topLogprobs > 0 ? options.topLogprobs : 0;
-  const capture = captureSelLp || captureTopK > 0;
+  const stepSampler = makeStepSampler(options, {
+    tokenRepresentation: "device",
+    grammarWait: "external",
+    historyUpdate: "manual",
+    captureSelectedLogprob: captureSelLp,
+    captureTopLogprobs: captureTopK,
+  });
+  const needsTokenHistory = stepSampler.needsHistory;
 
+  const closePrefill = diagnostics.trace?.begin("prefill.total", {
+    mechanism: diagnostics.mechanism ?? "serial",
+    promptTokens: promptTokens.length,
+    cachedTokens: options.cache?.[0]?.offset ?? 0,
+  });
+  const closeBatchSetup = diagnostics.trace?.begin("prefill.batch_setup", {
+    mechanism: diagnostics.mechanism ?? "serial",
+  });
   const ownsCache = !options.cache;
   const cache = options.cache ?? model.makeCache();
   const cachedTokens = cache[0]!.offset;
@@ -572,22 +614,10 @@ async function* generateInner(
   // prompt + maxTokens − 1 (step maxTokens−1's forward), so this bound
   // makes pool exhaustion an accounting-bug tripwire, not a request limit.
   maybePageKv(cache, options, promptTokens.length + maxTokens);
-  /** device-side token history (only maintained when processors need it) */
-  let history: MlxArray | null = null;
-
+  closeBatchSetup?.();
   /** 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) without appending casts behind the next decode step. */
   const readExtras = (e: StepExtras | null): TokenLogprobs | undefined => {
@@ -601,7 +631,7 @@ async function* generateInner(
         (a, b) => b.logprob - a.logprob,
       );
     }
-    disposeExtras(e);
+    disposeStepExtras(e);
     return out;
   };
 
@@ -611,52 +641,11 @@ async function* generateInner(
     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));
-    // Grammar mask (src/grammar.ts): apply AFTER the standard processors so the
-    // grammar has the final say on validity (a -inf'd token stays -inf through
-    // log-softmax). Off (no grammar) = zero cost — the branch is unset.
-    if (options.grammar) logits = disposing(logits, options.grammar.applyMask(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, extras };
+    const result = stepSampler.sample(logits3d, step);
+    return { tok: result.token, extras: result.extras };
   };
 
-  const pushHistory = (tok: MlxArray) => {
-    if (!needsTokenHistory) return;
-    if (!history) {
-      history = ops.reshape(tok, [1]);
-    } else {
-      const prev = history;
-      history = ops.concatAxis([prev, tok], 0);
-      prev.dispose();
-    }
-  };
+  const pushHistory = (tok: MlxArray): void => stepSampler.commitDevice(tok);
 
   // Decode-loop state lives at function scope so the finally can still
   // report stats and dispose in-flight arrays when the consumer
@@ -673,6 +662,7 @@ async function* generateInner(
   let nextExtras: StepExtras | null = null;
   let finished = false;
   let threw = false;
+  let closeTokenZero: (() => void) | undefined;
   const makeStats = (): GenerateStats => ({
     promptTokens: promptTokens.length,
     cachedTokens,
@@ -686,14 +676,19 @@ async function* generateInner(
 
   try {
     // ---- prefill ----
+    const closeInitialKv = diagnostics.trace?.begin("prefill.kv_maintenance", {
+      mechanism: diagnostics.mechanism ?? "serial",
+      boundary: "initial",
+    });
     maybeQuantizeKv(cache, options);
+    closeInitialKv?.();
     const tPrefill = performance.now();
     let h0: MlxArray;
     if (options.promptEmbeddings) {
       if (cachedTokens !== 0)
         throw new Error("promptEmbeddings cannot be combined with a pre-warmed cache");
       if (needsTokenHistory)
-        history = ops.fromInt32(promptTokens, [promptTokens.length]);
+        stepSampler.seedHistory(promptTokens);
       // e2b/e4b need the spliced token ids to build per-layer inputs
       // (multimodal soft-token positions zeroed inside forwardEmbeddings).
       const embedIds = ops.fromInt32(promptTokens, [1, promptTokens.length]);
@@ -715,6 +710,11 @@ async function* generateInner(
       if (snapAt > pos && snapAt < promptTokens.length) {
         while (pos < snapAt) {
           const chunk = promptTokens.slice(pos, Math.min(pos + prefillChunkSize, snapAt));
+          const closeChunk = diagnostics.trace?.begin("prefill.chunk", {
+            mechanism: diagnostics.mechanism ?? "serial",
+            startToken: pos,
+            tokens: chunk.length,
+          });
           const ids = ops.fromInt32(chunk, [1, chunk.length]);
           const h = await forwardHiddenForGeneration(model, ids, cache);
           ids.dispose();
@@ -722,6 +722,7 @@ async function* generateInner(
           evalCacheState(cache);
           maybeQuantizeKv(cache, options);
           clearCache();
+          closeChunk?.();
           pos += chunk.length;
           // Macrotask yield between chunks (runtime-isolation.md Phase 1):
           // each chunk is a synchronous multi-hundred-ms FFI eval; without
@@ -749,6 +750,11 @@ async function* generateInner(
           ? Math.min(prefillChunkSize, promptTokens.length - pos - 1)
           : prefillChunkSize;
         const chunk = promptTokens.slice(pos, pos + n);
+        const closeChunk = diagnostics.trace?.begin("prefill.chunk", {
+          mechanism: diagnostics.mechanism ?? "serial",
+          startToken: pos,
+          tokens: chunk.length,
+        });
         const ids = ops.fromInt32(chunk, [1, chunk.length]);
         const h = await forwardHiddenForGeneration(model, ids, cache);
         ids.dispose();
@@ -762,12 +768,15 @@ async function* generateInner(
         // measured, scripts/decode-split.ts; the context-scaling decode
         // gap's main term).
         clearCache();
+        closeChunk?.();
         pos += n;
+        if (runtimeValue("MLX_BUN_PREFILL_MEM_LOG") === "1")
+          console.error(`[prefill-mem] ${pos} active ${(activeMemory() / 2 ** 30).toFixed(2)} peak ${(peakMemory() / 2 ** 30).toFixed(2)}`);
         // Macrotask yield between chunks (runtime-isolation.md Phase 1).
         await new Promise<void>((r) => setImmediate(r));
       }
       if (needsTokenHistory) {
-        history = ops.fromInt32(promptTokens, [promptTokens.length]);
+        stepSampler.seedHistory(promptTokens);
       }
       // Under tailSplit this is exactly [promptTokens[len-1]] — the L=1
       // step-0 forward (uncompiled L=1 is bit-exact with compiled decode,
@@ -776,12 +785,27 @@ async function* generateInner(
       // PromptCache.put alignment) is unchanged: after step 0 the caches
       // cover exactly the prompt, as before.
       const lastChunk = promptTokens.slice(pos);
+      const closeChunk = diagnostics.trace?.begin("prefill.chunk", {
+        mechanism: diagnostics.mechanism ?? "serial",
+        startToken: pos,
+        tokens: lastChunk.length,
+        final: true,
+      });
       const ids0 = ops.fromInt32(lastChunk, [1, lastChunk.length]);
       h0 = await forwardHiddenForGeneration(model, ids0, cache);
       ids0.dispose();
+      closeChunk?.();
     }
     if (options.snapshotAt === undefined || options.snapshotAt >= promptTokens.length)
       options.onPrefillDone?.();
+    // The trace's prefill span is the uncached model-forward portion only.
+    // Token-0 lm-head, sampling, and readback are a separate additive stage.
+    if (diagnostics.trace && runtimeValue("MLX_BUN_P2R_SYNC") === "1")
+      synchronize(gpuStream);
+    closePrefill?.();
+    closeTokenZero = diagnostics.trace?.begin("token_zero.total", {
+      mechanism: diagnostics.mechanism ?? "serial",
+    });
     const [, L0, H] = h0.shape as [number, number, number];
     const hLast = h0.slice([0, L0 - 1, 0], [1, L0, H]);
     h0.dispose();
@@ -799,7 +823,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, ...extrasArrays(pendingExtras)]);
+    ops.asyncEvalAll([pending, ...stepExtrasArrays(pendingExtras)]);
 
     // ---- decode (pipelined) ----
     // Compiled decode (docs/design/optimization_plan.md Phase A): replay the per-step
@@ -886,13 +910,7 @@ async function* generateInner(
         // following iteration (supports() re-checks the grown caches).
         maybeQuantizeKv(cache, options);
         pushHistory(cur);
-        if (needsTokenHistory) {
-          const jt = ops.fromInt32(jumpEmit, [jumpEmit.length]);
-          const prev = history!;
-          history = ops.concatAxis([prev, jt], 0);
-          prev.dispose();
-          jt.dispose();
-        }
+        stepSampler.commitNumbers(jumpEmit);
         const chunk = [grammarTok, ...jumpEmit];
         const ids = ops.fromInt32(chunk, [1, chunk.length]);
         const h = await forwardHiddenForGeneration(model, ids, cache);
@@ -910,7 +928,7 @@ async function* generateInner(
           nextPending = sn.tok;
           nextExtras = sn.extras;
           logits.dispose();
-          ops.asyncEvalAll([nextPending, ...extrasArrays(nextExtras)]);
+          ops.asyncEvalAll([nextPending, ...stepExtrasArrays(nextExtras)]);
         } else {
           h.dispose(); // burst ends the generation (max_tokens or grammar done)
         }
@@ -945,7 +963,7 @@ async function* generateInner(
         nextPending = sn.tok;
         nextExtras = sn.extras;
         logits.dispose();
-        ops.asyncEvalAll([nextPending, ...extrasArrays(nextExtras), ...evalWith]);
+        ops.asyncEvalAll([nextPending, ...stepExtrasArrays(nextExtras), ...evalWith]);
       }
 
       // sync-read step n's token while n+1 computes
@@ -957,6 +975,8 @@ async function* generateInner(
         // (mlx-lm stream_generate's n==0 clock swap; the first token is
         // "free" on the decode clock there too)
         prefillMs = performance.now() - tPrefill;
+        closeTokenZero?.();
+        closeTokenZero = undefined;
         tDecode = performance.now();
       }
       cur.dispose();
@@ -967,10 +987,10 @@ async function* generateInner(
       if (nextPending !== null && !jumpEmit) forwarded.push(token);
 
       if (eosTokenIds.includes(token)) {
-        disposeExtras(curExtras);
+        disposeStepExtras(curExtras);
         nextPending?.dispose();
         nextPending = null;
-        disposeExtras(nextExtras);
+        disposeStepExtras(nextExtras);
         nextExtras = null;
         stop = true;
       } else {
@@ -1013,8 +1033,8 @@ async function* generateInner(
     if (!finished) {
       pending?.dispose();
       nextPending?.dispose();
-      disposeExtras(pendingExtras);
-      disposeExtras(nextExtras);
+      disposeStepExtras(pendingExtras);
+      disposeStepExtras(nextExtras);
     }
     if (ownsCache) for (const c of cache) c.dispose();
     // The grammar controller owns native WASM state (GrammarMatcher +
@@ -1022,7 +1042,7 @@ async function* generateInner(
     // throw). TokenizerInfo itself is process-cached (vocab-structural), so
     // only the per-request matcher/compiled are freed here.
     options.grammar?.dispose();
-    history?.dispose();
+    stepSampler.dispose();
     if (!finished && !threw) {
       // forced early return (consumer break at a yield): still report
       // stats — `forwarded` only lists tokens whose KV actually entered
@@ -1032,8 +1052,3 @@ async function* generateInner(
     }
   }
 }
-
-function disposing(old: MlxArray, next: MlxArray): MlxArray {
-  old.dispose();
-  return next;
-}

package/src/grammar.ts

@@ -33,6 +33,7 @@ import xgrammar, {
 import { readFileSync } from "node:fs";
 import { MlxArray } from "./mlx/array";
 import * as ops from "./mlx/ops";
+import { runtimeValue } from "./runtime-config";
 import type { LoadedTokenizer } from "./tokenizer";
 
 /** The request shape — the union of every structured-output field name across
@@ -581,5 +582,5 @@ function resolveGrammarRequest(req: GrammarRequest):
  *  (requests degrade to prompt injection). Matches the project's envelope-gated
  *  kill-switch discipline (MLX_BUN_COMPILED_DECODE, MLX_BUN_GRAMMAR_BATCH, …). */
 export function grammarEnabled(): boolean {
-  return process.env.MLX_BUN_GRAMMAR !== "0";
+  return runtimeValue("MLX_BUN_GRAMMAR") !== "0";
 }

package/src/harness-pi.ts

@@ -1,5 +1,5 @@
 // mlx-bun harness pi — register mlx-bun as a provider in the user's pi
-// install (Phase 16 P1, docs/investigations/pi-builtin-investigation.md).
+// install (Phase 16 P1, docs/archive/investigations/pi-builtin-investigation.md).
 //
 // Mechanism: a generated extension file in ~/.pi/agent/extensions/ that
 // calls pi.registerProvider() with dynamic discovery — it fetches