mlx-bun 0.0.8 → 0.0.10

package/docs/design/orpo-dynamic-lambda.md

@@ -1,9 +1,50 @@
 # Dynamic λ controller for ORPO (adaptive preference pressure)
 
 **Status:** proposed (design only — not built). **Opened:** 2026-06-20.
+**Reaffirmed + reframed by Josh 2026-07-01** (see "The band framing" below).
 **Surface:** `src/train/loss.ts` (orpoMetrics/orpoLoss), `src/train/trainer.ts`
 (orpoLoop, the λ config), the segmented ORPO classes in `src/train/segmented.ts`.
 
+## The band framing (Josh, 2026-07-01 — the product statement)
+
+> "Vary the lambda term based on the loss or another signal — maintain the
+> interest inside a band: learn preference HARDER while you are still able to
+> stay coherent. A self-stabilizing AND self-discovering hyperparameter, which
+> optimizes the training to get as much preference as possible while not
+> resulting in degeneracy."
+
+That is the controller's contract in one sentence: **λ is not a hyperparameter
+to tune, it is a process variable to control** — push preference pressure to
+the maximum the model can currently bear (self-discovering), retreat on
+confirmed coherence damage (self-stabilizing), and hold the system inside the
+healthy band rather than at a fixed point. Everything below (phases, AIMD/PID,
+displacement decomposition) is the mechanism for that contract.
+
+**Interplay with `sft_scope` (landed 2026-07-01):** the default chosen-NLL is
+now full prompt+response. The controller's displacement signal must key on the
+**response-only `lw`** (which `orpoMetrics` computes for the odds ratio in both
+modes) — full-scope NLL mixes prompt modeling into the level and would mask a
+crushed `logp(chosen-response)`. Instrument both, control on response-only.
+
+## Paper lens (2026-07-01): is this citable?
+
+Closest prior: **β-DPO (arXiv 2407.08639)** — dynamic β for DPO, calibrated
+per-batch from the reward gap. Ours differs on three axes worth claiming:
+(1) **closed-loop trajectory control** (a controller tracking a non-stationary
+optimum across the run) vs per-batch heuristic calibration; (2) the control
+signal is **displacement-aware** — the (dLw, dLr) decomposition targets the
+known ORPO/DPO failure mode directly (cf. likelihood displacement,
+arXiv 2410.08847 "Unintentional Unalignment"); (3) **objective-agnostic**
+actuator (ORPO λ, SimPO β/γ). The incremental P→PD→PID build IS the ablation
+section. The claim to stake: *"the optimal preference pressure is
+non-stationary, so a controller strictly dominates any fixed λ a sweep can
+find — and a displacement-aware control law finds the coherence cliff without
+a priori knowledge of it."* Referees will demand: ≥2 datasets, ≥2 model sizes,
+the static-λ sweep as the baseline family (the controller must beat the BEST
+constant, not λ=0.1), seeds, and the positive-control sanity run this doc
+already specifies. Workshop/arXiv-preprint tier as a single-model study;
+main-track needs the multi-model sweep.
+
 ## Motivation
 
 ORPO's loss is `L = NLL(chosen) + λ·OR`. The two terms learn on different

package/docs/design/orpo-training.md

@@ -260,9 +260,38 @@ intra-layer MLP split, shared prompt-prefix.
 
 $$\mathcal{L}_{ORPO} = \mathcal{L}_{NLL}(y_w) + \lambda \cdot \mathcal{L}_{OR}$$
 
