mlx-bun 0.0.4 → 0.0.6

package/docs/design/parallel-slots.md

@@ -1,13 +1,17 @@
 # Parallel slots — configurable batched serving
 
-Status: **design + phasing** (S0 landed; S1+ not started)
+Status: **S0 + S1b SHIPPED** (S2 not started; S3+ deferred)
 Owner: serving layer
-Default: **off** (`--slots 1` = today's serialized path, untouched)
+Default: **off** (`--batch 1` = today's serialized path, untouched)
+
+> **Flag rename note:** This document was drafted with `--slots N` as the design name.
+> The shipped flag is `--batch N` (`--decode-concurrency` accepted as an mlx_lm.server alias).
+> Occurrences of `--slots` in pre-ship text below refer to today's `--batch N`.
 
 ## Goal
 
 Let the server process more than one generation request concurrently,
-gated entirely behind a `--slots N` flag. `N=1` runs exactly what we
+gated entirely behind a `--batch N` flag (originally designed as `--slots N`). `N=1` runs exactly what we
 have today (the serialized promise chain). `N>1` loads a batch
 scheduler that runs multiple sequences through one forward pass.
 
@@ -78,7 +82,7 @@ concurrent ≈ arrival_rate × mean_generation_time).
 
 When it does **not** help: traffic so light that requests never overlap
 in time. Then there's only ever one live sequence and batching is idle
-— which is exactly why `--slots 1` is and stays the default.
+— which is exactly why `--batch 1` is and stays the default.
 
 ## Current code: what's batch=1 and what isn't
 
@@ -98,7 +102,7 @@ Good news: the hot path is already written over a batch dim `B`.
 | `src/generate.ts:295,315,378` | prompt/decode tensors built `[1, L]` / `[1, 1]` |
 | `src/generate.ts:320,235` | `hLast` slice `[1, …]`; `sampleStep` reshape `[1, V]` |
 | `src/generate.ts:247` | sampler history `[1]`, one sampled token/step |
-| `src/model/gemma4.ts:647-653` | per-layer-input reshapes `[1, L, …]` |
+| `src/model/gemma4.ts:647-653` | per-layer-input reshapes `[1, L, …]` *(pre-S1b inventory — this entry was fixed in S1b; e4b per-layer-input is now B-generic)* |
 | `src/model/gemma4-base.ts:210` | **one scalar `offset` per cache, shared across rows** |
 | `src/model/gemma4.ts:174,206` | **scalar RoPE offset** for the whole batch |
 | `src/model/gemma4-base.ts:256` | `makeMask` from that single offset (returns empty mask at N=1) |
@@ -237,7 +241,7 @@ Degradation levers: `--perf-kernel off` / `--fused-* off` drops L3→L2; bf16 KV
 drops L2→L1. Each layer must stand alone and stay reachable (don't-delete-
 optionality as a fallback ladder).
 
-So a feature like batched decode (`--slots`) is "done" only when green at
+So a feature like batched decode (`--batch`) is "done" only when green at
 **every (model × layer)** cell — and each cell may need its own fix:
 
 | Model · path | L1 (bf16 / mlx-lm) | L2 (quant / optiq) | L3 (perf / KL) |
@@ -323,7 +327,7 @@ Batched-decode landmines found by the modes analysis (2026-06-14):
   required), so quant batched decode falls to `quantizedSdpaUnfused` — correct
   but the perf kernel is bypassed (perf debt; a batched-aware fused decode kernel
   is a later item). `quantizedSdpaUnfused` must be validated with a 4-D bool mask.
-- **All cells**: compiled-decode forced off under `--slots>1`; LoRA batches only
+- **All cells**: compiled-decode forced off under `--batch>1`; LoRA batches only
   within one adapter group (else drain).
 
 ## KV memory: dynamic allocation, not static partition
@@ -357,11 +361,12 @@ Budget accounting is total-bytes from day one.
 **Current state (2026-06-14):** the batched *primitive* is **rung 2** —
 one shared `[B, H, S, D]` buffer per layer (`KVCache`/`RotatingKVCache`),
 growing in 256-token steps, width = longest live sequence, left-pad
