mlx-bun 0.4.0 → 1.0.0

package/docs/archive/investigations/steel-flash-cce-handoff.md

@@ -1,489 +0,0 @@
-# Steel flash-CCE — handoff
-
-Porting MLX's quantized `steel` GEMM verbatim into our flash-CCE ORPO head, then
-fusing the ORPO-specific epilogue (softcap + online-softmax logp / softmax−onehot
-coeff + dh) on top — to get **MLX-GEMM speed AND the `[M,V]`-free memory property**
-in one kernel. MLX ships no fused "quantized-matmul + softcap + response-logp" op, so
-we build the composition; the heavy GEMM is theirs, verbatim.
-
-## HANDOFF — current state (2026-06-19): SHIPPED in v0.0.5
-
-The whole ORPO training stack is **merged (PR #16) and released as v0.0.5** — `npm
-i mlx-bun` / `brew install joshuarossi/tap/mlx-bun` (tap formula at 0.0.5, sha verified).
-What shipped: the flash-CCE head (fwd+bwd), prefix-sharing, segmented composition, warm-start,
-the `scripts/train-orpo.ts` launcher, and the IFEval/UltraFeedback eval pieces. Release notes:
-GitHub release v0.0.5. Supported scope = **OptiQ-quantized MiniCPM5 + Gemma-4 (e4b/12B/26B)**.
-
-**Post-merge PR-review fixes (CodeRabbit, commits 890f380 / 63759fa) — read these, they
-caught real bugs:**
-- 🔴 **Prefix-shared head-primal UAF** (`branchLogpMeanGathered`/`orpoLossPrefixShared`): the
-  flash/fused-head CustomVjp PRIMAL (`hResp`, and `h`) was disposed *before* the lazy backward
-  recomputes from it → use-after-free (same class as the segmented `headSink` one). Fix: push
-  primals into the sink (freed post-eval) on the fused/flash path. **Tests passed before only
-  by lazy-eval timing luck** — this was latent. Watch for this pattern anywhere a CustomVjp's
-  primal is freed inline.
-- **Segmented "flash implies fused"** (`trainer.ts` `segFusedChunk`): `orpoFlashCe` alone left
-  the segmented non-prefix fallback on the full `[M,vocab]` head → e4b memory blowup. Now
-  `orpoFusedCe || orpoFlashCe` routes through the bounded chunked head.
-- Warm-start made atomic (validate all targets before disposing any), job payload exposes
-  `warm_start_adapter`/`mlp_split`, script GPU inputs use `fromBytesCopy` (page-align footgun),
-  env-precondition enforcement + flag validation in the probes.
-
-**Warm-start** (`warmStartFromAdapter` / launcher `RESUME=<adapter-dir>`): continue from a
-checkpoint's LoRA weights (optimizer + LR schedule restart). Insurance for interrupted runs.
-
-**The training dress-rehearsal (honest):** a CPM5 UltraFeedback ORPO run (full stack, lr 5e-5)
-reached ~step 4820, **val 1.66 → ~1.50, then plateaued** — expected: UF is open-ended, its loss
-floor is ~1.2-1.5, NOT the 0.3-0.4 you get on a narrow task. **IFEval moved 22.5% → 22.5%
-(flat)** — a real null result: general preference data + a partial run is the wrong lever for
-IFEval. The pipeline learns (val dropped); the metric just isn't the right one. Checkpoints
-every 200 under `./adapters/cpm5-uf-8h/checkpoints/` (best-val ~`step-04200-val1.5008`).
-
-**Gotchas the runs taught us:**
-- **Agent-spawned background runs get REAPED** (~47 min observed; not a crash/OOM — confirmed
-  via logs + `pmset` showed no sleep). **Run long training detached from your OWN shell**
-  (`nohup … &`), never as a session background task.
-- e4b training **requires** `MLX_BUN_PERF_KERNEL=0` + `MLX_BUN_FUSED_GELU=0` (the launcher
-  auto-sets them; the e4b parity probe now enforces them).
-- Adapters live in **`~/.cache/mlx-bun/mlx-bun-finetunes/`** (launcher default), never the repo.
-
-**THE next goal — the chunk segmenter (load-bearing).** UF was the rehearsal; the real run is
-**distilling Opus/GPT-5.5 conversation-segmentation into a local model** (data at
-`~/Code/lucien/benchmark/finetune/chunk-v3/dpo/orpo-curated-*.fixed.jsonl`, `{prompt:[sys,user],
-chosen, rejected}` — already supported by the loader). Score by **boundary/label accuracy vs
-the gold (chunk-eval), NOT val loss** (loss is dominated by low-entropy JSON/UUID tokens). This
-is what localizes the Lucien synthesis pipeline's `chunk-recent` stage (the whole "100% local
-memory" promise). ORPO is ideal here: chosen/rejected share the conversation + are both
-well-formed JSON, so the contrast isolates the segmentation decision, and the gold makes
-pair-ranking + eval objective. The stack is built for it (prefix-share the shared 8k
-conversation once, segment the long context, flash the big vocab). See memory
-[[training-tracks-are-appliance-components]]. After that: the e4b overnight, and the welcome
-assistant (a UF-tuned greeter — system prompt carries the mlx-bun knowledge, not training).
-
-**Run commands:** see [orpo-quickstart](../reference/orpo-quickstart.md) +
-[training.md](../reference/training.md). The full stack is `scripts/train-orpo.ts` (defaults on);
-resume with `RESUME=<ckpt-dir>`; run it detached (`nohup`).
-
-## TL;DR status
-
-| Piece | State |
-|---|---|
-| **Steel GEMM port** (verbatim MLX BlockMMA/QuantizedBlockLoader/qmm_t) | ✅ done, compiles, matches `quantizedMatmul` 1e-5, **101ms == MLX 106ms** on e4b |
-| **Forward** (steel GEMM + softcap + online-softmax → logp) | ✅ **in production** (`flashCceForward` default), parity PASS e4b/cpm/M=200, **180ms** (5× the old 848ms) |
-| **Backward** (steel GEMM ×2 + softmax−onehot coeff + dh accumulate) | ✅ **IN PRODUCTION** (`flashCceBackward` default, `BWD_STEEL_SOURCE`/`bwdSteelKernel`). Parity PASS e4b (dh 0.40%) + cpm (dh 0.28%); **754ms** @ e4b M=512 — phase-2 W dequant now uses MLX's fused `QuantizedBlockLoader` (HTw=group_size, reduction_dim=0), **5× the old SG's 3687ms exact** (was 1851ms with the scalar dequant). Peak **0.928 GB flat @ M=8192** (`[M,V]`-free preserved). Old SG kept as `MLX_BUN_CCE_BWD_NOSTEEL=1` fallback. |
-| Old shipping backward | the **OLD hand-rolled SG kernel** (`[M,V]`-free, correct, slow 3687ms) — now the fallback only |
-
-**So: BOTH directions are now fast + live (steel fwd 180ms, steel bwd 754ms, both
-`[M,V]`-free + exact). M1 + the phase-2 dequant optimization DONE.** Backward is now
-~2.1× the ~360ms ideal (down from ~5×). Everything below (M2 prefix-share, M3 e4b probe,
-M3-composition segmented+prefix, M4A integration tests) is also DONE — see sections.
-
-## Why this exists (the journey, so you don't repeat it)
-
-- Our hand-rolled simdgroup_matrix GEMM runs at **~8% of M1 Max peak**; MLX's `steel`
-  quantized GEMM runs at **~46%**. We spent a long time micro-optimizing the
-  hand-rolled kernel (token tiling, h-coalescing, dequant staging, bank-conflict
-  padding) — all marginal, because the gap is the **BlockMMA register-tiling +
-  K-pipeline**, not staging. The decisive move was: **stop hand-rolling, copy MLX
-  verbatim.** It works and hits their speed.
-- **MLX's `quantizedMatmul` does the whole head ~3.8× faster than our old kernel AND
-  exact** (fwd 151 + bwd 330 = 481ms). That's the `orpoFusedCe` head, already in the
-  repo. The ONLY reason to keep flash-CCE: it's a **single kernel** → `[M,V]`-free,
-  no per-chunk host-loop allocator churn.
-- **Memory caveat (important, was over-claimed then corrected):** the head is NOT the
-  memory wall at long seq — the layer activations are. Avoiding the full `[M,V]`
-  (8.6 GB at e4b/8192) IS essential, but BOTH flash-CCE (+0.17 GB) and a properly
-  streamed fused head (+0.27 GB) achieve it. flash-CCE backward peak is **0.93 GB flat
-  at M=8192** (weights-dominated; kernel scratch ~5 MB; linear in M, flat in V). So
-  option 2 (steel backward) is justified for **SPEED + single-kernel allocation
-  flatness**, not because the fused head "blows up" (it doesn't, if streamed right).
-
-## Files
-
-- `src/train/steel-qmm-header.ts` — **the verbatim MLX steel port** (auto-assembled
-  from MLX `quantized.h` + `steel/{gemm/mma.h,gemm/loader.h,gemm/transforms.h,
-  utils/integral_constant.h,utils/type_traits.h}`). Exports `STEEL_QMM_HEADER` (the
-  MSL header string). Contains: `BlockMMA`, `MMATile`, `BaseMMAFrag`, `tile_matmad`,
-  `BlockLoader`, `QuantizedBlockLoader`, `dequantize`, `get_pack_factor`, `qmm_t_impl`.
-- `src/train/flash-cce.ts` — the head. `FWD_STEEL_SOURCE` + `fwdSteelKernel(gs,bits)`
-  (the **live forward**, dispatched by `useSteel` default). The old SG/scalar/lane
-  forward kernels + the current backward (`BWD_SG_SOURCE` / `bwdSgKernel`, the
-  hand-rolled persistent-acc[80] kernel) are still there.
-- `src/train/loss.ts` — `makeFlashCceHeadVjp` (CustomVjp: fwd→Σlogp stash lse+blockMax,
-  bwd→flashCceBackward). `ChunkCtx.flash`. Routed via `orpoFlashCe` config flag.
-- `src/train/trainer.ts` / `job.ts` — `orpoFlashCe` / `orpo_flash_ce` knob.
-- `tests/train-orpo-fused-ce.test.ts` — flash-head fwd-parity + end-to-end training
-  tests (MLX_BUN_TEST_TRAIN=1, MiniCPM5).
-- **Validated experiment scripts** (preserved, portable imports):
-  - `scripts/experiments/qmm-port-test.ts` — steel GEMM correctness vs quantizedMatmul.
-  - `scripts/experiments/qmm-port-speed.ts` — steel GEMM speed vs MLX (101 vs 106ms).
-  - `scripts/experiments/steel-fwd-test.ts` — forward epilogue (logp) vs reference.
-  - `scripts/experiments/steel-coeff-test.ts` — **backward phase-1** (logit→coeff) PASS.
-  - `scripts/experiments/steel-bwd-test.ts` — **full backward dh** (phase-1+2) PASS 1.6e-5.
-  - `scripts/experiments/bwd-peak.ts` — clean isolated backward peak (0.93 GB @ M=8192).
-  - `scripts/experiments/fwd-gemm-bench.ts` — MLX QM baseline (fwd 151 / bwd 330ms).
-
-## The architecture (what's proven)
-
-Tile constants: **BM=8 tokens, BN=32 vocab, BK=32 hidden, WM=1, WN=4** (4 simdgroups
-along N), `BK_padded = BK + 16/sizeof(T)`. The `WM=1,WN=4` choice is THE unlock:
-BlockMMA's output tile is then split so **each simdgroup owns a different N-slice** —
-which for the dh GEMM means each simdgroup owns a different **H-slice**, i.e. exactly
-the persistent `acc[80]` layout we already proved works (160 regs/lane). So steel's
-BlockMMA and the `[M,V]`-free persistent accumulator ARE compatible.
-
-**Forward** (`FWD_STEEL_SOURCE`, live): grid `[128, NBLK, ⌈M/32⌉]` (BM=32 there, not 8),
-per (token-block z, vocab-block y) threadgroup loop BN vocab-tiles → steel K-loop →
-`Ctile.store` to threadgroup `Ls` → softcap + online-softmax → `[M,NBLK]` partials →
-MLX merge → logp. `load_safe` for the last token-block (M%BM); vocab aligned for e4b/cpm.
-
-**Backward** (proven in `steel-bwd-test.ts`, BM=8/WM=1/WN=4): per (token-block, vocab-
-block) threadgroup, persistent `dmma = BlockMMA<float,float,8,H,32,1,4,false,false>`
-Ctile = the dh accumulator. Loop vocab-tiles: phase-1 `lmma` BlockMMA logit → `Ls` →
-coeff epilogue (`g·(onehot−softmax(logit−lse))·sech²`) → threadgroup `coeffS[8,32]`;
-dequant W non-transposed → `Wd[32,H]`; `dmma.mma(coeffS, Wd)` accumulates `coeff@W`.
-After the loop `dmma.store_result(dh,...)`.
-
-## Gotchas (hard-won — read before touching the header/kernels)
-
-1. **String.raw escaping**: when writing the header to a `.ts` string, do NOT escape
-   backslashes (mangles macro `\` line-continuations). Only fails if content has a
-   backtick or `${` (the MLX code has neither).
-2. **MLX preamble**: `mx.fast.metal_kernel` already provides `METAL_FUNC` + `using
-   namespace metal` (redef warning, harmless) but NOT `STEEL_CONST`/`MLX_MTL_CONST`/
-   `SIMD_SIZE`/`Int<>`/`pointer_element_t` — we add the macros + `type_traits.h` +
-   `integral_constant.h` ourselves.
-3. **`#pragma METAL internals`**: `type_traits.h` and `integral_constant.h` need it
-   (struct inheritance / `is_integral`). Include them **verbatim with their own
-   self-managed enable/disable pragmas** — do NOT strip + wrap manually (type_traits'
-   trailing `disable` will turn it off before integral_constant compiles → "derived
-   classes not supported").
-4. **W dequant convention** (in MANUAL dequant, e.g. phase-2 `Wd`): use
-   `perWord = 32/bits` (4 elems per uint32 at 8-bit), NOT `get_pack_factor<bits,8>=1`
-   (that's MLX's byte-based packing for the QuantizedBlockLoader's internal uint8
-   addressing). Mixing these gave a 100%-wrong dh.
-5. **`BlockMMA.mma` reads `Bs[BK,BN]` full-width** from threadgroup — so you can't wrap
-   it for a full-H persistent accumulator (would need `W[32,2560]`=320KB tg). Either
-   keep H small (test) or tile H with manual `tile_matmad` (production).
-6. Watchdog: per-threadgroup vocab slice must stay bounded (`VOCAB_BLOCK=8192`).
-
-## NEXT STEPS — the plan (in order)
-
-### 1. H-tiling in the steel backward phase-2 (the one real remaining piece)
-The test stages the full `W[32,H]`; at e4b H=2560 that's 320 KB > 32 KB threadgroup.
-Production must tile H:
-- dh accumulator = `MMATile<float, 1, H/32>` per simdgroup (= its H-quarter; 80 frags
-  @ e4b = the proven `acc[80]` register budget).
-- per vocab-tile: loop H in `BKh` chunks → dequant `W[32, BKh]` → threadgroup (small) →
-  load coeff A-frags `MMATile<1,4>` + W B-frags `MMATile<4, BKh/8>` → `tile_matmad(
-  D_sub, A, B_sub, D_sub)` into the right sub-frags of the persistent D.
-- The intricate bit is the manual sub-frag indexing into D (BlockMMA does this
-  internally; here we do `tile_matmad` by hand because D persists across vocab-tiles).
-- Validate against `steel-bwd-test.ts` (extend it to H=2560, compare to QM(coeff,W,t=F)).
-
-### 2. Vocab-blocking grid + atomic dh
-- Grid `[128, NBLK, ⌈M/BM⌉]`; each (token-block, vocab-block) computes its block's dh
-  contribution, **atomic-add** into device `dh` (across the NBLK blocks) — same pattern
-  the current `BWD_SG_SOURCE` already uses (`atomicOutputs`, `initValue:0`).
-- Need `lse` (have it, forward returns it) + `blockMax` (optional, for the block-skip)
-  + `gv` (cotangent) passed in — same inputs as the current backward.
-
-### 3. Wire into `flashCceBackward`
-- Add `BWD_STEEL_SOURCE` + `bwdSteelKernel(gs,bits)` (header = STEEL_QMM_HEADER + defines).
-- Dispatch: `useSteel` default when H,V,blockV %32 (mirror the forward's gate).
-- Keep the old `bwdSgKernel` as fallback (H%32!=0) + the float4 scalar as last resort.
-- The CustomVjp (`makeFlashCceHeadVjp`) already calls `flashCceBackward` — no change there.
-
-### 4. Validate + benchmark
-- Parity: `E4B=1 M=512 bun scripts/experiments/flash-cce-parity.ts` (+ cpm, + M=200).
-  Expect dh ~0.4% (bf16-class) vs autograd.
-- Bench vs baselines: fused head **481ms** (speed), current flash-CCE (memory parity).
-  Target: steel backward in the same class as the forward (~hundreds of ms, vs 3375).
-- Memory: `scripts/experiments/bwd-peak.ts` should stay ~0.93 GB @ 8192 (flat).
-- **CAP ANY MEMORY PROBE well under 32 GB** and eval+free per iteration — a naive
-  full-graph host-loop allocation swapped the machine to 23 GB once. Single-kernel
-  flash probes are safe by construction.
-
-### 5. (optional) softcap + filter/block-skip parity in the steel backward
-- Softcap sech² is in the coeff epilogue already (steel-bwd-test uses cap=0; cap>0 path
-  is written, validate it).
-- The Apple-CCE coeff filter + blockMax skip are **opt-in, default-off** approximations
-  (data-dependent; flat-softmax probes showed dh sensitivity). They live in the OLD
-  backward; decide whether to port them to steel (probably not needed — steel is fast
-  enough exact).
-
-## Decision context (for when you pick the shipping config)
-
-- **Now**: `orpoFlashCe` = steel forward (fast) + old SG backward (slow but `[M,V]`-free).
-  Trains, memory-flat. `orpoFusedCe` = MLX QM both ways (481ms, exact, `[chunk,V]`
-  transient).
-- **After option 2 lands**: `orpoFlashCe` = steel fwd + steel bwd = fast AND `[M,V]`-free
-  AND single-kernel — strictly best; removes the speed/memory tradeoff.
-- The one genuinely-unanswered question: does the fused head's host-loop chunked backward
-  churn MLX's allocator at real e4b/8192 training? An end-to-end training-step peak
-  measurement (model + activations + head, with segmented/checkpointed backward) decides
-  whether the single-kernel property matters in practice. Worth measuring once.
-
----
-
-# END-STATE ROADMAP — "CPM5 test run with everything → e4b overnight @ 8192"
-
-**The actual goal.** Two deliverables Josh runs himself (we wire, we do NOT kick off
-the overnight run — ground rule):
-- **(A)** A short **MiniCPM5** test training session using ALL optimizations — proves
-  the whole stack composes and trains (loss decreases).
-- **(B)** An **e4b overnight** ORPO run at seq **8192** — the real target.
-
-"ALL of our stuff" = flash-CCE head (fwd+bwd at MLX speed, `[M,V]`-free) + prefix-sharing
-(single fwd/bwd, prompt shared) + segmented/checkpointed backward (layer-activation
-memory at long seq) + LoRA.
-
-## What's already wired (so you don't redo it)
-- flash-CCE **forward**: live + default (`flashCceForward`).
-- flash-CCE head CustomVjp (`makeFlashCceHeadVjp`) ↔ `orpoFlashCe` flag ↔ trainer.
-- flash threaded into BOTH non-segmented (`fusedLogpMeanB1`) and **segmented**
-  (`fusedLogpMeanFromHidden`) head paths.
-- Segmented backward exists (`SegmentedBackwardOrpo` / `…Gemma4`), `segmentSize` flag.
-- ORPO B=1 path (`branchLogpMean`→fused/flash) + tests (MiniCPM5 flash training PASSES).
-- prefix-sharing math (`orpoLossPrefixShared`, MiniCPM5, bit-exact, parity script).
-
-## Milestones (dependency order)
-
-### M1 — flash-CCE backward → production (the steel backward) ✅ DONE (2026-06-18)
-H-tiled phase-2 (Design C: persistent `MMATile<1,H/32>` accumulator + per-H-tile temp
-BlockMMA over bounded `Wd[32,128]` + lane-local frag accumulation) → vocab-blocking grid
-+ atomic dh → `BWD_STEEL_SOURCE`/`bwdSteelKernel` wired into `flashCceBackward` (`useSteel`
-gate: H,V,blockV %32 + H%128). Parity PASS e4b (dh 0.40%) + cpm (dh 0.28%); 2.0× the old
-SG (1851 vs 3687ms e4b M=512 exact); peak 0.928 GB flat @8192. Design validated first in
-`scripts/experiments/steel-bwd-htile-test.ts` (H=2560, PASS 1.1e-4 rel). Both flash
-directions now fast. KEY INSIGHT that avoided the "intricate manual tile_matmad": a TEMP
-BlockMMA per H-tile whose Ctile is accumulated into the persistent D via `frag += frag`
-(lane-local — the per-lane element layout is identical, so no get_coord math needed).
-
-### M2 — prefix-sharing → trainer (single fwd/bwd) ✅ DONE (2026-06-18)
-- `orpoPrefixShared` trainer flag (+ `orpo_prefix_shared` job field) dispatches to
-  `orpoLossPrefixShared` (MiniCPM5) / `orpoLossPrefixSharedGemma` (Gemma4 — the sibling
-  already existed) via `splitPrefixBatch`, falling back to two-forward `orpoLoss` when a
-  row's chosen/rejected prompts differ. Guarded against `segmentSize>0` (composition is
-  M3). End-to-end CPM5 training test PASSES (loss decreases).
-- **Composed with the flash head**: new `branchLogpMeanGathered` gathers each branch's
-  `[M,H]` response hiddens → `fusedRespLogpMean(flash)` per branch (reuses the analytic
-  per-branch VJP — cleaner than splitting one call's non-uniform length-normalized
-  cotangent). Validated: flash-routed prefix-share matches whole-vocab prefix-share to
-  **0.018%** (bf16-class) in `prefix-shared-parity.ts`. So prefix-share is now `[M,V]`-free
-  end to end (one model fwd + per-branch flash fwd/bwd).
-- **e4b caveat unchanged**: e4b's bf16 length-sensitivity (~2–14%) is a numerical-parity
-  nicety, not correctness. MiniCPM5 is bit-exact (forward 0%, grads ~1% at P=512). NOTE:
-  for the e4b OVERNIGHT at seq 8192, prefix-share can't be used until M3 composes it with
-  the segmented backward — the overnight uses **segmented + flash** (validated config shape
-  below), with prefix-share an M3 follow-on.
-
-### M4A — CPM5 smoke test of the complete new system ✅ DONE (2026-06-18)
-Automated end-to-end tests in `tests/train-orpo-fused-ce.test.ts` (run with
-`MLX_BUN_TEST_TRAIN=1`), all PASS — each trains CPM5 20 iters, asserts finite +
-decreasing loss + no OOM/NaN:
-- **flash head** (steel fwd+bwd) + LoRA.
-- **prefix-share + flash** (the single-concat-forward + `[M,V]`-free head).
-- **segmented + flash** (the e4b-overnight config shape: gradient-checkpointed layer
-  activations + `[M,V]`-free head). These exercise every new piece composing in a real
-  ORPO step. (prefix-share + segmented together = M3, not yet composed.)
-
-### M3 — long-sequence memory @ e4b/8192 (the layer activations, NOT the head)
-- Compose segmented/checkpointed backward (`segmentSize>0`) + flash head + prefix-share.
-  The head is ~free in memory (0.93 GB flat); the wall is the transformer activations,
-  which segmented backward already handles (peak 10.9→3.3 GB @2048 per the notes).
-- **THE measurement that actually decides "does e4b train at 8192"**: a real end-to-end
-  training-step peak (model + activations + head) with segmented backward on. Do this
-  EARLY — it tells us the segment size needed to fit 8192 in 32 GB, and whether the head
-  choice (flash vs fused) matters at all. (Cap probes; never swap the machine.)
-
-### M3/M4B — e4b @ 8192 memory probe + overnight config ✅ PROBED (2026-06-18)
-`scripts/bench-orpo.ts` sweeps the new-system configs (`flash head` / `segmented + flash`
-/ `prefix + flash`); `ONLY_NEW=1` skips the legacy baselines, `CONFIG="<substr>"` runs one,
-and it now prints peak GB + a finite-loss check. The **`segmented + flash` row IS the
-overnight config shape**.
-
-**MEASURED on the M1 Max 32 GB (e4b, segmented+flash, SEG=2, `MLX_BUN_PERF_KERNEL=0
-MLX_BUN_FUSED_GELU=0`):**
-| SEQ | peak GB | ms/step |
-|---|---|---|
-| 1024 | 8.66 | 20.7k |
-| 2048 | 10.17 | 40.4k |
-| 4096 | 11.68 | 80.8k |
-| **8192** | **16.14** | **175k** (loss 0.093→0.091 finite, decreasing) |
-
-Footprint is **linear** in seq (~+1.5 GB / 1024 tokens — segmenting + flash attention keep
-it linear, not quadratic; the head is flat 0.93 GB). **e4b @ 8192 fits in 16.14 GB with
-~16 GB headroom — the historical "e4b OOMs ≥2048" ceiling is BROKEN.** SEG=2 is the sweet
-spot: SEG=8 @ 8192 hit **33.47 GB (swapped) and was slower** — bigger segments hold more
-activations, so small segments win here. The run command:
-```
-MODEL=<e4b-snapshot> CONFIG="segmented + flash" SEQ=8192 SEG=2 ITERS=3 \
-  MLX_BUN_PERF_KERNEL=0 MLX_BUN_FUSED_GELU=0 bun scripts/bench-orpo.ts
-```
-
-**VALIDATED overnight config (segmented + flash + LoRA):**
-```
-method:"orpo", orpoFlashCe:true (implies fused), segmentSize:2, orpoChunkSize:512,
-maxSeqLen:8192, env MLX_BUN_PERF_KERNEL=0 + MLX_BUN_FUSED_GELU=0, rank/scale/lr + real
-dpo data per the run, LoRA targets default. Peak ~16 GB, ~2.9 min/step on the M1 Max.
-```
-Prefix-share is NOT in the overnight (segmented + prefix-share = the M3-composition below);
-the overnight uses segmented + flash. **Josh starts the overnight** (ground rule).
-
-### M3-composition — segmented backward + prefix-share together ✅ LANDED MiniCPM5 (2026-06-18)
-New `SegmentedBackwardOrpoPrefix` (`src/train/segmented.ts`): streams ONE concat
-`[prompt; chosen; rejected]` through the segmented backward with `PrefixSharedCache`
-(block-sparse mask) + `setMiniCpmPrefixPlan` active during the streaming forward AND every
-segment's `mlx_vjp` recompute; the head VJP gathers chosen/rejected response positions
-(`branchLogpMeanGathered`, the `[M,V]`-free per-branch flash head) and scatter-adds its dh
-into the `[1,T,H]` boundary that seeds the segmented backward. Trainer composes it (old
-guard removed) with a two-forward-segmented fallback on prompt mismatch. Validated
-(`scripts/experiments/prefix-shared-segmented-parity.ts`): grads **2.27%** bf16-class
-(single-segment diagnostic 1.16% → it's segment-boundary reassociation + flash-head 0.4%
-stacking, not a bug), loss **0.052%**, peak **32% lower** (3.64→2.47 GB @ P=512; grows to
-59% lower @ P=1024) — prompt-encode-once now holds at long seq with flat ~2.5 GB activation.
-Test #6 in `tests/train-orpo-fused-ce.test.ts` (PASS). **e4b (Gemma4) DEFERRED** — trainer
-throws a clear error; the logical-position sliding-window prefix mask + donor-KV threading
-across segmented donor boundaries is the remaining follow-on (the overnight uses
-segmented+flash without prefix-share, which works).
-
-### M4 — runnable configs + the two runs
-- **(A) MiniCPM5 test config**: short seq, small iters, `orpoFlashCe` (or prefix-share
-  mode) + segmented on, LoRA. Run `trainLora`; assert loss decreases. This is the
-  "everything composes" gate — basically the existing flash training test scaled up.
-- **(B) e4b overnight config**: seq 8192, `segmentSize` tuned from M3, flash head,
-  prefix-share (or two-forward if e4b variance is a problem), LoRA, full dataset
-  (`scripts/curate-ultrafeedback.ts` / the dpo data). **Josh runs this** — we hand him a
-  validated config + command, we do not start it.
-
-## THE POINT: the runs are INTEGRATION TESTS of the NEW system
-Josh does NOT care about training quality/results. The CPM5 run and the e4b run exist to
-**exercise everything we just built — steel backward + flash head + prefix-sharing — end
-to end in a real ORPO training step** and confirm the new code path RUNS (loss is finite,
-decreases a bit, no crash/NaN/OOM). So we FINISH all the new pieces, wire them ALL the
-way through, then run. Do NOT ship the old slow backward as a crutch — the steel backward
-is part of what's being tested.
-
-## Order of execution (build the new system, then test it)
-1. **M1 — steel backward → production** (H-tile phase-2, vocab-block + atomic dh, wire
-   `bwdSteelKernel` into `flashCceBackward`, parity e4b/cpm). The new backward IS under
-   test; it must be in the path.
-2. **M2 — prefix-sharing → trainer + compose with flash head** (gather `[Rc+Rr]` → one
-   flash-head call → split by branch). MiniCPM5 first.
-3. **M4(A) — CPM5 test run = the smoke test of the complete new system**: flash fwd+bwd
-   (steel) + prefix-share + segmented + LoRA, all on. Run `trainLora` a handful of steps;
-   assert finite + decreasing loss + no OOM. THIS is the test — it validates every new
-   piece composes in a real step. If it passes, the new system works.
-4. **M3 — e4b @ 8192 dry-run**: same new system on e4b, seq 8192, segmented on. First a
-   memory/step-peak check (cap it; the head is flat 0.93 GB, the wall is activations →
-   tune `segmentSize`), then a few real steps to confirm it RUNS. (e4b prefix-share bf16
-   variance is fine — we're not checking quality; if it NaNs, fall back to two-forward.)
-5. **M4(B) — e4b overnight**: hand Josh a validated config + command for the longer run
-   (we wire + validate a few steps; **Josh starts the overnight** — ground rule).
-
-**Risk register**: M1 mechanical (scoped — architecture+correctness already proven). M2/e4b
-bf16 variance = non-blocking for a RUN test (two-forward fallback; quality irrelevant). M3
-e4b/8192 memory = the one real unknown, but it's an activations/segmentSize question, not
-a head question. The runs test integration; the pieces are individually validated.
-
----
-
-# GLOSSARY — every named piece, what it is, status, where
-
-Status legend: ✅ done+live · 🟢 done, not production-wired · 🟡 partial · 🔬 proven in a
-test only · 📦 pre-existing (not ours this arc) · 🪦 dead end.
-
-## A. The "head" = how we turn response hiddens `[M,H]` into ORPO log-probs
-(M = number of response tokens; V = vocab; H = hidden.) These are *alternative
-implementations of the same math* (`logp_t = target_logit − logsumexp_v`).
-- **Full-logits head** 📦 — the naive reference: materialize `[M,V]` logits, logsumexp.
-  Correct, huge memory. Only used as the parity oracle.
-- **Chunked head** (`orpoChunkSize`) 📦 — token-chunked + per-chunk `Checkpoint`; bounds
-  `[M,V]`→`[chunk,V]`. Older, autograd through the head.
-- **Fused linear-CE head** (`orpoFusedCe`, `fusedRespLogpMean`) 📦✅ — a `CustomVjp` with
-  an ANALYTIC softmax−onehot backward (no autograd through the head), using **MLX
-  `quantizedMatmul`** for both GEMMs, token-chunked. Fast (**481ms** e4b), exact,
-  `[chunk,V]` transient. The speed baseline we're trying to match.
-- **Vocab-blocked head** (`makeVocabBlockedHeadVjp`) 📦 — pure-MLX online-softmax over
-  vocab blocks; `[chunk, vocabBlock]` only. Slow-ish; the MLX-only CCE.
-- **flash-CCE head** (`orpoFlashCe`, `makeFlashCceHeadVjp`) — **OURS**, the Metal-kernel
-  head: the vocab loop lives INSIDE one Metal kernel, so neither `[M,V]` nor a
-  dequantized `[V,H]` ever touches HBM → `[M,V]`-free, single kernel. Forward ✅ live,
-  backward 🟡 (old slow kernel live; steel kernel 🔬 proven, not wired). `flashCceForward`
-  + `flashCceBackward` in `src/train/flash-cce.ts`.
-
-## B. flash-CCE kernel internals (`src/train/flash-cce.ts`) — OURS
-The flash head dispatches to one of several kernel implementations of fwd/bwd:
-- **`FWD_STEEL_SOURCE` / `fwdSteelKernel`** ✅ — the NEW steel-GEMM forward (default). 180ms.
-- `FWD_SG_SOURCE` / `fwdSgKernel` 🪦-ish — old hand-rolled simdgroup_matrix forward (848ms).
-- `FWD_SOURCE` / `fwdKernel` — scalar fallback (H%32≠0).
-- `FWD_LANE_SOURCE` 🪦 — lane-sliced (2.6× slower, dead end, kept gated).
-- **`BWD_SG_SOURCE` / `bwdSgKernel`** ✅(live) — CURRENT production backward: hand-rolled
-  simdgroup_matrix, persistent `acc[80]` dh accumulator, phase-1 (logit→coeff) +
-  phase-2 (dh=coeffᵀ@W). Correct, `[M,V]`-free, **slow (3375ms)**.
-- `BWD_SOURCE` / `bwdKernel` — float4 scalar backward fallback.
-- **STEEL BACKWARD** 🔬 — the NEW one: `scripts/experiments/steel-bwd-test.ts`, validated
-  dh 1.6e-5, NOT yet a production kernel (`BWD_STEEL_SOURCE` is M1's job).
-
-## C. The steel port (`src/train/steel-qmm-header.ts`) — OURS (verbatim from MLX)
-MLX's quantized `steel` GEMM, copied verbatim so we run at MLX speed (~46% peak vs our
-hand-rolled 8%). The `STEEL_QMM_HEADER` string contains:
-- **`BlockMMA`** — the tiled matmul (BM×BN output, BK contraction, WM×WN simdgroups, the
-  K-pipeline). The thing that makes MLX fast.
-- **`MMATile` / `BaseMMAFrag` / `tile_matmad`** — the 8×8 simdgroup-matrix fragment
-  abstraction BlockMMA is built on (used directly for the persistent dh accumulator).
-- **`QuantizedBlockLoader`** — loads packed quantized weights and dequantizes them
-  straight into the threadgroup matmul tile (fused dequant, no separate pass).
-- **`BlockLoader`** — loads the dense input (X / coeff) into the tile.
-- **`qmm_t_impl`** — MLX's full quantized-matmul orchestration (we use it for the
-  forward; the backward composes BlockMMA + tile_matmad ourselves).
-- Tile params we use: **BM=8/32, BN=32, BK=32, WM=1, WN=4** (the `WM=1,WN=4` is the key
-  that makes BlockMMA's output tiling match our `[M,V]`-free per-simdgroup H-slice accum).
-
-## D. The epilogue (what we FUSE on top of the steel GEMM) — OURS
-MLX has no "quant-matmul + softcap + response-logp" op, so we add:
-- **Softcap** — `cap·tanh(raw/cap)` (Gemma final-logit cap; e4b cap=30, cpm none) + its
-  `sech²` chain-rule term in the backward.
-- **Online softmax** — running `(max, sumexp)` over vocab blocks → `lse` (forward).
-- **Gather / target capture** — pick the target token's logit per row.
-- **softmax−onehot coeff** — the backward gradient `coeff = g·(onehot − softmax)·sech²`.
-- **MLX merge** — combine the `[M,NBLK]` per-vocab-block partials → `logp`/`lse` in MLX.
-
-## E. Backward approximations (opt-in, default-off, data-dependent) — OURS
-Both are Apple-CCE-style skips that proved unreliable on flat/synthetic softmax, so they
-are OFF by default and live in the OLD backward; probably NOT needed once steel is fast.
-- **Coeff filter** (`MLX_BUN_CCE_BWD_FILTER_EPS`) — skip phase-2 chunks whose whole coeff
-  block < eps.
-- **blockMax vocab-block skip** (`MLX_BUN_CCE_BWD_BLOCK_EPS`) — skip whole cold
-  (token-block, vocab-block) programs using the forward's `blockMax`.
-
-## F. Model-level optimizations (above the head)
-- **Prefix-sharing** (`src/train/prefix-shared.ts`, `orpoLossPrefixShared`,
-  `blockSparseMask`, `PrefixSharedCache`) 🟢 — OURS (earlier this arc). One forward over
-  `[prompt ; chosen ; rejected]` with a block-sparse mask (chosen & rejected each attend
-  the shared prompt but NOT each other) → prompt encoded ONCE, single fwd/bwd for the
-  pair. Bit-exact for MiniCPM5; e4b has a bf16 length-sensitivity caveat. Built + parity
-  script, NOT wired into the trainer.
-- **Segmented backward** (`src/train/segmented.ts`, `SegmentedBackwardOrpo*`,
-  `segmentSize`) 📦✅ — gradient checkpointing for the TRANSFORMER LAYER activations (the
-  real memory wall at long seq; orthogonal to the head). Composes with the flash head
-  (`fusedLogpMeanFromHidden` carries the `flash` flag).
-
-## G. Wiring / config layer
-- **`makeFlashCceHeadVjp`** (`loss.ts`) — wraps flashCceForward/Backward as a `CustomVjp`
-  (fwd→Σlogp, stash lse+blockMax; bwd→dh×cotangent). OURS.
-- **`ChunkCtx`** (`loss.ts`) — the head config struct: `{fused, flash, chunkSize,
-  vocabBlock, sink}`.
-- **`orpoLoss` / `branchLogpMean`** (`loss.ts`) — the ORPO loss entry; B=1 routes to the
-  fused/flash head, B>1 to full-logits.
-- **Trainer flags** (`trainer.ts` / `job.ts`): `orpoFusedCe`, `orpoFlashCe`, `segmentSize`,
-  `orpoChunkSize`, `orpoLambda`. CLI: `--orpo-flash-ce` etc.
-
-## H. Baselines / numbers to remember
-- MLX `quantizedMatmul` raw GEMM (the speed reference): e4b fwd **151ms**, bwd **330ms**.
-- Fused head total: **481ms** exact.
-- Our steel forward: **180ms** (raw GEMM 101ms; +epilogue). Old SG forward: 848ms.
-- Our backward: old SG **3375ms** (live); steel target ~hundreds of ms (M1).
-- flash-CCE backward peak: **0.93 GB flat @ M=8192** (weights-dominated; kernel ~5MB;
-  linear in M, flat in V).

package/docs/archive/investigations/the-dreaming-handoff.md

@@ -1,141 +0,0 @@
-# THE DREAMING — HANDOFF (2026-06-28)
-
-Live handoff for the local personal-wiki memory system (`src/memory/`). A fresh
-session can resume from this doc. **Companion docs:** `research-journal.md` (dated
-findings), `docs/design/memory-inference-path.md` (the in-progress templating/batching
-rework), and the per-phase `p*-judge.md` reports.
-
----
-
-## TL;DR — where we are
-
-The Dreaming turns conversations → chunks → entities → **subject articles** →
-cross-linked, self-healing personal wiki. It **works end-to-end on real data.**
-
-- **Full-corpus import RESUMED from cursor 720/2096** into `~/.mlx-bun/wiki-full`, on the
-  corrected templating, batching off (serial, ~50h).
-- **Committed + merged to `main`** (the Dreaming pipeline + inference-path rework). tsc 0.
-- Templating fix **VERIFIED live** on real convs (bare prompt + a 40-msg conv → valid chunks).
-
-## Resume the import
-```
-# resumable driver; continues from cursor 720. Writes to ~/.mlx-bun/wiki-full (NOT the
-# real vault ~/.mlx-bun/wiki). ~25 min / 20-conv batch, ~50h total, SERIAL.
-bun scripts/memory/dreaming-full-run.ts          # resume (clean-slate guarded by an init flag)
-# state files: ~/.mlx-bun/full-run-cursor.txt, ~/.mlx-bun/full-run-initialized, ~/.mlx-bun/full-run.log
-# after the import finishes: reindex, then connect the graph:
-mlx-bun memory link
-```
-**Gotcha:** background bash + detached processes get **reaped across turns** in this
-harness. Drive long GPU runs as foreground bounded batches, or `nohup … & disown` (still
-iffy). The pipeline is resumable, so just re-launch to continue.
-
----
-
-## Architecture (as-built, after this session's corrections)
-
-**Staged, chronological, resumable pipeline** (`src/memory/stages.ts`). Each stage is an
-independent DB-state-keyed worker, ordered `updated_at ASC` (chronological), resumable
-(skips done work). CLI: `mlx-bun memory segment|extract|route|synthesize-stage|link` (run
-separately/concurrently) — and `mlx-bun memory synthesize` runs the full DAG.
-
-```
-ingest → SEGMENT (our e4b-chunk-300 chunker) → EXTRACT (entities) → ROUTE (subject-based,
-surface-everything) → CREATE/PATCH (per-section bounded synthesis; self-healing) →
-RECONCILE (infobox-ground-truth consistency) → LINK (cross-link) → WIKIFY (editorial) → commit
-```
-
-- **Vault = truth** (markdown files in `articles/`); **DB = derived cache** (`reindex`
-  rebuilds it). The article FILES are the real artifacts. The `entities.article_stem`
-  column being unset mid-run is cosmetic (file-presence is the source of truth).
-- **Vaults:** `~/.mlx-bun/wiki` (the REAL vault — untouched), `~/.mlx-bun/wiki-full`
-  (import target), `wiki-smoke|slice|pattern|pattern2|link-test` (scratch).
-- **Model:** ONE base **Gemma-4-e4b** (`SNAPSHOT_E4B`). **Only ONE adapter trained so far:
-  `e4b-chunk-300`** (the chunker), symlinked `~/.cache/mlx-bun/adapters/memory-chunk`.
-  `callLocal(stage,…)` picks `memory-<stage>` adapter iff it exists, else base.
-- **Lucien is a READ-ONLY oracle/data-source** (`~/Code/lucien/.lucien/lucien.db` opened
-  readonly; `~/Dreaming`). We reuse NONE of its derived chunks/buckets — only the raw
-  conversation text (≡ a JSONL export / scraper). Never write there.
-
----
-
-## What we LEARNED / CORRECTED this session (load-bearing — don't re-derive)
-
-1. **NO notability/ownership filter.** Surface everything; notable = "anything that recurs
-   in your thinking" (Lucien `Editorial_Guidelines.md`). The plan's "owned/decided/≥3
-   chunks → PARK as trivia" was **fabricated by a prior agent and ripped out.** ROUTE now:
-   surface everything, strongly prefer an existing article, mint a new one only for a
-   substantive/recurring subject, `_captured` bucket (searchable) for the rest — **nothing
-   dropped.**
-2. **Subject-based articles = the one deliberate improvement over Lucien** (which drifted
-   into `X_and_Y` topic-bins). Verified: frameworks (Kaizen, TPS, LessWrong, Free energy
-   principle) get good articles; coverage resolves (`resolve("Kaizen")` → article).
-3. **Our chunker must be used.** The pipeline's `segment` was a no-op reusing Lucien's
-   chunks; fixed — `e4b-chunk-300` segments the raw conversations (better, topic-coherent).
-4. **Chronological ordering.** Was `ORDER BY id` (= conv UUID, random); fixed to
-   `updated_at ASC` everywhere. Load-bearing for self-healing.
-5. **Self-healing.** A later-dated correction flips the LEAD + INFOBOX + every body section
-   to the latest value; KEEPS the trajectory in a `## History` section **only for the
-   USER's evolving positions** — an AI *factual error* is silently overwritten (no
-   history). Date-aware; `reconcile.ts` is the consistency authority (infobox = ground
-   truth). Proven via a timestamp-swap demo (`scripts/memory/dreaming-selfheal-demo.ts`).
-6. **Parroting (recurred 5×).** The small model copies ANY concrete example value from
-   prompts/Meta pages into unrelated articles (phantom S5IIX, "Anamorphic Video", etc.).
-   **ALL prompts + Meta pages must be SCHEMATIC** (no real names/specs/dates). Purged;
-   defensive strips added. Keep this rule for any new prompt.
-7. **Cross-linking is WIKIFY's job, NOT synthesis** (synthesis updates ONE article, can't
-   know others). Built a dedicated `link` stage (`crosslink.ts`): mention-based (prose
-   names another article → `[[link]]`) + co-occurrence (shared folded chunks → See also).
-   Result on a 587-article snapshot: See-also 18%→96%, resolved edges 309→**7026 (22.7×)**.
-   Idempotent; fence-fix handles malformed input.
-8. **Metrics.** Do NOT measure articles-per-conversation (there is no conv→article link).
-   Measure **topic aggregation** (e.g. `El_Salvador` ← 30 conversations → 1 article).
-   Quality is judged by a **CLOUD JUDGE on real output**, never Lucien bucket-F1 (that's the
-   bin taxonomy we replace).
-9. **Inference path (IN PROGRESS — see `docs/design/memory-inference-path.md`):**
-   - **Template mismatch (fixed in working tree):** the chunk adapter was SFT'd with a
-     **system role** (Gemma-4 renders `<|turn>system…<turn|>` distinctly), but inference
-     dropped it (one user message). `callLocal` now passes `{system, user}` matching
-     training; the chunk stage uses `CHUNK_SYSTEM` (its exact trained system prompt). Also
-     noted: the live `CHUNK_PROMPT` drifted from training (only 2 of 8 Meta pages; abstracted
-     examples) — adapter still generalizes, so quality delta is **modest**.
-   - **Batching = DEAD END.** Measured ~**1.7× SLOWER** (warm, uniform: 0.607 vs 0.357
-     s/chunk). The `BatchScheduler` pads heterogeneous prefills and doesn't amortize
-     weight-loads. Default OFF (`MLX_BUN_MEMORY_BATCH=1`). Import is genuinely ~50h serial;
-     batching is NOT the lever (a *better* batched-decode that amortizes weights could help
-     in principle, but that's a real perf project, not worth it for a one-time import).
-   - **Load-once / persistent server = the REAL efficiency win, for DAILY use.** Each CLI
-     invocation today reloads the model + pays warmup; routing memory generation through the
-     **resident server** (model loaded once) kills that. Within the single-process import
-     the model is already loaded once, so it doesn't speed the import.
-
-## Product framing (from the user)
-- **Daily local conversations (pi sessions) folding in is the product** — watermark-gated
-  ingest, minutes/day, self-healing. The 4-year Claude/ChatGPT import is a one-time
-  **onboarding feature + stress test**. Throughput only matters for the import.
-
----
-
-## NEXT STEPS (prioritized)
-
-1. **Finish the inference-path rework.** Templating fix is in the working tree (verify it
-   end-to-end). **Wire memory generation through the persistent server (load-once)** — the
-   user's actual ask — keeping a no-server fallback; default batching OFF. Verify the chunk
-   stage decodes on-distribution (system role present).
-2. **Import is RUNNING** (resumed from 720, serial, corrected templating). Just let it finish.
-   Re-segmenting the first 720 is **low-urgency** (old chunks were valid).
-3. **After import:** `reindex` → `mlx-bun memory link` → browse `wiki-full` in Obsidian;
-   cloud-judge a broad, diverse sample (not just frameworks/gear).
-4. **Promote to the real vault** when satisfied: point `~/.mlx-bun/wiki` at the bootstrapped
-   content; wire the **nightly daily fold-in** from pi sessions (`schedule.ts` exists).
-5. **Open quality items** (low priority): ~<1% articles leak a prompt fragment; some leads
-   run on; the `owned/used/chosen` framing in `Article_Conventions` lead spec can put
-   `owned: yes` on abstract concepts (reframe).
-6. **Commit when ready** — everything is uncommitted; the user has not asked to commit.
-
-## Gotchas / facts
-- `bunx tsc` sometimes can't resolve typescript in this checkout; use global `tsc`. Repo = 0 errors.
-- The `mlx-bun train-watch orpo-e4b-chunk` process is a **stale idle daemon, not training** — GPU is free.
-- e4b training still OOMs on this 32GB M1 Max; P10 adapter training stays USER-ACTION (segmented-backward).
-- The `Reference/` folder in vaults symlinks the mlx-bun repo docs (`setupVault`) — non-portable; fine for now, drop for the real personal vault.
-- Methodology memories: cloud-judge not bucket-F1; Lucien read-only; surface-everything notability.