-- $\mathcal{L}_{NLL}(y_w)$ — standard SFT cross-entropy on the **chosen**
-  response (prompt positions masked out), i.e. the negative mean response
-  log-prob $-\ell_w$.
+- $\mathcal{L}_{NLL}(y_w)$ — the SFT cross-entropy on the **chosen** sequence.
+  Scope is selected by **`sft_scope`** (TrainConfig `sftScope`, default
+  **`full`**):
+  - **`full`** (default; paper/TRL-faithful): the token-mean cross-entropy over
+    **all non-pad supervised positions — prompt AND response** (label-shifted;
+    `nn.CrossEntropyLoss reduction="mean"`), exactly TRL ORPOTrainer's
+    `chosen_nll_loss` (labels = input ids with only padding masked to −100) and
+    xfactlab/orpo's `outputs_pos.loss`. Computed from the **same chosen
+    forward** (no extra forward): the head additionally runs over the prompt
+    span and the two span means combine as
+    $\text{NLL}_{full} = -\frac{(P{-}1)\,\text{pm} + M\,\ell_w}{(P{-}1)+M}$
+    (`combineFullNll`, `src/train/loss.ts`).
+  - **`response`**: $-\ell_w$, the negative mean response log-prob — the
+    pre-2026-07 mlx-bun behavior (the house convention of masking the prompt),
+    kept **bit-exact** for reproducing older runs.
+
+  > **Resolved (2026-07-01).** The review found the then-hardcoded
+  > response-only scope diverged from the reference implementations (TRL /
+  > xfactlab/orpo supervise the full prompt+response, only padding masked).
+  > Decision (Josh): `sft_scope: full | response`, **default `full`** — loss
+  > curves are apples-to-apples with TRL/paper ORPO by default; `response`
+  > reproduces old runs. The odds-ratio $\ell_w/\ell_r$ terms stay
+  > **response-only length-normalized means in BOTH modes** — that also matches
+  > TRL (its `get_batch_logps` uses prompt-masked labels for the odds ratio;
+  > only the NLL term is full-sequence). Implemented across every ORPO path
+  > (naive / chunked / fused / flash-CCE / prefix-shared / all four segmented
+  > classes): the span-parameterized heads run once over the prompt span and
+  > once over the response span (disjoint — no token headed twice), keeping
+  > $\ell_w$ bit-identical to the `response` path. Cross-path full-scope parity
+  > + an independent oracle (full NLL ≙ `sftLoss` with `promptLen=1`, bf16
+  > class) are wired in `tests/train-orpo-fused-ce.test.ts`; loss-math units +
+  > exact response-scope regression pins in `tests/train-orpo.test.ts`.
 - $\mathcal{L}_{OR} = -\log \sigma(\text{log\_odds})$, the odds-ratio term:
 
 $$\text{log\_odds} = (\ell_w - \ell_r) - \big(\text{log1mexp}(\ell_w) - \text{log1mexp}(\ell_r)\big)$$
@@ -983,8 +1012,9 @@ Metal path**. Two openings stand out:
 
 - **Correctness:** the current MLX SOTA, `mlx-lm-lora`, ships a **non-canonical
   ORPO** — it uses `log_sigmoid(chosen_logp − rejected_logp)` and **omits the
-  additive SFT-NLL term** ORPO requires. Our paper-faithful loss (true log-odds
-  + λ·NLL, length-normalized) is *already* more correct than anything shipping.
+  additive SFT-NLL term** ORPO requires. Our loss (true log-odds + λ·NLL,
+  length-normalized; one known scope divergence — see the note under "The
+  objective") is *already* closer to the paper than anything shipping.
   **Do not copy their loss.**
 - **Efficiency:** a **Metal fused/chunked CE + fused LoRA** would be the first
   fused preference loss in the MLX world, on a quantized base, with grad

package/docs/design/parallel-slots.md

@@ -18,6 +18,16 @@ scheduler that runs multiple sequences through one forward pass.
 This is a server-path feature — serving speed is the user metric.
 Direct/embedding mode stays batch=1.
 
+> **v2 execution plan:** this document remains the PRIMITIVE-design record
+> (masks, caches, lanes). The living gap/debt plan and sequencing for the
+> batching engine is [batching-v2-plan.md](batching-v2-plan.md). Landed from
+> it (2026-07-01, steps 1–3): cache-capability gate (hybrid/SSM models route
+> serial), slice-before-head admit, ONE lock domain (curve `/generate`,
+> `/signal`, adapter mount/unmount via `gateway.runExclusive`), per-row
+> onToken failure containment, drain-on-serial-waiter (no starvation), and
+> the engine-hygiene transplant (pipelined decode via asyncEval, clearCache
+> 1/256 steps, chunked admit prefill interleaved with decode steps).
+
 ## STATUS (2026-06-14) — LIVE for full-attention models
 
 **`--batch N` now serves B>1** for full-attention models (CPM). The numerics
@@ -266,7 +276,8 @@ Harness: `tests/batched-decode-parity.test.ts` has two paths — (a)
 NOT the gate); (b) **`realBatchedGreedy` = the REAL gate**: left-pad → one
 `BatchedDecodeMaskCache` (handles prefill at offset 0 AND per-step decode) →
 batch-prefill → greedy batch-decode, trajectory compared bit-for-bit to the
