mlx-bun 0.0.13 → 0.3.0

package/docs/design/turboquant-kv.md

@@ -1,235 +0,0 @@
-# TurboQuant KV — design (Phase 13)
-
-Status: v1 wired end-to-end (2026-07-06): codec + `TurboQuantKVCache` +
-`--kv-quant turbo[:k<bits>v<bits>]` (cli.ts, serve/generate/bench) +
-`maybeQuantizeKv` conversion + `GenerationGateway` solo-only refusal +
-`kv-store.ts` persistence + `scripts/eval-turboquant-curve.ts` (KL/ppl curve
-script; not yet RUN against a real model — that's the next stage's gate).
-Owner seam: L2 cache format, L1-unmodified attention (dequantize-on-fetch),
-eval-gated. Opt-in only — naked default stays L1.
-
-## What it is
-
-Rotation-based KV-cache quantization: signed int8 affine keys + FWHT-rotated
-3-bit Lloyd-Max values, per-32-group fp16 scales. More usable context at the
-same memory (2.56× KV compression at the k8/v3 default, head_dim 128), NOT a
-weight quantizer. Composes with (is orthogonal to) the per-layer mixed-KV
-allocation axis — TurboQuant is the QUANTIZER axis (frontier plan §7.2).
-
-## Oracle chain (decided 2026-07-06 from the research fan-out)
-
-Two unrelated implementations share the name; do not conflate:
-
-- **The paper** (arXiv:2504.19874, Google): dense QR random rotation, no
-  groups, no K/V asymmetry, optional QJL residual stage. It anchors the WHY
-  (rotate→Gaussian marginal→Lloyd-Max optimal; 4^-b distortion law) but not
-  the layout.
-- **vllm-metal** (`vllm-project/vllm-metal`, `vllm_metal/attention/caches/
-  turboquant.py` + `kernels_v2/turboquant.metal`): the engineering variant.
-  optiq's `runtime/mtp/turboquant.py` config layer is lifted verbatim from it
-  (byte-identical `CENTROIDS_3BIT`, matching op signatures in optiq
-  `cache_state.py`). **This is our oracle** — "the oracle is whoever ships
-  it." Its Python reference functions are pure MLX; we vendor them and
-  generate bit-exact goldens with the oracle venv.
-- mlx-vlm's 6.3k-line module is a THIRD design (whole-vector norms, MSE/Prod
-  codecs, keys rotated too) — noted, not our target. Its numeric test oracles
-  (MSE≈{0.36, 0.117, 0.03} at bits {1,2,3} on unit vectors) are still useful
-  sanity anchors for the rotation+Lloyd-Max math.
-- vllm-metal's rendered docs page contradicts its own source (claims keys are
-  rotated; the code never rotates keys). Trust the source.
-
-## Algorithm (vllm-metal reference semantics, exact)
-
-Per token, per kv-head, groups of 32 along head_dim (head_dim ∈ {64,128,256,
-512}, must divide by 32):
-
-**Keys — NO rotation, asymmetric affine, per-32-group:**
-- signed (q8_0/int8, bits=8): `max_val=127; scale=(max−min)/(2·max_val);
-  zero=round((max+min)/(2·scale)); idx=clip(round(x/scale − zero), −127, 127)`
-  stored as int8, one byte per element (no packing at 8 bits).
-- unsigned (q5_0/q4_0/int2…): `max_val=2^bits−1; scale=(max−min)/max_val;
-  zero=round(min/scale); idx=clip(round(x/scale − zero), 0, max_val)`, packed.
-- Dequant (both): `x̂=(idx+zero)·scale`. scale/zero fp16, one pair per group.
-- Reference routes intermediates through fp16 (and deliberately drops the
-  +1e-8 scale epsilon because it underflows in fp16) — match the Python
-  reference's ACTUAL behavior, verified by goldens, not its comments.
-
-**Values — FWHT + Lloyd-Max, per-32-group RMS scale:**
-- `y = FWHT(x · signs)` where `signs ∈ {±1}^d` from `mx.random.randint(0,2,
-  (d,), key=mx.random.key(42))`, and FWHT normalizes by 1/√d.
-- `scale = sqrt(mean(y²))` per 32-group (RMS, not min/max); `y_n = y/(scale+1e-8)`.
-- `idx = searchsorted(BOUNDARIES, y_n)` — nearest centroid. 3-bit table:
-  centroids ±{0.24509, 0.75601, 1.34391, 2.15195}, boundaries ±{0.50055,
-  1.04996, 1.74793}, 0. Other bit-widths (2/4/5/8): iterative Lloyd-Max on a
-  unit normal — we generate those tables ONCE from the vendored reference and
-  hardcode them (parity by construction; no runtime iteration).
-- Pack: flat little-endian bitstream, element i at bits [i·b, i·b+b), values
-  may straddle two bytes (3-bit: 8 values → 3 bytes).
-- Dequant: `x̂ = InvFWHT(centroids[idx] · scale · signs⁻¹)` (inverse = apply
-  FWHT then multiply signs; Hadamard is an involution up to normalization).
-
-**Per-token-per-head bytes (k8v3, d=128):** 128 (K int8) + 8+8 (K scale/zero)
-+ 48 (V packed) + 8 (V scale) = 200 vs 512 bf16 → 2.56×. The centroid table is
-a shared constant, not per-token. This reproduces optiq's `compression_ratio`.
-
-**Sign vectors are hardcoded constants per head_dim** (64/128/256/512),
-generated from the reference's key(42) — exactly what vllm-metal does in its
-.metal source. Deterministic across machines and spill/restore.
-
-## v1 architecture in mlx-bun
-
-Dequantize-on-fetch: `TurboQuantKVCache implements Cache` (does NOT subclass
-KVCache — that keeps it auto-excluded from batching), `updateAndFetch(k,v)`
-quantizes the new tokens into packed storage and returns the dequantized bf16
-active window, so the standard `ops.sdpa` path runs unmodified — the same
-pattern optiq's `runtime/kv/rotating.py` uses for its v1. No generated-file
-changes: a novel cache class fails every `#matches()` guard and falls back to
-the monolith by design. Quantization groups run along head_dim only —
-token-axis slicing stays byte-safe (standing invariant).
-
-Files (from the codebase-map seams):
-- `src/mlx/turboquant-ops.ts` — NEW: fwht (bind `mlx_hadamard_transform`,
-  confirmed exported by libmlxc), encode/decode keys + values, pack/unpack.
-- `src/mlx/turboquant-tables.ts` — NEW, GENERATED by regen script: sign
-  vectors per head_dim + Lloyd-Max centroid/boundary tables per bit-width.
-- `src/model/gemma4-base.ts` — NEW `TurboQuantKVCache` (full-attention only in
-  v1; rotating/sliding-window layers stay bf16 — documented limitation).
-- `src/generate.ts` `maybeQuantizeKv()` — the single conversion chokepoint
-  (serve serial lane + generate + eval all flow through it); keep the
-  `offset===0` skip and `quantizedKvStart` semantics.
-- `src/cli.ts` — `--kv-quant turbo[:k<bits>v<bits>]` (default k8v3).
-- `src/serve/generation-gateway.ts` — explicit solo-only refusal (belt +
-  braces on top of the automatic instanceof exclusion).
-- `src/kv-store.ts` — new `CacheKind "turboquant"`, branches in
-  snapshotCache/loadKvCache/cloneKvCaches; header carries kBits/vBits.
-- `src/model/fingerprint.ts` — NOT touched, by design symmetry: uniform
-  `--kv-quant 4|8` is a CLI-only runtime lever that never feeds
-  `configFingerprint()` either (only the per-layer `kv_config.json` scheme
-  does); TurboQuant is the same class of lever, so it gets the same
-  non-treatment. The SSD-cache compat key (`server.ts`'s `schemeKey`) DOES
-  fold in the TurboQuant scheme directly (`turbo-k<bits>v<bits>`), same as it
-  already does for uniform bits — that's the actual anti-collision mechanism
-  for this class of scheme, not the fingerprint.
-- Docs same-commit: server-config.md, cli.md, features-matrix.md.
-- Tests: `tests/turboquant-ops.test.ts` (golden bit-exactness vs vendored
-  reference + model-free math props), `tests/turboquant-cache.test.ts` (cache
-  invariants, growth/trim, persistence roundtrip — synthetic arrays, no
-  model/golden dependency), `scripts/regen-turboquant-goldens.ts`.
-- `scripts/eval-turboquant-curve.ts` — the Gate 3 quality-vs-bpw curve script:
-  bf16 baseline + {k8v8,k8v4,k8v3,k8v2,k4v3,k4v2} sequentially, reusing
-  `src/eval/kl.ts`'s `evaluateKlKvArm` (teacher-forced serving-decode KL,
-  real `TurboQuantKVCache.fromKVCache` conversion) and
-  `src/eval/perplexity.ts`'s `evalPpl` (bf16-forward anchor column).
-  Typechecks + `--help`/arg-parsing verified; not yet RUN against a real
-  model — that run (and reading the resulting curve against the paper's law)
-  is the next stage's job.
-- Vendored reference: `lab/repro/vllm-metal-turboquant/` (their Python file +
-  provenance note), driven only by the regen script via the oracle venv.
-
-## Gates — ALL PASSED 2026-07-06 (M1 Max 32 GB)
-
-Gate 3 measured curve (MiniCPM5-1B, `scripts/eval-turboquant-curve.ts`,
-teacher-forced serving-decode KL vs bf16, 8 prompts × 128 tokens, 32 decode
-steps; affine baselines via the same `evaluateKlKvArm` harness):
-
-| scheme | effective KV bits | mean KL vs bf16 |
-|---|---|---|
-| bf16 | 16.00 | 0 (control) |
-| uniform kv8 (affine g64) | 8.50 | 0.00246 |
-| turbo k8v8 | 8.75 | 0.00214 |
-| turbo k8v4 | 6.75 | 0.00936 |
-| **turbo k8v3 (default)** | **6.25** | **0.0325** |
-| turbo k8v2 | 5.75 | 0.175 |
-| uniform kv4 (affine g64) | 4.50 | 0.0516 |
-| turbo k4v3 | 4.25 | 0.0622 |
-| turbo k4v2 | 3.75 | 0.205 |
-
-Shape reproduces the paper's law: value bits dominate, ~4^-b decay, cliff at
-v2. Headline: turbo adds operating points between affine kv8/kv4 — the k8v3
-default (2.56× KV compression) beats uniform kv4's KL with room to spare.
-Live smokes: cpm5 + e4b generation coherent under turbo:k8v3 (e4b emits the
-one-time rotating-layers-stay-bf16 warning). The paper's own LongBench
-numbers were NOT rerun (different benchmark harness) — the reproduction is
-shape-level plus the codec-level D_mse≈0.03 anchor, which our 3-bit
-roundtrip test hits.
-
-1. **Quantizer parity (bit-exact, the hard gate):** our encode indices,
-   scales, zeros, packed bytes, and dequant outputs match the vendored
-   vllm-metal Python reference on golden vectors (head_dim 64/128; configs
-   k8v3/k4v3/k4v2/k8v4/k2v5/k5v8 — spanning every kBits/vBits value the CLI
-   accepts, per the 2026-07-06 integration review). This is an oracle-backed
-   scheme at the codec level.
-2. **Math sanity (model-free):** FWHT involution/orthogonality; roundtrip MSE
-   at 3-bit on unit-normal data ≈ paper's D_mse≈0.03 band.
-3. **Quality-vs-bpw curve (Phase 13 exit criterion):** teacher-forced KL
-   (src/eval/kl.ts) + perplexity vs bf16 KV on MiniCPM5-1B (head_dim 128,
-   pure full-attention) across {k8v8, k8v4, k8v3, k8v2, k4v3, k4v2}; curve
-   shape must match the paper's law (≈quality-neutral ≥3.5 effective bits,
-   marginal degradation at 2.5, cliff below). e4b spot check second. Model
-   runs SEQUENTIAL (32 GB box rule).
-4. Whole-repo `tsc --noEmit` = 0; fast suite green; no batching regression
-   (turbo requests route serial with a clear reason).
-
-## Deferred inverse FWHT (LANDED 2026-07-06, post-v1, Josh-directed)
-
-Attention is linear in V, so `InvFWHT(Σᵢ wᵢ·v̂ᵢ) = Σᵢ wᵢ·InvFWHT(v̂ᵢ)`:
-`updateAndFetchDeferredV` returns the V window still in the rotated domain
-and the attention site un-rotates the OUTPUT once per query row
-(`tq.unrotateValues`) — O(q·d log d) per step instead of O(T·d log d).
-Opt-in per attention site via `SharedKv.vRotated` (monolith gemma4.ts —
-KV-shared consumer layers inherit the flag through sharedIn — and
-minicpm5.ts); every other consumer keeps calling `updateAndFetch` (eager)
-and stays correct. Not bit-identical to eager (bf16 rounding in the rotated
-domain before the transform); measured k8v3 KL 0.0338 vs eager 0.0325 —
-same quality. Paired fetch-path A/B on a LOADED M1 Max: never slower,
-~10-20% off fetch cost at 4k ctx, noise-level at 8k (dequant gather
-dominates both paths) — no speed claim until a quiet-machine benchmark.sh
-run; the win is removing the O(T) transform, which matters more once the
-gather itself is fused.
-
-## Future: rotation-based WEIGHT quantization (noted 2026-07-06, Josh)
-
-The same rotate-then-quantize mechanism applies to weights, and we already
-ship a quantization tool (`convert`/`fuse` verbs) to hang it on — as a
-further optimization of quantization itself, the quantizer leg of the §7.1
-mixed-precision program (allocation × quantizer, same orthogonality as KV):
-
-- Rotation smears outlier channels so a low-bit codebook covers the signal
-  mass — the QuaRot/SpinQuant family (both cited by the TurboQuant paper).
-- Weights beat KV on one axis: the rotation FOLDS into adjacent weight
-  matrices offline (R into one layer, Rᵀ into the next) → zero runtime
-  cost, no online transform at all.
-- Weights lose on the other: they're static, so calibration-aware methods
-  (GPTQ/AWQ/imatrix, optiq sensitivity maps) are admissible and set a
-  higher baseline than any calibration-free scheme — the win to chase is
-  rotation COMPOSED WITH calibration/allocation, not instead of it.
-- A Lloyd-Max weight FORMAT needs custom quantized-matmul kernels — the
-  26B gather-qmv shelving (dispatch fixed-cost ate the prize) is the
-  cautionary precedent. Rotation-folding into mlx's EXISTING affine format
-  needs no new kernels and is the cheap first experiment.
-- Activations (W4A4/QuaRot act-quant): more involved (online rotation,
-  fused kernels) and mispriced for this hardware — no int4 tensor cores,
-  decode is weight-bandwidth-bound. See w4a16-compute-precision-spike.md.
-  Not queued.
-
-Gate when picked up: perplexity + frozen 6-task eval at equal bpw vs the
-plain affine convert output (the §7.1 gate, eval DB rows).
-
-## Non-goals for v1 (recorded so they don't creep)
-
-- Fused quantized-SDPA Metal kernel (the remaining fetch cost is the
-  unpack+gather; deferred InvFWHT — see above — already landed).
-- Rotating/sliding-window TurboQuant cache; batched TurboQuant; QJL residual
-  stage (Q_prod); entropy coding (paper explicitly declined it too).
-- TurboQuant on the speculative lane: `--kv-quant turbo` + `--draft-model`
-  keeps the KV scheme — those requests decode serially WITHOUT speculation
-  (the spec loop is bf16-KV-only; excluded in the server's spec-eligibility
-  gate with a startup warning, 2026-07-07 review fix — turbo was previously
-  routed into the spec lane and silently dropped).
-- Head dims beyond {64,128,256,512}: REFUSED AT SERVER START (2026-07-07
-  review fix — validation was lazy on the first cache append, so an
-  unsupported model 500'd every request mid-prefill; createServer now
-  checks the config's full-attention head dim against
-  TURBOQUANT_HEAD_DIMS, and the cache-level check stays as the backstop).
-- Speed claims of any kind — v1 dequant-on-fetch is expected slower per-step
-  at long context; this ships as a memory/context feature like uniform KV.

