mlx-bun 0.2.0 → 0.3.0

package/src/mlx/turboquant-ops.ts

@@ -1,8 +1,8 @@
-// TurboQuant KV-cache codec — pure mlx-op composition, ported op-for-op
+// TurboQuant KV-cache codec — reference operations ported op-for-op
 // from the vendored vllm-metal reference
 // (lab/repro/vllm-metal-turboquant/turboquant_reference.py). See
-// docs/design/turboquant-kv.md for the algorithm writeup and
-// goldens/turboquant.json (via tests/turboquant-ops.test.ts) for the
+// docs/design/turboquant.md for the algorithm writeup and
+// goldens/turboquant.json (via tests/research/turboquant-ops.test.ts) for the
 // bit-exactness oracle. Quantization groups run along the LAST axis
 // (head_dim) only — never the token axis, per the standing invariant.
 //
@@ -27,6 +27,10 @@ import {
   type TurboQuantHeadDim, type TurboQuantValueBits,
 } from "./turboquant-tables";
 
+import { tryJointDecodePackedKv, type PackedKvArrays } from "./turboquant-kv-decode";
+import { tryDecodePackedKvInverse256 } from "./turboquant-kv-inverse";
+import { isShapelessTracing } from "./compile";
+
 type S = MlxHandle;
 
 const BLOCK_SIZE = 32;
@@ -421,6 +425,32 @@ export function decodeValuesRotated(
   return outBf16;
 }
 
+/** Optional exact packed K/V decode. Inputs stay borrowed; the caller owns
+ * the returned pair. The eager path keeps the reference inverse rotation. */
+export function tryDecodePackedKv(
+  inputs: PackedKvArrays, kBits: number, vBits: number, headDim: number,
+  deferV = false, s: S = gpuStream,
+): [MlxArray, MlxArray] | null {
+  // Reject unsupported value widths before requesting their centroid table.
+  if (![2, 3, 4, 5, 8].includes(vBits)) return null;
+  if (!deferV && headDim === 256 && kBits === 8 && vBits === 3 &&
+      s === gpuStream && !isShapelessTracing() && inputs[0].shape[2]! >= 8192) {
+    const decoded = tryDecodePackedKvInverse256(
+      inputs, centroidsFor(vBits), signArray(256), kBits, vBits, headDim, s,
+    );
+    if (decoded) return decoded;
+  }
+  const decoded = tryJointDecodePackedKv(inputs, centroidsFor(vBits), kBits, vBits, headDim, deferV, s);
+  if (!decoded || deferV) return decoded;
+  const [key, rotated] = decoded;
+  try {
+    const value = fwht(rotated, false, s);
+    try { return [key, value.astype(Dtype.bfloat16, s)]; }
+    finally { value.dispose(); }
+  } catch (error) { key.dispose(); throw error; }
+  finally { rotated.dispose(); }
+}
+
 /** Undo the value rotation on an attention output computed against
  *  decodeValuesRotated values. f32 through the transform (matching the
  *  eager decode path's precision), returned in the input's dtype. */

package/src/mlx/turboquant-tables.ts

@@ -1,5 +1,5 @@
-// GENERATED by scripts/regen-turboquant-goldens.ts — DO NOT EDIT.
-// Regenerate:  bun scripts/regen-turboquant-goldens.ts
+// GENERATED by scripts/regen.ts turboquant — DO NOT EDIT.
+// Regenerate:  bun scripts/regen.ts turboquant
 //
 // Source: lab/repro/vllm-metal-turboquant/turboquant_reference.py (vendored
 // vllm-metal oracle, see PROVENANCE.md in that directory), executed under

package/src/model/batched-mask.ts

@@ -2,11 +2,11 @@
 //
 // Companion to src/train/forward.ts's buildBatchedPadMask, which builds the
 // PREFILL mask ([B,1,L,L], offset 0, right-padded — the training case, already
-// parity-proven by tests/train-batch-e2e.test.ts). This module covers the
+// parity-proven by tests/parity/train-batch-e2e.test.ts). This module covers the
 // DECODE case for batched serving: left-padded rows sharing one growing
 // [B,H,S,D] KV buffer, where a step forwards N new query tokens at a nonzero
 // offset. (The two builders should be consolidated when batched prefill is
