mlx-bun 0.0.9 → 0.0.11

package/docs/design/grammar-spec-batching-integration.md

@@ -0,0 +1,301 @@
+# Grammar × spec decode × batching — integration plan + feature-matrix benchmark
+
+Status: **ALL PHASES EXECUTED** (A/B/C/E 2026-07-03; D completed
+2026-07-04: KV-budget admission + vectorized homogeneous sampling +
+extend-join, each oracle/A-B gated — see the execution log). Remaining
+follow-ups live in "Later / parked" with their triggers.
+Execution log at the bottom of this doc. Successor plan tying together
+three workstreams so they compose deliberately instead of colliding:
+
+- **Structured output** — SHIPPED on main (serial + batched lanes;
+  docs/design/structured-output.md). Open debt: B2 model-gated batch
+  tests + bench, F4 compiler cache, F5 regex, F6 escaping, F7 structural
+  tags, U1/U2 engine upgrades.
+- **Batching perf path** — P5 done; P0–P4 open
+  (docs/design/batching-perf-path.md).
+- **Speculative decoding** — research proven (assistant drafter bit-exact
+  vs optiq, 12B γ=1 ≈ 1.09× win, docs/design/spec-decode-larger-targets.md;
+  DSpark architecture verified, docs/investigations/dspark-handoff.md) but
+  **not served**: no `--draft-model`. Build plan already written:
+  docs/design/mlx-lm-tool-parity-plan.md §7 (the `DraftSource` interface,
+  two-model L1 source, serve loop, admission, prompt-cache composition).
+
+This doc does NOT restate those plans — it sequences them, defines the
+composition rules where they meet, and specifies the benchmark that
+measures every combination. The per-feature designs stay canonical in
+their own docs.
+
+## The composition matrix (v1 target state)
+
+`GenerationGateway.willBatch` is the single routing point; every rule
+below is one predicate there. "Serial" routing is a v1 rule, not a
+limitation to hide — upstream mlx-lm routes spec serial too
+(`is_batchable = draft_model is None`, server.py:371).
+
+| Request shape | Lane | Mechanism | Status |
+| --- | --- | --- | --- |
+| plain | batch (or serial) | shipped engine | DONE |
+| grammar only | batch (or serial) | B1 per-row matchers / serial mask loop | DONE (B2 gates open) |
+| spec only | **serial** | `hasDraft` routes serial (upstream parity) | Phase B |
+| grammar + spec | **serial** | constrained verify walk (novel — §Phase C) | Phase C |
+| batch + spec | serial per request | `hasDraft` wins over batch | Phase B (routing) |
+| batch + grammar + spec | serial per request | `hasDraft` wins | falls out of B+C |
+
+Non-goals (explicitly out of scope, revisit only on measured demand):
+per-row spec inside the batch scheduler (variable accept-lengths break
+uniform-B; a research project, not a feature gap), DSpark as a served
+source before its 27B retrain/data-scale lands, grammar on the
+`/v1/responses` shim beyond what chat inheritance already gives.
+
+## Phases
+
+Ordering rationale: A is small and closes the trust gap on what already
+shipped. B is the flagship parity feature and lays the `DraftSource`
+seam C needs. C is the novel composition. D runs parallel to B/C
+(scheduler files vs serial-loop files — near-zero overlap). E needs all
+of them and produces the numbers.
+
+### Phase A — structured-output debt (B2 + F4)  [S, 2–3 days]
+
+1. **B2 model-gated scheduler tests** (spec'd in structured-output.md):
+   mixed batch mask isolation; **all-grammar B=4 with four different
+   schemas** (THE bug-class gate — row↔matcher misalignment after
+   eviction/join); early termination + churn + joiner; max_tokens
+   mid-JSON truncation under batch; prefill-terminated row never merges.
+2. **F4 compiler cache**: one `GrammarCompiler` per `TokenizerInfo`
+   (both vocab-structural; `ownsCompiler=false` path already exists).
+   Measure compile time cold vs cached on a real JSON schema — this is
+   the agentic-replay win and a benchmark row in Phase E.
+3. Opportunistic: F6 choice escaping (tiny), F3 explicit padded-vocab
+   comment. F5/F7 stay parked for the structural-tag pass (see Later).
+
+Exit: gated suites green on cpm5 + e4b; compile-cache A/B number
+recorded.
+
+### Phase B — `serve --draft-model` (two-model, serial lane)  [M, ~1 wk]
+
+Execute mlx-lm-tool-parity-plan.md §7 as written (it is already the
+hand-to-an-implementer plan). Integration deltas this plan adds:
+
+- **`RequestShape.hasDraft`** beside `hasGrammar`; `willBatch` returns
+  false for it. Lane-debug line includes it. (The B0-grammar routing
+  pattern, reused verbatim.)
+- **The serve verify loop lands grammar-aware from day one**: the loop
+  emits bursts of ≤ n+1 tokens; the stop-sequence/detokenizer audit in
+  §7 ("risky parts") must treat the grammar mask hook as a first-class
+  participant even while it's unused in this phase — i.e. the per-round
+  walk is structured as `for pos in verifyWindow: sample(logits[pos],
+  maskHook?)` so Phase C is a hook implementation, not a loop rewrite.
+- **Kill switch** `MLX_BUN_SPEC=0` (degrade: ignore the flag, log once),
+  house style.
+
+Oracle: L1 spec-vs-spec token equality with `mlx_lm.server
+--draft-model` from the oracle venv (§7.8). Acceptance telemetry
+(`drafted/accepted/targetCalls`) in stats — Phase E consumes it.
+
+Exit: greedy token-for-token vs mlx-lm spec on a fixed pair; rotating
+wrap falls back non-spec (no 500); `--batch 4` + `--draft-model` routes
+spec requests serial and batches the rest; admission accounts for draft
+weights + KV.
+
+### Phase C — grammar × spec: the constrained verify walk  [M, ~1 wk]
+
+Novel territory — **no oracle serves grammar+spec** (mlx-lm has no
+grammar; oMLX has no spec). No-oracle (Lab-class) gating per the
+parity-tier contract: equivalence + validity gates, not bit-parity to an
+ancestor.
+
+Design (v1 = verify-side masking, drafter free-running):
+
+1. **Drafter is unconstrained.** Drafts are proposals; correctness never
+   depends on them (a grammar-invalid draft is simply rejected at
+   verify). No matcher state touches the draft path.
+2. **The verify walk applies the mask sequentially.** One batched target
+   forward over `[pending, ...drafts]` (n+1 positions) as in Phase B;
+   then the host-side accept walk goes position by position:
+   `mask_i = matcher.bitmask` → apply to `logits[i]` → sample → if
+   sampled == draft_i: `matcher.accept(draft_i)`, continue; else emit
+   the masked correction token, `matcher.accept(correction)`, stop.
+   The matcher only ever advances on **emitted** tokens — no rollback
+   needed in v1. Sequential bitmask fills cost 0.004–0.19 ms each
+   (measured, structured-output.md); at γ≤3 that's ≤0.8 ms/round against
+   a 12B target step of ~40 ms — noise.
+3. **Grammar termination mid-burst** truncates the accepted prefix at
+   the terminating token (same all-`-inf` guarantee as serial/batched:
+   nothing is sampled past termination), finish_reason "stop".
+4. **The measured fork** ([[flag-and-try-both-on-forks]]): on
+   grammar-tight spans (e.g. `guided_choice`, closing braces) a
+   free-running drafter's acceptance may crater — every draft off-grammar
+   is a wasted round. If Phase E shows that, the alternative is built
+   behind `MLX_BUN_SPEC_GRAMMAR_DRAFT=1`: mask the drafter's logits too
+   (matcher advances per draft, `rollBack(n − kAccept)` on rejection —
+   the WASM API exposes rollBack; this is its use case). Build the flag
+   only if the default's acceptance number demands it; don't pre-build.
+
+Gates (no external oracle):
+- **Equivalence**: greedy grammar+spec long-prefix-equal to greedy
+  grammar-only serial on fixed prompts (exact-match acceptance means
+  spec must reproduce the non-spec masked-greedy stream, modulo the
+  documented batched-verify-head knife-edges — same caveat class as L1
+  spec-vs-spec).
+- **Validity**: every grammar+spec output parses / conforms to its
+  schema (100%, no exceptions), matcher never rejects an emitted token.
+- **Composition regression fence**: grammar-only and spec-only suites
+  byte-unchanged.
+
+### Phase D — batching P0 (+ admission slice of P3)  [M, ~1 wk, parallel with B/C]
+
+From batching-perf-path.md, the subset that the benchmark needs and that
+doesn't collide with B/C's serial-loop work:
+
+1. **P0**: `scripts/bench-serving-load.ts` (N concurrent vs live server,
+   aggregate tok/s + TTFT p50/p95 at the SSE first-token event) —
+   Phase E builds on this harness; `extend` join (kill the O(B·S)
+   re-merge); vectorized homogeneous-sampler fast path.
+   - `extend` + grammar rows: the B1 read-before-build step must hold
+     across an extend-join — add one B2-style churn test with a joiner
+     landing while a grammar row is live.
+2. **P3 admission slice only**: `projectKvBytes` + `--kv-budget`
+   queueing (the OOM backstop the combined benchmark needs before it
+   pushes B=4 × long prompts on 24 GB). Prompt-cache-under-batching,
+   adapters, and the default flip stay in batching-perf-path P3.
+
+P1 (quantized KV at B>1) is NOT in this plan's critical path — it's the
+next wave after the matrix benchmark exists, and the benchmark is exactly
+what will size its payoff. (P2 — perf kernel at B>1 — is OBSOLETE: the
+kernel was deleted 2026-07-05, docs/design/unified-engine-frontier-plan.md §6.)
+
+### Phase E — the feature-matrix benchmark  [S–M, 3–4 days]
+
+One harness answering "what does each feature — and each legal
+combination — cost or buy, per model?" Extends P0's
+`bench-serving-load.ts`; preflight-gated like benchmark.sh (dirty-machine
+numbers are garbage; quotable rows need the clean-machine protocol).
+
+**Cells** (skip-if-unavailable, like the model-gated suites):
+
+- Lanes/features: `serial`, `batch4`, `serial+grammar`, `batch4+grammar`,
+  `serial+spec`, `serial+grammar+spec` — the six legal shapes from the
+  matrix. (batch+spec cells route serial by design; the harness records
+  them as routing checks, not perf rows.)
+- Models: cpm5 + e4b (fast movers, both lanes), 12B (+ its assistant
+  drafter — the measured spec win), Qwen3.5 (batch + thinking-model
+  grammar canary; grammar cells expected to expose F7 until structural
+  tags land — record, don't hide).
+- Workloads: short-chat (128 tok), JSON-extraction (a real schema,
+  grammar cells), long-form (512 tok). 4 concurrent for batch cells,
+  1 for serial cells; N=3 median.
+
+**Metrics per cell**: aggregate + per-stream decode tok/s · TTFT
+p50/p95 (SSE first token, wall-clock) · spec acceptance
+(drafted/accepted/targetCalls) · grammar overhead vs the same cell
+without grammar (%) · schema-conformance rate (**must be 100% in every
+grammar cell — a conformance failure fails the run, it's a correctness
+gate wearing a benchmark costume**) · peak RSS.
+
+**Denominator discipline** (spec-decode-larger-targets.md): every
+feature cell is compared to the SAME model+lane baseline cell from the
+SAME run. No cross-model or cross-run speedup claims.
+
+**Output**: eval-DB rows tagged `feature-matrix[<host>]` + a dated
+markdown artifact (gitignored, like benchmark.sh's); promoted rows land
+in benchmarks/RESULTS.md under a new "composition" table. Josh-gated
+clean-machine run for anything quotable.
+
+Exit: the matrix runs end-to-end on this machine; a first
+(loaded-machine, unquotable) sweep sanity-checks every cell; the
+clean-machine run is queued in STATUS Josh-gated.
+
+## Later / parked (pointers, so they aren't lost)
+
+- **F7 + F5 structural-tag pass** (thinking-model grammar + real regex)
+  — try the installed WASM package's `fromStructuralTag`/`RegexFormat`
+  first; U1 rebuild only if its Dec-2025 core predates the semantics
+  (structured-output.md addendum). Qwen3.5's grammar cells in Phase E
+  are the forcing function.
+- **U2 native XGrammar-2 via TVM-FFI** — trigger conditions unchanged
+  (B≫8, spec+grammar draft-tree masking via `traverse_draft_tree`, or
+  compile time visible despite F4). Phase C's v1 walk deliberately
+  avoids needing it.
+- **Batching P1** (quant-KV at B>1) — next wave, sized by the matrix.
+  (P2, perf kernel at B>1, is obsolete — kernel deleted 2026-07-05.)
+- **P4 device-side step chaining** — orthogonal serial-decode work; the
+  matrix's serial baselines will show its headroom.
+- **Assistant/DSpark sources behind `--draft-model`** —
+  mlx-lm-tool-parity-plan §7.9; DSpark stays research until 27B
+  retrain + data scale (dspark-handoff.md).
+- **Spec inside the batch lane** — research question, not scheduled.
+
+## Execution log (2026-07-03)
+
+- **Phase A** ✅ — B2 gates in `tests/batch-grammar.test.ts` (all five
+  scenarios + a sixth found later, below); F4 compiler cache went through a
+  bring-up fix (single-flight rebuild — concurrent compiles after a
+  tokenizer switch double-disposed the previous compiler); F6, F3 landed.
+- **Phase B** ✅ — `serve --draft-model`/`--num-draft-tokens`;
+  `src/spec/{source,two-model,serve-loop}.ts`. **L1 gate: 48/48
+  token-for-token vs mlx-lm spec** (Llama 3B+1B greedy, 65% acceptance,
+  `scripts/oracle-spec-two-model.py`). Deviations as designed (ring-wrap
+  degrades pre-pollution via an offset+n+1 gate — upstream's post-hoc check
+  would already have polluted the ring; prompt cache bypassed v1).
+- **Phase C** ✅ — the constrained verify walk in `samplePos`; gates green
+  (100% validity; 12/12 token-identical to grammar-only serial; choice
+  terminates mid-spec). Drafter-masking flag NOT built (acceptance 6/15 on
+  a tight schema — usable; revisit only if Phase E's clean run says so).
+- **Phase E** ✅ (harness) — `scripts/bench-feature-matrix.ts` (the planned
+  bench-serving-load extension became its own script: that file already
+  exists as a client-only stack-vs-stack tool and stays one). Six cells
+  green end-to-end on Llama 3B+1B; `usage.speculation` telemetry on all
+  four API usage blocks. Clean-machine run for RESULTS.md is Josh-gated.
+- **The conformance gate paid for the whole phase in one smoke run** —
+  three real bugs: (1) snake_case wire fields never reached the grammar
+  resolver (structured output was DEAD over HTTP since it landed);
+  (2) `#flushPipeline` didn't advance matchers → stale masks on mid-decode
+  joins (regression test added); (3) **UniversalDenseModel batched RoPE
+  uses the scalar cache.offset** → uneven-row batches decode joiners at
+  wrong positions — latent for ALL Tier-0 archs since v0.0.9. **Fixed
+  same day**: `UniversalRope.applyDynamic` (+ `ops.ropeScaledDynamic`)
+  routes the wrapper's per-row `ropeOffsetArr`; gated **token-exact vs
+  mlx-lm B=2 on Llama-3.2-3B** — static uneven rows AND the dynamic
+  join/leave protocol (tests/batched-decode-parity.test.ts "Llama 3B
+  Tier-0"). Plain full-attention universal archs now batch; maskArray
+  (gemma2-family, pad-blind causal mask in forwardLayers) and
+  sliding-window universal archs remain serial (unvalidated cells).
+  Lesson, per [[per-model-quant-specialization]]: the gateway's
+  capability gate admitted a model family no batched-parity suite ever
+  validated.
+
+### Phase D completion (2026-07-04)
+
+- **KV-budget admission** — `--kv-budget <GB>`: aggregate projected-KV gate
+  in the scheduler (queue-don't-OOM, oversized-alone rejects,
+  `/stats.batch.{pending_rows,kv_bytes,kv_budget_bytes}`;
+  tests/batch-kv-budget.test.ts).
+- **Vectorized homogeneous sampling** — all-plain-greedy batches sample via
+  ONE log-softmax+argmax over [B,V]; gateway flags `plainGreedy`; BIT-equal
+  A/B vs the per-row path (`MLX_BUN_BATCH_VEC_SAMPLE=0` kill switch;
+  tests/batch-vec-sample.test.ts).
+- **Extend-join** — `extendKVRows` (mlx-lm `BatchKVCache.extend` semantics:
+  one pad + one B-axis concat, running rows untouched, pads grow) replaces
+  the O(B·S) re-merge for full-attention layers in `#mergeJoiner`
+  (`MLX_BUN_BATCH_EXTEND=0` = re-merge). **Own oracle** (extend's pad
+  layout ≠ re-merge's): scripts/gen-batched-extend-golden.py →
+  token-for-token on CPM AND Llama-3B (tests/batched-decode-parity.test.ts
+  "EXTEND-JOIN ORACLE"); full scheduler suite incl. the grammar-churn join
+  test green through the extend path AND the kill-switch path.
+  Rotating-layer extend (Gemma joins re-merge their sliding layers, extend
+  their full layers) is a follow-up.
+
+## Risks
+
+- **Grammar+spec has no oracle** — mitigated by the three-gate stack
+  (equivalence to grammar-only greedy, 100% validity, regression fence)
+  and by keeping the drafter out of the grammar state in v1.
+- **Multi-token bursts through stop-sequence/detokenizer state** — the
+  known §7 risky part, now load-bearing for two phases; land the burst
+  audit with Phase B tests, not after.
+- **Memory at 12B+drafter / B=4 on 24 GB** — Phase D's admission slice
+  is the backstop; benchmark cells must go through it, not around it.
+- **Scope creep via the matrix** — the benchmark measures what exists;
+  it must not become a reason to build P1/P2/drafter-masking
+  speculatively. Each parked item has its trigger written above.

package/docs/design/minicpm5-decode-megakernel.md

@@ -15,8 +15,9 @@ The inference forward composes per-op MLX nodes (`quantized_matmul`, `sdpa`, …
 We tried to insert our own fused `fast.metal_kernel`s (steel `BlockMMA` GEMM) into
 that forward to fuse+speed it up. **Confirmed dead end:** MLX's lazy buffer planner
 corrupts *chained* `BlockMMA` custom kernels (nondeterministic NaN/garbage/crash) —
-the pointwise `fused-gelu` chains fine, ours do not, and ≥2 of ours in one lazy
-graph corrupt. The only robust fixes are (a) a per-kernel eval boundary (24 syncs/
+the pointwise `fused-gelu` kernel chained fine (that kernel has since been deleted,
+2026-07-05, in favor of the compiled activations — unified-engine-frontier-plan.md
+§6), ours did not, and ≥2 of ours in one lazy graph corrupt. The only robust fixes are (a) a per-kernel eval boundary (24 syncs/
 forward → net **slower**, 0.5–0.8×), or (b) **don't give the planner anything to
 chain: one resident kernel for the whole forward.** This plan builds (b).
 
@@ -48,9 +49,11 @@ miracle. The *architectural* win (no corruption, fully ours) is the point.
    entire MiniCPM5 decode forward in **one** `MetalKernel.apply` dispatch.
 2. **Correctness gate (L1, bf16 KV):** teacher-forced greedy trajectory matches
    `goldens/minicpm5-parity.json` — argmax agreement **≥ 98/100**, mean per-step
-   logit maxDiff small (this is an L3 path: ~1–2 bf16 ULP/op accumulating; it will
-   NOT be bit-exact, gate by teacher-forced agreement + KL like
-   `tests/perf-kernel-oracle.test.ts`, threshold ≥ 56/64-style). DETERMINISTIC
+   logit maxDiff small (this is a no-oracle Lab-gated path: ~1–2 bf16 ULP/op
+   accumulating; it will NOT be bit-exact, gate by teacher-forced agreement + KL —
+   the frozen-trajectory envelope method of the old `tests/perf-kernel-oracle.test.ts`,
+   threshold ≥ 56/64-style; that test was deleted 2026-07-05 with its kernel, but
+   the methodology stands on its own). DETERMINISTIC
    across runs (no nondeterminism — that was the corruption signature).
 3. **Perf gate:** decode tok/s ≥ the serial baseline on a paired same-process A/B
    (`scripts/experiments/` style). Not slower. (Stretch: measurably faster.)
@@ -383,7 +386,7 @@ q-level / group-min-max boundary and flips a level → a full quant-step dequant
 amplified 9.7e-4 → 1.38e-2. **NOT a quant bug — the megakernel's intrinsic ~1-ULP GEMV
 difference amplified by quant discontinuity. 93/100 is the L2 ceiling (as 97/100 is L1's);
 bit-exact L2 would need a bit-exact GEMV, defeating qmv4.** Gate L2 by KL + agreement
-(L3 class), not the bit-exact golden. Scripts: `megakernel-kv-{cmpkv,cmpl1,perstep}.ts`.
+(no-oracle, Lab class), not the bit-exact golden. Scripts: `megakernel-kv-{cmpkv,cmpl1,perstep}.ts`.
 **Increment 2b — MEMORY win (later):** bit-pack `q` into uint32 (flat KPACK/VPACK,
 per-layer literal offsets) to actually shrink 4-bit KV. Same logits as 2a. Deferred.
 **Build on the GENERATED kernel** (`genSourceMt`, `MLX_BUN_MEGAKERNEL_GEN`): the
@@ -404,8 +407,10 @@ teacher-forced agreement + KL. Defer until L1 solid (L1 IS solid — parity, Pha
 per-step `forwardHidden→logitsFromHidden` at decode (M=1) only; prefill unchanged.
 Guard: bf16 KV (L1) only, no adapters, no vision. Add `--megakernel` CLI flag.
 **Gate:** `bun test tests/minicpm5-*.test.ts` green (off); a new
-`tests/minicpm5-megakernel.test.ts` (teacher-forced gate, env-gated like
-`tests/perf-kernel-oracle.test.ts`); paired decode tok/s A/B logged. Update
+`tests/minicpm5-megakernel.test.ts` (teacher-forced gate, env-gated in the pattern
+of the old `tests/perf-kernel-oracle.test.ts` — deleted 2026-07-05 with its kernel,
+unified-engine-frontier-plan.md §6; the envelope-gate methodology stands on its
+own); paired decode tok/s A/B logged. Update
 `STATUS.md` + `benchmarks/RESULTS.md` (preflight-gated `benchmark.sh` for the
 quotable number — `[[dirty-machine-numbers-are-garbage]]`).
 
@@ -425,7 +430,7 @@ catch any divergence the moment it's introduced, not 24 layers + a quant step do
 Steps: (1) copy mlx `quantized.h` GEMV kernels INCLUDING the per-(bits,M,N,K) dispatch
 (`qmv_fast`/`qmv`/`qmv_quad`), `rms_norm`, `rope`, `sdpa`; (2) per-op maxDiff-0 harness
 before fusing; (3) analyze the single fused file; (4) optimize with the harness as a
-regression gate. This megakernel was built L3-first (KL-gated, NOT bit-exact by design),
+regression gate. This megakernel was built no-oracle-first (Lab-style, KL-gated, NOT bit-exact by design),
 so its 1-ULP/op drift is within contract — bit-exact is a focused REBUILD, not a patch
 (and may cost qmv4's speed if mlx's real 4-bit kernel is the slower qmv_quad).
 
@@ -438,8 +443,8 @@ so its 1-ULP/op drift is within contract — bit-exact is a focused REBUILD, not
   param. Bake dims as `templateInts` / read `x_shape[0]`, don't pass a runtime
   `shape` input array (extra lazy nodes).
 - **Match the reference's bf16 rounding op-for-op** (round to `T` at each boundary) —
-  `metal::precise::exp`/`tanh` are the only residual vs the reference; that's L3 and
-  fine. Do NOT compute everything in f32 then round once (drifts).
+  `metal::precise::exp`/`tanh` are the only residual vs the reference; that's
+  no-oracle (Lab-gated) territory and fine. Do NOT compute everything in f32 then round once (drifts).
 - **`atomicOutputs:true` + `initValue:0`** for any atomic-accumulated output.
 - **GQA:** 16 q-heads, 2 kv-heads, `nRep=8` → q-head `h` reads kv-head `h/8`.
 - **`rope_theta=5e6`, not 10000.** scale `=128^-0.5`. No MiniCPM residual scaling.

package/docs/design/mlx-lm-tool-parity-plan.md

@@ -324,7 +324,17 @@ follow-on if we build the shared capture instrumentation anyway.
 
 ## 7. `serve --draft-model` / `--num-draft-tokens` (speculative decoding at serve time)
 
-This is the concrete, hand-to-an-implementer plan.
+> **BUILT 2026-07-03** as grammar-spec-batching-integration.md Phase B —
+> `src/spec/{source,two-model,serve-loop}.ts`, flags in cli.ts, `hasDraft`
+> routing, fit admission, `usage.speculation` telemetry. **L1 gate passed:
+> 48/48 token-for-token vs mlx_lm's speculative path** (Llama 3B+1B,
+> `scripts/oracle-spec-two-model.py`, `tests/spec-serve.test.ts`). Grammar
+> composes (Phase C constrained verify walk). v1 deviations: prompt-cache
+> reuse bypassed on the spec path; ring-wrap degrades to plain decode
+> (pre-pollution gate) instead of raising. §7.9's assistant/DSpark sources
+> remain open behind the same `DraftSource` seam.
+
+This was the concrete, hand-to-an-implementer plan (kept for the record).
 
 ### What mlx_lm.server does (read: `server.py`, `generate.py`)
 

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

@@ -0,0 +1,98 @@
+# 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); since: grammar under batching, `--kv-budget` admission, Tier-0 universal (Llama) per-row RoPE (2026-07-03) | 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. **LANDED 2026-07-02, MERGED to main 2026-07-03** (branch
+   `feat/structured-output` deleted after merge):
+   `@mlc-ai/web-xgrammar` (WASM) on `/v1/chat/completions` + `/v1/completions`;
+   full `response_format`/`guided_*` surface, serial AND batched lanes;
+   L2-verified vs oMLX (byte-identical
+   content). See [docs/design/structured-output.md](structured-output.md);
+   follow-up debt sequenced in
+   [grammar-spec-batching-integration.md](grammar-spec-batching-integration.md).
+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. *Note: serve-time speculative
+   decoding itself LANDED 2026-07-03 (`serve --draft-model`, two-model
+   serial lane, 48/48 token-for-token vs mlx-lm's spec path —
+   grammar-spec-batching-integration.md Phase B); DFlash/DSpark would slot
+   behind the same `DraftSource` seam.*
+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/orpo-training.md

@@ -23,11 +23,11 @@ runs half the forwards per step.
 
 > Oracle note: neither `mlx_lm` nor `optiq` ships ORPO, so — unlike DPO
 > ([`loss.ts`](../../src/train/loss.ts) ported `optiq/lora/dpo.py`) — ORPO is
-> an **L3 mlx-bun original**. Correctness is anchored to a standalone
+> a **no-oracle mlx-bun original (Lab-class)**. Correctness is anchored to a standalone
 > torch/numpy reference of the paper/TRL math (see [Parity](#parity)), not a
 > bit-exact upstream port.
 
-> **L3 verification (2026-06-21).** The flash-CCE head has no oracle either, so
+> **No-oracle (Lab-gated) verification (2026-06-21).** The flash-CCE head has no oracle either, so
 > the "0.28% dh" figure is **fp-reassociation vs the full-logits *proxy*, not an
 > error** (two correct reductions in different orders can't be bit-identical).
 > Correctness is proven by **finite-difference** (`flash-fd-check.ts`) + the math
@@ -35,7 +35,7 @@ runs half the forwards per step.
 > (0.66→2.7%, from `flash-cce-parity.ts`'s random hidden) is its **worst case**:
 > random logits give a flat softmax, but real outputs are sharply peaked, so
 > filtering the ≈0-softmax tail is near-free — **measure teacher-forced on real
-> hiddens before judging** (it's plausibly a free speedup to enable). Standing L3
+> hiddens before judging** (it's plausibly a free speedup to enable). Standing Lab
 > gates (filter-on-real-data, teacher-forced grad fidelity, end-to-end quality)
 > + the host-buffer pin-leak post-mortem:
 > [orpo-flash-cce-pin-leak.md](../investigations/orpo-flash-cce-pin-leak.md).
@@ -68,11 +68,13 @@ else validates against):**
   [`scripts/experiments/parity-orpo.ts`](../../scripts/experiments/parity-orpo.ts):
   10/10 checks pass — scalar loss, log1mexp stability, gradient checks (finite
   differences vs autograd) at B=1 and B=3, edge cases.
-- ✅ `fusedGegluDifferentiable` in
-  [`src/model/fused-geglu-kernel.ts`](../../src/model/fused-geglu-kernel.ts):
-  `CustomVjp`-backed version with hand-derived backward (grad\_b = dc·gelu(a),
-  grad\_a = dc·b·gelu'(a)); tested in
-  [`tests/fused-geglu-vjp.test.ts`](../../tests/fused-geglu-vjp.test.ts).
+- ✅ `fusedGegluDifferentiable` in `src/model/fused-geglu-kernel.ts`
+  (**kernel DELETED 2026-07-05** in the Phase-1 deletion pass —
+  [unified-engine-frontier-plan.md §6](unified-engine-frontier-plan.md);
+  the compiled activations are the only geglu now, bit-exact, and training
+  needs no flag sanitization): `CustomVjp`-backed version with hand-derived
+  backward (grad\_b = dc·gelu(a), grad\_a = dc·b·gelu'(a)); was tested in
+  `tests/fused-geglu-vjp.test.ts` (deleted with it).
 - ✅ **Fused GeGLU wired into the e4b training path.** A training-mode flag
   (`setFusedGeluTraining`, set by the trainer around the loop, cleared in its
   finally) makes BOTH Gemma GeGLU sites — the MLP (`gelu(gate)·up`) and the
@@ -92,10 +94,11 @@ else validates against):**
   to the closed form at tol 0.05 in the vjp test). Composes with the segmented
   path (the fused `CustomVjp` nests inside the segmented `mlx_vjp`: segmented ORPO
   still PASSES with fused on, peak 11.25 GB full → 7.58 GB segmented). Unit guard
-  ([`tests/fused-geglu-vjp.test.ts`](../../tests/fused-geglu-vjp.test.ts)): the
+  (`tests/fused-geglu-vjp.test.ts`, deleted with the kernel): the
   flag toggles and the differentiable forward bit-matches the plain kernel.
-  Training still needs `MLX_BUN_PERF_KERNEL=0` (a separate no-vjp perf kernel —
-  its own follow-on).
+  (Historical: training at the time also needed `MLX_BUN_PERF_KERNEL=0` — that
+  perf kernel was likewise deleted 2026-07-05, so no training flag sanitization
+  remains.)
 - ✅ `responseOnlyLogpMean` helper in `loss.ts` (post-finalNorm h → mean response
   logp; used by the ORPO head VJP).
 - ✅ `SegmentedBackwardOrpo` in
@@ -520,7 +523,7 @@ tier below it, bottoming out at the math reference above:
 
 ## Memory & speed optimization
 
-ORPO has no upstream bit-exact oracle (L3 original), which gives us latitude to
+ORPO has no upstream bit-exact oracle (a no-oracle original, Lab-class), which gives us latitude to
 optimize the runtime aggressively as long as we stay self-consistent with a
 naïve reference (see [Parity](#parity)). This section is the result of a
 dedicated research spike across three areas — activation/segmentation memory,
@@ -568,7 +571,7 @@ GB** at this vocab scale. This **subsumes our current `responseOnlyCe` trick**
     LoRA-factored head backward (`∇A`, `∇B`) falls out per chunk — Apple's CCE
     has no LoRA path, so even adopting it would mean deriving this ourselves.
     Exact in forward and backward; no Metal kernel required. Reserve a
-    `mx.fast.metal_kernel` flash-CCE as an **L3 speed tier** only if per-chunk
+    `mx.fast.metal_kernel` flash-CCE as a **Lab speed experiment** only if per-chunk
     logits are still too large in practice.
 
 > ✅ **LANDED + validated (2026-06-18) — `fusedLogpMeanB1` in
@@ -640,7 +643,7 @@ GB** at this vocab scale. This **subsumes our current `responseOnlyCe` trick**
 > token-chunks (which free fine). Forcing an eval mid-chain breaks autodiff (a
 > 35 GB blowup). **Conclusion:** the CCE residency bound (`[chunk,Vblock]` peak) is
 > unreachable in pure MLX; it requires a `mx.fast.metal_kernel` with the vocab loop
-> INSIDE the kernel (the L3 flash-CCE tier). The pure-MLX implementation is kept as
+> INSIDE the kernel (the flash-CCE Lab kernel). The pure-MLX implementation is kept as
 > a **validated correctness reference** for that kernel and is **gated off** (not
 > wired to a trainer knob — it would only raise peak). So: token-chunking
 > (`orpo_chunk_size`) + the analytic fused head (`orpo_fused_ce`) are the usable
@@ -649,8 +652,8 @@ GB** at this vocab scale. This **subsumes our current `responseOnlyCe` trick**
 > 🔬 **flash-CCE Metal kernel — PROTOTYPED + working, scaling-limited at largest
 > vocab (2026-06-18).** [`src/train/flash-cce.ts`](../../src/train/flash-cce.ts):
 > `mx.fast.metal_kernel` forward (`logp` + `lse`) and backward (`dh`) that compute
-> logits in-kernel from the quantized head (in-Metal 4/8-bit affine dequant via the
-> qdot pattern from [`fused-decode-kernel.ts`](../../src/model/fused-decode-kernel.ts)),
+> logits in-kernel from the quantized head (in-Metal 4/8-bit affine dequant via
+> mlx's `quantized.h` qdot pattern — the deleted fused-decode kernel used the same),
 > online softmax across vocab, target capture, softcap + sech², so the vocab loop
 > lives INSIDE the kernel — neither `[M,V]` nor a dequantized `[V,hidden]` head
 > touches HBM. **Validated** ([`scripts/experiments/flash-cce-parity.ts`](../../scripts/experiments/flash-cce-parity.ts)):
@@ -710,8 +713,8 @@ GB** at this vocab scale. This **subsumes our current `responseOnlyCe` trick**
 >   knob (the existing `orpo_chunk_size` stays the token-axis tier; the two
 >   compose — token-chunk the M response rows, vocab-chunk each row's head). It
 >   replaces the per-chunk `logitsFromHidden` (which materializes `[Cc, V]`) with
->   the tiled engine below, wrapped in a single `CustomVjp` (our
->   [`fused-geglu-kernel.ts`](../../src/model/fused-geglu-kernel.ts) pattern: a
+>   the tiled engine below, wrapped in a single `CustomVjp` (the pattern the
+>   since-deleted `fused-geglu-kernel.ts` established: a
 >   `mx.custom_function` with a hand `.vjp`, validated against `mx.vjp`/autograd —
 >   never a hand formula, per the standing rule).
 > - **Explicit VJP inputs** `(hC [M,hidden], baseW, scales, biases, [loraA, loraB
@@ -952,7 +955,7 @@ for custom fused Metal kernels. Both apply to **forward and backward**.
    already fused both directions. (Note: `sdpa` forward only — its backward is
    the custom-kernel target in Tier 1.)
 
-**Tier 1 — custom fused fwd+bwd Metal kernels (the cutting-edge; L3 originals):**
+**Tier 1 — custom fused fwd+bwd Metal kernels (the cutting-edge; no-oracle originals, Lab-gated):**
 3. **Fused linear + log-softmax + gather loss (CCE/Liger-style) — THE moat.**
    Lever 1 of the memory plan as a single kernel: fuse the (quantized,
    vocab-chunked) head matmul + `logsumexp` + target gather so the `[T, 262k]`
@@ -968,14 +971,15 @@ for custom fused Metal kernels. Both apply to **forward and backward**.
    memory instead of today's O(L²). High value because it removes the term
    segmentation is currently compensating for. (FA3's extra tricks are
    Hopper-specific; mirror FA2 on Metal.)
-5. **Fused GeGLU/SwiGLU with a vjp.** We already have a fused-GeGLU *forward*
-   kernel ([`fused-geglu-kernel.ts`](../../src/model/fused-geglu-kernel.ts))
-   **disabled in training because it has no gradient** (the
-   `MLX_BUN_FUSED_GELU=0` constraint). The backward math is known (Liger): for
+5. **Fused GeGLU/SwiGLU with a vjp.** (The custom fused-GeGLU/SwiGLU kernels
+   were **DELETED 2026-07-05** — compiled activations own the fusion win
+   bit-exactly; see
+   [unified-engine-frontier-plan.md §6](unified-engine-frontier-plan.md). Kept
+   here as the backward-math record.) The backward math is known (Liger): for
    `c = act(a)⊙b`, save `a`,`b`, recompute the activation — `db = dc·act(a)`,
-   `da = dc·b·act'(a)`. Writing this `.vjp` re-enables the kernel in training —
-   doubly valuable since the MLP intermediate is our **dominant resident tensor**
-   (the intra-layer split target, lever 4).
+   `da = dc·b·act'(a)`. This `.vjp` was written and validated before the
+   deletion; the memory angle still matters since the MLP intermediate is our
+   **dominant resident tensor** (the intra-layer split target, lever 4).
 6. **Fused dequant-GEMM + LoRA epilogue (+ manual backward).** MLX's
    `quantized_matmul` already fuses dequant into the *base* GEMM; the missing
    piece is folding the LoRA epilogue `scale·(x@A)@B` into that kernel and a
@@ -1021,9 +1025,9 @@ Metal path**. Two openings stand out:
   checkpointing / segmented backward. Other gaps nobody fills: cached DPO
   reference log-probs, a concatenated/prefix-shared chosen+rejected forward.
 
-These fused kernels are **L3 originals** (mlx-bun's own tier) — validated
+These fused kernels are **no-oracle originals (Lab-gated)** — validated
 against the naïve path via the [oracle ladder](#parity), consistent with the
-three-level fidelity model.
+fidelity model (external-oracle tiers L1/L2 plus the no-oracle Lab class).
 
 ## Model coverage, quantization, batching, dropout
 
@@ -1181,8 +1185,8 @@ Runbook (the generate/train steps are yours to run — they're long / download):
 2. **Baseline:** `MODEL=<e4b> bun scripts/run-ifeval.ts ifeval.jsonl` → record.
 3. **Fit check:** `MODEL=<e4b> CHUNK=512 bun scripts/bench-orpo.ts` to confirm
    the config fits 24 GB (e4b's 262k vocab makes the chunked head matter).
-4. **Train ORPO** (e4b is Gemma → set the no-vjp kernels off):
-   `MLX_BUN_FUSED_GELU=0 MLX_BUN_PERF_KERNEL=0` + `method:"orpo"`,
+4. **Train ORPO** (no env-flag sanitization needed — the no-vjp kernels were
+   deleted 2026-07-05): `method:"orpo"`,
    `learning_rate 1e-5`, `orpo_lambda 0.1`, `orpo_chunk_size 512`,
    `max_seq_length ~1024`, `rank 16`. ⚠️ Segmented backward for ORPO is
    **MiniCPM5-only** today (`SegmentedBackwardOrpoGemma4` is a follow-on), so