package/docs/design/w4a16-compute-precision-spike.md

@@ -1,232 +0,0 @@
-# Research spike — W4A16 head GEMM (compute precision, not storage)
-
-**Status:** proposed (not started) · **Opened:** 2026-06-19 · **Owner:** Josh
-**Surface:** `src/train/flash-cce.ts` (steel + SG kernels) · **Gate:** the existing
-1e-5 CCE parity gate + greedy-prefix logit parity.
-
-> One-line thesis (Josh): *we've maxed **storage** quant (4-bit weights), but we
-> dequant into fp32 and run the matmul in fp32 — so the head GEMM is **W4A32**.
-> The untouched lever is **compute** precision: dequant 4-bit → half, matmul in
-> `half` operands with an fp32 accumulator (**W4A16**), gate on parity. Same 4-bit
-> storage, untouched. The only change is the float width we dequant into and
-> multiply in.*
-
----
-
-## TL;DR
-
-- The flash-CCE head (forward logit GEMM `E@Wᵀ` and backward phase-2 `coeffᵀ@W`)
-  currently runs **W4A32**: 4-bit weights are dequantized into **fp32** threadgroup
-  tiles and multiplied with **fp32** simdgroup matrices. Verified in source (below).
-- The 4-bit packing bought the cheap HBM read; the fp32 matmul throws away the
-  compute advantage. Apple GPUs run **fp16 FMA at ~2× fp32** on the regular SIMD
-  ALUs (no tensor cores). Both head GEMMs are **compute-bound** (see roofline note),
-  so half operands are the lever.
-- **Industry-standard for the forward** (MLX `quantized_matmul`, llama.cpp, every
-  on-device 4-bit stack do exactly W4A16). Low risk; do it first.
-- **Backward needs the mixed-precision caveat** — fp32 accumulate is non-negotiable
-  (already planned), and the **operands** (`coeffᵀ`, `h`) need a dynamic-range probe
-  before trusting fp16, specifically the near-zero coeff tails after the skips fire.
-- **M1 Max constraint:** no `bf16` simdgroup_matrix (M2+ only), so **fp16 is the
-  only half option** — which makes the range-probe matter more here than on M3.
-- Gate exactly like the existing skip epsilons: **fp32 stays the default-exact
-  parity mode**; half is the training default; the 1e-5 gate test runs against fp32.
-
----
-
-## Current state — verified W4A32
-
-The default production path for **both** forward and backward is the **steel**
-kernel (MLX's verbatim `BlockMMA` GEMM, `steel-qmm-header.ts`), selected when
-`H,V,blockV0 % 32 == 0` (`flash-cce.ts:1034`, `:1130`). The `simdgroup_float8x8`
-SG kernels and the scalar/lane kernels are fallbacks for non-32-divisible shapes or
-explicit `MLX_BUN_CCE_*` overrides.
-
-Every one of these instantiations is fp32-operand:
-
-| path | site | instantiation | operand width |
-|---|---|---|---|
-| steel fwd (default) | `flash-cce.ts:950` | `BlockMMA<float,float,…>` | fp32 |
-| steel bwd phase-1 (default) | `flash-cce.ts:799` | `BlockMMA<float,float,…>` (lmma_t) | fp32 |
-| steel bwd phase-2 (default) | `flash-cce.ts:800` | `BlockMMA<float,float,…>` (dmma_t) | fp32 |
-| steel staging tiles | `:596–598`, loaders `:801` | `threadgroup float` / `BlockLoader<float,…>` | fp32 |
-| SG fwd (fallback) | `:251–254` | `simdgroup_float8x8`, `Csh`/`Esh` = `threadgroup float` | fp32 |
-| SG bwd (fallback) | `:617`, `:684–704` | `simdgroup_float8x8`, `Wsh` = `threadgroup float` | fp32 |
-
-The dequant itself (`Csh[…] = sc * q + bi`, `flash-cce.ts:244`) computes into fp32
-and stores fp32. The dh accumulator is already fp32 (`MMATile<float,…>`,
-`flash-cce.ts:828`) — i.e. **the accumulator is already where it should be**; only
-the matmul *operands* are wider than they need to be.
-
----
-
-## The lever — what actually changes
-
-W4A16 = **same 4-bit storage, half-width operands, fp32 accumulate.** Concretely:
-
-1. **Steel operand type** — the `BlockMMA<float,float,…>` first template param is the
-   GEMM input/operand type; flip operand to `half` while keeping the
-   accumulator/fragment in `float`. (**Spike task 0: confirm the exact MLX steel
-   template param that selects operand vs accumulator width** in
-   `steel-qmm-header.ts` before editing — don't assume the param order.)
-2. **Staging tiles** — `Esh`, `Cstg`, `Wsh`, `Csh` go `threadgroup half`; dequant
-   arithmetic can stay fp32 then narrow on store, or compute in half.
-3. **Loaders** — `BlockLoader<float,…>` → `BlockLoader<half,…>` to stage half tiles.
-4. **SG fallback** — `simdgroup_float8x8` → `simdgroup_half8x8` for the operands,
-   keep an `simdgroup_float8x8` accumulator (Apple supports the mixed `half×half→float`
-   `simdgroup_multiply_accumulate`).
-
-Everything reduction-shaped (the online-softmax `m`/`d` recurrence, the dh accum,
-logsumexp) **stays fp32**.
-
----
-
-## First thoughts (mine) — what to add to Josh's analysis
-
-The analysis is right and the change is well-targeted. Five additions, each grounded
-in the current source:
-
-**1. Aim at the steel kernel, not the SG snippet.** Josh's note led with the SG path
-(`simdgroup_float8x8`, `Csh`). That path is a *fallback*. The lever that moves the
-measured 180 ms fwd / 754 ms bwd (e4b) numbers is the **steel `BlockMMA<float,float>`**
-instantiations (`:950`, `:799`, `:800`) and their `BlockLoader<float>`/staging tiles.
-Do the SG path too for completeness, but the steel path is where the production time is.
-
-**2. The 2× is an upper bound gated by the FMA fraction — get the roofline split
-first.** The whole flash-CCE design is `[M,V]`-free and amortizes dequant across
-`BLOCK_B` tokens specifically because the **dequant ALU** was "the dominant cost in
-v1" (`flash-cce.ts:59`). The runtime splits into (a) nibble-unpack + `sc·q+bi`
-dequant and (b) the FMA matmul. **W4A16 only speeds up (b).** The kernel is
-unambiguously compute-bound, not BW-bound — **confirmed live** (see the measured box
-below), and Josh flagged it from `mactop` before this doc existed. *But* the win
-ceiling is the FMA share of that ALU time. The dequant-amortization work was done
-precisely to shrink (a)'s share, which *should* leave (b) dominant and favorable — but
-**measure the (a)/(b) split before predicting 2×.** The live numbers prove the
-*whole kernel* is ALU-bound; they do **not** resolve the intra-kernel (a)/(b) split,
-which is what actually bounds the reclaim. If the kernel is still dequant-ALU-bound,
-half operands move it less.
-
-> **Live roofline — measured (M1 Max 32 GB, MiniCPM5-1B ORPO, `--no-segment`, seq 4096):**
-> GPU **100% @ 1296 MHz, 81 °C, Nominal** (sustained, not throttling) · memory
-> **BW 129.2 GB/s ≈ ⅓ of the M1 Max's ~400 GB/s ceiling** (bus two-thirds idle) ·
-> mem **17.37 / 32 GB (54%)**, swap 0 · CPU ~5%, ANE 0%. The training `bun` proc is
-> 92.6% GPU. **GPU pegged + memory bus mostly idle = ALU-bound.** This is the regime
-> where W4A16 pays; if it were BW-bound (BW near 400, GPU spiky) the spike would be
-> moot. The 4-bit storage already did its job — *that's why BW is low* — and the GPU
-> is now spending its cycles on the fp32 FMAs.
-
-**3. Threadgroup memory is a free secondary win — possibly a compounding one.** The
-staging tiles (`Esh`/`Cstg`/`Wsh`/`Csh`) are the kernel's threadgroup-memory pressure,
-and tile sizing is *already* occupancy-constrained — see the note at `flash-cce.ts:71`
-("the 8KB `Csh` tile measured SLOWER"). Half-width staging **halves that footprint**,
-which can either raise occupancy or admit larger tiles. So W4A16 may buy more than the
-raw 2× FMA: it could relieve the exact constraint that capped the current tiling.
-Worth re-sweeping `BLOCK_B`/tile dims *after* the precision change, not just measuring
-the drop-in.
-
-**4. SDPA in the segmented forward is probably already half — don't chase it.** Josh's
-aside ("if your SDPA tiles are also fp32, same logic"). The segmented forward uses
-`ops.sdpa` — MLX's own fused flash kernel ([[opssdpa-dk-vjp-bug]]) — which runs in the
-model's dtype (bf16), **not** our fp32 tiles. So that lever is likely already pulled by
-delegating to MLX. One-line confirm, then drop it; the head GEMM is the real target.
-
-**5. Free instrumentation hooks already exist for the backward range-probe.** The
-backward already computes `cmax = simd_max(coeff)` for the FILTER_EPS skip
-(`flash-cce.ts:682`). That same reduction is the natural hook to log the coeff dynamic
-range (min nonzero |coeff|) across a real run *before* committing the operand to fp16 —
-the underflow-to-zero case Josh flagged. Pair it with a max-abs probe on `h`: at e4b's
-`scale=1.0` peaked attention the residual stream can carry outliers
-([[e4b-forward-length-sensitivity]], [[siglip-vision-parity-cross-build]]), and with no
-bf16 on M1 Max we can't lean on a wider exponent to eat them.
-
----
-
-## Risk analysis (preserving the forward/backward split)
-
-**Forward — low risk, boring industry default.** Logits are well-scaled; the
-online-softmax max-subtraction bounds the `exp` args. W4A16 here is what the dense
-path already assumes. Validate against the 1e-5 gate (fp16 operands won't hold 1e-5 —
-that's expected; the **fp32 path stays the exact-parity gate**, half is graded by
-greedy-prefix logit agreement, same as the vision tier-a bar).
-
-**Backward — fp16 matmul is mainstream, but as *mixed precision*, with eyes on
-operand range.** What makes the forward safe (well-scaled, max-subtracted) is **not**
-automatically true here:
-- `coeff_v = g·(onehot − softmax)·sech²` — after the skips fire, the survivors are the
-  **near-zero tails**. Small magnitudes are exactly where fp16's ~3-digit mantissa and
-  narrow range (subnormals < ~6e-5) bite. **`coeffᵀ` as an fp16 operand is the highest
-  risk** — tail coeffs underflowing to zero silently drop gradient contribution. Probe
-  the range first (task 5 above).
-- `h` (activations) as fp16 operand — residual outliers; the classic bf16-beats-fp16
-  case, but bf16 isn't available on M1 Max. Range-check before trusting.
-- **fp32 accumulate is non-negotiable** (already planned; `MMATile<float>` already is).
-- **Validate grads against autograd, not a hand-derived backward** ([[opssdpa-dk-vjp-bug]]).
-
----
-
-## Spike plan
-
-Sequenced so each step is independently measurable and revertible (Josh: *"Don't flip
-both at once — you'll learn more, and if parity drifts you'll know which GEMM did it."*).
-
-- [ ] **0. Confirm the steel template semantics.** In `steel-qmm-header.ts`, identify
-  which `BlockMMA`/`BlockLoader` template params are operand vs accumulator width.
-  Exit: a one-line mapping, no code change.
-- [x] **1a. Whole-kernel roofline — DONE (live).** GPU 100% / BW 129 GB/s (~⅓ of 400)
-  on M1 Max confirms the kernel is ALU-bound, not BW-bound (see the measured box above).
-  The spike is worth chasing.
-- [ ] **1b. Intra-kernel split (still open).** Profile the steel fwd (180 ms) and bwd
-  (754 ms, e4b) into dequant-ALU vs FMA fractions — 1a says the GPU is the wall, this
-  says *which part* of the GPU work. Exit: a number that bounds the expected win. *If
-  FMA < ~40% of either kernel, reset expectations before coding.*
-- [ ] **2. Add the precision gate.** New codegen `#define` (e.g. `CCE_A_HALF`) +
-  env flag (e.g. `MLX_BUN_CCE_A16`), mirroring `CCE_BWD_FILTER_EPS`/`CCE_BWD_BLOCK_EPS`
-  (`flash-cce.ts:732`,`:739`). Default: fp32 (exact) for the parity suite; half is the
-  training default. Per [[three-level-fidelity-tree-model]] the exact path must stay
-  selectable.
-- [ ] **3. Forward → half operands.** Steel fwd `BlockMMA`/loader/staging to half.
-  Validate: 1e-5 gate (fp32 path) + greedy-prefix logit parity (half path) + measure
-  the fwd GEMM delta. Land as its own change.
-- [ ] **4. Backward range probe.** Instrument coeff (via the `simd_max` hook) and `h`
-  dynamic range across a real MiniCPM5/e4b ORPO step. Exit: min-nonzero-coeff and
-  max-|h| numbers + a go/no-go on fp16 operands.
-- [ ] **5. Backward → half operands** (only if 4 says go). `coeffᵀ`/`h`/`W` operands
-  half, dh accumulator stays fp32. Validate grads **vs autograd** + measure bwd delta.
-- [ ] **6. Re-sweep tiling.** With half-width staging, re-check `BLOCK_B`/tile dims and
-  occupancy (first-thought #3) — the precision win may unlock a tiling win.
-
-**Exit criteria for the spike:** a measured fwd/bwd speedup on e4b (V=262144) and
-MiniCPM5 (V=130560) on the bench machine (M4 Pro 24 GB; numbers verified against
-`~/.cache/mlx-bun/evals.sqlite`, not docs — [[verify-benchmarks-against-evaldb-not-docs]]),
-with the fp32 exact path still passing the 1e-5 gate and the half training path inside
-the greedy-prefix parity bar. Either it lands as the training default behind the gate,
-or the roofline/range data explains why it didn't pay and the spike closes with that.
-
----
-
-## Open questions
-
-- Steel `BlockMMA` mixed operand/accum: does MLX's steel already support
-  `half`-operand / `float`-accum, or does it need a local template tweak? (task 0)
-- Is the dequant cheaper to do **in half** (fewer ALU bits) or **in fp32 then narrow**
-  (preserves the scale/bias precision)? A/B inside task 3.
-- Does the half staging change the `+4-float` padding trick (`flash-cce.ts:197`, the
-  BK_padded bank-conflict avoidance)? Half-width may re-introduce or shift bank
-  conflicts — re-derive the pad.
-
----
-
-## References
-
-**Code:** `src/train/flash-cce.ts` — steel fwd `:937–1028`, steel bwd `:763–935`, SG
-fwd `:174–270`, SG bwd `:561–726`, gates `:732`/`:739`/`:1034`/`:1130`; staging tiles
-`:596–598`; dequant `:244`; coeff `simd_max` hook `:682`; dh fp32 accum `:828`.
-`src/train/steel-qmm-header.ts` (the verbatim MLX BlockMMA). `src/mlx/metal-kernel.ts`
-(kernel infra).
-
-**Cross-links:** [[flash-cce-metal-kernel-prototype]] · [[three-level-fidelity-tree-model]]
-· [[opssdpa-dk-vjp-bug]] · [[e4b-forward-length-sensitivity]] ·
-[[verify-benchmarks-against-evaldb-not-docs]] · `docs/investigations/steel-flash-cce-handoff.md`.
-
-**Sibling deferred lever** (queued together, post-run): `docs/design/orpo-future-enhancements.md`
-→ "Deferred levers" (expose `--grad-accum`).