-mlx-lm B=N oracle fixture (`tests/fixtures/batched-golden-*.json`, generated by
+mlx-lm B=N oracle golden (`goldens/batched-golden-*.json`, machine-layered via
+`tests/goldens.ts` — per-GPU overrides in `goldens/<machine>/` — generated by
 `scripts/gen-batched-golden.py` in the oracle venv).
 
 **CPM L1 oracle-verified 2026-06-14d:** mlx-bun's real batched greedy trajectory

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

@@ -61,16 +61,24 @@ Worked examples:
 | 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) |
+| `perf-kernel` | **envelope-gated, no bit-exact oracle** — mlx-bun-original kernel; gate is ≥56/64 teacher-forced argmax vs *our own* frozen compat trajectory (`perf-kernel-oracle.test`, `freeze-perf-oracle.ts`), and the optiq kv-quant goldens track the **stock unfused L=1 decode** (`regen-kvq-goldens.ts`) | faster | **L3 / explicit opt-in only** (off in bare `--l2`; on in `--l3` and the no-tier default) |
+| `fused-sdpa` | matches optiq bit-for-bit, not mlx-lm (`fused-sdpa.test` tier-a goldens) | 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
+each tier are the **performance** axis (the fastest kernel that still holds it). On
+the *decode path* L2 and L3 now differ by exactly one node — the perf kernel (bare
+`--l2` = the optiq-golden composition: fused N-tiled prefill + stock unfused L=1
+decode; `--l3` adds the envelope-gated fused decode kernel). L3's remaining
+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.
 
