mlx-bun 0.0.9 → 0.0.10

package/docs/design/batching-perf-path.md

@@ -0,0 +1,232 @@
+# Batching perf path — promote continuous batching out of compat mode
+
+Status: P5 + two gate fixes LANDED (2026-07-02, same day); rest PLANNED.
+
+**Headline result (M1 Max 32 GB, loaded machine, 4 concurrent × 128 tok,
+median of 3, wall-clock aggregate; oMLX 0.4.5.dev1, same OptiQ snapshots):**
+
+| Model | oMLX | mlx-bun `--batch 4` | mean TTFT (theirs/ours) |
+| --- | --- | --- | --- |
+| MiniCPM5-1B | 339.0 | **345.4** | 397 / **119 ms** |
+| gemma-4-e4b | **89.8** | 87.1 (−3%) | 815 / **397 ms** |
+| Qwen3.5-4B | **100.9** | 99.5 (−1%) | 848 / **368 ms** |
+
+What landed (all bit-gated, `tsc` 0, suites green):
+1. **P5 — SSM batched path** (`SSMCache.mergeRows`/`filter` in
+   `src/model/qwen3-delta.ts`; "ssm" kind in `batch-scheduler.ts` merge/wrap/
+   filter; per-row RoPE via `ropeOffsetArr` in `qwen3_5.ts`; gate relaxed in
+   `generation-gateway.ts`, kill switch `MLX_BUN_BATCH_SSM=0`). Oracle:
+   `tests/qwen35-batched-parity.test.ts` — solo-prefill+merge batched greedy
+   is **token-for-token identical** to mlx-lm B=2 left-padded batch
+   (`goldens/apple-m1-max/batched-golden-qwen35.json`). No ssm_mask needed:
+   rows solo-prefill unpadded, decode feeds one real token per row.
+2. **Logits processors batch** (`generation-gateway.ts`): the per-row sampler
+   folds `makeLogitsProcessors` over a per-row device-side token history
+   (generate()'s pushHistory), so repetition/presence/frequency penalties +
+   logit_bias + min_p/XTC no longer route serial. Load-bearing discovery:
+   **Qwen3.5 ships a default repetition penalty in generation_config.json**,
+   which routed EVERY request serial under `--batch N` — the SSM port alone
+   changed nothing over HTTP until this landed.
+3. Diagnostic breadcrumb: `MLX_BUN_LANE_DEBUG=1` logs per-request lane
+   decisions + shape (server.ts) — how (2) was found.
+
+Bench-method traps found while measuring — **both FIXED 2026-07-02**:
+- `serve`/`bench` accepted `--model` but IGNORED it (positional-only fuzzy
+  query; auto-pick silently loaded the wrong model). Now: `--model
+  <path-or-query>` is an explicit override — a directory with config.json
+  loads via `scanSnapshot` (no registry dependency), anything else resolves
+  like the positional. Precedence: `--model` > positional > `--query` >
+  auto-pick.
+- Event-loop starvation during generation was the SERIAL lane, not batch
+  (batch's drive loop yields per step — `/stats` answered in ~8 ms mid-
+  generation once measured properly): serial decode is an unbroken microtask
+  chain, so `/stats`//health/new accepts stalled for the whole generation
+  (measured 2.5 s on a 512-token run). Fixed with a ≥25 ms rate-limited
+  macrotask hop wrapped around onToken in `GenerationGateway.run`'s serial
+  branch — covers all four API surfaces, streaming and not. After: `/stats`
+  10–44 ms mid-generation, decode unchanged (cpm5 serial 193 tok/s). Motivated by the oMLX head-to-head (M1 Max,
+2026-07-02): at 4 concurrent requests oMLX's continuous batching gets
+1.46–1.89× our aggregate throughput (Qwen3.5-4B 97.3 vs 66.5 tok/s;
+MiniCPM5 328 vs 173; e4b 85 vs 52) and our mean TTFT queue-inflates to
+1.1–3.8 s. Single-stream we match or beat them (e4b 55.8 vs 53.7,
+Qwen3.5 66.3 vs 63.9, TTFT 3–7× better everywhere); the gap is entirely
+the serial lane going flat under parallel load. Full comparison + numbers:
+the 2026-07-02 oMLX report (Claude artifact) and `~/.omlx` bench scripts.
+
+Design stance: the batch lane grows the **same parity-tier DAG as the
+serial lane** (docs/design/parity-tier-dag.md). `--l1 --batch N` stays
+bit-exact mlx-lm-B=N compat mode forever (the oracle checkpoint); the
+default/`--l3` batch path becomes the optimized one. `--batch` remains a
+mode switch, not a load fallback.
+
+## 1. Current state
+
+`src/serve/batch-scheduler.ts` (`BatchScheduler`) is a real
+continuous-batching engine: iteration-level scheduling, pipelined decode
+(`#step` via `ops.asyncEvalAll`, mlx-lm `GenerationBatch._step`
+semantics), chunked joiner prefill (2048-token chunks) interleaved with
+decode, per-row eviction (`filterKVRows` / `BatchedRotatingCache.filter`),
+per-row failure containment, drain-on-serial-waiter, `clearCache` every
+256 steps. Bit-parity with mlx-lm B=N verified at B=2 across all 4 model
+paths (`tests/batched-decode-parity.test.ts`, `tests/batched-rotating.test.ts`,
+`tests/batch-scheduler.test.ts`).
+
+`src/serve/generation-gateway.ts` (`GenerationGateway.willBatch`) routes;
+one `AsyncMutex` is the GPU/`loraState` exclusion domain. The batch lane
+is deliberately compat mode: it calls `model.forwardHidden` /
+`logitsFromHidden` directly, bypassing `generate()` and every perf lever.
+
+### Levers the batch lane lacks, and why
+
+| Lever | Serial mechanism | Why batch can't use it today |
+| --- | --- | --- |
+| Quantized KV (`kv_config`) | `maybeQuantizeKv` (`src/generate.ts:101`) | `QuantizedKVCache` (`src/model/gemma4-base.ts:419`) has no dynamic-B ops (no merge/filter/extend/temporalView over the packed/scales/biases triples); `willBatch` hard-gates `kvQuant` → serial |
+| Perf-kernel fused decode SDPA | dispatched from `quantizedSdpa` (`gemma4-base.ts:1560`) | kernel M=1-specialized: `fused-decode-kernel.ts:262` rejects `B !== 1 \|\| L !== 1`, and requires empty mask (batched decode needs per-row left-pad validity) |
+| Fused N-tiled SDPA (`--fused-decode`) | `quantizedSdpaTiled` | needs 2-D mask; batched masks are 4-D `[B,1,1,S]` (`buildBatchedDecodeMask`); quantized-only anyway |
+| Compiled decode | `CompiledDecode.for()` (`generate.ts:518-525`) | trace is `[1,1]`; B changes every join/evict |
+| Prompt-cache prefix reuse | `PromptCache.take/put` in `runGeneration` | bypassed by design in v1; every batched row solo-prefills (`cachedTokens: 0`) |
+| Vectorized sampling | n/a | `#step` samples per row (B slice+sample+concat graphs per token) |
+
+Two more gap sources specific to the benchmark:
+
+- **Qwen3.5 never batches at all**: `SSMCache` (`src/model/qwen3-delta.ts`)
+  has no merge/filter, so the capability gate routes every request serial.
+  The Qwen3.5 concurrency loss is pure serial queueing. mlx-lm batches
+  Mamba-family state via `ArraysCache` — oracle exists.
+- Joins re-merge the whole batch (O(B·S) slice storm, `#mergeJoiner`)
+  instead of mlx-lm's `extend`.
+
+## 2. Gap analysis: quantized KV + fused kernels at B>1
+
+- `QuantizedKVCache` already stores `[B,H,S,D']` triples; missing is
+  dynamic-B surgery. New `src/model/batched-quant.ts`:
+  `BatchedQuantizedKVCache` = quantized twin of `BatchedRotatingCache`
+  (merge/extend/filter/extract/temporalView = the `mergeKVRows`/`filterKVRows`
+  B-axis ops from `src/model/batched-mask.ts:171,196`, applied to all three
+  triples). Left-pad must align to quantization group boundaries — pad in
+  whole 256-token STEP granules (both caches already grow in STEP=256).
+- Conversion point mirrors `maybeQuantizeKv`: row solo-prefills bf16,
+  `toQuantized` per layer per `kv_config` before merge.
+- **Scope cut**: shipped `kv_config.json` quantize full-attention layers
+  only ⇒ P1 = quantized full layers + bf16 `BatchedRotatingCache` sliding
+  layers (the serial `config` composition). Batched rotating-quant deferred.
+- Correctness fallback: `quantizedSdpaUnfused` (`gemma4-base.ts:1273`) is
+  already B-generic with `mask.mode === "array"`; needs the batched mask
+  broadcast as `[B,1,1,1,N]` for the 5-D GQA score reshape.
+- Perf kernel at B>1: (1) drop the `B !== 1` gate (keep `L !== 1`);
+  (2) grid gains a B axis (`[TG_THREADS, H·B, G]`), `SCORE_ROW` offsets
+  index `((b·KV + kv)·N + n)`; (3) replace the empty-mask requirement with
+  a per-row `leftPad` int32 `[B]` input — each row's n-loop starts at
+  `max(blockStart, leftPad[b])`, so pad bytes are never read.
+- RoPE per-row positions already solved (`ropeOffsetArr`, the S1b.1 fix).
+
+## 3. Phases
+
+### P0 — Baseline + engine quick wins (no new numerics)  [S–M, 2–4 days]
+- Write `scripts/bench-serving-load.ts` (N concurrent vs live server,
+  aggregate tok/s + TTFT p50/p95). Re-measure vs oMLX **with `--batch 4`
+  actually on** to size the true engine gap vs the serial-queueing gap.
+- Land `extend` join (full-attention twin + `BatchedRotatingCache.extend`);
+  oracle: extend `scripts/gen-batched-dynamic-golden.py` (mlx-lm `extend`).
+- Vectorize the homogeneous-sampler fast path in `#step` (all-greedy →
+  one `ops.argmaxAxis` on `[B,V]`); mixed rows keep the loop.
+- Exit: gated suites unchanged; B=4 baseline recorded; joins not O(B·S).
+
+### P1 — Quantized KV at B>1  [M–L, 1–1.5 wk]
+- `src/model/batched-quant.ts` + scheduler wiring (layer kind
+  `"full" | "rot" | "qfull"`); solo-prefill bf16 → `toQuantized` → extend.
+- Batched mask 5-D broadcast; validate `quantizedSdpaUnfused`.
+- `willBatch`: `kvQuant` stops being a serial gate when quant layers are
+  full-attention-only.
+- Exit (parity ladder; no external oracle — mlx-lm quant batching NYI):
+  (1) B=1-row batched-quant **bit-exact vs serial `--kv-quant config`**;
+  (2) row-isolation within the calibrated bf16-reduction envelope;
+  (3) KL vs batched-bf16 ≤ serial quant-vs-bf16 KL; 6-task eval unchanged;
+  (4) `--l1 --batch N` suites untouched.
+
+### P2 — Perf kernel at B>1 + fallback matrix  [M, ~1 wk]
+- Kernel edits per §2; dispatch when every row supported, else unfused.
+- Fallback matrix documented in server-config.md.
+- Batched frozen-oracle envelope (B=2/4 `freeze-perf-oracle` variant,
+  ≥56/64 per row, `tests/perf-kernel-oracle.test.ts` pattern).
+- Exit: dispatch counter proves engagement; envelope green; measured
+  uplift at B=4 vs P1 unfused on a clean machine; if <5%, ship the
+  fallback matrix only and close honestly.
+
+### P3 — Admission, prompt cache, adapters, defaults  [M–L, 1–1.5 wk]
+- KV-budget admission: `projectKvBytes(model, promptLen, maxTokens)`
+  beside `fit()`; `--kv-budget`, `--prompt-concurrency`;
+  `/stats.batch.{pending_rows,kv_bytes,kv_budget_bytes}`.
+- Prompt-cache reuse under batching: `take` at admit, forward uncached
+  tail; on finish/evict `extract(row)` → solo caches → `put`. Quantized
+  entries are non-trimmable mid-group — match on full-prefix or STEP
+  boundaries. Real `prefillMs/decodeMs`; `BatchStats.cachedTokens` real.
+- Single-adapter batches: adapter-spec = batch compatibility key
+  (join if match, drain if not; `loraState.active` set once per epoch).
+- Default review: if batch lane at 1 live row is within ~2% of serial,
+  propose flipping the shipped default to `--batch 4` (Josh decision,
+  numerics change vs serial default); else README recommends `--batch 4`
+  for agent workloads.
+- Exit: TTFT p50 <500 ms / p95 <1.5 s at 4 concurrent short prompts;
+  `cached_tokens>0` on repeats under `--batch 2`; oversized-request
+  queueing test (no OOM); B=4 aggregate within 15% of oMLX on cpm5 + e4b.
+
+### P4 — Burst decode  [REVISED after reading + porting oMLX's source]
+- **oMLX-style event-loop bursting: BUILT AND REFUTED (2026-07-02).** Read
+  their actual implementation (`omlx/engine_core.py _step_burst` — full
+  source ships in `/Applications/oMLX.app/Contents/Resources/omlx/`): it is
+  NOT a GPU trick — it runs K `scheduler.step()` calls per
+  `run_in_executor` hand-off with adaptive time budgets (100 ms single /
+  30 ms concurrent, 64-step cap) because each per-token hand-off ping-pongs
+  the GIL with asyncio/uvicorn (~1 ms/token; their comment: 80 vs 74 tok/s).
+  Ported faithfully to `BatchScheduler.#drive` and measured: **regression**
+  — cpm5 B=4 345→289, batch-lane B=1 149→121, TTFT +~100 ms (SSE flushes
+  wait out the budget). Bun's `setImmediate` hop costs microseconds — no
+  GIL to dodge — so the port kept the cost and lost the win. Reverted;
+  numbers restored (cpm5 B=4 349.1). Breadcrumb comment left at the drive
+  loop. This also explains their ~21-token SSE bursts (measurement trap in
+  the comparison report).
+- **Still open (different mechanism, not refuted): device-side step
+  chaining** — generalize `generateInner`'s pipeline to depth k (k chained
+  step graphs fed device-side, one host sync per k tokens). That attacks
+  OUR real per-token cost (FFI graph-build + sync readback), which is not
+  what oMLX's burst attacks. Exit unchanged: bit-exact vs depth-1, cpm5
+  single-stream gap vs oMLX (178 vs ~214) is the target; no regression at
+  k=1.
+
+### P5 — SSM batched path (Qwen3.5)  [DONE 2026-07-02 — see header]
+- Landed in one day, much cheaper than the ~1 wk estimate: `ssm_mask` under
+  left-padding turned out unnecessary (our join model solo-prefills unpadded,
+  so SSM state never sees a pad token; mlx-lm's masked-pad prefill is
+  equivalent — proven by the token-exact oracle match). Batch ops are plain
+  B-axis concat/take on the two state slots.
+- Bonus fix that the exit benchmark forced out: logits processors now batch
+  (per-row fold; Qwen3.5's default repetition penalty had been routing every
+  request serial). Exit numbers in the header table.
+
+Ordering: P1→P2 strictly ordered; P4 and P5 independent (P5 first if the
+Qwen3.5 headline matters most). Total ~5–6 wk serial, ~4 wk with P4/P5
+parallel.
+
+## 4. Risks
+
+- **No external oracle for quant-batched** — mitigated by the P1 parity
+  ladder; `--l1 --batch N` stays the always-reachable bit-exact checkpoint.
+- **Metal kernel at B>1**: leftPad start-index never touches pad bytes;
+  no atomics (deterministic); gate = batched frozen-oracle envelope +
+  dispatch assertions + A/B vs unfused.
+- **Memory**: B× quant caches + prompt cache + weights on 24 GB — P3
+  admission is the backstop; mirror mlx-lm's
+  `trim_to(total − activeBatchBytes)` so LRU and live batch share one
+  ceiling. GPU OOM is uncatchable — queueing test before default changes.
+- **Ring-wrap at B>1 short-context-verified only**: land a long-context
+  bf16 batched golden before any P3 default flip. Quantized rotating out
+  of scope.
+- **Default flip changes numerics** vs today's serial default — explicit
+  Josh decision with doc trail.
+- **Benchmark hygiene**: quotable numbers via `bench-serving-load.ts` on a
+  clean machine; TTFT at the SSE first-token event, not server-side
+  (oMLX's own server log over-reports by excluding ~180 ms/req cache
+  bookkeeping — measure them wall-clock only).
+- **Regression fence**: all existing gated batch suites + the serial suite
+  green unchanged at every phase.