-// wired into the serving path — phase S1a, docs/design/parallel-slots.md.)
+// wired into the serving path — phase S1a, docs/design/batching.md.)
 //
 // Why left-padding for decode: with rows right-aligned in the KV buffer, every
 // row's next write lands in the SAME column, so one advancing offset serves
@@ -15,6 +15,7 @@
 // thing the mask must encode is (a) causality, (b) each row's left padding,
 // and (c) the sliding window, if any.
 
+import { cacheSignature } from "./gemma4-base";
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { createCausalMask, KVCache, type Cache, type Mask } from "./gemma4-base";
@@ -92,6 +93,11 @@ export class BatchedDecodeMaskCache implements Cache {
   #ropeArr: MlxArray | null = null;
   #ropeForOffset = -1;
 
+  /** A decorator, not a storage kind: never claims the inner cache's
+   *  identity (isPlainKvCache(wrapper) must stay false — it is not a
+   *  KVCache), but records what it wraps for diagnostics/persistence. */
+  signature(): string { return `kv:batched-mask:${cacheSignature(this.inner)}`; }
+
   constructor(
     private readonly inner: Cache,
     private readonly B: number,

package/src/model/batched-quant.ts

@@ -17,7 +17,7 @@
 // (maybeQuantizeKv at chunk boundaries — see BatchScheduler.#quantizeSolo),
 // so a row's quantized bytes are bit-exact vs serial `--kv-quant config` by
 // construction; this module only re-arranges those bytes across the batch
-// axis. Gates: tests/batched-kv-quant-parity.test.ts.
+// axis. Gates: tests/parity/batched-kv-quant-parity.test.ts.
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";

package/src/model/batched-rotating-quant.ts

@@ -22,9 +22,9 @@
 // every step; it self-tracks per-row offset/leftPad and the scalar ring
 // state. N=1 decode updates only (rows solo-prefill serially, then merge).
 //
-// Gates: tests/batched-rotating-quant.test.ts (model-free byte-identity vs
+// Gates: tests/unit/batched-rotating-quant.test.ts (model-free byte-identity vs
 // the serial oracle per row, through ring wrap, B=1 and B=2) + the gemma
-// scheduler gate in tests/batched-kv-quant-parity.test.ts.
+// scheduler gate in tests/parity/batched-kv-quant-parity.test.ts.
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";

package/src/model/batched-rotating.ts

@@ -20,7 +20,7 @@
 // temporal coordinates then ROLLED to the ring's physical layout — see
 // buildBatchedRotatingMask (a column-by-column port of mlx-lm make_mask,
 // including the `roll(shift=idx+1)`). Gated model-free against mlx-lm
-// (tests/batched-rotating.test.ts) + end-to-end vs a long-context Gemma oracle.
+// (tests/unit/batched-rotating.test.ts) + end-to-end vs a long-context Gemma oracle.
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";

package/src/model/compiled-decode.ts

@@ -1,4 +1,4 @@
-// Compiled decode step (Phase A of docs/design/optimization_plan.md).
+// Compiled decode step (Phase A of docs/archive/investigations/optimization_plan.md).
 //
 // The single-token decode graph is identical token-to-token except for
 // integer state (RoPE offset, cache write position, active length). This

package/src/model/diffusion-gemma.ts

@@ -1,3 +1,4 @@
+import { cleanupFailure, disposeResources } from "../engine/resources";
 // DiffusionGemma-26B-A4B-it (model_type "diffusion_gemma") — the first
 // NON-autoregressive model in the codebase. Instead of a left-to-right AR loop
 // it runs an ENCODER prefill over the prompt (builds a KV cache) and then a
@@ -10,7 +11,7 @@
 // Ported verbatim from optiq/vlm/_mlxvlm/models/diffusion_gemma/language.py.
 // Reuses mlx-bun's quantized primitives (QuantizedLinear/Embedding/SwitchLinear,
 // RMSNorm, KVCache/RotatingKVCache) from gemma4-base. Architecture deltas vs
-// gemma4 (verified against the checkpoint, see docs/design/diffusion-gemma-port.md):
+// gemma4 (verified against the checkpoint, see docs/design/generic-model-support.md):
 //   - attention scale=1.0, NO attention softcap (only final-logit softcap 30 fp32)
 //   - QK/V-norm POST-proj / PRE-RoPE; v_norm is RMSNormNoScale; V gets no RoPE
 //   - full-attention layers (5,11,17,23,29) reuse k as v (no v_proj), hd=512 kv=2,
@@ -626,8 +627,8 @@ export class DiffusionGemmaModel {
   /** Prefill from a pre-spliced image+text prompt → populated cache. */
   prefillVision(splicedIds: number[], pixels: MlxArray): Cache[] {
     const cache = this.makeCache();
-    this.encodeVisionPrefill(splicedIds, pixels, cache);
-    return cache;
+    try { this.encodeVisionPrefill(splicedIds, pixels, cache); return cache; }
+    catch (error) { cleanupFailure(error, () => disposeResources(cache)); }
   }
 
   /** One KVCache (full) / RotatingKVCache (sliding) per layer — the encoder
@@ -899,16 +900,15 @@ export class DiffusionGemmaModel {
    *  loop reuses it across all steps of a canvas block). */
   prefill(promptIds: number[]): Cache[] {
     const cache = this.makeCache();
-    this.encodePrompt(promptIds, cache);
-    return cache;
+    try { this.encodePrompt(promptIds, cache); return cache; }
+    catch (error) { cleanupFailure(error, () => disposeResources(cache)); }
   }
 
   /** One denoising-step decoder pass -> softcapped logits [1, canvas, vocab]. */
   decoderLogits(canvasArr: MlxArray, cache: Cache[], scEmbeddings: MlxArray | null): MlxArray {
     const h = this.decodeCanvasArr(canvasArr, cache, null, scEmbeddings);
-    const logits = this.logitsFromHidden(h);
-    h.dispose();
-    return logits;
+    try { return this.logitsFromHidden(h); }
+    finally { h.dispose(); }
   }
 
   /** Dequantized embedding table (once per call), for self-conditioning soft

package/src/model/factory.ts

@@ -24,16 +24,55 @@ import {
 } from "./glm52-memory";
 import { UniversalDenseModel } from "./universal/dense";
 import { genericArgsFor } from "./universal/archs";
+import {
+  ModelImplementationRegistry,
+  type ModelImplementation,
+  type ModelImplementationProvider,
+} from "./implementation";
 import {
   assertResolvedModelProfile,
   resolveModelProfile,
   type ResolvedModelProfile,
+  type ResolveModelProfileOptions,
+  type ModelGraph,
 } from "./profile";
 
 export type RuntimeModel =
   | Gemma4Model | MiniCPM5Model | Qwen35Model | Qwen3Model | Qwen3MoeModel
   | DiffusionGemmaModel | Glm52Model | UniversalDenseModel;
 
+function residentImplementation(
+  id: string,
+  graph: ModelGraph,
+  create: ModelImplementation<Weights, RuntimeModel>["create"],
+): ModelImplementation<Weights, RuntimeModel> {
+  return { id, graph, loader: "safetensors",
+    loop: graph === "diffusion-gemma" ? "diffusion" : "autoregressive", create };
+}
+
+/** Engine-owned registrations. Exact quant profiles can name additional
+ * implementations in a composed registry, without changing sessions or files. */
+export const MLX_MODEL_IMPLEMENTATIONS = new ModelImplementationRegistry<Weights, RuntimeModel>([
+  residentImplementation("diffusion-gemma", "diffusion-gemma", (weights, config) => new DiffusionGemmaModel(weights, config)),
+  residentImplementation("minicpm5", "minicpm5", (weights, config) => new MiniCPM5Model(weights, config)),
+  residentImplementation("qwen3.5", "qwen3.5", (weights, config) => new Qwen35Model(weights, config)),
+  residentImplementation("qwen3-moe", "qwen3-moe", (weights, config) => new Qwen3MoeModel(weights, config)),
+  residentImplementation("qwen3", "qwen3", (weights, config) => new Qwen3Model(weights, config)),
+  residentImplementation("gemma4", "gemma4", (weights, config) => new Gemma4Model(weights, config)),
+  residentImplementation("gemma4-generated", "gemma4", (weights, config) => {
+    const Graph = GENERATED.get(configFingerprint(config));
+    if (!Graph) throw new Error("no generated Gemma graph for this config; refusing to fall back");
+    return new Graph(weights, config);
+  }),
+  residentImplementation("universal-dense", "universal-dense", (weights, config) =>
+    new UniversalDenseModel(weights, config, genericArgsFor(config)!)),
+]);
+
+export interface ModelOpenOptions<Model = RuntimeModel> extends Glm52RuntimeOpenOptions {
+  readonly profiles?: ResolveModelProfileOptions;
+  readonly implementations?: ModelImplementationProvider<Weights, Model>;
+}
+
 export interface Glm52RuntimeOpenOptions {
   /** Whole-process ceiling. Defaults to the smaller of the validated 25 GiB
    * preset and physical RAM. */
@@ -100,47 +139,51 @@ export async function openGlm52RuntimeModel(
  * model.safetensors.index.json, so it must bypass Weights.open and use its
  * dedicated header catalog/tensor source.
  */
-export async function openModel(
+export function openModel<Model>(
   modelDir: string,
-  options: Glm52RuntimeOpenOptions = {},
-): Promise<RuntimeModel> {
+  options: ModelOpenOptions<Model> & { readonly implementations: ModelImplementationProvider<Weights, Model> },
+): Promise<Model>;
+export function openModel(modelDir: string, options?: ModelOpenOptions): Promise<RuntimeModel>;
+export async function openModel<Model>(
+  modelDir: string,
+  options: ModelOpenOptions<Model> = {},
+): Promise<Model | RuntimeModel> {
   const config = await loadModelConfig(modelDir);
-  const profile = resolveModelProfile(config);
-  if (profile.profile.execution.loader === "colibri")
+  const profile = resolveModelProfile(config, options.profiles);
+  if (profile.profile.execution.loader === "colibri") {
+    if (options.implementations || profile.profile.execution.implementation !== undefined)
+      throw new Error("custom Colibri implementations require a streamed loader binding; refusing to fall back");
     return (await openGlm52RuntimeModel(modelDir, options)).model;
-  return createModel(await Weights.open(modelDir), config, profile);
+  }
+  // Resolve before opening weights; missing or incompatible code allocates nothing.
+  const implementation = (options.implementations ?? MLX_MODEL_IMPLEMENTATIONS).select(config, profile);
+  const weights = await Weights.open(modelDir);
+  try { return implementation.create(weights, config, profile); }
+  catch (error) { weights.dispose(); throw error; }
 }
 
-/** Construct the graph named by a previously validated profile. */
+/** Construct the binding named by a validated profile. A supplied registry can
+ * return any backend-owned interface; the compatibility default returns the
+ * legacy model union. The caller retains ownership of the supplied weights. */
+export function createModel<Model>(
+  weights: Weights, config: ModelConfig, resolved: ResolvedModelProfile | undefined,
+  implementations: ModelImplementationProvider<Weights, Model>,
+): Model;
 export function createModel(
+  weights: Weights, config: ModelConfig, resolved?: ResolvedModelProfile,
+  implementations?: ModelImplementationProvider<Weights, RuntimeModel>,
+): RuntimeModel;
+export function createModel<Model>(
   weights: Weights,
   config: ModelConfig,
   resolved: ResolvedModelProfile = resolveModelProfile(config),
-): RuntimeModel {
+  implementations: ModelImplementationProvider<Weights, Model> | ModelImplementationProvider<Weights, RuntimeModel> = MLX_MODEL_IMPLEMENTATIONS,
+): Model | RuntimeModel {
   assertResolvedModelProfile(config, resolved);
-  const execution = resolved.profile.execution;
-  switch (execution.graph) {
-    case "glm5.2":
-      throw new Error(
-        "glm_moe_dsa uses the direct Colibri container; construct it with " +
-        "openModel(modelDir) or Glm52Model.open(modelDir)",
-      );
-    // DiffusionGemma is non-autoregressive. The profile pairs this graph with
-    // the denoising loop; generate() recognizes the constructed model surface.
-    case "diffusion-gemma": return new DiffusionGemmaModel(weights, config);
-    case "minicpm5": return new MiniCPM5Model(weights, config);
-    case "qwen3.5": return new Qwen35Model(weights, config);
-    case "qwen3-moe": return new Qwen3MoeModel(weights, config);
-    case "qwen3": return new Qwen3Model(weights, config);
-    case "gemma4": {
-      // Generated specializations remain exact config matches. Output-changing
-      // request methods stay independently resolved and cannot be fallbacks.
-      const cls = execution.specialization === "generated"
-        ? GENERATED.get(configFingerprint(config))!
-        : Gemma4Model;
-      return new cls(weights, config);
-    }
-    case "universal-dense":
-      return new UniversalDenseModel(weights, config, genericArgsFor(config)!);
-  }
+  if (resolved.profile.execution.loader === "colibri")
+    throw new Error(
+      "glm_moe_dsa uses the direct Colibri container; construct it with " +
+      "openModel(modelDir) or Glm52Model.open(modelDir)",
+    );
+  return implementations.select(config, resolved).create(weights, config, resolved);
 }

package/src/model/fingerprint.ts

@@ -1,4 +1,4 @@
-// Config fingerprint for generated-model dispatch (docs/design/optimization_plan.md
+// Config fingerprint for generated-model dispatch (docs/archive/investigations/optimization_plan.md
 // Phase C): a stable hash over every config field that changes the
 // decode graph structure, plus the kv_config quant layout. A generated
 // specialization is used only when the fingerprint of the loaded config

package/src/model/flash-attention.ts

@@ -545,7 +545,7 @@ export function flashBackward(
     // (dK_j[d] → [kv_row][d]); no post-transpose. (An earlier "fix" transposed
     // it on the false belief the buffer was swapped — that transpose is only the
     // identity at Tkv==D and corrupted dK for Tkv≠D; removed after
-    // scripts/experiments/flash-dkv-debug.ts showed the raw output is correct.)
+    // flash-dkv-debug.ts (deleted 2026-08-23; git history) showed the raw output is correct.)
     return [dq[0]!, dkv[0]!, dkv[1]!];
   } finally {
     Dvec.dispose();
@@ -621,8 +621,8 @@ export function flashAttention(
 //  - This flash kernel is the L2 path: a port of mlx-optiq's
 //    flash_attention_metal (O(L) backward memory). Its parity oracle is OPTIQ,
 //    not finite differences. Confirmed: optiq's flash dK == ops.sdpa to f16
-//    (scripts/flash-optiq-check.py, rel 0.0%), and mlx-bun's flash == ops.sdpa
-//    (tests/flash-attention.test.ts) ⟹ mlx-bun flash == optiq flash (L2 parity).
+//    (scripts/oracle/flash-optiq-check.py, rel 0.0%), and mlx-bun's flash == ops.sdpa
+//    (tests/unit/flash-attention.test.ts) ⟹ mlx-bun flash == optiq flash (L2 parity).
 //
 // Two real port bugs were fixed to reach that parity (the original diverged
 // ~100% from optiq on dK): a spurious dK transpose in flashBackward (corrupted

package/src/model/gemma4-base.ts

@@ -1,5 +1,5 @@
 // gemma4-base — the config-INDEPENDENT machinery shared by every
-// architecture variant (docs/design/optimization_plan.md Phase B: pure code
+// architecture variant (docs/archive/investigations/optimization_plan.md Phase B: pure code
 // movement out of the gemma4.ts monolith; every definition here is
 // verbatim from it). Cache classes, quantized-KV SDPA, masks, the
 // quantized primitives + LoRA machinery, and small graph helpers live
@@ -13,6 +13,7 @@ import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
+import { quantizedMatmulRows } from "../mlx/quantized-rows";
 import { expertOffloadArray } from "../expert-offload";
 import * as tq from "../mlx/turboquant-ops";
 import { runtimeValue } from "../runtime-config";
@@ -155,8 +156,10 @@ export class QuantizedLinear {
     return this._scales.shape[0]!;
   }
 
-  forward(x: MlxArray): MlxArray {
-    let out = ops.quantizedMatmul(
+  forward(x: MlxArray, independentRows = false): MlxArray {
+    let out = independentRows ? quantizedMatmulRows(
+      x, this._w, this._scales, this._biases, this._spec,
+    ) : ops.quantizedMatmul(
       x,
       this._w,
       this._scales,
@@ -263,8 +266,13 @@ export type SharedKv =
 
 export interface Cache {
   offset: number;
-  /** Stable storage identity for compatibility guards and persistence. */
-  signature?(): string;
+  /** Stable storage identity for compatibility guards and persistence.
+   *  REQUIRED: every capability predicate (isPlainKvCache, isRotating*,
+   *  kv-store codecs) is a string compare on this value, so a missing
+   *  override used to fail OPEN into wrong routing (PRs #42/#43:
+   *  BatchedRotatingCache shipped without one and batched joins dropped
+   *  running rows). Wrappers report their own kind plus the inner kind. */
+  signature(): string;
   /** Physical cache storage consumed by one additional token for one row.
    *  Recurrent caches return 0 because their state does not grow with the
    *  sequence. Optional for stateless/training adapters that are never
@@ -303,7 +311,7 @@ export interface Cache {
 }
 
 export function cacheSignature(cache: Cache | undefined): string {
-  return cache?.signature?.() ?? "unknown";
+  return cache ? cache.signature() : "unknown";
 }
 
 export interface RowBatchCache extends Cache {
@@ -1499,7 +1507,7 @@ export class RotatingQuantizedKVCache implements Cache {
 /** One TurboQuant-encoded (K, V) storage tuple — the 5 arrays kv-store.ts
  *  and toQuantized/fromKVCache pass around. Not `ops.QuantizedTensor`
  *  (mlx's affine int4/int8 scheme): this is the rotation + Lloyd-Max
- *  layout (docs/design/turboquant-kv.md), asymmetric-affine for keys,
+ *  layout (docs/design/turboquant.md), asymmetric-affine for keys,
  *  FWHT+Lloyd-Max for values. Field order is the kv-store.ts tensor-slot
  *  contract — do not reorder without updating snapshotCache/loadKvCache. */
 export interface TurboQuantTensor {
@@ -1518,7 +1526,7 @@ export const disposeTurboQuant = (t: TurboQuantTensor): void => {
   t.vScales.dispose();
 };
 
-/** TurboQuant KV cache — v1 (docs/design/turboquant-kv.md): dequantize-
+/** TurboQuant KV cache — v1 (docs/design/turboquant.md): dequantize-
  *  on-fetch. Deliberately does NOT subclass KVCache/RotatingKVCache (or
  *  QuantizedKVCache) — a novel class fails every generated-file
  *  `#matches()` guard and every batching `instanceof` allow-list, so it
@@ -1548,6 +1556,8 @@ export class TurboQuantKVCache implements Cache {
   offset = 0;
   /** Validated lazily on first update, once head_dim is known. */
   #headDim: number | null = null;
+  /** Experimental operation selection, captured once for this cache. */
+  readonly #fusedDecode = process.env.MLX_BUN_TURBOQUANT_FUSED_DECODE === "1";
 
   constructor(readonly kBits: number, readonly vBits: number) {}
 
@@ -1626,6 +1636,13 @@ export class TurboQuantKVCache implements Cache {
       const [B, H, , D] = a.shape as [number, number, number, number];
       return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
     };
+    if (this.#fusedDecode) {
+      const inputs = [cut(t.kIdx), cut(t.kScales), cut(t.kZeros), cut(t.vPacked), cut(t.vScales)] as const;
+      try {
+        const decoded = tq.tryDecodePackedKv(inputs, this.kBits, this.vBits, headDim, deferV);
+        if (decoded) return decoded;
+      } finally { for (const input of inputs) input.dispose(); }
+    }
     const kIdxCut = cut(t.kIdx);
     const kIdxUnpacked = this.kBits < 8 ? tq.unpackBits(kIdxCut, this.kBits, headDim) : kIdxCut;
     if (kIdxUnpacked !== kIdxCut) kIdxCut.dispose();
@@ -2080,7 +2097,7 @@ export function quantizedSdpaTiled(
 function fusedSdpaSupported(q: MlxArray, mask: Mask, groupSize: number, bits: number): boolean {
   // Escape hatch mirroring optiq serve's --no-fused-kv: forces the
   // stock unfused path everywhere. Also the A/B lever for
-  // scripts/bench-fused-prefill.ts. Read per call (cheap next to the
+  // scripts/bench-levers.ts fused-prefill. Read per call (cheap next to the
   // FFI work) so tests and paired A/B harnesses can flip it in-process.
   if (runtimeValue("MLX_BUN_NO_FUSED_SDPA") === "1") return false;
   if (bits !== 4 && bits !== 8) return false;

package/src/model/gemma4.ts

@@ -23,7 +23,7 @@ import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
-import { isExpertTracing, recordRouting } from "../expert-trace";
+import { isExpertTracing, recordRouting } from "../lab/expert-trace/expert-trace";
 
 // The shared, config-independent machinery lives in gemma4-base.ts
 // (Phase B extraction); this file keeps the architecture-specific
@@ -56,7 +56,7 @@ import { Checkpoint } from "../mlx/checkpoint";
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
 import { unrotateValues as tqUnrotateValues } from "../mlx/turboquant-ops";
 import { CompiledFunction } from "../mlx/compile";
-import { flagOn } from "../flags";
+import { flagOn } from "../runtime-config";
 
 /** Verbatim port of mlx_lm/models/gemma4_text.py:
  *    `@partial(mx.compile, shapeless=True) def geglu(gate, x): return nn.gelu_approx(gate) * x`
@@ -242,7 +242,7 @@ class Attention {
   /** A number offset takes the static fast::rope; an array offset (set by
    *  compiled-decode trace adapters) takes the dynamic variant — same
    *  kernel, offset read from the array (bit-exactness asserted in
-   *  tests/compile.test.ts). */
+   *  tests/unit/compile.test.ts). */
   rope(x: MlxArray, offset: number | MlxArray): MlxArray {
     return typeof offset === "number"
       ? ops.rope(x, this.headDim, this.ropeBase, offset, this.ropeFreqs)
@@ -370,7 +370,7 @@ class Attention {
       // verified route: its dQ/dK/dV vjps are exact (finite-difference checked;
       // it's what mlx-lm's tuner differentiates). The two historical bugs were
       // in THIS flash kernel (spurious dK transpose, divergent dQ barrier),
-      // both fixed and gated by tests/flash-attention.test.ts vs the ops.sdpa
+      // both fixed and gated by tests/unit/flash-attention.test.ts vs the ops.sdpa
       // oracle — see flash-attention.ts's routing comment. Flash also remains
       // Gemma-guarded in the trainer (e4b >=2K SIGTRAP not yet re-validated).
       // Full layers → window 0 (pure causal); sliding layers → their window.
@@ -885,7 +885,7 @@ export class Gemma4Model {
    *  which attend bidirectionally among themselves (use_bidirectional_attention:
    *  "vision" — text stays causal; audio prompts pass NO bidir mask, and a
    *  mixed image+audio prompt drops the image overlay too, §3.3 Q1 of
-   *  docs/design/audio-input-plan.md). `ids` ([1, L], the spliced token ids)
+   *  docs/design/generic-model-support.md). `ids` ([1, L], the spliced token ids)
    *  is required for per-layer-input models (e2b/e4b): multimodal soft-token
    *  positions get token 0's per-layer embedding, and their tower content
    *  enters only through the projection term (the merged hidden). Matches
@@ -1015,7 +1015,7 @@ export class Gemma4Model {
     return h;
   }
 
-  // --- Segmented-backward support (docs/design/segmented-backward-training.md
+  // --- Segmented-backward support (docs/design/orpo-training.md
   // §4 Phase B). These are ADDITIVE — forwardLayers is untouched. The segmented
   // driver builds masks + per-layer inputs ONCE, then drives runLayerRange per
   // segment, threading the KV-shared donor K/V across segment boundaries.

package/src/model/generated/gemma4-12b.ts

@@ -4,8 +4,8 @@
 // fingerprint: 9f812d2eb461fcbe
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 

package/src/model/generated/gemma4-26b.ts

@@ -4,8 +4,8 @@
 // fingerprint: c9dd67ed5a525231
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 

package/src/model/generated/gemma4-e4b.ts

@@ -4,8 +4,8 @@
 // fingerprint: 418e9adc386ea67c
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 

package/src/model/glm52-dsa.ts

@@ -291,7 +291,7 @@ function getDsaRankKeyKernel(): MetalKernel {
           ? ~bits
           : (bits ^ 0x80000000u);
         const uint descending = ~ascending;
-        keys[i] = (ulong(descending) << 32) | ulong(i);
+        keys[i] = descending;
       `,
     });
   }
@@ -308,7 +308,8 @@ function getDsaContractOrderKeyKernel(): MetalKernel {
         const uint i = thread_position_in_grid.x;
         const uint position = selected_positions[i];
         const uint threshold_class = selected_scores[i] > threshold ? 0u : 1u;
-        keys[i] = (ulong(threshold_class) << 32) | ulong(position);
+        // MLX axis lengths are signed int32; positions leave the high bit free.
+        keys[i] = (threshold_class << 31) | position;
       `,
     });
   }
@@ -318,8 +319,8 @@ function getDsaContractOrderKeyKernel(): MetalKernel {
 /**
  * Deterministic on-device equivalent of `selectDsaThresholdTiesF32`.
  *
- * The first uint64 key sorts by score descending and then position ascending,
- * making the top-k set deterministic even at the threshold. The second key
+ * Stable sorting of the first uint32 key orders scores descending and keeps
+ * equal-score positions ascending, including threshold ties. The second key
  * restores Colibri's observable two-scan order: all scores strictly above the
  * threshold in position order, followed by threshold ties in position order.
  * The caller owns the returned arrays and must call `dispose()` once the last
@@ -340,18 +341,21 @@ export function selectGlm52DsaDevice(
     ? scores
     : scores.astype(Dtype.float32);
   const [rankKeys] = getDsaRankKeyKernel().apply([scoresF32], {
-    outputs: [{ shape: [contextLength], dtype: Dtype.uint64 }],
+    outputs: [{ shape: [contextLength], dtype: Dtype.uint32 }],
     grid: [contextLength, 1, 1],
     threadGroup: [Math.min(256, contextLength), 1, 1],
   });
-  const partition = ops.argpartitionAxis(rankKeys!, topK - 1, 0);
-  const selectedUnordered = partition.slice([0], [topK]);
+  // Stable argsort retains ascending positions for equal score bits, so the
+  // ranking key needs no position field. MLX 0.31.2's GPU argpartition also
+  // sorts fully, but does not promise this tie order in its public contract.
+  const rankedPositions = ops.argsortAxis(rankKeys!, 0);
+  const selectedUnordered = rankedPositions.slice([0], [topK]);
   const selectedScores = ops.takeAlongAxis(scoresF32, selectedUnordered, 0);
   const threshold = ops.minAxis(selectedScores, 0, false);
   const [orderKeys] = getDsaContractOrderKeyKernel().apply(
     [selectedScores, selectedUnordered, threshold],
     {
-      outputs: [{ shape: [topK], dtype: Dtype.uint64 }],
+      outputs: [{ shape: [topK], dtype: Dtype.uint32 }],
       grid: [topK, 1, 1],
       threadGroup: [Math.min(256, topK), 1, 1],
     },
@@ -360,7 +364,7 @@ export function selectGlm52DsaDevice(
   const positions = ops.takeAlongAxis(selectedUnordered, order, 0);
 
   rankKeys!.dispose();
-  partition.dispose();
+  rankedPositions.dispose();
   selectedUnordered.dispose();
   selectedScores.dispose();
   orderKeys!.dispose();

package/src/model/glm52-memory.ts

@@ -4,7 +4,7 @@ import {
   type Glm52Config,
 } from "./glm52-config";
 import { ColibriGlm52Container } from "./glm52-container";
-import type { MemoryPlan } from "../memory-plan";
+import type { MemoryPlan } from "../admission-plan";
 
 export const GLM52_G5_GIB = 1024 ** 3;
 export const GLM52_G5_DEFAULT_MACHINE_BYTES = 32 * GLM52_G5_GIB;