mlx-bun 0.0.9 → 0.0.11

package/docs/investigations/decode-roofline-lookagain.md

@@ -1,5 +1,11 @@
 # Decode perf "look again" — adversarial re-exam of the M=1 floor claim (2026-07-01)
 
+> **2026-07-02 CORRECTION — the host-graph-tax claim is itself overturned; see
+> §7.** The fix-plan items 1a–1d were built on "host graph build is serial with
+> the GPU"; the overlap spike (decode-overlap-probe.ts) shows the pipelined loop
+> ALREADY hides the build. The §3 wall times and %-of-roofline stand; the
+> §4(b)/(c) decomposition does not — the whole gap is GPU-side.
+
 **Machine:** Apple M1 Max, 32 GB (MacBookPro18,2), `sysctl` verified. All eval-DB
 rows are also M1 Max — same box. **All numbers below are directional (session,
 loaded-machine caveat); quotable numbers remain benchmark.sh-gated.** GPU idle
@@ -225,12 +231,33 @@ is suspect, profile on the other laptop before building anything.**
    1d. **12B short-context: compile the concat phase too** (or accept 1b as
        the fix). The current segmented form compiles nothing below the 1024
        window and never compiles the 8 full-attention layers. ~+9% @600.
-2. **26B GPU overhead (~4 ms/tok over floor):** profile gather-qmm expert
-   reads (each expert slice ~3.3 MiB → the ~180 GB/s small-buffer regime);
-   candidates: expert-contiguous layout so 8 experts read as fewer/larger
-   spans, the open perf-kernel enablement (review #7, +1–3%), router-chain op
-   fusion. Potential **+15–25% 26B decode** if expert reads reach big-GEMV
-   bandwidth. Effort M (profile first).
+2. **26B GPU overhead (~4 ms/tok over floor):** ~~profile gather-qmm expert
+   reads~~ **PROFILED 2026-07-02** (moe-expert-read-profile.ts, exact 26B
+   shapes; real decode @600 re-measured 21.0–22.6 ms/tok ✓). The gap is real
+   and located, but BOTH cheap candidates are refuted: expert-contiguous
+   layout does nothing (a chain with ONE expert repeated ×8 — 1/8 the unique
+   bytes — runs at the same speed: not bandwidth/locality-bound), and merging
+   gate+up stacks to cut dispatches is SLOWER (9.6 vs 8.1 ms). Root cause:
+   **mx.gather_qmm at M=1 is compute-bound** (~99 GB/s effective vs the qmv
+   kernel's ~180 on identical bytes; upstream python mlx 0.31.2 reproduces
+   6.18 vs 4.01 ms) — it lacks a qmv-class fast path, so each per-expert
+   [1,2816]→[704] product runs on mostly-idle GEMM tiles. The router chain
+   adds a secondary ~1.6–2.6 ms of tiny-op dispatches.
+   **Custom-kernel attempt (same day): built, correct, SHELVED.** A gather-qmv
+   Metal kernel (scripts/experiments/moe-qmv-kernel.ts + moe-qmv-parity.ts)
+   passes bounded-divergence on all three 26B dispatch patterns, but five
+   structural variants (occupancy/staging/bank-conflict/mask-trick sweeps —
+   post-mortem table in the kernel header) all measure SLOWER than gather_qmm
+   (best 12.5 vs 8–10 ms on the 30-layer chain), and the killer measurement is
+   structural: **mx.fast.metal_kernel dispatches cost ~60–95 µs each in a
+   dependent chain, nearly independent of the work inside** — at 3
+   dispatches/layer × 30 layers that fixed cost alone eats the entire ~4 ms
+   prize. Remaining routes, in order: (a) UPSTREAM — give mlx's gather_qmm an
+   M=1 qmv-class path in C++ (no dispatch overhead; natural companion to the
+   quantized_matmul M=2/3 bug report), (b) a fused whole-MLP kernel (gate+up+
+   gelu in one dispatch, down in a second → 60 dispatches/step) with a
+   qmv-class body — a dedicated fused-decode-v2-style session starting from
+   the post-mortem map.
 3. **e4b GPU overhead (1.8–3.7 ms):** dispatch-count reduction on the
    per-layer-input machinery (the ~1–2 MiB gate/proj GEMVs and altup
    elementwise chains; batch the 42 per-layer gate GEMVs into one [42·width]
@@ -247,7 +274,66 @@ is suspect, profile on the other laptop before building anything.**
    bandwidth wall (weight read shared across K tokens) — unchanged priority,
    now with the roofline to prove why nothing else can.
 
+## 7. The overlap spike (2026-07-02) — fix 1b REFUTED; §4(c) decomposition corrected
+
+Fix 1b asked: can the host graph build hide under the GPU (build-two-ahead /
+off-thread dispatch)? `decode-overlap-probe.ts` answers two ways — the
+decisive one is DIRECT REMOVAL: take the per-step host build out of the loop
+entirely (pre-build all N step graphs up front, chained through the lazy
+argmax outputs — identical graph structure to production — then run
+dispatch+read with zero interleaved build) and see if decode gets faster.
+
+M1 Max, greedy, uncompiled, all arms token-IDENTICAL to baseline:
+
+| arm | CPM5 @132 (spin 2 ms) | e4b @600 (spin 4 ms) |
+|---|---|---|
+| baseline (pipelined) | 4.05 ms/tok (graph 0.86, dispatch 3.13) | 19.02→18.06 across passes (graph ~3.3–3.6, dispatch ~14) |
+| **prebuilt (host build REMOVED)** | **4.11 (+0.07)** | **19.03 (+0.01)** |
+| prebuilt-1buf (also no per-step dispatch boundaries) | 4.33 | 19.73 |
+| serial (evalAll, no pipeline) | **+1.75 ms** | **+3.36 ms** |
+| spin-build / spin-pre / spin-post (+W injected host work) | **+0.4 ms** (not +2) | **+1.0–1.4 ms** (not +4) |
+| chain-2 / chain-3 (deeper pipeline) | ±0.01 ms | +0.02 / −0.96 (== pass drift) |
+
+- **Removing the work gains nothing**: with the entire 0.9 ms (CPM5) /
+  3.6 ms (e4b) per-token build lifted out of the loop, decode speed is
+  unchanged (+0.07 / +0.01 ms, inside pass-to-pass drift). Even the one-buffer
+  floor — no host work AND no per-step dispatch boundaries — is no faster.
+  **The build already overlaps; wall = GPU step time.**
+- The injection arms agree from the other direction: extra host work is
+  ABSORBED (dispatch shrinks 3.13→1.6 exactly as `gpu − host` predicts) — the
+  GPU is busy with step n while the host builds step n+1.
+- The serial anchor costs MORE than the build itself (CPM5 +1.75 vs 0.86;
+  e4b +3.36 vs 3.6): the pipelined loop hides the build AND the per-step
+  launch gap. The loop shape is already right; mlx-lm's identical shape gets
+  the identical hiding.
+- Deeper pipelining (chain-2/3) buys exactly nothing — there is no residual
+  host term to hide.
+
+**Corrections to this doc:**
+- §4(c) "host JS loop — the single biggest NON-physics term" is WRONG as an
+  additive claim. The "host graph ms" column was measured as time-in-JS-calls,
+  but that window runs concurrently with the GPU; "GPU ms" (= the dispatch
+  wait) is `gpu − graph`, so summing them reconstructs the wall without any
+  recoverable host term. The true GPU step time ≈ the measured wall itself.
+- §6.1a (CPM5 CompiledDecode "+20–25%"): premise gone. Any compiled-decode win
+  is GPU-side (fewer/denser dispatches from C++ replay), not host removal —
+  re-bill as profile-first, expected small. This also re-explains §4(d):
+  e4b @600 compiled wall 18.9→18.1 means compiled made the GPU ~0.8 ms faster
+  (net of concat tax), not "host −2.4, GPU +1.7".
+- §6.1b (this spike): answered, REFUTED — nothing to build.
+- §6.1c/1d (SharedKv segmented-compile plumbing ~L effort, 12B concat-phase
+  compile): the host-tax share of their billing is gone; only GPU-side effects
+  remain. De-prioritize pending a GPU-side profile.
+- §6 items 2 (26B expert reads), 3 (e4b dispatch count), 4 (backlog #4
+  contiguity), 5 (CPM5 KV-path), 6 (spec decode) are UNAFFECTED — they were
+  always GPU-side, and they are now the entire recoverable gap.
+- 12B/26B not re-probed (structurally the same loop; 12B was at-wall under
+  either reading; 26B's 3.7 ms graph ≪ its ~18 ms step has the same slack).
+
 ## Probes added (scripts/experiments/, tsc-clean)
+- `decode-overlap-probe.ts` — direct-removal (prebuilt / prebuilt-1buf) +
+  spin-injection + serial-anchor + chain-depth arms for the overlap question
+  (§7).
 - `decode-roofline-bw.ts` — machine bandwidth: raw read, decode-shaped qmv
   GEMV (multi-copy, one-eval chaining), bf16 GEMV, dispatch chain.
 - `decode-roofline-step.ts` — factory-generic pipelined decode split

package/docs/investigations/dspark-handoff.md

@@ -1,11 +1,176 @@
 # DSpark — handoff
 
-**State:** faithful DSpark (DFlash KV-injection + Markov + confidence) is **built
-and proven architecturally correct** (overfit τ=3.24, paper-range). It does not
-yet net a speedup on e4b — that's a data-scale + target-speed problem, not an
-architecture problem. Full design + measured numbers:
+**State (2026-07-06): the DSpark paper is CODE-COMPLETE and serve-integrated.**
+Architecture proven correct (overfit τ=3.24, paper-range); serve-integrated
+behind `--draft-model` (Phase 1 — the `DraftSource` seam is KV-borrowing-ready;
+sibling AssistantSource verified bit-lossless on real weights,
+`tests/spec-serve-assistant.test.ts`); and every remaining paper component
+landed in Phase 2 the same day: Alg-1 confidence-scheduled draft-length pruning
+(checkpoint-driven, zero change for uncalibrated checkpoints), STS calibration
+(§3.2.1, `calibration.ts` + `scripts/dspark-calibrate.ts`), the RNN sequential
+head (Eq 6, `--seq-head rnn`, init-equivalent to Markov), the tightened draft
+loop (on-device token chaining + deferred conf reads + `collectLogits:false`
+on serve — bit-identity pinned), and the `dspark` variant rename + central
+loader. Suites: smoke 21/21 · dspark tests 17/17 · adjacent server suites
+green. **What remains is exactly the GPU work below** — data scale + 12B
+retarget + train + calibrate + measure; e4b not netting a speedup is a
+data/target-speed problem, not architecture. Full design:
 [docs/design/dspark-speculative-decoding.md](../design/dspark-speculative-decoding.md).
 
+## PATH A (preferred, no training) — DeepSeek's published 12B drafter
+
+DeepSeek released a TRAINED drafter for our exact target:
+`deepseek-ai/dspark_gemma4_12b_block7` (6.86 GB bf16, γ=7; DeepSpec repo is
+the MIT reference impl + the oracle — see the "DeepSpec ground truth" section
+of the design doc). Once the deepspec-variant port lands:
+
+```sh
+# 1. download the trained drafter (~6.9 GB)
+mlx-bun get deepseek-ai/dspark_gemma4_12b_block7
+# 2. oracle gate: dump their reference trace (temp 0 = RNG-free deterministic)
+#    in a torch venv, then compare ours round-for-round (scripts staged)
+# 3. measure + serve (auto-detected via config.json architectures field)
+mlx-bun serve --model gemma-4-12B-it-OptiQ-4bit --draft-model <snapshot-of-dspark_gemma4_12b_block7>
+```
+
+### Oracle protocol (step 2, expanded)
+
+Two staged scripts implement the temp-0 deterministic trace protocol against
+DeepSpec's own reference implementation (github.com/deepseek-ai/DeepSpec, MIT).
+At temperature 0 their leaky-rejection verify degenerates to exact argmax
+token-match — the whole round trace is RNG-free and reproducible, so a
+plain round-for-round diff is a real bit-exact gate (same discipline as the
+D2 diffusion oracle, minus RNG parity entirely — see the design doc's
+"DeepSpec ground truth" section).
+
+```sh
+# a. dump the reference trace — needs its OWN torch venv (NOT the mlx-lm
+#    oracle venv at /Users/joshrossi/Code/mlx-lm/.venv, which has no
+#    torch/deepspec and shouldn't); see the script header for the exact
+#    venv-setup + pip install steps.
+.venv-deepspec/bin/python scripts/oracle-dspark-deepspec.py \
+    --target <snapshot-dir-of-google/gemma-4-12B-it> \
+    --drafter <snapshot-of-dspark_gemma4_12b_block7> \
+    --data prompts.jsonl --n 8 --max-new-tokens 128 \
+    --temperature 0 --confidence-threshold 0 \
+    --out goldens/dspark-deepspec/trace-thr0.jsonl
+
+# b. compare our port round-for-round
+bun scripts/dspark-deepspec-compare.ts \
+    --fixture goldens/dspark-deepspec/trace-thr0.jsonl \
+    --target gemma-4-12B-it-bf16 --drafter <snapshot-of-dspark_gemma4_12b_block7>
+```
+
+Run the pair twice — once at `--confidence-threshold 0` (no pruning), once at
+`0.5` — to gate both the base draft loop and the confidence-truncation path.
+
+Fixture location: `goldens/dspark-deepspec/*.jsonl` follows the house
+manifest convention (goldens/README.md) — these are machine-independent
+**JSON manifests** (prompts + token-id traces, no tensor blobs), so they are
+the **tracked** kind, not the `.bin` machine-specific kind. Regenerate by
+rerunning step (a); there is no `scripts/regen-*` wrapper for this family yet
+since the oracle run needs its own GPU-hosting torch venv, not the mlx-lm
+oracle venv this repo otherwise standardizes on.
+
+**The OptiQ-4bit caveat is load-bearing, not a footnote**: the fixture above
+is generated against the bf16 HF target. A `dspark-deepspec-compare.ts` run
+with `--target` pointing at our bf16 12B is the TRUE bit-exact gate. A run
+with `--target` pointing at the OptiQ-4bit-quantized 12B we actually serve is
+a **separate acceptance-rate measurement** — tapped hiddens differ
+numerically at 4-bit, so round-for-round divergence there is expected, not a
+bug; pass `--acceptance-only` on that arm so the script reports an aggregate
+accept-length reading instead of a PASS/FAIL assert.
+
+Caveat: trained against the bf16 HF target; our serving 12B is OptiQ-4bit —
+tapped hiddens differ numerically, so measure acceptance rather than assuming
+their 60–85%.
+
+**FIRST LIVE RUN (2026-07-07, loaded machine — directional only):** τ ≈ 2.8
+committed/target-forward (200 tok / 72 verify rounds; per-token acceptance
+26–33% — the predicted OptiQ-4bit degradation), output coherent, losslessness
+holding. **But wall-clock LOST ~3.4×** (spec 14.6 vs serial 49.8 agg tok/s;
+the 20 s bench TTFT is conc-4 queueing on the serial-forced spec lane, not
+prefill). The tax is the drafter itself. **The full step-by-step program for
+the next pass is
+[docs/design/dspark-serving-program.md](../design/dspark-serving-program.md)**
+(phases, gates, dependency graph; pickup at TurboQuant merge, only its
+Phase 5 waits for it). Lever summary:
+1. **Quantize the drafter** — the big one, and the ideal quantization target
+   PERIOD: drafter numerics only move ACCEPTANCE, never correctness
+   (losslessness is structural), so aggressive bpw is safe and the gate is a
+   single acceptance A/B. 6.9 GB bf16 → ~1.8 GB @4-bit kills both the
+   ~3 GB/round bandwidth tax (2 GB lm_head + 7×134 MB markov_w2 reads) and
+   the 14 GB two-model memory pressure on 24 GB. Uniform 4-bit (mlx-native)
+   first as the baseline; **TurboQuant (Phase 13), when it lands, competes
+   on acceptance-per-byte against that baseline** — the drafter is its
+   lowest-risk first customer. Enabler either way: quantized-weight support
+   in deepspec-module's forward (quantized_matmul, the house pattern).
+2. **Tighten deepspec draftBlock** — ~14 host syncs/round (7 argmax + 7
+   confidence reads); replay the on-device chaining + deferred-conf
+   tightening already done for our own module (threshold=0 needs NO conf
+   reads).
+3. **captureLayer in the GENERATED gemma forwards** — tapped verifies
+   currently fall back to the monolith (correct-but-unspecialized); teach
+   scripts/gen-model.ts to emit the tap so spec verifies keep the fast path.
+4. **γ sweep** — γ=7 at ~30% acceptance wastes drafts;
+   `--num-draft-tokens 2..3` may net better wall-clock free of any code.
+
+## PATH B — train our own module (custom targets / research)
+
+The payoff target is **12B** (27B is memory-infeasible to train on 24 GB; the
+pipeline stays dim-generic so 27B remains a future recipe). All three scripts
+take `--model` and (new, 2026-07-06) `--tap-layers` — the retarget is now a
+pure parameterization, no code edits. **`--tap-layers` MUST be identical across
+regen and train** (the shard feature dim is `m*H`, `m = tapLayers.length`).
+
+12B tap set: use **`5,17,29,41,46`** — DeepSeek's actual trained
+`target_layer_ids` for this model (their layer-output convention matches our
+tapLayers indexing; audited 2026-07-06). Our earlier `24,37,47,48` guess is
+superseded.
+
+```sh
+# 0. one-time: cache the 12B target + a prompt/topic corpus (thousands of
+#    on-distribution topics — the data-scale lever; 160 only reached per-pos ~0.17).
+mlx-bun get mlx-community/gemma-4-12B-it-OptiQ-4bit
+# (12B target may need the serve process stopped to free memory for regen.)
+
+# 1. regen multi-layer training shards from the 12B target's own generations
+bun scripts/dspark-regen-dflash.ts \
+  --model gemma-4-12B-it-OptiQ-4bit --tap-layers 5,17,29,41,46 \
+  --topics <thousands-of-topics.txt> --out data/dspark-12b --max-resp 320
+
+# 2. train the drafter against the 12B (warm-start survivable via --resume;
+#    GPU runs kept getting killed ~step 5000 pre-resume). Optional A/B:
+#    --seq-head rnn trains the Eq-6 RNN head instead of Markov (Eq 5).
+bun scripts/dspark-train-dflash.ts \
+  --model gemma-4-12B-it-OptiQ-4bit --tap-layers 5,17,29,41,46 \
+  --data data/dspark-12b --out ckpt/dspark-12b \
+  --iters 8000 --batch 8 [--seq-head rnn] [--ddraft 2560] [--resume ckpt/dspark-12b]
+
+# 3. STS-calibrate (§3.2.1): fit per-position confidence thresholds on real
+#    verify outcomes and stamp them into dspark.json — this ACTIVATES the
+#    Alg-1 confidence scheduler at decode (uncalibrated = fixed-γ, no change).
+bun scripts/dspark-calibrate.ts \
+  --model gemma-4-12B-it-OptiQ-4bit \
+  --drafter ckpt/dspark-12b --data <prompts.jsonl> [--target 0.5]
+
+# 4. measure live τ + tok/s vs vanilla 12B decode (clean machine — dirty
+#    numbers are garbage). Run BOTH pre- and post-calibration for the
+#    scheduler's own A/B (MLX_BUN_DSPARK_MINCONF=<c> = manual uniform knob).
+bun scripts/dspark-measure-dflash.ts \
+  --model gemma-4-12B-it-OptiQ-4bit \
+  --drafter ckpt/dspark-12b --data <prompts.jsonl>
+
+# 5. serve it (auto-detected as kind=dspark via dspark.json; γ pinned to cfg.gamma)
+mlx-bun serve --model gemma-4-12B-it-OptiQ-4bit --draft-model ckpt/dspark-12b
+```
+
+Decision gate (write it down, RESULTS.md only if clean-machine): does 12B live
+τ net a wall-clock win? If the draft loop dominates, do the loop-tightening
+(per-position host syncs + double 262K LM-head) first — see the open items in
+the design doc. Memory guard (`server.ts`) caps context to fit; a client
+sending huge `max_tokens` trips admission — lower it.
+
 ## What happened (short version)
 - Built v1 first as a single-vector shortcut — it dropped the paper's core KV
   injection and was never really DSpark. Chunk (τ≈1.26) and article (τ≈1.16)
@@ -31,20 +196,21 @@ architecture problem. Full design + measured numbers:
 - Training `--resume` (warm-start from checkpoint) — the GPU runs kept getting killed ~step 5000; resume survives that.
 
 ## Next (in priority order)
-1. **Retarget to a slow model (27B or 12b).** The drafter reads dims from the
-   model + tapLayers is config, so it ports — but needs regen+train on that
-   model. 27B on 32GB is memory-tight (17.75GB weights); may need the serve
-   process stopped to free the ~14GB for regen. This is where a real speedup lives.
-2. **Scale training data** — thousands of on-distribution generations (not 160)
-   to lift generalizing τ toward the overfit ~0.75.
-3. **Tighten the draft inference loop** — per-position host syncs + double 262K
-   LM-head make the draft heavy; needed for τ to translate to wall-clock.
-4. **Finish the paper:** RNN head (Eq 6), STS calibration (§3.2.1), hardware-aware
-   prefix scheduler (Alg 1 — single-user = confidence-scheduled draft-length pruning).
-5. **Rename** `dflash`→`dspark` (faithful module IS DSpark; v1 = legacy single-vector).
+1. **Retarget to 12B + scale data + train + calibrate + measure** — the GPU
+   recipe above is the complete remaining work. (27B stays a future recipe:
+   memory-infeasible to train on the 24 GB M4 Pro.)
+2. ~~Tighten the draft inference loop~~ **DONE 2026-07-06** (on-device token
+   chaining, deferred conf reads, `collectLogits:false`; bit-identity pinned).
+3. ~~Finish the paper~~ **DONE 2026-07-06**: RNN head (Eq 6, `--seq-head rnn`),
+   STS calibration (§3.2.1), Alg-1 confidence-scheduled draft-length pruning
+   (single-user form). ⚠ Eq 6 + §3.2.1 shapes are design-doc-faithful readings
+   (paper PDF absent) — verify against the paper when available.
+4. ~~Rename `dflash`→`dspark`~~ **DONE 2026-07-06** (canonical variant +
+   central loader `src/spec/dspark/loader.ts`; legacy stamps still load).
 
 ## Key files
-`src/spec/dspark/{module,data,generate}-dflash.ts`, `{loss,sample}.ts`;
+`src/spec/dspark/{module,data,generate}-dflash.ts`, `{loss,sample,calibration,loader}.ts`;
+`src/spec/{source,serve-loop,dflash-source,assistant-source}.ts` (the serve seam);
 `src/model/gemma4.ts` (hiddenTap), `gemma4-base.ts` (trim bypass);
 `scripts/dspark-{regen,train,measure}-dflash.ts`, `dspark-dflash-smoke.ts`.
 

package/docs/investigations/kernel-perf-review-2026-07.md

@@ -12,7 +12,7 @@ area reviews + full mechanisms): `reports/kernel-perf-review-2026-07-01.json`
   vendored MLX steel machinery, exact-by-default backward with autograd + FD
   parity — marred by decision drift: the PLAN-decided 1.35× filter win was
   never landed and the in-code comments argued three contradictory positions
-  (comments fixed 2026-07-01; the FLIP itself is still backlog #1).
+  (comments fixed 2026-07-01; the FLIP LANDED 2026-07-02 — see backlog #1).
 - **Training flash-attention kernel**: faithful port with exemplary post-bug
   discipline, correctly opt-in — but it copies a naive oracle (~30× slower)
   and its tests covered none of the production configs (extended 2026-07-01:
@@ -47,35 +47,109 @@ evalSftLoss error swallowing · tokens_per_sec metric · inverted comments
 (gemma4.ts:298, flash-cce filter trio) · silent steel fallback (now warns;
 H%128 throw narrowed) · flash-attn test blindness.
 
-STILL OPEN: **FUSED_DECODE=1 inside whole-graph compiled decode** bakes the
-tile loop at trace-time N — q-cat adapters grow N per step under shapeless
-replay, so the newest KV rows are silently never attended (gemma4-base.ts
-quantizedSdpaTiled; deferred for file-ownership at fix time — guard the combo
-or make the tile loop trace-safe).
+FIXED 2026-07-02: **FUSED_DECODE=1 inside whole-graph compiled decode** baked
+the tile loop at trace-time N — q-cat adapters grow N per step under shapeless
+replay, so the newest KV rows were silently never attended. REPRODUCED on e4b
+(42 layers traced tile=true at N=17; diverges within 3 tokens then
+repetition-loops). The 12B was never affected — its segmented form runs concat
+layers as uncompiled js layers. Fix: generate() refuses to compile under
+MLX_BUN_FUSED_DECODE=1 (explicit opt-in wins, like LoRA/MoE) + quantizedSdpa
+throws if the tiled path is ever selected inside a trace (lands in the
+transactional fallback, never silent). Regression in
+tests/compiled-decode.test.ts, whose e4b block also now uses the dynamic
+snapshot resolver — the old hard-coded hash silently skipped it on this
+machine, which is how the combo went unexercised.
+
+NEW BUG (found 2026-07-02 by the backlog-#8 gate test; FIXED in mlx-bun,
+UPSTREAM in mlx): **mx.quantized_matmul(transpose=false) returns wrong values
+at row count 2 or 3** — relnorm ~0.33-0.39 on our dylib, ~0.98-1.14 on stock
+mlx 0.31.2 (oracle venv repro); M=1 and M>=4 exact, transpose=true exact at
+every M. mlx's autograd computes dh for a transpose=true head qmm via an
+internal transpose=false call at the cotangent's row count, so EVERY training
+head silently produced wrong dh for 2-3-token responses, 2-3-token tail
+chunks, and tiny prompt spans (f32 ground truth: dh relnorm up to 0.99 at the
+affected positions, cos ~0.2). Serving/decode unaffected (transpose=true).
+Fix: (a) loss.ts `logitsFromHiddenPadM` pads differentiated head spans to 4
+rows (pad cotangents are exactly zero) — applied at all seven head sites;
+(b) ops.quantizedMatmul pads direct transpose=false calls (the fused head's
+analytic backward). Verified vs f32-dequant ground truth (chunk-4 dh
+0.47→0.006) + M-sweeps. Worth filing upstream with the 10-line python repro.
 
 ## The ranked optimization backlog (task #12)
 
-1. **[S] Land the coeff filter at eps~1e-5** (+ measure blockMax skip on real
-   data) — 1.35× ORPO head backward, ALREADY MEASURED
-   (flash-cce-filter-realdata.ts, PLAN.md "filter-on-real-data"); gates:
-   E4B=1 rerun + teacher-forced grad fidelity + eps=0 byte-identical.
-2. **[M] planSegments full-attention isolation** — e4b@8K ~17.5→10 GB;
-   unblocks 8K training on the 24 GB M4 Pro.
-3. **[S] Auto-dispatch the training head by M** — exact fused QM head for
-   short-M (~1.9×: 481 ms vs 934 ms at e4b M=512), flash only when memory
-   demands; both heads already share the fusedRespLogpMean interface.
-4. **[S] Fused decode full-capacity KV buffers + activeN** — kills 6
-   ensureRowContiguous copies/layer/step; 1-3% @8k, ~4× that at 32k.
+1. ~~**[S] Land the coeff filter at eps~1e-5**~~ **LANDED 2026-07-02** — BOTH
+   skips default 1e-5 (coeff filter + blockMax block skip; env=0 restores exact).
+   All gates run on the M1 Max with real chunk-ORPO data
+   (flash-cce-filter-realdata.ts, extended: E4B=1 + blockEps sweep + Vjp
+   full-logits fidelity): filter CPM5 0.343%/1.41×, e4b 0.158%/1.70×; blockMax
+   skip ~lossless (≤0.004%) and real text DOES go cold (the synthetic probe was
+   the artifact) — **combined 1.71× (CPM5) / 3.16× (e4b)**; teacher-forced
+   fidelity cosine ≥0.99993, filter adds ≤0.07% over the eps=0 reassociation
+   floor (standing test in tests/train-orpo-fused-ce.test.ts). The "eps=0
+   byte-identical" gate resolved structurally: eps=0 compiles the identical
+   pre-flip kernel, but run-to-run dh is never byte-stable (atomic-add
+   reassociation) — byte-replay was the wrong spelling of that gate.
+2. ~~**[M] planSegments full-attention isolation**~~ **REFUTED BY MEASUREMENT
+   2026-07-02** — built, A/B'd @8K on real chunk rows: ZERO peak win (18.09 vs
+   18.02 GB; scripts/experiments/seg-isolation-smoke.ts + MLX_BUN_SEG_MEM_LOG
+   phase probes). The §5 cost-model premise is false: mlx's sdpa BACKWARD
+   materializes O(L²) scores for EVERY layer (~3.5 GB/layer @8K e4b — the
+   sliding window is just an additive mask; sliding pair +7.1 GB ≈ full+sliding
+   +7.15). Worst segment = layer count alone, so **segment_size is the whole
+   knob: seg1 measures 14.59 GB @8K (+3% step time, loss identical) — which
+   DOES fit the 24 GB M4 Pro today, no code needed.** Inter-segment clearCache
+   also measured no-effect (watermark is live memory). Real next levers: the
+   head vjp's ~3 GB [M,V] spike (backlog #8) and an O(L)-memory attention
+   backward (the flash-attention track).
+3. ~~**[S] Auto-dispatch the training head by M**~~ **LANDED 2026-07-02** —
+   fusedRespLogpMean routes requested-flash batches with M <
+   MLX_BUN_FLASH_MIN_M (default 1024) to the exact fused head. Fresh crossover
+   sweep (head-dispatch-sweep.ts, M1 Max — the review's 1.9× @M=512 was
+   pre-filter-flip): fused wins e4b M=64 1.6× / M=512 1.13×, tie @1024; the
+   #1-flip made flash FASTER at long M too (e4b ≥2048; 8K 10.7 vs 13.2 s) on
+   top of its memory win. CPM5: fused faster at every M — flash stays a pure
+   memory lever there (≥1024 honors the caller's intent). Dispatch proven by
+   bit-identical-to-fused loss in tests/train-orpo-fused-ce.test.ts.
+4. ~~**[S] Fused decode full-capacity KV buffers + activeN**~~ **REFUTED BY
+   MEASUREMENT 2026-07-02** — the variant was built (an int32 activeN input,
+   capacity-strided rows, fetch attaches full buffers), gated BYTE-IDENTICAL
+   on all three recorded dispatch shapes and perf-oracle-green — and then
+   measured SLOWER end-to-end: interleaved in-process A/B on the 12B uniform
+   kv4 (tokens identical across arms, M1 Max) gave 24.64 vs 26.18 tok/s @8k
+   (0.94) and 23.07 vs 23.83 @22k (0.97) — the "~4× at 32k" scaling is
+   refuted too. Two holes in the estimate: the 12B's full-attention dispatch
+   shape is KV=1 (H16 D512), where the fetch view is a CONTIGUOUS PREFIX and
+   the six-copies model never applied; and kernel-only micro-A/Bs proved
+   unstable (±50% run-to-run) — the numbers behind "1-3%" were never
+   end-to-end. REVERTED; evidence + re-application recipe:
+   scripts/experiments/fused-decode-activen-ab.ts.
 5. **[M] Collapse per-token host syncs in the spec/DSpark loop** — ~2γ+1
    round-trips → ~2 per verify cycle; directly attacks DSpark's fixed draft
    overhead.
 6. **[M] Pipeline the batched decode loop + clearCache cadence** — folded into
    the batching-v2 wave (docs/design/batching-v2-plan.md).
 7. **[S] Enable the 26B perf kernel** (open PLAN Phase E box; oracle frozen).
-8. **[S] Bound the SFT segmented head** — replace full-[M,V] responseOnlyCe in
-   the segmented head vjp; ~2-4 GB peak on long-response e4b SFT.
-9. **[S] Segmented-step overhead pass** — detachLeaf single-copy, one evalAll
-   barrier/segment, memoize the CPM5 prefix block-sparse mask.
+8. ~~**[S] Bound the SFT segmented head**~~ **LANDED 2026-07-02** — both SFT
+   segmented classes now use boundedSftCe (token-chunked Checkpoint head, the
+   proven ORPO plumbing; MLX_BUN_SEG_HEAD/MLX_BUN_SEG_HEAD_CHUNK knobs).
+   Head-only A/B at e4b M=6000: whole-vocab head 16.60 GB peak vs bounded
+   6.60 GB (**~10 GB head-vjp transient eliminated**), loss identical, dh
+   relnorm 0.00000; head time +45% (Checkpoint recompute, small share of a
+   step). The @8K smoke's short responses were never head-bound (peak set by
+   the per-layer sdpa backward — see #2), so its wall peak moved little; the
+   win is exactly the long-response regime the review named. Gate test:
+   tests/train-orpo-fused-ce.test.ts "bounded SFT head".
+   **Landing this exposed a NEW upstream mlx bug (see ledger below).**
+9. ~~**[S] Segmented-step overhead pass**~~ **LANDED 2026-07-02** — all three
+   fixes: MlxArray.detachCopy (ONE mlx→mlx copy; the old rawBytes+
+   fromBytesCopy route was two copies via the JS heap), detachLeaves (ONE
+   evalAll barrier per segment's whole vjp set, applied in all six segmented
+   backward loops), and a step-shared block-sparse-mask memo for the CPM5
+   prefix classes (prefixSharedCaches — makeMask hands out slice views).
+   Gate: loss + grads BYTE-IDENTICAL on CPM5-SFT / CPM5-ORPO-prefix / e4b-SFT
+   fixed steps. Step time (short-seq fixtures, M1 Max): CPM5 ORPO-prefix
+   238→158 ms (−34%), e4b SFT 2065→1272 ms (−38%); e4b @8K flat (GPU-bound,
+   as predicted) with peak unchanged (14.49 GB).
 10. **[M] e4b prefill gap** — profile-first; NOTE: the decode look-again could
     NOT reproduce RESULTS.md's 304-vs-373 row on the M1 Max (no eval-DB
     backing) — re-bench on the M4 Pro before building anything.
@@ -85,8 +159,13 @@ or make the tile loop trace-safe).
 The review's framing assumed the megakernel post-mortem's "M=1 decode is at
 the floor". The follow-up investigation (docs/investigations/
 decode-roofline-lookagain.md) OVERTURNED that: only the 12B is at the wall
-(~92-93% of measured roofline); CPM5 58-64%, e4b 64-70%, 26B 60-62%, with the
-host JS graph build as the #1 recoverable term. Its ranked fixes (CPM5
-compiled-decode port; build-ahead graph overlap spike; 26B gather-qmm
-bandwidth; e4b dispatch batching) COMPOSE with the backlog above and should be
-sequenced first where they overlap.
+(~92-93% of measured roofline); CPM5 58-64%, e4b 64-70%, 26B 60-62%.
+**2026-07-02 second correction (lookagain §7):** the "host JS graph build =
+#1 recoverable term" part of that verdict was itself wrong — the overlap
+spike (decode-overlap-probe.ts, spin-injection + serial anchor) proved the
+pipelined loop ALREADY hides the build; wall = GPU step time. The lookagain's
+host-side fixes (1a–1d: CPM5 compiled-decode port, build-ahead spike, SharedKv
+segmented-compile, 12B concat-phase compile) are refuted or de-prioritized;
+its GPU-side fixes (26B gather-qmm bandwidth, e4b dispatch batching, CPM5
+KV-path, backlog #4 contiguity, spec decode) are now the ENTIRE recoverable
+gap and still compose with the backlog above.

package/docs/investigations/lab-build-journal.md

@@ -1,7 +1,7 @@
 # Build journal — mlx-bun Lab (web UI + native quantize/train/dataset + pi web chat)
 
 *Branch `web-ui-and-native-lab`, 2026-06-13/14. Visual summary:
-[archive/mlx-bun-lab-report.html](../../archive/mlx-bun-lab-report.html). Evidence
+[archive/mlx-bun-lab-report.html](../../docs/archive/mlx-bun-lab-report.html). Evidence
 artifact: `benchmarks/lab-verification.json` (regenerable, gitignored — produced by
 `scripts/verify-lab.ts`).*
 
@@ -22,7 +22,7 @@ page. Compute is **native Bun/MLX (no Python)** per Josh's decision.
 
 2. **Gate the riskiest unknown first.** Native LoRA needs autograd across the
    FFI boundary. Rather than discover mid-build that it's infeasible, the first
-   thing built was `spikes/phase-train-vag.ts`: a tiny LoRA-shaped loss
+   thing built was `lab/spikes/phase-train-vag.ts`: a tiny LoRA-shaped loss
    (`mean((x + (x@A)@B − y)²)`) differentiated by `mlx_value_and_grad` through a
    `bun:ffi` JSCallback closure, checked against central finite differences.
    **PASS (<1e-2), exactly 2 grads returned (base frozen).** Only then was the