package/docs/design/omlx-adoption-map.md

@@ -0,0 +1,90 @@
+# oMLX adoption map — every feature, where it slots, status
+
+Status: LIVING (started 2026-07-02). oMLX (github.com/jundot/omlx, Apache
+2.0) is the systematic port source: read its actual implementation before
+building (full source ships locally in
+`/Applications/oMLX.app/Contents/Resources/omlx/`), port the IDEA into our
+architecture, and keep it only if OUR benchmarks improve (the burst-decode
+refutation below is the cautionary tale). Attribution: note derived code in
+THIRD_PARTY_LICENSES.md; idea ports get a source comment.
+
+## Where oMLX sits in the fidelity-tier model
+
+It doesn't slot at L1 or L2 — it decomposes. The tiers are NUMERICS
+contracts (parity-tier-dag.md); oMLX is a Python appliance ON mlx-lm:
+
+- Its stock forwards are mlx-lm's numerics = our existing **L1 oracle**
+  (no new oracle value; mlx-lm itself stays the reference).
+- Its own numeric inventions (oQ, custom kernels, DFlash, MTP) are
+  oracle-less = **L3-class** here, and get our L3 gating (KL/eval/envelope
+  + kill switch) that they don't have upstream.
+- Its product surface (SSD cache, EnginePool, menu bar, MCP, grammar) is
+  **tier-agnostic serving layer** — keyed by the effective tier/scheme,
+  never part of the decode-numerics DAG.
+
+## Scoreboard
+
+### Ported, kept (benchmarked wins)
+| Feature | Theirs | Ours | Result |
+| --- | --- | --- | --- |
+| Continuous batching parity | scheduler on mlx-lm BatchGenerator | `--batch N` lane + P5 SSM port + per-row logits processors (2026-07-02) | cpm5 345 vs 339 (win), e4b −3%, Qwen3.5 −1%, TTFT 2–3× better (batching-perf-path.md) |
+| SSD KV cold tier | `cache/paged_ssd_cache.py` (content-hashed blocks) | `src/ssd-cache.ts` + kv-store v2 (whole-entry spill, zero-copy mmap) | restart TTFT 12.1s→0.24s vs their 1–3s; 0% decode overhead vs their ~20% (ssd-kv-cold-tier.md) |
+
+### Ported, refuted (do not re-add without new evidence)
+| Feature | Why it doesn't transfer |
+| --- | --- |
+| Adaptive burst decode (`engine_core.py _step_burst`) | Amortizes Python GIL/asyncio ping-pong (~1ms/token). Bun has no GIL; faithful port REGRESSED cpm5 B=4 345→289 + TTFT +100ms. Reverted with breadcrumb (batching-perf-path.md P4). |
+
+### Queue (roughly by leverage; ★ = Josh explicitly wants)
+1. **Structured output / JSON-schema constrained decoding** (`api/grammar.py`)
+   — serving layer. The biggest API-surface hole vs both oMLX and agent
+   frameworks. We own the sampler loop; a token-mask grammar engine slots
+   beside the existing sampler extensions. Serial lane first, batch lane
+   via the per-row processor fold (already built).
+2. **★ Menu bar app** (`apps/omlx-mac/`, native SwiftUI + Sparkle) — product
+   layer. Their repo is the structural reference; our signed/notarized
+   single binary is the sidecar. Supersedes tauri-desktop-app.md's Electron
+   ambivalence: native SwiftUI is the received-well shape.
+3. **Multi-model serving / EnginePool** (`engine_pool.py`: LRU + pinning +
+   load/unload API) — serving layer. Fits our single-USER framing as model
+   SWITCHING (one active, LRU-evict the rest), not concurrent tenants.
+4. **oQ-style sensitivity-driven quantization** (`oq.py`,
+   docs/oQ_Quantization.md) — L3-class, lands in `convert` beside
+   `--target-bpw`: calibration-measured per-layer sensitivity, boosts on
+   non-expert tensors only, batched expert GPTQ. Gate: perplexity + 6-task
+   eval vs our knapsack at equal bpw. (arXiv-lens candidate.)
+5. **DFlash serving wiring** (`engine/dflash.py` + their published -DFlash
+   drafts) — L3-class; unblocks DSpark (docs/design/dspark-speculative-decoding.md).
+   Their integration pattern: separate engine class, ctx-length fallback
+   routing, tape-replay cache rollback.
+6. **Vision feature cache** (`cache/vision_feature_cache.py`) — serving
+   layer; encoder features for repeated-image agent turns. Natural sibling
+   of the SSD tier (same keying discipline).
+7. **MCP tool execution in serve** (`--mcp-config`, `mcp_routes.py`) —
+   serving layer. Note overlap: Pi already owns tool orchestration for our
+   UIs — scope this to bare-API consumers; don't duplicate Pi.
+8. **Prefill progress observability** (`prefill_progress.py`) — /stats +
+   web UI surfacing for long prompts. Pairs with the serial-lane hop fix.
+9. **Rerank endpoint** (`/v1/rerank`) — needs a reranker model family
+   first; park until a Qwen3-Reranker-class port is wanted.
+10. **Document ingestion (MarkItDown)** — product layer; ours would go
+    through the chat UI / Pi attachments, not the server.
+11. **Admin one-click benchmark with prefix-cache-hit testing** — fold into
+    bench-serving-load.ts as a `--cache-hit-ratio` mode instead of a UI.
+
+### Explicitly not porting
+- **Audio (STT/TTS/STS)** — audio tower deferred (project scope).
+- **Their SSE burst streaming** — side effect of the GIL workaround;
+  per-token streaming is strictly better UX at zero cost for us.
+- **Python-side scheduler details** (GIL executors, collector reaping) —
+  runtime-specific.
+
+## Porting discipline (learned this week)
+1. Read their real source first (local app bundle) — the README lies less
+   than the code.
+2. Port the idea into OUR architecture; never transliterate Python.
+3. Benchmark before/after on THIS machine; keep only wins (burst decode
+   died here). Wall-clock metrics only — their own logs over-report.
+4. Numeric features get L3 gates + kill switches; serving features get
+   the effective-scheme keying; nothing touches L1/L2 contracts.
+5. Verify the served model id (/v1/models) before trusting any bench.