-"waste" on short rows. But the **budget/admission half of rung 2 is NOT
-built** (no `B × S_max` projection, no KV budget enforcement) — that's
-scheduler work (S2), and the scheduler itself isn't built (`--slots`
-inert). And **rung 3 (paged) is a deferred spike**. So today: rung-2
-*allocation shape*, no *budget control*, not wired into serving.
+"waste" on short rows. Rung-2 allocation IS wired into serving via
+`BatchScheduler` + `GenerationGateway` (S1b SHIPPED). But the
+**budget/admission half of rung 2 is NOT built** (no `B × S_max`
+projection, no KV budget enforcement) — that's scheduler work (S2, NOT
+STARTED). And **rung 3 (paged) is a deferred spike**. So today: rung-2
+*allocation + scheduling wired*, no *budget/admission control* yet.
 
 KV quantization is a force multiplier: 4-bit KV (`generate.ts:67`,
 already supported) is ~4× smaller, so the same budget buys ~4× the slots
@@ -462,34 +467,37 @@ predicate (full-attention model + no vision/adapter/repetition-penalty/user-seed
 
 ## Config surface
 
-- `ServerOptions.slots` (default 1); later `kvBudgetBytes`,
+- `ServerOptions.batch` (default 1); later `kvBudgetBytes`,
   `promptConcurrency`.
-- CLI: `--slots N` (S0); later `--kv-budget <GB>`.
-- `slots === 1` → today's serialized promise chain, untouched.
-  `slots > 1` → scheduler.
+- CLI: `--batch N` (shipped; `--slots N` was the original design name, `--decode-concurrency` accepted as mlx_lm.server alias); later `--kv-budget <GB>`.
+- `batch === 1` → today's serialized promise chain, untouched.
+  `batch > 1` → scheduler.
 
 ## Phasing (incremental, parity-gated)
 
-- **S0 — config seam (DONE).** `--slots N` / `ServerOptions.slots`
+- **S0 — config seam (DONE).** `--slots N` (now shipped as `--batch N`) / `ServerOptions.slots`
   plumbed, validated, surfaced (ready card + `/stats`). `N>1` warns that
   batched execution lands in S1 and runs serially. No behavior change.
-- **S1 — static 2-wide, base model only.** Split by the reuse finding:
+- **S1 — static 2-wide, base model only. (DONE — S1b LIVE as `--batch N`)** Split by the reuse finding:
   - **S1a (prefill)** — reuse the training machinery (`buildBatchedPadMask`
     / `BatchedMaskCache`); wire the serving path to prefill B prompts in one
     forward. Mostly generalization + consolidating the two mask builders.
-  - **S1b (decode)** — the new work: a batched KV cache that grows with
-    per-row offsets, per-row RoPE via the array-offset path, and the per-row
-    `[B,1,1,S]` decode mask (`buildBatchedDecodeMask`, landed 2026-06-14).
-    Plus the B-token-per-step decode loop + per-row stream fan-out.
+  - **S1b (decode) — SHIPPED 2026-06-14.** `BatchScheduler`
+    (`src/serve/batch-scheduler.ts`) + `GenerationGateway`
+    (`src/serve/generation-gateway.ts`) wired into `createServer` behind
+    `--batch N`. L1 oracle-verified for CPM + all 3 Gemmas (full-attn +
+    sliding-window via `BatchedRotatingCache`). Per-row SSE fan-out live.
+    See STATUS block above for details.
   **Teacher-forced gate**: a 2-row batch must produce per-row logits matching
   two solo runs within bf16 tolerance (the train-batch-e2e methodology, now
   applied to decode).
-- **S2 — N-wide + continuous injection/eviction.** Rows retire on EOS;
-  queued requests prefill into freed rows. Dynamic byte-budget admission.
-- **S3+ — paged KV** (custom paged-attention kernel + block manager),
+- **S2 — N-wide + continuous injection/eviction. (NOT STARTED)** Rows retire on EOS;
+  queued requests prefill into freed rows. Dynamic byte-budget admission
+  (`B × S_max` projection, KV budget enforcement) still TODO.
+- **S3+ — paged KV (DEFERRED)** (custom paged-attention kernel + block manager),
   KV-quant under batch, LoRA-group batching.
 
-Every phase ships default-off behind `slots=1`; the serialized path is
+Every phase ships default-off behind `batch=1`; the serialized path is
 never removed (only bypassed).
 
 ## Open questions / to measure

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

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

package/docs/design/runtime-isolation.md

@@ -1,6 +1,6 @@
 # Runtime isolation — the AI may crash, the UI never may
 
-Status: **design** (not started)
+Status: **design** (not started — deferred future work; no scheduled PLAN.md phase)
 Owner: serving layer / runtime
 Default today: inference runs **in-process**, on the server's event-loop
 thread. This doc proposes moving it to a **subprocess**.

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

@@ -15,8 +15,10 @@ it holds all 42 layers' recompute activations at once, so e4b's training backwar
 spikes to **23 GB live @2048** (resting is only 6.95 GB) and **crashes at 4096+**.
 We proved a fix — **segmented backward** — that forces per-segment streaming:
 forward saves detached boundary activations at segment edges, then backprop runs
-segment-by-segment via a surrogate-loss `value_and_grad`. Only one segment's
-activations are ever live. On an activation-dominated toy it beat per-layer
+segment-by-segment via `mlx_vjp` (the shipped production implementation — see §9.4;
+an earlier surrogate-loss `value_and_grad` approach was tried but leaked ~32 MB/seg/step
+and was replaced). Only one segment's activations are ever live.
+On an activation-dominated toy it beat per-layer
 checkpointing (4.26 vs 5.97 GB) with **bit-exact gradients**. Extrapolated to e4b
 this lands at **~10 GB @8K** — fitting on the 32 GB M1 Max, where the optiq/mlx-lm
 reference cannot (it crashes on e4b @2048+). Next: wire it into the trainer,
@@ -97,12 +99,16 @@ Algorithm:
    small (one [B,L,hidden] per edge) and cheap.
 2. **Compute the loss** from the final hidden (LM head + response-only CE) and its
    gradient w.r.t. the final hidden — this `dh_out` seeds the last layer segment.
-3. **Backward, reverse over segments.** For each segment k (last→first), build a
-   `value_and_grad` over `[boundary_leaf_k, ...segment_LoRA_params]` whose loss is
-   the surrogate `sum(stop_gradient(dh_out) ⊙ segment_forward(boundary_leaf, params))`
-   (the optiq chunked-CE surrogate, generalized to the whole backward). It returns
+3. **Backward, reverse over segments.** For each segment k (last→first), call
+   **`mlx_vjp`** (bound as `Vjp` in `src/mlx/autograd.ts`) over
+   `[boundary_leaf_k, ...segment_LoRA_params]` with cotangent `dh_out`. It returns
    `[dh_in, ...dLoRA_segment]`. `dh_in` becomes `dh_out` for segment k−1; accumulate
-   `dLoRA`. Only segment k's activations are live during its `value_and_grad`.
+   `dLoRA`. Only segment k's activations are live during its vjp.
+   **NOTE:** an earlier design used a surrogate-loss `value_and_grad` here
+   (`sum(stop_grad(dh)⊙segment_forward(...))`). That proved to leak ~32 MB per
+   segment per step (linear, no plateau) and was replaced by `mlx_vjp` in §9.4.
+   The production code uses `Vjp`/`mlx_vjp` throughout — not surrogate-loss
+   `value_and_grad`.
 
 Why it's exact: `d/dx sum(stop_grad(dh)⊙y) = (dy/dx)ᵀ dh = dLoss/dx`, and likewise
 for params — this is reverse-mode AD done in segments. Verified **0.000%** vs the
@@ -140,9 +146,14 @@ quant. The trivial case to validate the machinery on a real model.
 2. **Trainer:** add a `segmented` training path (flag `TrainConfig.segments` or auto;
    gate alongside the existing `gradCheckpoint`). It:
    - runs the segmented forward (save detached boundaries),
-   - computes response-only CE loss + `dFinalHidden` (already memory-bounded),
-   - runs the reverse per-segment surrogate `value_and_grad`, threading the LoRA
-     params for each segment as that `value_and_grad`'s argnums, accumulating grads,
+   - computes response-only CE loss and `dFinalHidden` via a `Vjp` of the loss head
+     with scalar cotangent 1.0 (already memory-bounded),
+   - runs the backward per-segment using **`mlx_vjp`** (`Vjp` in
+     `src/mlx/autograd.ts`): for each segment k (last→first), calls
+     `vjp.apply([boundary_k, ...LoRA_k], [dh_out])` to get `[dh_in, ...dLoRA_k]`,
+     threads `dh_in` as `dh_out` for segment k−1, and accumulates `dLoRA` — no
+     surrogate scalar, no `value_and_grad` (see §9.4 for why that approach leaked
+     ~32 MB/seg/step and was replaced),
    - also handles the embedding + (e4b) per-layer-input grads from the first
      segment's `dh_in`.
 3. **Validate:** grads bit-match the current trainer (response-only CE, same seed);
@@ -232,17 +243,24 @@ returns the split indices + which donor K/V to save, computed from the above.
 - `scripts/ft-chunk-smoke.ts` — e4b train/fwd/infer smoke (GRAD_CKPT, MEM_LIMIT_GB, MLX_BUN_MEM_LOG).
 - `scripts/sdpa-vs-manual.ts`, `scripts/flash-dk-debug.ts`, `scripts/flash-grad-test.ts` — gradient correctness.
 - `scripts/chunk-{eval,finetune,filter}.ts` — the chunk experiment pipeline.
-- `src/train/trainer.ts` — training loop (combined-eval, mem-log, gradCkpt ctx). Add the segmented path here.
+- `src/train/trainer.ts` — training loop (combined-eval, mem-log, gradCkpt ctx); contains the segmented path (`TrainConfig.segmentSize`).
 - `src/train/forward.ts`, `src/train/loss.ts` — training forward + response-only CE.
-- `src/model/gemma4.ts` — e4b model: `forwardLayers`, `runCheckpointedLayer`, `computePerLayerInputs`, KV-sharing (`numDonors`, `previousKvs`, `cacheIndex`), `layer_types`. Add `runLayerRange`.
-- `src/model/minicpm5.ts` — MiniCPM5 (Phase A target).
-- `src/mlx/autograd.ts` (`ValueAndGrad`), `src/mlx/checkpoint.ts` (`Checkpoint`), `src/mlx/ffi.ts` (memory fns, `stopGradient`, `evalAll`).
+- `src/model/gemma4.ts` — e4b model: `forwardLayers`, `runCheckpointedLayer`, `computePerLayerInputs`, KV-sharing (`numDonors`, `previousKvs`, `cacheIndex`), `layer_types`, `runLayerRange` (added for segmented backward).
+- `src/model/minicpm5.ts` — MiniCPM5 (Phase A).
+- `src/mlx/autograd.ts` (`ValueAndGrad`, `Vjp`), `src/mlx/checkpoint.ts` (`Checkpoint`), `src/mlx/ffi.ts` (memory fns, `stopGradient`, `evalAll`, `mlx_vjp`).
 - Memory: `[[e4b-lora-training-seqlen-ceiling]]` (full findings), `[[opssdpa-dk-vjp-bug]]`, `[[three-level-fidelity-tree-model]]`.
 
-## 8. First action on resume
-Phase A step 1: add `runLayerRange` to `MiniCPM5Model`, then a segmented training
+## 8. Original first action (COMPLETED)
+~~Phase A step 1: add `runLayerRange` to `MiniCPM5Model`, then a segmented training
 path in `trainer.ts`, validate bit-exact grads + measure peak on
-`scripts/ft-chunk-smoke.ts`-style MiniCPM5 run. Then Phase B (e4b) with the §5 plan.
+`scripts/ft-chunk-smoke.ts`-style MiniCPM5 run. Then Phase B (e4b) with the §5 plan.~~
+
+**Both Phase A and Phase B are done — see §9 and §10.** The genuine remaining work
+is `planSegments` full-attention isolation for the ~10 GB @8K e4b target: put each
+O(L²) full-attn layer (5,11,17,23,29,35,41) in its own short segment so it isn't
+materialized alongside sliding-layer activations. Currently 17.5 GB @8K (seg2);
+resting is ~6.6 GB, so ~3.4 GB of transient budget remains. The mechanism and tunable
+segmentation knob are proven; this is a peak-optimization step. See §10 close.
 
 ---
 
@@ -339,11 +357,11 @@ memory win confirmed, no leak, trains end-to-end through `scripts/chunk-finetune
    baseline (11.89). json/schema/labels/nonempty 100%, anchors 86%. Segmented
    training produces a quality adapter. Adapter:
    `~/.cache/mlx-bun-finetunes/minicpm5-chunk-segmented`.
-2. **Phase B (e4b)** with the §5 plan (KV-sharing, per-layer-input, full-attn O(L²)
-   isolation). e4b adds `runLayerRange` to `Gemma4Model` with donor-KV /
-   per-layer-input boundary threading; the `SegmentedBackward` + vjp machinery and
-   the §5 segmentation strategy carry over. This is the path to the original goal:
-   e4b @8K under ~10 GB where the optiq/mlx-lm reference crashes.
+2. ✅ **DONE — see §10.** **Phase B (e4b)** with the §5 plan (KV-sharing,
+   per-layer-input, full-attn O(L²) isolation). e4b `runLayerRange` added to
+   `Gemma4Model` with donor-KV / per-layer-input boundary threading;
+   `SegmentedBackwardGemma4` + vjp machinery wired into trainer, validated on real
+   e4b. Full details in §10.
 
 ---
 
@@ -467,42 +485,43 @@ are proven; this is a peak-optimization, not a blocker.
   (Config note: mlx-lm's default LoRA scale is 20.0; mlx-bun's `DEFAULT_TRAIN_CONFIG`
   is 1.0 — a config difference, not correctness. The chunk runs used 1.0.)
 
-**Trainer side — `SegmentedBackwardGemma4` (to build + validate):** same skeleton as
-the MiniCPM5 `SegmentedBackward`, with TWO additions.
+**Trainer side — `SegmentedBackwardGemma4` (BUILT + VALIDATED, `src/train/segmented.ts`):**
+same skeleton as the MiniCPM5 `SegmentedBackward`, with TWO additions — both are
+implemented and verified in the Phase B results above.
 
 (a) **per-layer-input is a pure constant boundary — NO gradient threading.** The
 `per_layer_input_gate` / `per_layer_projection` LoRA live INSIDE the layer (applied
 to the sliced `pls`), so they're differentiated within each segment's vjp normally.
 The `perLayer` tensor itself feeds in as a constant (its grad would flow to the
-embed / `per_layer_model_projection`, neither a LoRA target) — so detach it once and
-slice per segment; discard its cotangent. Easy.
+embed / `per_layer_model_projection`, neither a LoRA target) — detached once and
+sliced per segment; its cotangent is discarded.
 
 (b) **the KV-shared donor K/V is a SECOND boundary stream WITH its own cotangent.**
 The donor's K/V depend on the donor layer's q/k/v LoRA and are attended by every
 sharer, so `dLoss/d(donor LoRA)` gets contributions from the donor's own attention
 AND all sharers. Threading:
 - **Forward (boundary-saving):** when a segment produces a reused donor's K/V
-  (`donorKvOut`), detach it (rope-applied K + V, tiny — ~1 KB/tok) and save as a
-  boundary; later sharer segments pass it in via `donorKvIn`.
+  (`donorKvOut`), it is detached (rope-applied K + V, tiny — ~1 KB/tok) and saved as
+  a boundary; later sharer segments pass it in via `donorKvIn`.
 - **Backward (reverse), two extra pieces:**
-  - keep a cotangent accumulator `dKV[d]` per reused donor d (init zero).
+  - a cotangent accumulator `dKV[d]` per reused donor d (init zero).
   - a SHARER segment's vjp differentiates `[boundary_h, donorK, donorV, ...LoRA]`
     with cotangent `[dh]`; its vjps for `donorK/donorV` are ACCUMULATED into `dKV[d]`.
   - a DONOR segment (produces donor d's K/V) has a MULTI-OUTPUT forward
     `[h_out, donorK, donorV]`; its vjp takes cotangents `[dh, dKV[d].k, dKV[d].v]`
-    (the `Vjp` class already supports nOut>1 + multiple cotangents) and returns
+    (the `Vjp` class supports nOut>1 + multiple cotangents) and returns
     `[dh_in, ...donor LoRA]`. This folds the sharers' gradient back into the donor.
   - donors 22 (sliding) and 23 (full) are each reused by all sharers of their type;
     if both live in one segment, that segment's forward outputs 5 arrays
     `[h, k22, v22, k23, v23]` with 5 cotangents.
-- **Segmentation (the §5 plan):** respect the donor/sharer split at layer 24 and
-  ISOLATE the 7 full-attention layers (5,11,17,23,29,35,41) — each is O(L²) in the
-  backward under ops.sdpa (~1 GB bf16 @8K), so put each full layer in a short
-  segment and group the sliding layers. The natural first cut is the 6-layer
-  `sssssF` period → 7 segments of 6 (one full layer each). `planSegments` from §5
-  returns the split indices + which donor K/V to save.
+- **Segmentation (the §5 plan):** the donor/sharer split at layer 24 is respected and
+  the 7 full-attention layers (5,11,17,23,29,35,41) — each O(L²) in the backward
+  under ops.sdpa (~1 GB bf16 @8K) — are candidates for isolation into short segments;
+  the natural first cut is the 6-layer `sssssF` period → 7 segments of 6 (one full
+  layer each). `planSegments` from §5 returns the split indices + which donor K/V to
+  save. Full-attention ISOLATION for the ~10 GB @8K target remains the next
+  optimization step (currently 17.5 GB @8K, seg2 — see §10 close above).
 
 Everything else (forward boundaries via `detachLeaf`, head vjp with cotangent 1.0,
-reverse `dh` chain, `mlx_vjp` not surrogate) is identical to the MiniCPM5 path.
-Target: e4b @8K ≤ ~10 GB, then `chunk-eval` quality on e4b (the original goal —
-chunk data is mostly 3.4–8.2K, useless at the ≤2K the reference is capped to).
+reverse `dh` chain, `mlx_vjp` not surrogate) is identical to the MiniCPM5 path and
+is verified end-to-end: peak 10 GB @2048, active flat (no leak), adapter saved.

package/docs/design/tauri-desktop-app.md

@@ -0,0 +1,127 @@
+# Design: mlx-bun desktop app (Tauri)
+
+A native macOS desktop app that wraps the existing mlx-bun chat UI in a
+window — double-click an icon, get local AI, no terminal. The app is a
+**thin shell**: it spawns `mlx-bun serve` as a sidecar and points a
+WKWebView at the chat UI the server already serves. Almost no new code;
+the value is packaging + UX, not new functionality.
+
+Companion docs: [embedding.md](../reference/embedding.md) (the sidecar
+layout this builds on), [distribution.md](../reference/distribution.md)
+(the sign/notarize pipeline this reuses).
+
+## Where it lives — separate repo (recommended)
+
+Ship the app from its **own repo** (`mlx-bun-desktop`), not inside this
+one.
+
+- **Toolchain isolation.** This repo is Bun/TS. Tauri drags in Rust,
+  Cargo, and `src-tauri/target/` (heavy build artifacts). Keeping it out
+  preserves the core repo's clean dependency surface and scope.
+- **Different product, cadence, audience.** The app's artifact is a
+  signed `.dmg` / Homebrew **Cask** for end users; the core repo ships a
+  CLI tarball + npm package for developers. They version independently.
+- **The sidecar is already a published artifact.** The app pulls a
+  pinned, signed, notarized `mlx-bun` bundle from a core-repo release
+  (the versionless `mlx-bun-arm64.tar.gz`, or a tagged
+  `mlx-bun-v<ver>-arm64.tar.gz`) — no source coupling needed.
+- **Precedent.** The Homebrew tap is already a separate repo; this fits.
+
+In-repo alternative (a `desktop/` subdir) is viable if tight version
+lockstep ever matters more than toolchain cleanliness — the mechanics
+below are identical, only the sidecar source changes (local `dist/`
+instead of a downloaded release). Default to separate.
+
+## Architecture
+
+```
+mlx-bun-desktop.app
+ ├─ src-tauri (Rust shell)
+ │    on launch  → pick a free port, spawn sidecar:
+ │                 mlx-bun serve <model> --port <p> --no-open
+ │    readiness  → poll GET http://127.0.0.1:<p>/v1/models until 200
+ │    window     → WKWebView → http://127.0.0.1:<p>/#/chat
+ │    on quit    → SIGKILL the sidecar (restart→ready is ~0.4s; safe)
+ └─ Contents/Resources/
+      mlx-bun  + libmlxc.dylib + libmlx.dylib + libjaccl.dylib
+      + mlx.metallib + pi assets   (the bundle build-binary.sh produces)
+```
+
+The window renders the **same page** `mlx-bun serve` serves today
+(`/#/chat`, the embedded pi web chat). No bespoke UI.
+
+## Components
+
+1. **Sidecar binary.** Tauri `externalBin` (target-triple-suffixed, e.g.
+   `mlx-bun-aarch64-apple-darwin`); the dylibs + metallib + pi assets go
+   in `bundle.resources` beside it. mlx-bun resolves its libs via
+   `dirname(realpath(execPath))`, so the flat Resources layout works
+   unchanged (verified: Bun realpaths execPath through symlinks).
+2. **Spawn / teardown** (`main.rs`). Use Tauri's shell/sidecar API (or
+   `std::process::Command`) to launch on `setup`, kill on window close /
+   app exit. One generation runs at a time server-side, so no client
+   concurrency to manage.
+3. **Port.** Bind an ephemeral port in Rust, pass it to both the spawn
+   (`--port`) and the webview URL, to avoid 8090 collisions with a
+   user's own `mlx-bun serve`.
+4. **Readiness gate.** Show a lightweight "starting…/downloading model…"
+   splash until `/v1/models` returns 200. Warm start ~0.4s; first-ever
+   run downloads the starter model (~0.9 GB) — the existing onboarding
+   page already narrates this, so the splash can hand off to it.
+
+## Distribution
+
+A `.app` is a better Gatekeeper citizen than the flat CLI binary:
+
+- **Stapleable.** Unlike a bare Mach-O, a `.app`/`.dmg` accepts a stapled
+  notarization ticket → works **offline**, no online check needed.
+- **Channels:** signed `.dmg` direct download, and/or a Homebrew **Cask**
+  (`brew install --cask mlx-bun`) — a natural 5th install path.
+- **Not App Store.** JIT + a spawned child process + the Metal model
+  don't fit the sandbox. Developer-ID + notarized + DMG/Cask outside the
+  store is the same flow this project already uses.
+
+### Signing / entitlements
+
+Reuse the core pipeline. The nested sidecar `mlx-bun` still needs the
+same entitlements as the standalone build
+([packaging/entitlements.plist](../../packaging/entitlements.plist)):
+
+- `com.apple.security.cs.allow-jit` + `allow-unsigned-executable-memory`
+  (JavaScriptCore JIT) — without these the sidecar is killed on launch.
+- `com.apple.security.cs.disable-library-validation` — the sidecar
+  dlopens the bundled MLX dylibs.
+
+Tauri's bundler signs nested binaries during the `.app` signing pass;
+point these entitlements at the sidecar specifically. Notarize the
+`.app` (or `.dmg`) with the existing `AC_PROFILE` notary credentials,
+then `xcrun stapler staple`.
+
+## Milestones
+
+1. **Spike** — `tauri init` in a new repo; hardcode port 8090; spawn a
+   locally-built `mlx-bun serve --no-open`; load `/#/chat` in the window;
+   `cargo tauri dev`. Proves the wrap end-to-end.
+2. **Robust runtime** — ephemeral port; readiness poll + splash;
+   clean teardown on quit; surface sidecar stderr to a log.
+3. **Bundle + sidecar** — wire `externalBin` + `bundle.resources` to a
+   pinned core-repo release tarball (download in a build script, or check
+   in a fetch step); first-run model-download UX.
+4. **Sign + notarize + package** — Developer-ID sign with nested
+   entitlements, notarize, staple, emit `.dmg`.
+5. **Cask** — formula in a tap pointing at the `.dmg` release asset;
+   `brew install --cask`.
+
+## Open questions / risks
+
+- **Tauri sidecar entitlement plumbing** — confirm Tauri applies the
+  per-binary entitlements to `externalBin` (may need a custom signing
+  hook in the bundler config). Validate early in milestone 4.
+- **Window vs. menubar** — a normal window is the MVP; a menubar/tray app
+  ("always-on local model") is a possible later mode.
+- **Model picker** — MVP serves one model (the starter). Switching models
+  means restarting the sidecar; defer a UI for that.
+- **Version pinning** — decide whether the app tracks `latest` or pins a
+  specific `mlx-bun` release per app version (pin is safer for repro).
+- **Universal binary** — arm64-only today; Tauri can target both, but
+  MLX is Apple-Silicon-only, so x86 is moot.