+> History: commit `f1bf5cc` put the perf kernel in the L2 preset claiming
+> "perf-kernel-oracle.test asserts the fused kernel == the frozen optiq goldens".
+> Both halves were wrong — that golden is frozen from mlx-bun's own compat engine,
+> and the gate is argmax-agreement, not equality — so 2026-07-01 restored the tier
+> rule: the bare tier is the guarantee; perf kernels are opt-ins within it.
+
 ## Why this matters for CLI flags
 
 We are accumulating a **lot** of flags (`--grad-accum`, `--grad-clip`, `--seg`,

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/spec-decode-larger-targets.md

@@ -0,0 +1,229 @@
+# Spec decode on larger targets — does the Gemma drafter pay on 12B / 26B?
+
+Status: **12B γ=1 is a WIN after the verify fix; matches optiq bit-exactly** (2026-06-14; 26B not run)
+Owner: decode / spec
+Default: **off** (speculation is opt-in; non-spec greedy is untouched)
+
+## Result (2026-06-14, v2) — fixed the wrong-oracle bug; γ=1 now wins
+
+The first 12B measurement (v1, below) was made through a verify path that
+targeted the **wrong oracle**. Fixing it changed the conclusion.
+
+### The wrong-oracle bug
+
+Assistant-drafter speculation is an **optiq** feature — mlx-lm has only
+generic standalone-`draft_model` two-model spec and structurally cannot
+drive a KV-borrowing drafter. So the correctness oracle is optiq's
+`spec_generate`. **Both optiq and mlx-lm BATCH the verify lm-head** (one
+matmul over the γ+1 window, then argmax per position); neither is bit-
+exact to stock token-at-a-time decode. Our port instead verified
+**per-position** to be bit-exact to *stock decode* — a target no real
+implementation has, and one that re-reads the ~1 GB tied lm-head matrix
+γ+1× per iteration instead of once. So we were exact to the wrong thing
+*and* paying a γ-scaling tax for it. Fix: `picksBatched` in
+`src/spec/generate.ts` — one batched verify lm-head, matching optiq.
+
+### Do we match optiq? YES — bit-exact (e4b)
+
+Cross-checked our `specGenerate` against optiq's `spec_generate` on the
+identical prompt ids (`scripts/oracle-spec.py` + `scripts/spec-dump.ts`),
+e4b γ=2: **identical 48-token output AND identical accept/reject trace**
+(drafted 60, accepted 17, target calls 31 in both). Before the fix we
+matched *stock*, not optiq, at knife-edges; now we match optiq.
+
+### Is it still a loss? Not at γ=1 (12B)
+
+12B pair, batched verify (loaded machine — ratios paired/trustworthy,
+absolutes depressed & not quotable):
+
+| γ | v2 speedup (batched) | v1 (per-position) | acceptance |
+|---|----------------------|-------------------|-----------|
+| 1 | **1.09× (WIN)** | 0.96× | 42% |
+| 2 | 0.91× (near break-even) | 0.77× | 29% |
+| 3 | 0.72× | 0.67× | 23% |
+| 4 | 0.56× | 0.52× | 17% |
+
+Acceptance is unchanged (the fix only touches verify cost, not which
+tokens are accepted). The per-position verify was a large self-inflicted
+tax: removing it flips γ=1 from loss to **win** and moves γ=2 to the edge.
+γ≥2 still loses because the *drafter's own* cost dominates there — the
+12B drafter is heavy (hidden 1024 vs e4b 256, 16 heads, and a full 262k
+tied-embedding head **every draft step**) and acceptance falls with depth.
+
+**So the v1 "hypothesis refuted / net loss at every γ" was itself partly
+an artifact of the wrong-oracle verify.** Corrected conclusion: γ=1
+speculation on 12B is a real ~9% decode win (pending clean-machine
+confirmation before quoting); γ≥2 doesn't pay with the shipped heavy
+drafter.
+
+### Bonus finding: we're ahead of optiq on the unified drafters
+
+optiq's *installed* runtime can't actually drive the 12B/26B
+`gemma4_unified_assistant` drafters: its head detection is config-based
+(`use_ordered_embeddings = num_centroids and top_k`), but those artifacts
+declare `num_centroids` while shipping **no** centroid tensors and **no**
+`lm_head` (tied `embed_tokens` instead). optiq loads the centroid path
+with zero-init weights → garbage drafts → **0% acceptance** (measured:
+drafted 46, accepted 0 on 12B). Our `drafter.ts` detects the head by
+**tensor presence** and uses the tied-embedding path → 29% acceptance on
+the same artifact.
+
+### Remaining levers / next
+
+- **Clean-machine rerun** before promoting any number to RESULTS.md
+  (baseline jumped 15.4→24.4 between runs under memory pressure — exactly
+  why only paired ratios are trusted here).
+- **Cheaper drafter head**: the 262k tied-embed argmax every draft step is
+  the dominant remaining γ≥2 cost; capping/approximating it (e4b-style)
+  could extend the win past γ=1.
+- **26B-MoE**: still unlikely — MoE decode reads only top-8/128 experts
+  (~2.4 GB active), so per-token target cost is *lower* than its size →
+  harder to beat. Measure if curious.
+- **Optional strict mode**: the old per-position verify gave "spec output
+  identical to spec-off" — a product nicety no oracle provides. Could be
+  kept as a default-off `strictVerify` flag rather than deleted
+  ([[dont-delete-optionality-on-one-measurement]]) — not implemented; ask.
+
+## v1 (superseded) — per-position verify, "net loss at every γ"
+
+> Kept for the record. The numbers (0.96/0.77/0.67/0.52×) were real but
+> measured through the wrong-oracle per-position verify; v2 above is the
+> optiq-faithful measurement.
+
+## TL;DR
+
+The Gemma `-assistant-bf16` drafters are Gemma's substitute for an MTP
+head: a 4-layer Q-only module that reuses the target's KV + hidden state
+to draft tokens (`src/spec/drafter.ts`), verified by the target's own
+greedy pick (`src/spec/generate.ts`). True MTP (a trained
+`mtp.safetensors` head + an `optiq/runtime/mtp` backend) **does not exist
+for Gemma** and never will — that's the Qwen path (Phase 14). The
+assistant drafter is the functional equivalent we already have.
+
+It was measured a **net loss on e4b** (Phase 6): γ=2 ran at 0.78× the
+54 tok/s baseline, ~23% acceptance — because e4b decode is already too
+fast for a low-acceptance drafter to beat. **This doc is the experiment
+to find out whether the same drafter wins on a slower target**, where
+each skipped target-forward is worth more.
+
+## The hypothesis
+
+The drafter's per-step cost is roughly fixed. The *value* of each
+accepted draft = one skipped target forward, which grows as the target
+gets slower per token. e4b is the worst case for speculation (fastest
+target). So the SAME acceptance rate can flip net-positive as the target
+slows:
+
+| target | non-spec decode | per-accepted-draft value | drafter | status |
+|--------|-----------------|--------------------------|---------|--------|
+| e4b              | ~54 tok/s | low  | 160 MB (cached) | **measured net loss** (0.91/0.78/0.63/0.51× @ γ=1–4) |
+| 12B              | ~25 tok/s | ~2× e4b | 846 MB (cached) | **measured net loss** — see Result above (hypothesis refuted) |
+| 26B-A4B (MoE)    | ~32 tok/s | high (MoE forward) | 839 MB (HF, **not downloaded**) | untested (follow-on; expect loss) |
+
+> The reasoning below is the *pre-measurement* hypothesis — kept for the
+> record. It predicted the crossover assuming a fixed-cost drafter; the
+> 12B drafter is not fixed-cost (it scales with the target), which is
+> exactly why the Result section above refutes it.
+
+Crossover intuition: a drafter is a net win when
+`acceptance × (target_step_cost − draft_step_cost) > verify_overhead`.
+On e4b the target step is cheap, so the bar isn't cleared at 23%. On 12B
+the target step is ~2× more expensive while the drafter cost barely
+moves — same acceptance, very different sign. **Measure, don't assume**
+(it could still lose; this is exactly the kind of thing the e4b result
+warned against assuming).
+
+Acceptance itself may also differ on 12B (a stronger target predicts
+itself differently). We don't know the direction — that's why we run it.
+
+## What is already wired (this change)
+
+- **`scripts/spec-bench.ts`** — now multi-pair, **one pair per process**
+  (memory-safe; 12B+drafter ≈ 8 GB, 26B ≈ 17 GB, never co-resident):
+  - `bun scripts/spec-bench.ts e4b` (default) — reproduces the baseline.
+  - `bun scripts/spec-bench.ts 12B` — the experiment.
+  - `bun scripts/spec-bench.ts 26B` — follow-on.
+  - Auto-detects whether each `(target, drafter)` snapshot pair is
+    present; exits with a clear "download with: hf download …" message
+    otherwise. Records baseline + γ-sweep (acceptance + tok/s) to the
+    eval DB tagged `spec-bench[<id>]`.
+- **`tests/spec-decode.test.ts`** — added a `(12B)` parity block,
+  auto-skipped until the drafter is downloaded. Gate = long-prefix
+  agreement (greedy spec vs greedy non-spec). It guards the shared
+  accept/reject/**rollback + cache-trim** machinery against 12B's
+  distinct layer/cache shapes — NOT drafter correctness (see the
+  important nuance below). The calibrated e4b block (exact `toEqual` +
+  knife-edge) is unchanged.
+
+The port (`drafter.ts` + `generate.ts`) is already generic over the
+target's `backbone_hidden_size` — drafters are target-bonded artifacts
+(e4b 2560, 12B 3840, 26B 2816) but there's no per-pair code, just the
+matching weights. Expect to *validate* each cell per the project's
+per-model-specialization reality, not assume it's free.
+
+## Important nuance: a broken drafter still produces correct tokens
+
+Speculation can't corrupt output. The target verifies every position and
+replaces every rejected draft with its own greedy pick; drafts only ever
+*skip* target forwards. So a mis-wired drafter (wrong KV borrow, wrong
+projection) shows up as **low/zero acceptance — a perf signal**, not bad
+tokens. This splits the validation cleanly:
+
+- **Correctness** (parity test): does the shared accept/reject/rollback
+  + cache-trim logic still hold on 12B's cache shapes? A trim bug *would*
+  corrupt output, shortly after the first rejection — that's what the
+  long-prefix gate catches. (Shape-incompatible KV borrow throws, it
+  doesn't silently corrupt.)
+- **Payoff** (bench): is acceptance high enough, and the target slow
+  enough, that γ>0 beats the non-spec baseline? This is the actual
+  question.
+
+## Procedure (once the 12B drafter is downloaded)
+
+1. `hf download mlx-community/gemma-4-12B-it-assistant-bf16`
+   (846 MB; `HF_HUB_DISABLE_XET=1`). The 12B OptiQ-4bit target is
+   already cached.
+2. **Correctness first** — `bun test tests/spec-decode.test.ts`
+   (the `(12B)` block now runs). Long-prefix gate must pass for γ=1,2,3.
+   If it fails early → rollback/trim bug on 12B caches, fix before
+   trusting any number.
+3. **Calibrate + promote (optional)** — find a tie-free prompt
+   on-device, promote the 12B gate from long-prefix to exact `toEqual`
+   like e4b, for a stronger standing guard.
+4. **Measure** — `bun scripts/spec-bench.ts 12B` on a clean machine
+   (preflight matters: dirty-machine numbers are garbage). Read
+   acceptance + speedup-vs-baseline per γ.
+5. **Record + decide** — numbers land in the eval DB. Promote into
+   `benchmarks/RESULTS.md` only if quotable (clean machine).
+6. **26B follow-on** — repeat with `26B` if 12B is encouraging; MoE
+   target adds a validation cell (donor KV borrow through the MoE
+   decoder), and a slower/heavier forward that should favor speculation
+   even more — but the MoE forward also reads only top-8/128 experts at
+   decode, so the per-token cost advantage is smaller than weight size
+   suggests. Measure.
+
+## Exit criteria
+
+- 12B spec parity gate green (rollback/trim correct on 12B caches).
+- 12B acceptance + tok/s-vs-baseline recorded in the eval DB across
+  γ=1–4.
+- A decision written down: ship speculation as the **default** only for
+  the (model, context) cells where it measurably wins; otherwise record
+  it as a characterized default-off lever (same disposition as the e4b
+  net-loss result — a documented option, never a deletion).
+
+## Caveats / interactions
+
+- **Batched serving (Phase 18).** Speculation needs trimmable caches for
+  partial-accept rollback — `generate.ts` throws once a `RotatingKVCache`
+  ring wraps past the sliding window — and variable accept-lengths per
+  slot break a uniform-B batch. Speculation is naturally a **B=1 latency**
+  optimization; batched-throughput serving is a different regime. Layering
+  both is real work, not free composition.
+- **One serve path, two backends.** The assistant-drafter loop and the
+  Qwen MTP runtime are two backends of the *same* "draft γ → verify in one
+  forward → accept-longest-prefix → rollback" abstraction. As Phase 14
+  (MTP) and Phase 18 (slots) land, build one verify/accept executor, not
+  two — the drafter vs MTP-head difference is just what fills the draft.
+- **Denominator discipline.** Speedups are only ever vs the *same model's*
+  non-spec baseline. e4b's 54 tok/s is the wrong denominator for 12B.

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/design/tauri-desktop-app.md

@@ -62,7 +62,7 @@ The window renders the **same page** `mlx-bun serve` serves today
    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
+   (`--port`) and the webview URL, to avoid 8080 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
@@ -99,7 +99,7 @@ then `xcrun stapler staple`.
 
 ## Milestones
 
-1. **Spike** — `tauri init` in a new repo; hardcode port 8090; spawn a
+1. **Spike** — `tauri init` in a new repo; hardcode port 8080; 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;