mlx-bun 0.0.5 → 0.0.6

package/docs/design/orpo-dynamic-lambda.md

@@ -0,0 +1,179 @@
+# Dynamic λ controller for ORPO (adaptive preference pressure)
+
+**Status:** proposed (design only — not built). **Opened:** 2026-06-20.
+**Surface:** `src/train/loss.ts` (orpoMetrics/orpoLoss), `src/train/trainer.ts`
+(orpoLoop, the λ config), the segmented ORPO classes in `src/train/segmented.ts`.
+
+## Motivation
+
+ORPO's loss is `L = NLL(chosen) + λ·OR`. The two terms learn on different
+timescales, and a **static λ** can't serve both:
+
+- **The SFT/NLL half learns easily** — the model quickly raises `logp(chosen)`
+  (fits the response *style*). Observed live: val loss falls steadily while
+  preference accuracy/margin barely move.
+- **The preference/OR half is hard** — on subtle data (UltraFeedback on a 1B) the
+  model doesn't generalize "which of two plausible answers is better."
+- **Over-weighting OR backfires** via **likelihood displacement**: the model widens
+  the *relative* margin by **crushing `logp(rejected)`** (and sometimes dropping
+  `logp(chosen)` too) rather than genuinely preferring chosen — degrading fluency.
+
+A fixed λ either under-pressures the preference (λ too low → never leans on OR) or
+over-pressures it (λ too high → displacement/degeneration). The hypothesis: an
+**adaptive, displacement-aware λ** — low while SFT learns, ramped once SFT
+saturates, backed off the moment displacement appears — can find a generalization
+regime a static λ sails past. (Static λ tuning was the blunt version we rejected;
+this is the targeted version.)
+
+## Control signals (what we measure)
+
+`orpoMetrics` (`loss.ts:1249`) already computes everything; we just don't surface it:
+
+- **`lw`** = mean `logp(chosen)` (currently folded into `nll = mean(−lw)`).
+- **`lr`** = mean `logp(rejected)` (currently only used inside the odds ratio).
+- **`margin`** = mean log-odds (chosen vs rejected) — already emitted.
+- **`accuracy`** = fraction `lw > lr` — already emitted.
+
+The controller keys on the **decomposition of margin growth**:
+- margin ↑ because **`lw` ↑** → *healthy* (genuinely preferring chosen).
+- margin ↑ because **`lr` ↓** while `lw` flat/falling → *displacement* (the cheat).
+
+So the load-bearing new instrumentation is **surfacing `lw` and `lr` separately**
+per eval (not just `nll`/`margin`).
+
+## The controller (state machine)
+
+Runs at **each val eval** (every `stepsPerEval` steps) — val (`lw`, `lr`) is the
+stable signal; per-step B=1 train metrics are too noisy to control on. Maintains a
+short window of recent val (`lw`, `lr`, margin, accuracy).
+
+```
+state: { phase, lambda, window[] }   # lambda starts at LAMBDA_MIN
+on each val eval (lw, lr, margin, acc):
+  dLw = lw - lw_prev ; dLr = lr - lr_prev
+
+  PHASE 1 — SFT warmup (lambda = LAMBDA_MIN, low):
+    stay while the NLL is still improving (dLw > EPS_SFT, i.e. chosen logp rising).
+    when dLw plateaus (|dLw| < EPS_SFT over K evals) → SFT saturated → PHASE 2.
+
+  PHASE 2 — preference ramp:
+    each eval, lambda += RAMP (cap LAMBDA_MAX), applying OR pressure —
+    AS LONG AS chosen holds: dLw >= -EPS_DISP.
+    if accuracy/margin respond (margin widening with dLw >= 0) → keep ramping.
+    if displacement detected (dLw < -EPS_DISP, i.e. chosen logp dropping, or
+       dLr << 0 with dLw <= 0) → PHASE 3.
+
+  PHASE 3 — displacement backoff:
+    lambda -= BACKOFF (floor LAMBDA_MIN) until chosen recovers (dLw >= 0),
+    then hold (hysteresis) before considering another ramp.
+```
+
+Knobs (all tunable): `LAMBDA_MIN≈0.05`, `LAMBDA_MAX≈0.5`, `RAMP/BACKOFF≈0.05`,
+window `K≈3` evals, `EPS_SFT`/`EPS_DISP` set from the observed val-`lw` noise scale.
+Hysteresis (require K consecutive evals before a phase flip) prevents oscillation.
+
+**Degeneration guard (hard stop):** if `lw` falls below its phase-1 starting value
+by more than a threshold, freeze λ at `LAMBDA_MIN` and log — the model is paying for
+margin with chosen's likelihood, which a backoff alone isn't recovering.
+
+## Control-theory framing (PID / AIMD) — added 2026-06-23
+
+The phase state-machine above is one instance of a more general and more honest
+frame: **this is closed-loop control.** The **process variable** is the SFT signal
+(NLL / `lw`), the **actuator** is the preference weight (ORPO's λ, *or* SimPO's
+β/γ — the controller is objective-agnostic), and the **control law** is what maps the
+measured degradation back to the actuator. Two laws worth naming:
+
+- **AIMD** (additive-increase / multiplicative-decrease — TCP congestion control):
+  ramp the preference up *gently* while the SFT is healthy, **slam it down hard** the
+  moment degradation ("congestion") appears. Famously robust and forgiving of its own
+  knobs. The PHASE-2 ramp / PHASE-3 backoff is essentially AIMD.
+- **PID** decomposes the response into three terms, each fixing a failure the others
+  can't:
+  - **P (proportional · the level):** how degraded is `lw` *right now* → back off in proportion. (Leaves a steady-state offset.)
+  - **D (derivative · the slope):** is degradation *accelerating* → a **leading indicator** that catches the helping→hurting **inflection before the level rises**, so you pull back *earlier* and leave less damage to heal. **Gotcha:** differentiating a noisy signal amplifies noise → use a *robust* slope (slope of the EMA, or a windowed regression), never raw step deltas. The slope is where most of the value is.
+  - **I (integral · accumulated error):** the "missing letter." P leaves a chronic offset and D ignores it (flat slope at steady-state) → the model can settle hovering with a *small persistent* NLL elevation = quiet permanent SFT damage. I accumulates that residual and drives `lw` **back to its floor**. **Gotcha — integral windup:** if the actuator is pinned (λ at floor) while error persists, the integral piles up and **overshoots** on recovery → needs **anti-windup** (clamp/stop accumulating when saturated).
+
+**Build it incrementally — P → PD → PID** — each term earning its place against an
+*observed* failure of the simpler controller (P oscillates/offsets → add D; chronic
+offset remains → add I with anti-windup). That's both good control practice and an
+**ablation that writes the experimental section**, and it stops us from tuning three
+gains blind.
+
+**Why a controller beats a sweep (the real argument):** the optimal preference
+pressure is **non-stationary** — fragile early (push less), robust late (push more) —
+so **no fixed λ is optimal.** A sweep finds the best *constant*; the controller tracks
+the moving target and finds the best *trajectory*, which strictly dominates. The pitch
+isn't just "one run instead of six," it's "**beats any fixed value** a sweep could
+find." (The displacement backoff is also why this is *more* than auto-tuning — it
+targets the degradation signal directly.)
+
+**Degenerate failure to guard (beyond the §degeneration hard-stop):** if the law backs
+off too eagerly it **collapses the preference term to ~0 → you've trained pure SFT.**
+The AIMD/PID balance must **favor pushing** (aggressive-ish increase, retreat only on
+*confirmed* degradation), and target a small **safety margin below the cliff** rather
+than the absolute max-bearable (the edge is fragile).
+
+**Synergy with the diffusion frontier:** all of this is noise-robustness machinery for
+controlling off a noisy process variable — so it transfers *especially* well to a
+diffusion LM, where the likelihood is a stochastic ELBO (even noisier). See
+[orpo-base-uf-experiment-and-directions.md](../investigations/orpo-base-uf-experiment-and-directions.md) §6.7.
+
+## Wiring
+
+1. **Instrument (small):** add `lw`, `lr` to `orpoMetrics`'s return + the trainer's
+   `metric` emit (train and val). `pref-control.ts` / `metrics.json` log them so the
+   controller (and our plots) can see the decomposition.
+2. **Make λ dynamic:** today `cfg.orpoLambda` is a constant threaded into
+   `orpoLoss(model, batch, cfg.orpoLambda, …)` and **baked into the segmented ORPO
+   classes' constructors** (`new SegmentedBackwardOrpo(model, lora, ranges,
+   cfg.orpoLambda, …)`). Change to a **mutable λ source**:
+   - simplest: a `{ value: number }` box (or a getter) passed instead of the scalar;
+     the loss / segmented step reads `lambdaRef.value` each step.
+   - the segmented classes expose a `setLambda(x)` (or read the ref) so the
+     constructor-baked constant becomes a live field.
+3. **The controller** lives in `orpoLoop` (`trainer.ts`): after each val eval, feed
+   the val (`lw`, `lr`, margin, acc) to `updateLambda(state)` and write the new λ
+   into the ref. Log every λ change (`· λ 0.10 → 0.15 (SFT saturated, ramping)`).
+4. **Off by default:** gate behind a config flag (`orpoLambdaSchedule: "static" |
+   "adaptive"`); static = today's behavior, exact.
+
+## Experiment design (how we validate it)
+
+Head-to-head, **same data / seed / everything else**, on the from-base UF setup:
+- **Arm A:** static λ=0.1 (the current run — the baseline).
+- **Arm B:** adaptive λ (this controller).
+
+Compare:
+- **val accuracy + margin** trajectories (does B clear the noise floor / widen
+  margin where A stalls?);
+- **`lw` trajectory** (does B avoid displacement — `lw` non-decreasing — where a
+  high static λ would crater it?);
+- the **downstream** numbers: Exp-1 capability x2 and Exp-2 win-rate (does B's model
+  actually win more head-to-head?).
+
+Success = B generalizes measurably better (val acc clears ~0.55 with margin lifting,
+and/or Exp-2 win-rate > A's) **without** degeneration (`lw` holds). Null result is
+also informative: "even adaptive λ can't beat the 1B capacity ceiling on UF."
+
+## Open questions / risks
+
+- **Cadence:** controlling on val every `stepsPerEval` may be too coarse (few control
+  steps per run). Option: a cheaper, more frequent "control eval" on a small fixed
+  probe set distinct from the reported val set.
+- **Confounds the replication claim:** an adaptive-λ result is no longer "standard
+  ORPO" — it's our method. Keep Arm A (static 0.1) as the honest paper-replication;
+  Arm B is the "can we do better" research arm. Report both.
+- **Tuning the controller's own knobs** risks the same trap (we'd be tuning a tuner).
+  Mitigate: set EPS thresholds from observed noise scale, not by hand-fitting to one
+  run; validate the controller on the *positive control* (prefer-uppercase) first —
+  it should ramp λ early (SFT saturates fast) and never trip displacement.
+
+## References
+
+- `src/train/loss.ts:1249` `orpoMetrics` (computes `lw`,`lr`,`nll`,`or`,`accuracy`,`margin`).
+- `src/train/trainer.ts` `orpoLoop` (λ config, val emit), `DEFAULT_TRAIN_CONFIG.orpoLambda`.
+- `src/train/segmented.ts` segmented ORPO classes (λ baked in constructor).
+- `scripts/experiments/pref-control.ts` (trajectory logging — extend with `lw`,`lr`).
+- Context: [[opssdpa-dk-vjp-bug]] (validate grads vs autograd), the dynamic-λ
+  discussion + the live from-base UF val trajectory (2026-06-20).

package/docs/design/orpo-future-enhancements.md

@@ -14,6 +14,48 @@ comfortable.
 > [steel-flash-cce-handoff.md](../investigations/steel-flash-cce-handoff.md). The
 > remaining sections (§2, §4–§9, the objective variants) are still genuinely future.
 
+## Deferred levers — queued, do NOT start mid-run (2026-06-19)
+
+Two optimizations identified during the MiniCPM5 UltraFeedback ORPO run. **Both are
+deferred until the active run finishes** — the GPU is fully occupied, so these get
+*measured* later, not "quickly tested" now (a quick probe steals the training GPU).
+
+1. **Expose `--grad-accum` (gradient accumulation).** The machinery already exists —
+   `accumulateStep` (`src/train/trainer.ts:975`) loops `gradAccumSteps` B=1 micro-steps,
+   **averages** their grads (`scale = 1/accumSteps`, line 989), and evals+frees each
+   micro-graph so **peak memory stays flat** (line 1015). It's wired into the ORPO loop
+   (`trainer.ts:837`) but `gradAccumSteps` defaults to 1 and has **no CLI flag**.
+   - **Why it helps us:** ORPO here is hard **B=1** (segmented + prefix-share assume a
+     single stream — `segmented.ts:636` throws on B>1), so every optimizer update is
+     computed from exactly **one** preference pair — the noisiest possible gradient (the
+     per-step loss swinging 0.4↔3.3 is that noise). Grad-accum gives an effective batch of
+     G — averaging G pairs/update, variance ↓ ~1/G, steadier descent, tolerates a higher
+     LR — at **B=1 memory**. It's the *only* way to a bigger effective batch when B>1 is
+     impossible.
+   - **What it is NOT:** not a speedup (G× passes per optimizer update → same total
+     passes, same wall-clock, same data seen); and not a guaranteed win (the prior run
+     converged fine at B=1, val 1.66→1.50). It's a **try-and-measure** lever — expose the
+     flag, A/B G=4/8 judged on **val-margin**, don't assume.
+   - **Note on prefix-share:** the existing `orpo_prefix_shared` is *already* a hand-built
+     batch-of-2-with-shared-prefix (1F + 1B over `[P; Rc; Rr]`, prompt deduped in **both**
+     directions — strictly cheaper than a generic 2-batch's `2F + 2B`). Generic N-way
+     batching would re-duplicate per-pair prefixes and need a block-diagonal mask layer
+     *on top of* the prefix-share mask — high friction, and the payoff is dominated by
+     grad-accum (effective batch, free) + W4A16 (the real GPU bottleneck). Park batching
+     unless a profile ever shows real short-row GPU idle.
+
+2. **W4A16 compute-precision spike.** The head GEMM dequants 4-bit weights into **fp32**
+   and matmuls in fp32 (W4A32); dropping the operands to `half` (fp32 accumulate) attacks
+   the actual bottleneck. Confirmed **compute-bound live** (GPU 100% / BW ~⅓ of 400 on
+   M1 Max). Full plan + parity gating: [w4a16-compute-precision-spike.md](w4a16-compute-precision-spike.md).
+
+**Context — seg vs no-segment A/B (MiniCPM5, seq 4096, same-machine M1 Max):**
+`--no-segment` measured **~1.25 s/step steady @ peak 9.26 GB** vs `--seg 2`
+**~1.58 s/step @ 2.62 GB** — ~20% faster for 3.5× peak memory, **bit-identical output**
+(segmenting is just gradient checkpointing). Dropping the segs removed the recompute (real
+GPU work) → GPU busy but no longer pinned. Fine when memory is spare; the current 3-epoch
+run uses `--no-segment`.
+
 Product relevance: ORPO is not only a researcher feature. It is the path to
 small, downloadable task adapters that make the default local assistant better
 without forcing every user onto a larger base model. The performance work below

package/docs/design/orpo-training.md

@@ -27,6 +27,19 @@ runs half the forwards per step.
 > torch/numpy reference of the paper/TRL math (see [Parity](#parity)), not a
 > bit-exact upstream port.
 
+> **L3 verification (2026-06-21).** The flash-CCE head has no oracle either, so
+> the "0.28% dh" figure is **fp-reassociation vs the full-logits *proxy*, not an
+> error** (two correct reductions in different orders can't be bit-identical).
+> Correctness is proven by **finite-difference** (`flash-fd-check.ts`) + the math
+> audit — not by matching a reference. The **coeff filter's** synthetic-data cost
+> (0.66→2.7%, from `flash-cce-parity.ts`'s random hidden) is its **worst case**:
+> random logits give a flat softmax, but real outputs are sharply peaked, so
+> filtering the ≈0-softmax tail is near-free — **measure teacher-forced on real
+> hiddens before judging** (it's plausibly a free speedup to enable). Standing L3
+> gates (filter-on-real-data, teacher-forced grad fidelity, end-to-end quality)
+> + the host-buffer pin-leak post-mortem:
+> [orpo-flash-cce-pin-leak.md](../investigations/orpo-flash-cce-pin-leak.md).
+
 ## Implementation status
 
 **Landed (naïve correctness path — the shippable unit + the oracle everything

package/docs/design/parity-tier-dag.md

@@ -0,0 +1,143 @@
+# Parity-tier DAG → meaningful CLI flags
+
+**Status:** vision / roadmap (2026-06-21). The first artifact exists — a zoomable,
+tier-tagged map of the training + inference DAGs at
+[`docs/dag/training-inference-map.html`](../dag/training-inference-map.html). This
+doc records where it goes and *why*: to turn a sprawling pile of CLI flags into a
+small, coherent, tier-organized surface.
+
+## The framework (recap)
+
+Every node in the compute graph sits in the **lowest parity tier whose oracle it
+matches bit-for-bit**:
+
+- **L1** — drop-in for **mlx-lm** (its op is the oracle; bit-exact or it's a bug).
+- **L2** — drop-in for **mlx-optiq** (optiq's behavior is the oracle).
+- **L3** — **our originals** (flash-CCE head, ORPO loss, prefix-share, segmented
+  backward, HLG sampler, expert offload). **No oracle ships**, so correctness is
+  gated by finite-difference / KL / quality, *not* parity.
+
+**The gate is the oracle.** A node can only live in a lower tier if that oracle
+exists AND it matches bit-exactly; otherwise it bubbles up. Note the subtlety the
+map already shows: an *optimization* that still matches the oracle stays low —
+`compiled-decode` is ours but is L1 because it replays the same ops bit-for-bit.
+Only nodes with **no** oracle (flash-CCE) are *forced* to L3. (Background +
+post-mortem: [orpo-flash-cce-pin-leak.md](../investigations/orpo-flash-cce-pin-leak.md).)
+
+## Two axes: parity gates performance
+
+Every optimization has **two independent coordinates**, and conflating them is what
+made the flag surface unreadable:
+
+1. **Parity** — *which reference does it reproduce bit-for-bit?* This sets the
+   **lowest level** it may live in: matches mlx-lm → L1; matches mlx-optiq but not
+   mlx-lm → L2; matches neither → L3 (gated by KL + tests, no oracle).
+2. **Performance** — *is it the fastest correct way?* This sets whether it's the
+   **default (on)** within its level, or a kept-but-**off** opt-in.
+
+**Parity gates performance.** A faster kernel that breaks mlx-lm parity is simply
+*not in L1* — it bubbles up to whatever level its parity allows. You never trade the
+guarantee for speed silently; you change levels, and the user opts in. The level is
+your floor of trust; the default within it is the best way to honor that floor.
+
+### Decision procedure for any new optimization
+
+1. **Measure its parity** (run it both ways, compare — "one flag, one measurement
+   per lever"). The lowest oracle it matches bit-for-bit is its level ceiling.
+2. **Measure its performance** vs that level's current default:
+   - **faster AND holds the guarantee** → it becomes the **default (on)** for that level.
+   - **slower but still correct** → **keep it as a documented default-off flag**
+     (optionality for A/B; never the default, never deleted).
+   - **breaks the guarantee but offers a real tradeoff** (e.g. 2× tok/s at KL
+     0.0015) → an **L3 opt-in**; the user chooses exact vs fast.
+
+Worked examples:
+- *5% less memory, still bit-for-bit with mlx-lm* → **default-on in L1.**
+- *2× tok/s, but KL 0.0015 to mlx-lm* → **L3 opt-in** (user picks exact or 2×).
+- *another way, 3% slower, still correct* → **kept, off in every level.**
+
+### Our decode levers under this rule (today)
+
+| lever | parity | perf | placement |
+|---|---|---|---|
+| `compiled-decode` | bit-exact w/ uncompiled (`compiled-decode.test`) | faster | **default-on, every level** |
+| `perf-kernel` | matches optiq, not mlx-lm (`perf-kernel-oracle.test`) | faster | **default-on from L2** (the L2 bridge) |
+| `fused-sdpa` | matches optiq, not mlx-lm (`fused-sdpa.test`) | faster | **default-on from L2** (the L2 bridge) |
+| `fused-decode` | (would be L3) | ~3% slower (`optimization_plan.md`) | **kept, off everywhere** |
+
+So `--l1/--l2/--l3` is the **parity** axis (the guarantee), and the defaults within
+each tier are the **performance** axis (the fastest kernel that still holds it). For
+the *decode path* L2 == L3 today (every decode kernel matches mlx-lm or optiq); L3's
+real territory is the no-oracle **features** — HLG sampler, expert offload, batched
+mixed-precision — which live off the decode axis and are gated by KL + quality.
+
+## Why this matters for CLI flags
+
+We are accumulating a **lot** of flags (`--grad-accum`, `--grad-clip`, `--seg`,
+`--no-flash`, `--no-prefix`, `--no-segment`, `--lambda`, `--val-size`, kv-quant
+knobs, sampler knobs, …). The problem isn't any one flag — it's the **combinatorics**
+and the fact that **some are never really a choice** (always-on guards, or the only
+viable route). The tier-tagged DAG makes each flag legible by asking: **what route
+does this flag push the computation down, and what tier is that route?**
+
+A flag falls into exactly one of three buckets:
+
+| flag selects… | tier effect | what it IS | action |
+|---|---|---|---|
+| the only viable route (no alternative node) | — | not a choice | **always-on default; drop the flag** |
+| between two routes at the SAME tier / same oracle | none | cosmetic / safety guard | **default it; expose only if it ever fails** |
+| between a **lower-tier** route and an **L3** route | **L1/L2 ↔ L3** | a real **parity ⇄ optimization** knob | **keep, document as such, group by tier** |
+
+So the DAG lets us answer "**what flags push down THAT route**" for any path, and
+from that:
+
+1. **Prune always-on flags to defaults** — anything whose alternative route is
+   never selected, or whose toggle doesn't change the tier, shouldn't be a flag.
+2. **Surface the real axis.** The flags that genuinely matter all toggle
+   *bit-exact-lower-tier ⇄ our-L3-optimization*. Those could collapse into a small
+   number of intent-level switches (e.g. a single `--exact` that forces every node
+   down to its lowest tier, vs the optimized default) instead of N independent
+   kernel toggles the user has to understand individually.
+3. **Group the rest by what they trade** — memory (`--seg`, all bit-exact → not a
+   parity choice, a memory↔compute knob), hyperparameters (`--lambda`, `--lr` →
+   training knobs, orthogonal to tier), etc.
+
+### Worked classification (current training flags)
+
+- `--grad-clip 1.0` → standard guard, L1, **always-on default** (not a route choice).
+- `--seg N` / `--no-segment` → memory↔compute, **bit-exact either way** → a *memory*
+  knob, **not** a parity choice. Should read as "memory budget," not "exactness."
+- `--grad-accum N` → changes the accumulation route (L1; smoother grads) → a real
+  *training-quality* knob.
+- `--no-flash` → **L3 flash-CCE → L2 MLX fused head** → a genuine *parity ⇄ perf*
+  fallback (the one place the head leaves L3).
+- `--no-prefix` → **L3 prefix-share → bit-exact two-forward** → a *perf* knob with no
+  numerical cost.
+- `--lambda` → ORPO hyperparameter (pure L3) → a *training* knob.
+
+The pattern: most flags are **memory/training knobs or always-on guards**; only a
+couple are true **parity ⇄ optimization** toggles. Today they're a flat list the
+user must reason about individually; tier-tagging lets us present them as what they
+actually are.
+
+## Roadmap to make it executable
+
+The map is hand-authored today. To make the framework *enforce* itself:
+
+1. **Derive the graph from code** — extract the call graph + per-node op/oracle tags
+   so it stays in sync and is **queryable** ("every L3 node", "paths that cross into
+   L3", "the pure-L1 subgraph", "flags that select route X").
+2. **CI-enforce the gate** — a node tagged L1 must have a passing bit-exact parity
+   test vs mlx-lm; if it can't, it's a bug or must be re-tagged L3. "The oracle is
+   the gate" becomes a check, not a principle.
+3. **Shrink the L3 surface** — audit whether anything in the L3 island is actually
+   bit-exact-able and just unproven (e.g. the prefix-shared *forward* is already
+   bit-exact — `prefix-shared-parity.ts` — so only the encode-once *fold* is truly
+   L3). Minimizing L3 minimizes what needs KL/quality gates.
+4. **Generate the flag surface from the graph** — once flags map to routes, the CLI
+   help + defaults can be derived: always-on routes have no flag, parity⇄opt routes
+   collapse to intent switches, knobs group by what they trade.
+
+End state: instead of a growing flat list of flags the user memorizes, a small set
+of intent-level switches whose meaning is exactly "which tier route does the compute
+take," backed by a DAG that's checked in CI.

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

@@ -0,0 +1,232 @@
+# 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`).