package/docs/design/segmented-backward-training.md

@@ -184,6 +184,21 @@ quant. The trivial case to validate the machinery on a real model.
 
 ## 5. The segmentation strategy (the part to get right)
 
+> **2026-07-02 CORRECTION — this section's cost model is WRONG, measured.**
+> The full-attention isolation planner was built and A/B'd @8K on real chunk
+> rows (scripts/experiments/seg-isolation-smoke.ts, MLX_BUN_SEG_MEM_LOG=1
+> phase probes): ZERO peak win (18.09 vs 18.02 GB seg2). The premise "sliding
+> layers are O(L·window) in the backward" is false — mlx's sdpa BACKWARD
+> materializes the O(L²) scores for EVERY layer (~3.5 GB/layer @8K e4b; the
+> sliding window is only an additive mask; a sliding PAIR costs +7.1 GB ≈
+> sliding+full's +7.15). The worst segment is therefore set by LAYER COUNT
+> alone and `segment_size` is the whole knob: **seg1 = 14.59 GB @8K, +3% step
+> time, loss identical — fits the 24 GB M4 Pro today.** The head vjp adds
+> ~3 GB (full-[M,V] responseOnlyCe — kernel-review backlog #8). Going below
+> ~14 GB needs that head fix plus an O(L)-memory attention backward (flash
+> track), not smarter segmentation. Section kept for the (still-correct)
+> donor/boundary constraints; ignore its sliding-vs-full peak arithmetic.
+
 NOT a naive "divide by N". Compute the split points from the model:
 
 **Constraints & inputs (all readable from the model/config):**

package/docs/design/ssd-kv-cold-tier.md

@@ -0,0 +1,158 @@
+# SSD cold tier for the prompt/KV cache
+
+Status: **LANDED P1–P3 (2026-07-02, same day)** — P4 hardening extras remain.
+
+**Measured (M1 Max, MiniCPM5, 13.7k-token prefix, kv_config quant):**
+- Restart TTFT: **12,083 ms (full prefill) → 236 ms** restored — 2% of full
+  prefill, 13,675 tokens back via a 3 ms zero-copy mmap. Beats oMLX's 1–3 s
+  claim and both exit criteria (<2.5 s, ≤10%).
+- Steady-state decode: **0% overhead** (192.0 vs 192.3 tok/s, flag off/on) —
+  nothing runs on the token loop, as designed (vs oMLX's measured ~20%).
+- Correctness: SSD-restored + trimmed continuation is **identical to the RAM
+  tier's cache-hit output** (control experiment); divergence from a fresh
+  full prefill is the pre-existing bf16 prefix-reuse property both tiers
+  share, not introduced by persistence.
+
+What landed: kv-store v2 (`src/kv-store.ts` — all five cache kinds incl.
+quantized + SSM, invalidation metadata, per-tensor hashes, atomic writes;
+`restoreState` added to both quantized cache classes), `src/ssd-cache.ts`
+(SsdCacheStore per D4), tier integration (`src/prompt-cache.ts` spill hook +
+retain thunks; `src/server.ts` tiered take with post-restore trim,
+debounced write-behind snapshots under the gateway lock, `/stats.ssd_cache`),
+flags `--ssd-cache <dir>` / `--ssd-cache-max <GB>` / `--ssd-cache-verify`.
+Tests: `tests/kv-store.test.ts` (fork-compare token-identical for bf16,
+quantized, SSM), `tests/ssd-cache.test.ts` (restart recovery, quarantine,
+corruption, eviction, supersede). One deviation from the plan below: find()
+returns PARTIAL prefix matches and the server trims after restore
+(divergent-tail conversations — full-entry-only missed the real pattern). Inspired by oMLX's paged SSD cache
+(Apache 2.0 — `omlx/cache/paged_ssd_cache.py`, `boundary_snapshot_store.py`;
+we port ideas, not code — add a `THIRD_PARTY_LICENSES.md` note if any code
+ends up derived). Their headline: coding-agent TTFT 30–90 s → 1–3 s by
+restoring long prefixes from SSD across requests AND restarts. Their
+measured costs (independent tests): ~2–3× cold TTFT and ~20% steady-state
+decode overhead — this design exists to get the win without those costs.
+Flag-gated, off by default.
+
+**Head start**: `src/kv-store.ts` already implements page-aligned KV
+persistence (`MLXBUNKV1`, 16 KiB-aligned tensors) with zero-copy COW mmap
+reload (`MmapFile.open(path, "cow")` → `MlxArray.fromPointer` — exactly
+the Metal page-alignment constraint, proven by expert-offload), including
+mid-rotation `RotatingKVCache` restore (`restoreState`,
+`src/model/gemma4-base.ts:880`). `scripts/experiments/cold-start.ts` is a
+working sub-1 s cold-restore harness. The SSD tier is wiring +
+quantized-cache support + an index/eviction/recovery layer.
+
+## Current state / gaps
+
+- `src/prompt-cache.ts`: byte-capped LRU RAM tier (default 2 GB,
+  `--prompt-cache`), prefix-entry based, adapter-namespaced, serial lane
+  only, race-free under the gateway lock. **A restart loses everything.**
+- `kv-store.ts` gaps: (1) no quantized kinds — `saveKvCache` throws for
+  `QuantizedKVCache`/`RotatingQuantizedKVCache` and neither has
+  `restoreState`, yet serving defaults to the model's `kv_config.json`
+  scheme, so the default production cache can't persist today; quantized
+  state is 6 arrays/layer (packed u32 + scales/biases bf16 × K/V);
+  (2) no invalidation metadata / checksums / atomic writes; (3) latent
+  bug: `loadKvCache` checks `caches.length !== model.layers.length` but
+  `makeCache()` returns donor layers only — wrong for KV-shared e4b
+  (verify in Phase 1).
+- Versioning primitive exists: `configFingerprint`
+  (`src/model/fingerprint.ts`) — sha256 over graph-shaping config incl.
+  `kvQuant`.
+
+## Design decisions
+
+**D1 — Whole-prefix-entry spill, not paged blocks.** oMLX pages for
+multi-user shared-prefix dedup; we are single-user with a prefix-entry
+cache that grows monotonically per conversation. Paging is exactly where
+their overheads come from: per-block SHA256 + save-queue bookkeeping
+*during decode* (their ~20% tax) and block-granular safetensors restore
+with copies (their 2–3× cold TTFT). Ours: **zero work on the token loop**
+(prefix tokens ARE the key), spill/persist only at request boundaries and
+evictions, restore = one zero-copy mmap with lazy page fault-in (~5 GB/s
+NVMe; 2 GB entry ≈ 0.4 s amortized into the first forward). Partial-prefix
+reuse works without paging: restore + `Cache.trim()`. Revisit paging only
+for multi-user batch dedup (v2, out of scope).
+
+**D2 — One store, two tiers.** `PromptCache` grows a cold tier behind the
+same `take`/`put` API:
+- RAM→SSD: `put()`'s evict loop spills (snapshot bytes → async writer →
+  dispose GPU arrays) instead of just disposing.
+- SSD→RAM: `take()` prefers a strictly-longer SSD prefix; restored entry
+  carries its `MmapFile`, unmapped on entry dispose (extensions copy into
+  fresh anonymous buffers via concat-grow anyway).
+- **Restart persistence needs write-behind at request completion**, not
+  just spill-on-evict: after `promptCache.put(...)` in `runGeneration`
+  (`src/server.ts:1019`), a debounced, coalesced-per-ns async snapshot.
+  Byte extraction (the only GPU-thread work) runs on the idle serial lane
+  between requests; file I/O is async. (oMLX's boundary-snapshot idea at
+  whole-entry granularity.)
+
+**D3 — Extend the existing format to `MLXBUNKV2`**, not safetensors
+(safetensors offsets aren't page-aligned; oMLX pays a copy via `mx.load`).
+Add: `"qkv"`/`"rotating-qkv"` kinds (six tensor slots + groupSize/bits;
+slice quantized buffers to `[.., :offset, :]` — don't persist 256-step
+slack), header fields `formatVersion`/`modelId`/`configFingerprint`/`ns`/
+`tokenizerHash`/`createdAt`, per-tensor checksums + header checksum.
+Integrity: header checksum always verified on scan; full verification
+opt-in (`--ssd-cache-verify`) because eager reads destroy lazy fault-in;
+default trusts atomic writes (tmp + fsync + rename). Corrupt/incompatible
+→ delete + log, never fatal. `loadKvCache` becomes model-generic
+(accept `{ makeCache(): Cache[] }`, donor-count check).
+
+**D4 — Files are the database.**
+`<ssd-dir>/<configFingerprint>/<nsHash>/<uuid>.mlxkv`. Startup recovery =
+header-only directory scan → in-memory `SsdIndex`
+(`{ns, tokens, path, bytes, mtimeMs}[]`); tokens live in the header JSON.
+LRU = file mtime (`utimes()` on hit); byte cap enforced at write time by
+oldest-mtime eviction; disk-full → drop write, warn once, keep serving.
+
+**D5 — Flags & interactions.** `--ssd-cache <dir>` +
+`--ssd-cache-max <GiB>` (default 32) → `ServerOptions.ssdCacheDir/MaxBytes`.
+Requires the RAM tier (`--prompt-cache 0` + `--ssd-cache` = startup
+error). `--memory-budget` unchanged (restored caches are clean file-backed
+pages until dirtied; admission already assumes worst-case bf16). `/status`
+gains `ssd: { entries, bytes, max_bytes, restores, spills, restore_ms_last }`.
+v1 is serial-lane only (where `PromptCache` lives); batch-lane interplay is
+docs/design/batching-perf-path.md P3's problem.
+
+## Phases
+
+- **P0 — Baseline bench (½ d)**: `scripts/experiments/ssd-cache-bench.ts`
+  grown from `cold-start.ts`: ≥16k-token prefix on the standard snapshot,
+  kv_config quant active; measure full-prefill TTFT / warm RAM-hit TTFT /
+  steady-state decode. Exit: baseline table committed.
+- **P1 — kv-store v2 (2–3 d, highest uncertainty)**: `restoreState` on
+  both quantized cache classes; all four kinds in save/load; header v2;
+  atomic writes; model-generic load. Exit: round-trip tests all kinds
+  (incl. wrapped ring) + **greedy decode parity** — N tokens continued
+  from a restored cache bit-exact vs never-persisted; `tsc` 0.
+- **P2 — `src/ssd-cache.ts` store (1–2 d)**: `SsdCacheStore
+  { scan, find(prompt, ns), store (async, extraction split from I/O),
+  evictToCap, remove, stats }`. Exit: unit tests — restart recovery,
+  incompatible-fingerprint quarantine, corrupt-header deletion, mtime cap
+  eviction, disk-full soft-fail.
+- **P3 — Tier integration + flags (2–3 d)**: cold tier in `PromptCache`,
+  debounced end-of-request snapshots, CLI/`/status`. **Exit (headline)**:
+  (1) restart TTFT on a 16k prefix **< 2.5 s and ≤ 10% of full-prefill**
+  (beat oMLX's 1–3 s; a true-cold miss runs the unchanged prefill path);
+  (2) steady-state decode overhead **< 2%** vs flag off (vs their ~20%);
+  (3) flag absent ⇒ zero behavior change.
+- **P4 — Hardening + docs (1–2 d)**: kill-during-write crash test,
+  adapter-ns isolation, kv-quant-scheme-change invalidation, README/help,
+  license note. Exit: suite green; on/off × cold/warm/restart matrix
+  recorded.
+
+Total ~7–10 working days.
+
+## Risks
+
+- **mmap lifetime vs mlx laziness**: unmap only in entry `dispose()` after
+  successor state is evaluated; COW mapping already covers buffer donation.
+- **Byte-extraction stall on multi-GB entries**: runs between requests on
+  the idle lane; measure in P3; chunk with yields if needed.
+- **COW dirty-page accounting**: extended restored caches migrate to
+  anonymous memory — probe with `phys_footprint` (expert-offload
+  discipline) before making footprint claims.
+- **Stale-token poisoning**: `configFingerprint` + `tokenizerHash` + ns;
+  chat-template drift degrades to a shorter matched prefix (safe).

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