mlx-bun 0.4.0 → 1.0.0

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

@@ -1,157 +0,0 @@
-# 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](../../docs/archive/mlx-bun-lab-report.html). Evidence
-artifact: `benchmarks/lab-verification.json` (regenerable, gitignored — produced by
-`scripts/verify-lab.ts`).*
-
-## Goal
-
-Re-implement the whole mlx-bun web surface as one beautiful, keynote-grade SPA
-with five sections — a **full pi-agent chat**, plus **Quantize**, **Fine-tune**,
-and **Build Dataset** ported from the OptIQ "lab" — while keeping the status
-page. Compute is **native Bun/MLX (no Python)** per Josh's decision.
-
-## The decisions that shaped it
-
-1. **The big fork (asked up front).** Quantize + fine-tune compute could be (a)
-   native via mlx-c, (b) orchestrate the installed OptIQ Python, or (c) proxy
-   the running `optiq lab`. Josh chose **(a) native** — truest to mlx-bun's "the
-   buggy layer doesn't need Python" thesis. Chat = **full pi agent** (all tools).
-   Verification = **real tiny runs, pausing for go-ahead** before heavy jobs.
-
-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 `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
-   trainer worth writing. This single check de-risked the entire subsystem.
-
-3. **One SPA, hash-routed, text-imported.** Kept the exact single-binary
-   embedding (`with { type: "text" }`) the status page already used — no
-   bundler, no asset-resolution to patch in `build-binary.sh`. Design tokens
-   lifted from `status-page.html` so the five sections feel like one product.
-
-4. **Mixed job-execution model by safety class.** Async Metal command-buffer
-   OOM is **uncatchable** (it `std::terminate`s). So GPU jobs (quantize,
-   finetune) run as **crash-isolated child `mlx-bun` subprocesses** behind a GPU
-   lease; dataset jobs (pure JS + loopback HTTP) run **in-process**. Both write
-   the same NDJSON event log + SQLite row, so the SSE/tail path is uniform.
-
-5. **Open-book ports, not inventions.** Each engine follows a readable source:
-   dataset → `optiq/lab/dataset_templates.py` (13 templates, field-for-field);
-   trainer → `optiq/lora/{trainer,dpo,config}.py` + mlx-lm `default_loss`;
-   quantize → mlx-lm `convert.py` rules + the on-disk OptIQ model format; jobs →
-   `optiq/lab/{jobs,db}.py`; chat → pi SDK `examples/sdk/12-full-control.ts`.
-
-## How it was built — parallel orchestration
-
-After laying the shared job contract (`src/jobs/types.ts`), six subsystems were
-fanned out to concurrent sub-agents on **non-overlapping files**, then
-integrated centrally:
-
-| subsystem | files | how verified |
-|---|---|---|
-| shared jobs | `src/jobs/*` | 12 tests (lifecycle, GPU lease, crash isolation, zombie, SSE) |
-| dataset | `src/dataset/*` | 32 tests (generators, CSV, 90/10 split, 13 templates) |
-| quantize | `src/quantize/*` | 6 tests + real-model e2e |
-| trainer | `src/train/*`, `src/mlx/autograd.ts` | 19 fast + real-model e2e |
-| web SPA | `src/web/app.html` | 7 in-process route tests |
-| pi embed | `src/pi-web.ts` | 9 unit + live chat turn |
-
-Shared/owned-by-me files were edited centrally to avoid conflicts:
-`src/server.ts` (websocket + Lab routes + serve the SPA), `src/mlx/{ffi,ops}.ts`
-(autograd/random/save bindings), `src/cli.ts`, `package.json` (pinned
-`@earendil-works/pi-coding-agent@0.79.3`).
-
-## Key technical findings
-
-- **`mlx_save_safetensors` exists** (mlx-c `io.h`) — the quantizer writes real
-  HF-layout shards natively; no hand-rolled byte writer, no dtype/alignment bugs.
-- **`mlx_value_and_grad` + `mlx_closure_new_func_payload`** are drivable from a
-  `bun:ffi` JSCallback (the repo already uses one for the mlx error handler).
-  The closure swaps the differentiated primals into `LoraWeights.{a,b}`, builds
-  the loss graph, and restores; mlx's graph keeps its own refs, so only A/B
-  differentiate and the quantized base stays frozen.
-- **No optimizer in mlx-c** → AdamW hand-rolled as array ops (m/v EMA, bias
-  correction, decoupled weight decay), each step swapping the leaf handle back
-  into the `LoraWeights` so the next forward sees the update.
-- **pi tool-approval gate** is a real pre-execution hook: an inline extension
-  via `DefaultResourceLoader.extensionFactories` registers `pi.on("tool_call")`
-  returning `{block:true}` to deny — read-only tools auto-allow, bash/edit/write
-  await a browser decision (120s auto-deny).
-
-## Bugs found and fixed (during the real-model runs)
-
-- **LoRA init key shape** — `mlx_random_normal` got a `[1,2]` key (a row slice of
-  the split-keys array) but mlx wants `[2]`. Reshaped the per-module subkey.
-- **Quantize e2e path** — the gated test passed a model *id* to
-  `quantizeModelDir` instead of a resolved snapshot *path* (the job runner
-  resolves via the registry; the direct call didn't). Pointed it at
-  `SNAPSHOT_MINICPM5`.
-- **Flaky chat test** — the agent sometimes takes a tool-call turn before
-  emitting text; the test resolved on the first `turn_end` and saw no tokens.
-  Fixed to settle on the first turn that actually streams text. The embed was
-  always correct (a debug run showed 25+ text deltas + a tool execution).
-
-## Measured evidence (M4 Pro, MiniCPM5-1B, on-disk)
-
-- **Quantize:** 170 modules → **4.50 bits/weight**, **0.62 GB** dir, **7.0 s**,
-  reopened correctly through the loader (quantized embed_tokens as `uint32`).
-- **Train:** 20 SFT iters, **loss 4.48 → 0.68** in **4.6 s**, 168 adapters
-  mounted. Behavior changed — on "Tell me about your day": base replies normally,
-  the UPPERCASE-fixture adapter replies `"MY DAY AHEAD. APPROVE AND PROPERLY…"`.
-- **Chat:** a real WebSocket turn streamed text + executed a tool to completion.
-- **Tests:** 85+ new, all green; `tsc --noEmit` clean; existing server suite
-  17/17 (no regressions).
-
-## Scope: shipped vs designed-seam
-
-- **Shipped (v1):** full SPA (5 sections), pi web chat with tools + approval
-  gate, native uniform-affine quantize, native SFT + DPO LoRA, 12/13 dataset
-  templates, shared job system, real tiny-run verification.
-- **Designed seams (follow-ups):** ~~OptIQ sensitivity+knapsack mixed-precision
-  quantization; batched (B>1) training; `hf_dataset_import`; HF push-to-hub;
-  folding pi's assets into the compiled single binary.~~ **ALL DONE in Round 2 —
-  see below.**
-
-## Round 2 — the leftover seams, ported 1:1 (2026-06-14)
-
-Josh: "all of those, port, one to one functionality." Built the five remaining
-items as faithful ports of the OptIQ / mlx-lm source (open book), in parallel on
-non-overlapping files, integrated centrally (server routes + one `app.html` UI
-pass), and verified.
-
-- **Mixed-precision quantization** (`src/quantize/{allocator,sensitivity,
-  calibration}.ts`) — `optimize_mixed_precision` ported **1:1** from
-  `optiq/core/optimizer.py` (greedy knapsack + block-aware floor + protected
-  layers + block-run guard + the uniform-vs-bf16 `_kl_reduction` sign detection).
-  Verified **bit-identical to OptIQ** on shared synthetic inputs by spawning the
-  venv Python in the test (4/4 cases). Per-layer KL sensitivity
-  (`analyze_sensitivity_exact`) runs on our graph: dequantize each layer's weight
-  to bf16, re-quantize at each candidate bit, swap into the running
-  `QuantizedLinear`, forward on a WikiText-2 calibration set, KL vs reference.
-  Real e2e: MiniCPM5-1B at target 5.0 bpw → 80×4-bit + 90×8-bit, **5.19 bpw, 32 s**.
-- **Batched (B>1) training** (`src/train/{dataset,forward,loss}.ts`) — ragged
-  padding + a `[B,1,L,L]` causal-AND-key-validity mask via a cache wrapper (no
-  model/op edits); B=1 stays a separate bit-identical path. Bug found: comparing a
-  length-4 reference to a length-6 batched row conflated mask correctness with
-  bf16 sequence-length rounding — the corrected shape-held parity test passes,
-  confirming the mask isolates rows.
-- **`hf_dataset_import`** (`src/dataset/generators.ts`) — the 13th template via
-  the HF datasets-server REST API (pagination, config resolve, label/min-chars/
-  cap filters, three output formats).
-- **HF push-to-hub** (`src/hf-push.ts`) — native Hub upload (create repo → LFS
-  preupload → S3 PUT → NDJSON commit) for model/adapter/dataset repos, token at
-  `~/.mlx-bun/hf.json`. Mock-server tested end-to-end (11 tests); real uploads
-  untested by design (would publish under the user's account).
-- **Single-binary web chat** (`scripts/build-binary.sh`) — pi's JS already
-  bundles under `bun build --compile`; only `photon_rs_bg.wasm` is sidecarred.
-  The compiled binary's headless pi session was verified to construct.
-
-UI: one `app.html` pass enabled the Uniform/Mixed quantize toggle + target-bpw
-slider, live merge/export/push on the fine-tune done step, push on quantize &
-dataset, and an HF-token settings gear. New server routes:
-`/api/finetune/{merge,export}`, `/api/{quantize,finetune,dataset}/push`,
-`/api/settings/hf-token`, + mixed-precision params on quantize submit.

package/docs/archive/investigations/minicpm5-decode-megakernel.md

@@ -1,489 +0,0 @@
-# MiniCPM5 Decode Megakernel — end-to-end build plan (handoff)
-
-**Read this top to bottom, then execute the phases in order. Each phase has an
-exact validation gate; do not advance until it passes. At the end you have one
-resident Metal kernel that computes the entire decode forward (all 24 layers →
-logits) in a single dispatch, never handing MLX's planner a graph to mis-plan.**
-
-This document is self-contained. You should not need to ask anything to execute it.
-
----
-
-## 0. Why this exists (the problem, settled)
-
-The inference forward composes per-op MLX nodes (`quantized_matmul`, `sdpa`, …).
-We tried to insert our own fused `fast.metal_kernel`s (steel `BlockMMA` GEMM) into
-that forward to fuse+speed it up. **Confirmed dead end:** MLX's lazy buffer planner
-corrupts *chained* `BlockMMA` custom kernels (nondeterministic NaN/garbage/crash) —
-the pointwise `fused-gelu` kernel chained fine (that kernel has since been deleted,
-2026-07-05, in favor of the compiled activations — unified-engine-frontier-plan.md
-§6), ours did not, and ≥2 of ours in one lazy graph corrupt. The only robust fixes are (a) a per-kernel eval boundary (24 syncs/
-forward → net **slower**, 0.5–0.8×), or (b) **don't give the planner anything to
-chain: one resident kernel for the whole forward.** This plan builds (b).
-
-The training path already proved the shape: its forward is MLX ops (which chain
-fine) + **one** terminal fused kernel (flash-CCE loss), `detachLeaf` at coarse
-segment boundaries (`src/train/segmented.ts`). Our megakernel is the inference
-analogue — but because inference has no terminal heavy-reduction, the "one kernel"
-must be the *whole forward*, resident, coordinating its own stages.
-
-**Scope = DECODE (single new token, M=1).** This is the only case physically
-expressible as one dispatch: a single Metal dispatch cannot barrier across
-threadgroups, so cross-token prefill attention can't be one kernel — but one token
-attending to the *cached* KV has no cross-token dependency. Decode is also the
-latency path that *is* serving speed (`[[serving-speed-is-the-user-metric]]`).
-Prefill stays on the existing (correct, fast, parity) MLX-ops path.
-
-**Honest perf expectation:** decode is weight-bandwidth-bound (reading ~1 GB of
-4/8-bit weights per token is the floor for any method). The megakernel's win is
-(1) eliminating ~240 per-op dispatch launches/token, (2) no intermediate device
-round-trips (hidden stays resident), (3) zero MLX-planner involvement. Expect a
-modest but real decode-latency win (overhead elimination), not a bandwidth
-miracle. The *architectural* win (no corruption, fully ours) is the point.
-
----
-
-## 1. Success criteria (definition of done)
-
-1. A function `decodeMegakernel(model, tokenId, position) → logits[V]` that runs the
-   entire MiniCPM5 decode forward in **one** `MetalKernel.apply` dispatch.
-2. **Correctness gate (L1, bf16 KV):** teacher-forced greedy trajectory matches
-   `goldens/minicpm5-parity.json` — argmax agreement **≥ 98/100**, mean per-step
-   logit maxDiff small (this is a no-oracle Lab-gated path: ~1–2 bf16 ULP/op
-   accumulating; it will NOT be bit-exact, gate by teacher-forced agreement + KL —
-   the frozen-trajectory envelope method of the old `tests/perf-kernel-oracle.test.ts`,
-   threshold ≥ 56/64-style; that test was deleted 2026-07-05 with its kernel, but
-   the methodology stands on its own). DETERMINISTIC
-   across runs (no nondeterminism — that was the corruption signature).
-3. **Perf gate:** decode tok/s ≥ the serial baseline on a paired same-process A/B
-   (`scripts/experiments/` style). Not slower. (Stretch: measurably faster.)
-4. Default **OFF** behind `MLX_BUN_MEGAKERNEL=1`; the existing forward untouched;
-   `bun test tests/minicpm5-*.test.ts` stays green with it off.
-5. Wired into the real decode loop (`src/generate.ts`) for `--megakernel` / env.
-
----
-
-## 2. Hard facts (MiniCPM5-1B-OptiQ-4bit) — do not re-derive
-
-**Snapshot:** `/Users/joshrossi/.cache/huggingface/hub/models--mlx-community--MiniCPM5-1B-OptiQ-4bit/snapshots/664aabaed233c653f82716d8dc822234d0091f78`
-
-**Config:** `hidden_size H=1536`, `num_hidden_layers=24`, `num_attention_heads=16`,
-`num_key_value_heads=2` (GQA, `nRep=8`), `head_dim=128`, `intermediate_size I=4608`,
-`vocab_size V=130560`, `rms_norm_eps=1e-6`, **`rope_theta=5_000_000`** (NOT 10000 —
-the code reads `t.ropeParameters.full_attention?.ropeTheta`), `tie_word_embeddings=false`.
-`attention scale = head_dim^-0.5 = 128^-0.5`. Plain Llama: **no MiniCPM residual
-scaling** (verified — `LlamaLayer.forward` is `inputNorm→attn→+x→postAttnNorm→mlp→+h`).
-
-**Per-layer attention output dims:** q_proj out = `16*128=2048`; k_proj/v_proj out =
-`2*128=256`; o_proj in = `2048`, out = `1536`. MLP: gate/up out = `4608`, down in =
-`4608`, out = `1536`. lm_head out = `130560`, in = `1536`. embed = quantized 8-bit.
-
-**Per-TENSOR mixed precision (weight quant):** every linear has its OWN `{bits,
-group_size}` baked in `config.quantization` (group_size always 64; bits ∈ {4,8}).
-This is why the megakernel must be **generated/de-branched per model** — bits is a
-compile-time literal per (layer, matrix), not a runtime value. Full table is in
-`config.json` under `quantization`; read it with `quantFor(config.quantization,
-path)` (`src/model/gemma4-base.ts`). Examples: L0 all 8-bit; L1 q4/k4/v8/o8/g4/u4/d8;
-L7 q8/k4/v4/o8/g4/u8/d4; L23 all 8-bit. **Read the table programmatically per build —
-do not hardcode.**
-
-**Quantized weight layout** (`QuantizedLinear`, `src/model/gemma4-base.ts:76`):
-- `.w`: packed `uint32` `[N_out, K_in*bits/32]`. `.scales`,`.biases`: `bf16`
-  `[N_out, K_in/group_size]` (affine: `deq = scale*q + bias`).
-- Dequant math is in `STEEL_QMM_HEADER` (`dequantize<T,N,bits>`, `get_pack_factor`,
-  `get_bytes_per_pack`) — `src/train/steel-qmm-header.ts`. REUSE IT.
-
-**KV cache:** L1 path = `KVCache` (bf16, `src/model/gemma4-base.ts:243`); its
-`updateAndFetch(k,v)` appends `[1,heads,1,D]` at `cache.offset` and returns the full
-`[1,heads,offset+1,D]`. `cache.offset` is the position. L2 path = per-layer
-`QuantizedKVCache` (mixed bits from `config.kvQuant`) — DEFER to Phase 4.
-
-**RoPE** (`ops.rope`, `src/mlx/ops.ts:100`): `rope(x[B,heads,L,D], dims=head_dim,
-base=rope_theta, offset, scale=null)`. For decode L=1, offset=position. The
-in-kernel formula (standard NeoX/Llama interleaved-half): for dim pair `i<D/2`,
-`theta_i = position / base^(2i/D)`; rotate `(x[i], x[i+D/2])` by `theta_i`. **Match
-`ops.rope`'s exact convention** — confirm against `src/mlx/ops.ts` rope binding and
-a 1-token parity probe before trusting it (Phase 2 gate catches mismatch).
-
-**Validation goldens (ALREADY EXIST, no regen needed):**
-- L1: `goldens/minicpm5-parity.json` (prompt_ids `[0,608,4894,304,6918,357]`,
-  greedy_ids[100], vocab 130560) + `goldens/minicpm5-logits-step{0..99}.bin`
-  (f32 last-position logits). Produced by `scripts/regen.ts minicpm5`.
-- L2: `goldens/minicpm5-kv-parity.json` + `…-kv-logits-step*.bin`.
-
-**Reference forward to copy (the source of truth for correctness):**
-`src/model/minicpm5.ts` — `LlamaAttention.forward` (95–193), `LlamaMLP.forward`,
-`LlamaLayer.forward`, `MiniCPM5Model.forwardHidden/forwardLayers/logitsFromHidden`.
-Transcribe THIS op-for-op into the kernel.
-
----
-
-## 3. Already-built, validated building blocks (reuse, don't rebuild)
-
-All default-OFF, `bun test tests/minicpm5-*.test.ts` green. From this session:
-- `src/train/steel-qmm-header.ts` — `STEEL_QMM_HEADER`: MLX's verbatim quantized
-  `BlockMMA` GEMM + `dequantize`/`QuantizedBlockLoader`/`BlockLoader`. The GEMM
-  building blocks for every projection. `qmm_t_impl` (line 1334) is the full GEMM
-  reference (takes dims as `constant int&` — for our metal_kernel we inline the
-  `BlockMMA` loop with templateInt dims instead, as `FWD_STEEL_SOURCE` does).
-- `src/model/fused-mlp-kernel.ts` — `fusedMlp`: whole MLP (gate+up+silu+down +
-  absorbed residual) in ONE dispatch, `hidden` never materialized, per-loader
-  mixed-precision bits baked, atomic-accumulated f32 output. **1.9–3.4× standalone
-  @ M=8–256, ~1 ULP correct.** This is a working prototype of one megakernel STAGE
-  and the dequant/GEMV patterns to copy. (Note: it currently corrupts when CHAINED
-  in the lazy graph — that's the whole reason for the megakernel; the kernel math
-  itself is correct, proven by `fused-mlp-check.ts (deleted; git history)`.)
-- `src/model/steel-linear-kernel.ts` — `steelLinear`: plain quantized linear as our
-  dispatch (q/k/v/o/down), ~1 ULP vs `ops.quantizedMatmul`. The projection GEMV
-  pattern.
-- `src/mlx/metal-kernel.ts` — `MetalKernel`: `apply(inputs, {outputShapeFn|outputs,
-  grid, threadGroup, templateInts, templateDtypes, atomicOutputs, initValue})`.
-  `*_shape` helpers available in-kernel. `metalCapture` for Xcode GPU traces.
-- `src/train/segmented.ts:81` `detachLeaf` (contiguous→eval→`fromBytesCopy`) — the
-  graph-break primitive (replicated in `minicpm5.ts`).
-
-**Probe scripts (templates for new validation):** `scripts/experiments/swiglu-*.ts`,
-`fused-mlp-*.ts`, `steel-linear-check.ts`, `swiglu-teacherforced.ts` (the
-teacher-forced gate harness — copy it).
-
----
-
-## 4. Architecture of the megakernel
-
-### 4.1 The shape
-
-One resident dispatch. State lives in **device** buffers (shared across
-threadgroups); the kernel walks layer 0→23 as a sequence of **stages**, with a
-**grid-wide barrier** between stages so stage N+1 sees all of stage N.
-
-Per layer, the stages (each a grid-wide-parallel step, barrier after):
-1. `xn = RMSNorm(hidden, inputNorm_w)`  → device `xn[H]`
-2. `q = Wq·xn` (`[2048]`), `k = Wk·xn` (`[256]`), `v = Wv·xn` (`[256]`) — quantized
-   GEMVs, outputs split across threadgroups → device
-3. `rope(q, pos)`, `rope(k, pos)` (per head, D=128)
-4. write `k,v` into KV cache at `pos` (device)
-5. **attention:** for each of 16 q-heads (GQA: head h uses kv-head `h/8`), online
-   softmax of `q_h · K[0..pos]` over the `pos+1` cached keys (scale `128^-0.5`),
-   weighted sum of `V[0..pos]` → `attn[2048]`
-6. `o = Wo·attn` (`[1536]`); `hidden += o` (residual)
-7. `hn = RMSNorm(hidden, postAttnNorm_w)` → device `hn[H]`
-8. MLP: `gate = Wg·hn`,`up = Wu·hn` (`[4608]`), `silu(gate)·up`, `down = Wd·(…)`
-   (`[1536]`); `hidden += down` (residual)  — the `fusedMlp` body, inlined
-After 24 layers: `RMSNorm(hidden, finalNorm_w)` → `lmHead·hidden` → `logits[V]`.
-
-### 4.2 The grid-wide barrier (the crux)
-
-A single dispatch has no built-in cross-threadgroup barrier. Implement a
-**persistent-kernel software barrier** via two device atomics (`atomic_uint`):
-`arrive` (count) + `sense` (phase). Each barrier: `tg_leader` does
-`atomic_fetch_add(arrive,1)`; if it was the last (`==G`), reset `arrive=0` and flip
-`sense`; all threadgroups spin-read `sense` until it flips; `threadgroup_barrier`
-to fan back into the threadgroup. **HARD REQUIREMENT:** all `G` threadgroups must be
-**co-resident** or this deadlocks (a threadgroup waiting on one never scheduled).
-So `G` must be ≤ the GPU's max concurrent threadgroups.
-
-Pick `G` conservatively for the M4 Pro (~16–20 GPU cores). Phase 3 determines the
-safe `G` empirically (start `G=1`, raise until just before the largest occupancy;
-if any deadlock/timeout, lower). Use `dispatchThreads`/grid = `G * TGN`. Provide an
-env `MLX_BUN_MEGAKERNEL_G` to sweep. Document the chosen `G`.
-
-> **De-risking:** Phase 2 builds `G=1` (single threadgroup, `threadgroup_barrier`
-> only — NO software barrier, hidden in threadgroup memory) to get CORRECTNESS
-> first. Phase 3 generalizes to `G>1` with the software barrier for SPEED. Keep the
-> per-stage math identical between them.
-
-### 4.3 Weight access — the repack (Phase 1)
-
-A `metal_kernel` can't take ~500 tensors. Repack at load time into a SMALL fixed set
-of device buffers + an offset/spec table:
-- `WBYTES`: one `uint8` buffer = every quantized weight's packed bytes concatenated
-  in a fixed order (L0.q, L0.k, … L23.lm_head). Plus `SCALES` (bf16 concat),
-  `BIASES` (bf16 concat). Plus `NORMS` (bf16 concat of all RMSNorm weights).
-- `META`: a `uint32`/`int32` table the kernel indexes by `(layer, matrix)` giving:
-  byte-offset into WBYTES, scales offset, biases offset, N, K. **bits/group_size are
-  baked as template literals per (layer,matrix) in the GENERATED kernel**, not in
-  META (they drive the dequant template instantiation).
-- Because bits varies per matrix, the kernel is **generated** (like
-  `scripts/gen-model.ts`): emit a `forwardLayer_L<i>` inline body per layer with that
-  layer's 7 matrices' bits baked, OR a single body that `switch`es on a baked
-  per-layer-constant array. Prefer generation: produces
-  `src/model/generated/minicpm5-megakernel.ts` with all bits as literals.
-
-KV cache buffers: `KBUF`,`VBUF` device `[24, max_seq, 2, 128]` bf16 (L1). `pos` is a
-scalar input. Logits out: `[V]` f32.
-
----
-
-## 5. Phased plan (execute in order; gate each)
-
-### Phase 1 — Weight repack (foundation)
-**File:** `src/model/megakernel-pack.ts`.
-**Do:** `packMiniCpm5(model) → { wbytes, scales, biases, norms, meta, layout }`.
-Walk `model.layers[i]` (`.attn.{qProj,kProj,vProj,oProj}`, `.mlp.{gate,up,down}`,
-`.inputNorm`,`.postAttnNorm`) + `model.finalNorm`,`model.lmHead`,`model.embed`. For
-each quantized linear concat `.w` bytes → WBYTES, `.scales`→SCALES, `.biases`→BIASES
-(record offsets, N, K, bits, group in `layout`). Concat norm weights → NORMS. Use
-`MlxArray.rawBytes()` (eval first) to get bytes; build with `Bun`'s buffers; upload
-as `MlxArray.fromBytesCopy`.
-**Gate:** unit test — for 3 sample (layer,matrix), slice WBYTES/SCALES/BIASES at the
-recorded offset, rebuild a `QuantizedLinear`, run `ops.quantizedMatmul` on a fixed
-input, assert **byte-identical** to the original `.forward`. (`scripts/experiments/
-megapack-check.ts`.) This catches layout bugs before any kernel.
-
-### Phase 2 — Single-threadgroup correct megakernel (G=1)
-**Files:** `src/model/megakernel-kernel.ts` (the `MetalKernel` + dispatch),
-`gen-minicpm5-megakernel.ts (deleted; git history)` (emits the per-layer-bits-baked SOURCE into
-`src/model/generated/minicpm5-megakernel.ts`).
-**Design:** ONE threadgroup (`G=1`, `TGN=256` threads). `hidden[H]` in threadgroup
-memory (6 KB). Loop layers; each stage uses `threadgroup_barrier`. GEMVs: the 256
-threads cover the N outputs in chunks (each thread does `ceil(N/256)` dot-products
-over K, dequantizing weight via `STEEL_QMM_HEADER`'s `dequantize`). RMSNorm: parallel
-sum-of-squares reduction in threadgroup. Attention: loop heads; per head, online
-softmax over `pos+1` cached keys (threads split the key positions, threadgroup
-reduce). Inline the `fusedMlp` math for stage 8 (gate/up into threadgroup, silu·up,
-down). KV read/write from `KBUF`/`VBUF`. Final: finalNorm + lmHead GEMV → logits[V]
-(threads split V).
-**Transcribe** `src/model/minicpm5.ts` op-for-op; match dtypes: accumulate in f32,
-round to bf16 at each op boundary to track the reference (`(T)` casts, as `fusedMlp`
-does). RoPE: match `ops.rope` exactly (verify with a single-layer probe).
-**Gate (THE correctness gate):** `megakernel-teacherforced.ts (deleted 2026-08-23; git history)`
-(copy `swiglu-teacherforced.ts`): drive 100 decode steps via the megakernel,
-teacher-forced on `goldens/minicpm5-parity.json`, KV cache in `KBUF`/`VBUF`. Require
-**≥98/100 argmax agreement, deterministic across 3 runs, no NaN.** If an early step
-mismatches, bisect by stage: add a debug output for `hidden` after layer L (env
-`MLX_BUN_MEGA_DUMP_LAYER=L`) and compare to the reference model's per-layer hidden
-(capture via the wrapper trick in `swiglu-layerN.ts (deleted; git history)`).
-
-### Phase 3 — Persistent multi-threadgroup (G>1) for SPEED
-**Do:** generalize Phase 2 to `G` threadgroups. Move `hidden`,`xn`,`hn`,`q/k/v/attn/
-gate/up` to DEVICE scratch (shared). Implement the §4.2 software barrier (`arrive`/
-`sense` atomics, leader-flips-phase, spin). Split each stage's outputs across
-threadgroups by `tid.x` block (GEMV: threadgroup `g` does N-rows `[g*N/G,
-(g+1)*N/G)`). Determine safe `G`: sweep `MLX_BUN_MEGAKERNEL_G` from 1 up; the largest
-that (a) never deadlocks/timeouts and (b) maximizes tok/s. Add a 1-second host-side
-watchdog around the dispatch during the sweep to catch deadlock (kill + report).
-**Gate:** same teacher-forced ≥98/100 + deterministic, AND paired tok/s ≥ baseline.
-If `G>1` corrupts (barrier race), the barrier impl is wrong — verify the leader
-reset ordering and that EVERY thread (not just leaders) waits on `sense` before
-reading shared state, with a `threadgroup_barrier` + `mem_flags::mem_device` after.
-
-### Phase 3 — RESULTS (2026-06-24)
-
-Multi-threadgroup persistent kernel built + correct. Findings, in order discovered:
-
-- **Occupancy sweep:** TGN must be a **power of 2** (the RMSNorm tree reduction
-  `for(s=TGN/2;s>0;s>>=1)` silently drops lanes otherwise — TGN=192 gave KL 0.54).
-  Smaller TGN packs more co-resident threadgroups (TGN=128 scales past where TGN=256
-  collapses). Peak total threads ≈ 12–16k; G beyond co-residency **deadlocks** the
-  software barrier (G=144 @ TGN=128 hung). Sweep needs a host watchdog.
-- **Barrier cost** (measured, `megakernel-barrier-time.ts`): ~3.8 µs/barrier @ G=48,
-  ~145 barriers/token after fusion → ~0.5 ms = ~10% of decode. Not the dominant cost.
-- **bits-literal specialization REJECTED for the *naive* GEMV.** Templating on BITS
-  and/or K bloated code/registers → lowered max co-resident G → regressed peak
-  (164→102–156; full-K-unroll collapsed to 16). The kernel was occupancy/bandwidth-
-  bound, not branch-bound. **Lesson: feed the bottleneck (DRAM-latency-hiding =
-  occupancy×ILP), not instruction count.**
-- **THE win — port mlx's `qmv_fast` decode GEMV verbatim** (`quantized.h`
-  qmv_fast_impl/load_vector/qdot → `qmv4`/`load_vec_tg`/`qdot_t`). Three things ours
-  lacked: (1) `load_vector` pre-scales x so `qdot` is mask-only (no per-nibble
-  shifts); (2) affine folds once per group as `scale*accum + sum*bias`; (3) **4 rows
-  per simdgroup reusing a register-resident x_thread** → 4× arithmetic per x-load +
-  4 independent accumulators hiding FMA/load latency. Our weight layout already
-  matches mlx's, all N,K divide the fast-path blocks → no safe-tail path. This is the
-  M=1 decode primitive that was never ported (steel-qmm-header is the M>1 GEMM/BlockMMA
-  path). **Jumped 0.70× → ~0.93–0.95× baseline; KL improved to 9.5e-4; removed the
-  G≥24 down-tiling constraint (any G now correct).**
-- **Barrier fusion:** RMSNorm computed in-place per-threadgroup (`rms_norm_local`,
-  every TG normalizes the full hidden — tiny, removes the rms→GEMV grid barrier +
-  d_nrm device traffic) and SwiGLU `silu(gate)·up` folded into the down-proj's input
-  staging. ~145 barriers/token (from ~194). Perf-neutral at G=32 (barriers already
-  hidden) but removes device round-trips.
-- **Generated (layer-unrolled, constants-baked) kernel built** (`genSourceMt`, env
-  `MLX_BUN_MEGAKERNEL_GEN=1`): per-(layer,matrix) bits/offsets/N/K + layer index L as
-  compile-time literals, no META lookups, no runtime bits dispatch. Layer loop
-  unrolled (sequential → registers reused, NO occupancy hit, unlike the inner-K
-  unroll). **Correct (KL identical) but perf-NEUTRAL** (202.4 vs 202.9 tok/s) —
-  confirms the kernel is bandwidth/barrier-bound, not lookup-bound. Kept behind the
-  flag as the architecturally-correct path.
-- **Peak config:** TGN=256, G=32. Paired same-process A/B: megakernel **~202 tok/s**
-  vs baseline **~216 tok/s = ~0.94×**. Correctness 97/100 teacher-forced (3 argmax
-  near-ties; KL 9.7e-4), deterministic, no NaN.
-- **Remaining gap (~5%) is the resident-kernel structural tax** (~6 grid barriers/
-  layer for the row-split GEMV producer→consumer chain + per-stage staging), which is
-  largely fundamental to multi-threadgroup decode. Candidate lever to cross baseline:
-  fold rope-q into attention to drop one barrier/layer. Needs a CLEAN-machine paired
-  measurement (Josh-gated) — the session machine is loaded, only the ratio is trustworthy.
-
-### Phase 4 — L2 quantized KV (IN PROGRESS 2026-06-24, spec ready)
-Swap `KBUF/VBUF` for the per-layer quantized-KV format. **Facts gathered:**
-- CPM5 `kv_config.json`: per-layer `bits ∈ {4,8}`, `group_size=64`, quantized along
-  **head_dim=128** (so 2 groups per head per k/v). Loaded as `config.kvQuant`.
-- Layout (`QuantizedKVCache`, gemma4-base.ts:407): per (layer,pos) store quantized
-  triple along the KVDIM=256 axis: `packed` uint32 `[256*bits/32]`, `scales`/`biases`
-  bf16 `[256/64=4]`. Affine `deq = scale*q + bias` (== our weight dequant).
-- **mlx quantize formula to replicate EXACTLY** (`affine_quantize`, quantized.h:2432;
-  this is what the golden used via optiq `updateAndFetchQuantized`):
-  ```
-  n_bins=(1<<bits)-1; eps=1e-7
-  wmin=min(group), wmax=max(group)
-  scale=max((wmax-wmin)/n_bins, eps); side=|wmin|>|wmax|; scale = side? scale : -scale
-  edge = side? wmin : wmax;  q0=round(edge/scale)
-  scale = q0==0? scale : edge/q0;  bias = q0==0? 0 : edge
-  q = min(round((w-bias)/scale), n_bins)     // dequant: scale*q+bias
-  ```
-**Implementation steps:**
-1. ✅ **DONE — quantize formula validated** (`kv-quant-formula-check.ts (deleted 2026-08-23; git history)`):
-   the affine_quantize transcription matches `ops.quantize` to bf16 ULP (scale/bias diff
-   1.8e-3 @4-bit, 6e-5 @8-bit; dequant residual = expected bf16-storage rounding). The
-   riskiest piece is de-risked. Compute scale/bias in f32, store bf16 (as mlx does).
-2. **Host:** `MegakernelRunner` holds per-layer quantized KBUF/VBUF (packed/scales/
-   biases) sized to kvSeq; bits per-layer baked (literal in the generated kernel).
-3. **Kernel write:** after rope, each head's roped k/v (128 dims) → 2 groups of 64 →
-   compute wmin/wmax (threadgroup or per-head reduce) → scale/bias → pack q → store at
-   current pos. Replaces the bf16 `knewOut/vnewOut` emit.
-4. **Kernel read (attend_simd):** dequant EVERY position incl. current (optiq quantizes
-   then reads back, so the current entry must also round-trip) from packed+scales+biases.
-
-**Increment 1 (DONE, bf16-storage round-trip): 93/100, KL 1.49e-2, deterministic, no NaN**
-(`kv_rt_group<KVBITS>` after rope, `megakernel-kv-teacherforced.ts`,
-`MLX_BUN_MEGAKERNEL_KVQUANT=1`). Quantize formula correct. **Gap diagnosed:** optiq's
-`quantizedSdpaUnfused` uses `quantizedMatmulQT`, which dequants K/V **to f32 on-the-fly
-inside the matmul** — never storing a bf16 dequantized value. Increment 1 stores
-`bf16(scale·q+bias)` in the bf16 cache → ONE extra bf16 rounding optiq lacks → KL
-15× the L1 path. NOT a formula bug; a storage-precision artifact of the shortcut.
-**Increment 2a — PARITY fix (store-int-q, the cleaner path):** the quantized int `q`
-(0..15 / 0..255) is EXACTLY representable in bf16 (ints ≤256 are exact). So store `q`
-itself in the existing bf16 kcache/vcache (lossless) + per-group scale/bias in small
-side buffers `[NLAYERS,KVSEQ,nKvHeads·headDim/64=4]` bf16; `attend_simd<KVBITS>` computes
-`scale·q+bias` in f32 on read = optiq's exact f32 dequant, ZERO extra rounding → should
-reach the attend_simd ceiling (~97, == L1). Much smaller than bit-packing: kv_rt_group
-stores `q` (not dequant) to d_k/d_v + emits scale/bias new-rows; host adds 4 scale/bias
-buffers + merges; attend dequants (current-pos scale/bias from device scratch, cached
-from the side-buffer inputs). NO memory win (q stored loosely as bf16), but full optiq
-parity — the Phase 4 correctness goal.
-**Increment 2a — DONE, store-int-q: 93/100, KL 1.38e-2, deterministic, no NaN.**
-Stores int q in kcache/vcache (exact in bf16) + bf16 scale/bias; `attend_simd_q<KVBITS>`
-dequants to f32 on read. **Found + fixed a real bug:** the current-pos scale/bias were
-first written to a NON-ATOMIC output (`sbNew`) read cross-threadgroup by attention →
-incoherent (stale) per [[metal-persistent-kernel-coherence]] → 78/100. Fix: route
-current-pos scale/bias through ATOMIC scratch `d_sb` (A_LD/A_ST), copy to `sbNew` only
-for the host merge; cached scale/bias read from the `sb` INPUT (coherent). → 93/100.
-**Surprise / honest correction:** 2a removed increment-1's extra bf16-rounding but
-landed at the SAME 93/100 — so **dequant precision was NOT the gap.** The residual
-(~14× L1's KL) is a genuine divergence between our quantized KV and optiq's, cause
-still OPEN: formula validated vs `ops.quantize` (bf16 ULP), all 24 layers' bits match
-kv_config, coherence fixed. Candidate causes to investigate: optiq may keep recent
-token(s) unquantized (a bf16 staging window) rather than quantize-on-write; or its
-fused quantized-SDPA accumulates differently than our dequant-then-online-softmax.
-**ROOT CAUSE FOUND (decisive, in-process vs our bit-exact reference):** the reference
-`QuantizedKVCache`+`quantizedSdpaUnfused` is bit-exact with the golden (maxDiff 0,
-`tests/minicpm5-kv-parity.test.ts`). Comparing the megakernel's dequantized layer-0 K
-to the reference's: **max diff 0.166 (≈ one quant step).** Comparing the *pre-quant*
-bf16 roped K (L1): **max diff 1 bf16 ULP** (9.0 vs 9.0625; 425/1536 elems > 1 ULP) —
-the `qmv4`-vs-`ops.quantizedMatmul` GEMV rounding difference, i.e. exactly what gives
-L1 its 9.7e-4. **Quantization is DISCONTINUOUS**, so that 1-ULP K difference lands on a
-q-level / group-min-max boundary and flips a level → a full quant-step dequant error →
-amplified 9.7e-4 → 1.38e-2. **NOT a quant bug — the megakernel's intrinsic ~1-ULP GEMV
-difference amplified by quant discontinuity. 93/100 is the L2 ceiling (as 97/100 is L1's);
-bit-exact L2 would need a bit-exact GEMV, defeating qmv4.** Gate L2 by KL + agreement
-(no-oracle, Lab class), not the bit-exact golden. Scripts: `megakernel-kv-{cmpkv,cmpl1,perstep}.ts`.
-**Increment 2b — MEMORY win (later):** bit-pack `q` into uint32 (flat KPACK/VPACK,
-per-layer literal offsets) to actually shrink 4-bit KV. Same logits as 2a. Deferred.
-**Build on the GENERATED kernel** (`genSourceMt`, `MLX_BUN_MEGAKERNEL_GEN`): the
-per-layer KV `bits` is fixed on disk (`kv_config.json`), so bake it as a compile-time
-literal per layer (same "pre-determined" principle as the weight bits) — the quantize/
-dequant helpers are `template<int KVBITS>`, and each unrolled layer emits its literal
-`KVBITS` (and `group_size=64`). No runtime kv-bits table lookup. This is why Phase 4
-makes the generated path the primary one (the looped kernel would need a runtime kv
-table). Phase 4 therefore promotes `genSourceMt` from flag-gated to the L2 default.
-**Gate:** `goldens/minicpm5-kv-parity.json` + `…-kv-logits-step*.bin` (already exist),
-teacher-forced agreement + KL. Defer until L1 solid (L1 IS solid — parity, Phase 3).
-
-### Phase 5 — Integration + ship
-**Do:** `MiniCPM5Model.decodeStep(tokenId, pos) → logits` using the megakernel when
-`MLX_BUN_MEGAKERNEL=1`; allocate the packed buffers + KV scratch once at model load
-(or first decode), reuse across steps. Wire into `src/generate.ts` decode loop
-(n23–n44 in `docs/dag/training-inference-map.html`) — the megakernel REPLACES the
-per-step `forwardHidden→logitsFromHidden` at decode (M=1) only; prefill unchanged.
-Guard: bf16 KV (L1) only, no adapters, no vision. Add `--megakernel` CLI flag.
-**Gate:** `bun test tests/minicpm5-*.test.ts` green (off); a new
-`tests/minicpm5-megakernel.test.ts` (teacher-forced gate, env-gated in the pattern
-of the old `tests/perf-kernel-oracle.test.ts` — deleted 2026-07-05 with its kernel,
-unified-engine-frontier-plan.md §6; the envelope-gate methodology stands on its
-own); paired decode tok/s A/B logged. Update
-`STATUS.md` + `benchmarks/RESULTS.md` (preflight-gated `scripts/bench-serve.ts all` for the
-quotable number — `[[dirty-machine-numbers-are-garbage]]`).
-
----
-
-## 5.5 Methodology for the next megakernel (learned 2026-06-24)
-
-**Copy mlx's exact code verbatim into one file → prove bit-identical (per-op harness,
-maxDiff 0 vs `ops.*`) → static-analyze that one file → THEN optimize incrementally,
-re-running the bit-exact harness after each step.** Why: a hand-port of `qmv_fast` was
-1-ULP-correct but NOT bit-identical to what `ops.quantizedMatmul` actually dispatches
-for **4-bit M=1** (it uses a different kernel/accumulation order — `qmv4-bitexact-check.ts`
-shows 8-bit bit-exact, 4-bit off by exactly 1 bf16 ULP on 12/256 rows, e.g. 1552 vs 1544
-= adjacent bf16 values). Invisible at L1 (9.7e-4) but quantization's DISCONTINUITY
-amplifies it into the L2 gap. A provably-identical starting copy + per-step gating would
-catch any divergence the moment it's introduced, not 24 layers + a quant step downstream.
-Steps: (1) copy mlx `quantized.h` GEMV kernels INCLUDING the per-(bits,M,N,K) dispatch
-(`qmv_fast`/`qmv`/`qmv_quad`), `rms_norm`, `rope`, `sdpa`; (2) per-op maxDiff-0 harness
-before fusing; (3) analyze the single fused file; (4) optimize with the harness as a
-regression gate. This megakernel was built no-oracle-first (Lab-style, KL-gated, NOT bit-exact by design),
-so its 1-ULP/op drift is within contract — bit-exact is a focused REBUILD, not a patch
-(and may cost qmv4's speed if mlx's real 4-bit kernel is the slower qmv_quad).
-
-## 6. Pitfalls (learned this session — do not relearn)
-
-- **Chained custom kernels corrupt under MLX's planner.** The ENTIRE reason for one
-  kernel. Don't reintroduce intermediate custom-kernel boundaries.
-- **`templateDtypes:{T:x.dtype}` is required** or `T` is undeclared (Metal compile
-  error). **`.template store<…>`** needed for dependent names when `T` is a template
-  param. Bake dims as `templateInts` / read `x_shape[0]`, don't pass a runtime
-  `shape` input array (extra lazy nodes).
-- **Match the reference's bf16 rounding op-for-op** (round to `T` at each boundary) —
-  `metal::precise::exp`/`tanh` are the only residual vs the reference; that's
-  no-oracle (Lab-gated) territory and fine. Do NOT compute everything in f32 then round once (drifts).
-- **`atomicOutputs:true` + `initValue:0`** for any atomic-accumulated output.
-- **GQA:** 16 q-heads, 2 kv-heads, `nRep=8` → q-head `h` reads kv-head `h/8`.
-- **`rope_theta=5e6`, not 10000.** scale `=128^-0.5`. No MiniCPM residual scaling.
-- **Software barrier deadlocks if `G` > co-resident threadgroups.** Start G=1.
-- **`detachLeaf` ≠ `evalAll`:** eval alone doesn't break the graph; only a fresh
-  `fromBytesCopy` leaf does. (Relevant only if you keep any boundary.)
-- **`MlxArray` has `.astype`, not `ops.astype`; `ops.reshape`, not `.reshape`.**
-- Oracle venv for any golden regen: `/Users/joshrossi/Code/mlx-lm/.venv/bin/python`,
-  `MLX_BUN_ORACLE_VENV=…`, `HF_HUB_DISABLE_XET=1`. (Goldens already exist; only
-  regen if you change the prompt.)
-
----
-
-## 7. File manifest (what you create)
-
-```
-src/model/megakernel-pack.ts            # Phase 1: weight repack + layout table
-gen-minicpm5-megakernel.ts (deleted; git history)      # Phase 2: emit per-layer-baked kernel SOURCE
-src/model/generated/minicpm5-megakernel.ts   # (generated) the kernel SOURCE + bits
-src/model/megakernel-kernel.ts          # Phase 2/3: MetalKernel + decodeMegakernel()
-megapack-check.ts (deleted 2026-08-23; git history)        # Phase 1 gate
-megakernel-teacherforced.ts (deleted 2026-08-23; git history)  # Phase 2/3 correctness gate
-megakernel-perf.ts (deleted 2026-08-23; git history)       # Phase 3/5 paired tok/s A/B
-tests/minicpm5-megakernel.test.ts       # Phase 5 CI gate (env-gated)
-```
-Modify: `src/model/minicpm5.ts` (add `decodeStep` megakernel branch),
-`src/generate.ts` (decode-loop branch), `src/cli.ts` (`--megakernel`).
-
-## 8. Order of operations (the literal checklist)
-
-1. [ ] Phase 1: `megakernel-pack.ts` + `megapack-check.ts` → byte-identical gate green.
-2. [ ] Phase 2: generator + `G=1` kernel; `megakernel-teacherforced.ts` ≥98/100,
-       deterministic ×3, no NaN. Bisect by layer if it fails.
-3. [ ] Phase 3: `G>1` + software barrier; sweep `MLX_BUN_MEGAKERNEL_G`; same
-       correctness gate + tok/s ≥ baseline. Record chosen `G`.
-4. [ ] Phase 4 (optional): L2 quantized KV vs kv-parity golden.
-5. [ ] Phase 5: wire into `decodeStep`/`generate.ts`, env+CLI gate, CI test,
-       benchmark, update `STATUS.md`/`RESULTS.md`/memory
-       `[[cpm5-generate-and-fuse-inference]]`.
-
-**Done = one resident dispatch computes the MiniCPM5 decode forward, ≥98/100
-teacher-forced, deterministic, ≥ baseline tok/s, default-off, CI-gated.**