mlx-bun 0.0.10 → 0.0.12

package/docs/design/audio-input-plan.md

@@ -0,0 +1,356 @@
+# Audio input — design & implementation plan (gemma-4 audio tower)
+
+Status: **A0–A4 DONE** (2026-07-07: served end to end on e4b, oracle-EXACT
+greedy offline AND over HTTP; A5 bench + 12B coverage open).
+Scope: audio **understanding** — audio-in, text-out through the chat API.
+TTS / STT-as-an-endpoint / STS stay out of scope (see
+docs/design/omlx-adoption-map.md); §5 lists non-goals explicitly.
+
+## 0. TL;DR
+
+e4b ships a 12-block Conformer audio tower that we already have on disk —
+the `optiq_vision.safetensors` sidecar in the local e4b OptiQ-4bit snapshot
+contains 752 `audio_tower.*`/`embed_audio.*` tensors, and its `config.json`
+carries the full `audio_config` plus the audio token ids. mlx-lm strips all
+of this and rejects audio requests, so (as with vision) **the oracle is
+optiq's internal machinery**, which is complete end-to-end but not exposed
+by optiq's own serve frontend. The implementation is a fourth tower next to
+`src/vision/{siglip,embedder,diffusion-vision}.ts`, reusing every existing
+multimodal seam: sidecar lazy-load, `<|placeholder|>` splice, merged
+embeddings into `forwardEmbeddings`, serial-lane routing. New primitives
+needed: a `mlx_conv2d` binding (SSCP subsampler) and a small TS-side
+mel-spectrogram extractor (USM pipeline). Everything else is pattern-match
+on the vision port.
+
+## 1. What the ancestors do (surveyed 2026-07-06, pinned oracle venv)
+
+### mlx-lm 0.31.3 — nothing, and actively strips
+
+- `models/gemma3n.py:602` and `models/gemma4.py:55` `sanitize()` pop
+  `audio_tower`, `embed_audio` (and the vision keys) from the weights at
+  load. The upstream capability is deliberately discarded.
+- `server.py:118–150` `process_message_content()` raises
+  `ValueError("Only 'text' content type is supported.")` for any non-text
+  content part — `input_audio` requests hard-fail.
+- No mel/spectrogram/feature code anywhere in the package; no mlx-vlm /
+  mlx-audio / mlx-whisper installed in the oracle venv.
+
+**Parity-tree consequence:** there is no mlx-lm arm for audio. Like the
+vision port, correctness gates come from optiq's internals; anything no
+ancestor does (e.g. ever batching audio requests) is novel territory →
+KL + quality gates, not bit-parity (see three-tier doctrine in PLAN.md).
+
+### mlx-optiq — complete machinery, unexposed at the serve layer
+
+All under `optiq/vlm/_mlxvlm/models/gemma4/` in the oracle venv:
+
+- **Feature extraction** (`audio_feature_extractor.py`,
+  `Gemma4AudioFeatureExtractor`) — USM pipeline: 16 kHz mono input,
+  frame 20 ms / hop 10 ms, magnitude STFT, 128 HTK-scale mel bins over
+  0–8 kHz, `log(mel + 1e-3)` floor, optional per-bin mean/stddev
+  normalization; `preemphasis=0.0`, `fft_overdrive=False`, `dither=0.0`
+  by default. Runs CPU-side (numpy). Output
+  `input_features [B,T,128]` + validity mask.
+- **Tower** (`audio.py`, `AudioEncoder`) —
+  `SubSampleConvProjection`: two Conv2d stages (channels `(128, 32)`,
+  each 2× time stride) then linear to `hidden_size=1024`; then
+  **12 Conformer blocks**: FF (residual_weight 0.5) → chunked local
+  attention (chunk 12, left context 13 chunks, right 0, logit softcap
+  50.0, invalid logits −1e9) → depthwise light Conv1d (kernel 5) → FF →
+  clamp + RMSNorm; then output projection to `output_proj_dims=1536`.
+  `use_clipped_linears=True`: linears clamp activations to recorded
+  per-tensor `input_min/max`/`output_min/max` ranges (those stats ship in
+  the sidecar — confirmed present in our local e4b sidecar).
+- **Embedder** (`merge.py`, `MultimodalEmbedder`) — RMSNormNoScale +
+  bias-free Linear into text hidden size; embeddings scattered into the
+  input sequence at `audio_token_id` positions via `masked_scatter`
+  (`gemma4.py get_input_embeddings`).
+- **Tokens & splicing** (`processing_gemma4.py:264–299`) —
+  `<|audio|>` expands to `boa + audio_token×n + eoa` with
+  `n = min(ceil(duration_ms / 40), 750)` (40 ms/token, ≤30 s ≙ 750 soft
+  tokens). Ids (confirmed in our local e4b config.json):
+  `boa=256000`, `audio=258881`, `eoa=258883`.
+- **Request parsing** (`processing_gemma4.py:690`) — accepts OpenAI-style
+  parts `{"audio", "input_audio", "audio_url"}`.
+- **Sidecar & quant policy** — `vlm/sidecar.py` extracts
+  `audio_tower.*`/`embed_audio.*` into the same `optiq_vision.safetensors`
+  as vision, kept **bf16**; `backends/mlx_backend.py:491` skips
+  `audio_tower.` in the quantization predicate. Calibration recipe (and a
+  worked example of driving the audio path programmatically):
+  `calibration/datasets.py load_audio_calibration` (LibriSpeech clean).
+- **The gap:** `vlm/gemma4/frontend.py` only wires **vision** into optiq
+  serve; audio works only through the internal model class. Golden
+  generation therefore scripts the internal model directly (the
+  calibration loader shows exactly how) — same move as the vision port.
+
+## 2. What we already have (verified on this machine, 2026-07-06)
+
+- **e4b OptiQ-4bit** (snapshot `fcdb12…`): sidecar contains **752 audio
+  tensors** (`audio_tower.*` + `embed_audio.*`, incl. clipped-linear
+  min/max stats); `config.json` has `audio_config`
+  (`model_type=gemma4_audio`, 12 layers, hidden 1024, heads 8, conv
+  kernel 5, subsampling channels (128,32), logit cap 50, residual 0.5,
+  `output_proj_dims=1536`, `use_clipped_linears=true`) plus
+  `audio_token_id/boa_token_id/eoa_token_id`; `tokenizer_config.json` has
+  `audio_token: "<|audio|>"`. **No download needed — audio-ready today.**
+- **12B OptiQ-4bit**: `audio_config` present but the local sidecar holds
+  only **1** audio tensor → needs a sidecar rebuild via optiq's
+  `build_vision_sidecar` (its selective-download path pulls only the
+  audio shards, not the full base model).
+- **26B-A4B, DiffusionGemma, bf16 assistants**: no `audio_config` — no
+  audio for these, ever (architecture, not a porting gap).
+- **Vision pipeline is the template** — every seam already exists:
+  content-part detection `src/server.ts:2234`, extraction + splice + merge
+  `src/vision/prompt.ts`, lazy tower `src/server.ts:326–350`, sidecar
+  detection `src/registry.ts:487` / `src/config.ts:287`, embeddings
+  prefill `src/generate.ts:468` → `forwardEmbeddings`
+  `src/model/gemma4.ts:876–896` (per-layer-input id zeroing + bidir mask).
+- **FFI gaps**: `mlx_conv1d` is bound (`src/mlx/ffi.ts:173`, Qwen3.5
+  DeltaNet, depthwise) but there is **no conv2d and no FFT** binding.
+  SSCP needs conv2d; mel extraction needs an FFT (resolved TS-side, §3.1).
+
+## 3. Design
+
+### 3.1 New files (mirror of `src/vision/`)
+
+- **`src/audio/decode.ts`** — WAV parser (PCM16/24/32/float32) → mono
+  mix-down → resample to 16 kHz (windowed-sinc or linear; gate vs oracle
+  resampler on the fixture). Non-WAV formats (mp3/m4a/flac/ogg): shell out
+  to macOS `afconvert` (built into darwin, our only platform) to produce
+  WAV — zero new dependencies. **Duration (and thus soft-token count) is
+  computed from decoded samples**, so format never affects splicing.
+- **`src/audio/features.ts`** — USM mel extractor, verbatim-semantics port
+  of `audio_feature_extractor.py` (megakernel copy-verbatim methodology:
+  match first, optimize never — this is ~100×128 floats, not a perf
+  path). STFT frame 320 samples / hop 160 / FFT 512
+  (`fft_overdrive=False`), 128 HTK mel bins 0–8 kHz, `log(x + 1e-3)`,
+  per-bin normalization from shipped stats (open Q3). Implemented as a
+  small radix-2 real FFT in TS, f64 accumulate → f32 out, CPU-side —
+  exactly where optiq does it (numpy), keeping the mlx graph clean.
+  *Considered and rejected for v1*: binding `mlx_fft_rfft` and doing it
+  on-device — new binding surface for no perf need.
+- **`src/audio/conformer.ts`** — the tower, bf16 mlx graph:
+  - **needs `mlx_conv2d` binding** — read the full mlx.h signature first
+    (trailing-optional-param → shifted stream arg hazard, CLAUDE.md).
+  - SSCP: Conv2d(1→128) → Conv2d(128→32), each 2× time stride → linear
+    to 1024 (exact shapes read from the sidecar weights).
+  - 12 × ConformerBlock as in §1; chunked local attention implemented as
+    an explicit dense additive mask (≤750 frames pre-subsample, ~188
+    post — dense is trivially cheap at this size) with logit softcap 50
+    and −1e9 invalid fill.
+  - Clipped linears: clamp to shipped `input_min/max`/`output_min/max` —
+    parity depends on this, port exactly (toggle test in A2).
+  - Output proj → 1536, then `embed_audio` (RMSNormNoScale + Linear →
+    text hidden), **pre-divided by embed_scale** — same convention as
+    vision towers (`forwardEmbeddings` re-multiplies).
+- **`src/audio/prompt.ts` folded into a shared multimodal prompt
+  builder** — recommendation: generalize `src/vision/prompt.ts`
+  (`buildVisionPrompt` → `buildMultimodalPrompt`) rather than add a
+  parallel audio-only path, because one message can legally carry images
+  AND audio and both must splice in document order. Vision-only and
+  audio-only fall out as special cases; the vision goldens gate the
+  refactor (spliced ids must stay byte-identical).
+
+### 3.2 Server surface
+
+- Accept OpenAI content parts:
+  `{"type":"input_audio","input_audio":{"data":"<b64>","format":"wav"|"mp3"}}`
+  (OpenAI canonical) plus optiq's `{"type":"audio"}` / `{"type":"audio_url"}`
+  aliases. `hasAudio` detection lands next to `hasImages`
+  (`src/server.ts:2234`).
+- `getAudioTower(ctx)` lazy-loads next to `getVisionTower`
+  (`src/server.ts:326–350`) — same sidecar file, separate tower object;
+  text-only sessions never pay for it. Auto-enable on
+  `audio_config` + sidecar audio tensors present, mirroring vision —
+  **no new flags** (target surface: zero CLI delta).
+- Routing: audio requests take the **serial lane**, same rule as vision —
+  embeddings-prefill doesn't batch, and `promptEmbeddings` already
+  excludes the pre-warmed prompt cache (`src/generate.ts:468`). Batching
+  is a mode, not a fallback; batched audio has no oracle → if we ever
+  want it, it goes through KL + quality gates as new territory.
+- Anthropic-protocol endpoint: the Anthropic messages API has no audio
+  input block type → clean 400 with a pointer to the OpenAI endpoint;
+  documented.
+- **Docs land in the SAME commit** (standing rule):
+  `docs/reference/server-api.md` (content-part schema, ≤30 s / 750
+  soft-token cap, 16 kHz mono internal format),
+  `docs/reference/features-matrix.md` (audio row incl. serial-lane
+  note), README capability line, `cli.md` only if a flag appears
+  (target: none).
+
+### 3.3 Open questions — RESOLVED against the oracle sources (2026-07-07)
+
+1. **LM attention semantics for audio tokens: strictly CAUSAL.**
+   `language.py:484–545 _make_masks`: the blockwise bidirectional overlay
+   keys off `is_vision = (mm_token_type_ids == 1) | (== 2)` only — audio
+   (type 3) never enters a block. Stronger: `use_bidirectional_vision`
+   requires `not has_audio_tokens`, so **any audio token in the prompt
+   disables the vision overlay entirely** (oracle comment: "Audio spans
+   are sequential; keep mixed image+audio prompts causal"). mlx-bun rule:
+   audio runs get NO bidir mask, and a mixed image+audio request drops
+   the image bidir mask too — pass no bidir mask at all in that case.
+2. **Per-layer-input zeroing: yes, audio included.** `gemma4.py:92–99`
+   zeroes per-layer ids for the union mask
+   `(ids==image_token) | (ids==audio_token) | (ids==video_token)`.
+   Extend the existing image-mask zeroing in
+   `src/model/gemma4.ts:876–896` to the combined multimodal mask.
+3. **Per-bin mel stats: constructor defaults, no normalization.**
+   `processing_gemma4.py:919–934`: standard HF checkpoints ship no
+   `feature_extractor` key (our e4b snapshot has no
+   processor_config.json at all) → `Gemma4AudioFeatureExtractor()` with
+   defaults; "the USM parameters are fixed for all Gemma 4 models".
+   `per_bin_mean/stddev = None` → that branch is dead for us. Also
+   pinned there: `audio_seq_length` default 750, `audio_ms_per_token`
+   default 40.
+4. **boa/eoa are splice-side.** e4b `chat_template.jinja:301,334` emits
+   bare `<|audio|>` for audio parts (both content shapes); the
+   boa + audio×n + eoa expansion is processor-side — same ownership as
+   vision's boi/eoi. Splice in our prompt builder, not the template.
+
+### 3.4 Parity strategy (tree doctrine)
+
+Oracle = optiq internal gemma4 model in the pinned venv, driven by
+`scripts/regen-audio-goldens.py` (modeled on their calibration loader;
+remember `enable_thinking` pinning for any templated prompt). Tiers follow
+the e4b vision precedent exactly:
+
+- **T0 model-free** (CI, no weights): WAV decode + mel features vs
+  oracle-dumped golden (f32 tolerance ~1e-5 — DSP is deterministic);
+  spliced ids **exact**; soft-token count **exact**.
+- **T1 golden tower**: `audio_tower`+`embed_audio` output vs oracle,
+  rel-RMSE gate (expect vision-like ~1% from bf16 composition; calibrate
+  the threshold from the first honest run, don't invent it).
+- **T2 e2e greedy** (`tests/e4b-audio.test.ts`): first-N greedy prefix
+  match vs oracle + grounded-output check (speech fixture → transcript
+  keyword present; chirp fixture → tone/beep mentioned).
+- **T3 serve**: curl e2e with base64 `input_audio` against a running
+  server (manual/gated; Josh runs servers).
+
+Fixtures (tracked in `fixtures/`, tiny): (a) 1.6 s synthesized chirp WAV
+(~50 KB, generated by a checked-in deterministic script), (b) ~2 s public
+domain speech clip (LibriSpeech clean sample, ~64 KB). Goldens:
+`goldens/e4b-audio.json` manifest tracked, `.bin` blobs regenerable and
+untracked (goldens/README.md convention).
+
+### 3.5 Perf & memory
+
+Expectation (to be *measured*, not asserted): the tower is tiny —
+≤750 mel frames → ~≤188 post-subsample positions × 12 layers × hidden
+1024 bf16 — so added TTFT should be dominated by CPU feature extraction
+plus one extra graph eval; sidecar RSS delta is the 752 bf16 tensors.
+Every number that gets quoted comes from preflight-gated `benchmark.sh`
+cells (tower ms, TTFT delta vs text-only, peak RSS delta) and lands in
+`benchmarks/RESULTS.md`.
+
+## 4. Work plan
+
+**Phase A0 — groundwork**
+- [ ] Bind `mlx_conv2d` (full header signature; CPU/GPU sanity test vs
+      oracle conv output)
+- [ ] `scripts/regen-audio-goldens.py` + fixtures; dump mel golden,
+      tower-output golden, spliced ids, greedy ids
+- [ ] Resolve open questions §3.3 (write the answers into this doc)
+- Exit: goldens on disk; masking/zeroing semantics documented.
+
+**Phase A1 — decode + features**
+- [ ] `src/audio/decode.ts` (WAV + mono + resample; `afconvert` fallback)
+- [ ] `src/audio/features.ts` USM mel port
+- Exit: T0 green (mel within tolerance; token counts + spliced ids exact).
+
+**Phase A2 — tower**
+- [x] `src/audio/conformer.ts` loads `audio_tower.*`/`embed_audio.*` from
+      `optiq_vision.safetensors`; forward matches golden — rel-RMSE **2.4e-8**
+      (bit-exact: the tower runs f32 activations over bf16 weights because
+      the oracle feeds f32 mel and mlx promotes; the residual is 1 f32 ulp
+      from the test's /embed_scale·embed_scale roundtrip). Sidecar conv
+      weights are ALREADY MLX layout [C_out,kH,kW,C_in] — no transpose.
+      Golden reference point: embed_audio output BEFORE the /embed_scale
+      pre-division; features() returns pre-divided (vision convention).
+- [x] Clipped-linear toggle test (off → 90% rel-RMSE divergence, on → gate;
+      in tests/e4b-audio-tower.test.ts)
+- Exit: T1 rel-RMSE gate green. ✅ (gate 1e-6, ~40× margin)
+
+**Phase A3 — prompt + LM integration**
+- [x] Generalize prompt builder (images + audio, document order):
+      `buildMultimodalPrompt` in src/vision/prompt.ts; `<|audio|>` →
+      `boa + audio×n + eoa`; `buildVisionPrompt` kept as a byte-identical
+      wrapper (vision goldens green, e4b-vision 3/3).
+- [x] Merged embeddings + masks; per-layer id zeroing per Q2 decoupled from
+      the bidir overlay (`forwardEmbeddings` gained an optional `multimodal`
+      zeroing mask; generate() `multimodalMask` option). Merge numerics are
+      oracle-exact: the divide happens AFTER the bf16 cast with a weak
+      (bf16) embed_scale scalar — `AudioTower.features(preDivide=false)` +
+      astype + div, mirroring `features.astype(bf16) / embed_scale` (mlx
+      weak scalars take the array dtype; divide-in-f32-then-cast is up to a
+      bf16 ulp off and fails the exact gate).
+- Exit: T2 green — and stronger than planned: the FULL greedy stream +
+  decoded text match the oracle EXACTLY for both fixtures (the tower is
+  bit-exact, so no prefix allowance was needed; generate()'s
+  promptEmbeddings path prefills in one forward exactly like the golden
+  script — no tail-split gap). `tests/e4b-audio.test.ts`.
+
+**Phase A4 — serve + docs**
+- [x] `src/server.ts` audio branch, lazy tower, serial-lane routing,
+      mixed image+audio request handled. `hasAudio` detection over the
+      three part shapes; `makeAudioLoader`/`getAudioTower` mirror vision
+      (same sidecar, zero flags, auto-enable on `audio_config` +
+      sidecar) — but a request WITH audio on a towerless model is an
+      explicit 400 naming the model, never a silent text-only degrade.
+      The gateway's `Vision` payload generalized to
+      `{embeddings, imageMask?, multimodalMask?}` (bidirMask stays null
+      whenever audio is present, §3.3 Q1); `shape.hasVision = !!payload`
+      routes every media request serial and skips the prompt cache —
+      verified live: `/stats` batch.submitted_rows stays 0 for audio
+      while a concurrent text request advances it (batch default 8).
+      Non-WAV containers (mp3/m4a/flac/ogg/aiff) transcode via
+      `afconvert` (src/audio/transcode.ts, RIFF-sniffed content-based
+      passthrough, temp files under os.tmpdir(), clear 400 on failure);
+      >30 s clips truncate at 480 000 samples in decodeAudio — the
+      oracle feature extractor's own default (truncation=True,
+      max_length=480_000) verified in audio_feature_extractor.py:302.
+      /v1/messages (Anthropic) rejects smuggled audio blocks with a 400
+      pointing to the OpenAI endpoint (protocol has no audio block).
+- [x] server-api.md (schema + curl + caps + 400s + lane note) /
+      features-matrix.md (audio row + footnote: mlx-lm strips it, optiq
+      never serves it — we expose it) / README capability line, same
+      change set; Anthropic endpoint 400 documented.
+- [x] capability discovery (post-merge CodeRabbit follow-up, 2026-07-07):
+      `audio` flag mirrors `vision` on every surface — /v1/models served
+      row, /library rows (`audioCapable` in src/registry.ts:
+      `has_audio_config AND has_audio_tower`, header-only sidecar scan),
+      and the pi ws `ready` frame. The predicate checks TENSORS, not just
+      config: the 12B stub sidecar (audio_config + embed_audio-only) reads
+      audio=false, and `makeAudioLoader` header-checks the sidecar so the
+      runtime flag can't lie either. Discovery-only in pi land — pi's
+      `Model.input` type has no "audio" modality to declare.
+- Exit: T3 green — `tests/e4b-audio-serve.test.ts`
+      (MLX_BUN_TEST_AUDIO_SERVE=1): base64 WAV over
+      /v1/chat/completions returns the EXACT golden transcription
+      ("The quick brown fox jumps over the lazy dog.", 83 prompt /
+      11 completion tokens); m4a transcode path transcribes; malformed
+      part, undecodable bytes, and Anthropic audio block all 400.
+      docs-map hygiene green. ✅ 2026-07-07
+
+**Phase A5 — bench + coverage**
+- [ ] benchmark.sh cells → RESULTS.md (tower ms, TTFT delta, RSS delta)
+- [ ] 12B audio cell: rebuild its sidecar via optiq `build_vision_sidecar`
+      (selective download), own goldens — per-model doctrine, every
+      (model, quant) cell gets its own validation or an explicit deferral
+- [ ] features-matrix: audio×batching = serial documented
+- Exit: numbers curated; e4b cell validated; 12B validated or deferred
+  with reason.
+
+Sizing: A0–A4 ≈ 3–4 focused sessions (the tower is the only genuinely new
+model code; everything else is seam reuse); A5 ≈ 1.
+
+## 5. Non-goals (v1)
+
+- TTS, STS, and `/v1/audio/transcriptions`-style endpoints — a
+  transcription shim over this same path is a cheap follow-up, noted, not
+  committed.
+- Streaming/real-time audio input.
+- >30 s audio: cap at 750 soft tokens exactly as the oracle does.
+- Video (`mm_token_type_ids=2`) — same seams, explicitly later.
+- Audio for 26B-A4B / DiffusionGemma / bf16 assistants (no
+  `audio_config` in those snapshots — architectural, not deferred work).
+- Batched audio prefill (novel territory; serial lane only in v1).

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

@@ -1,6 +1,19 @@
 # Batching perf path — promote continuous batching out of compat mode
 
-Status: P5 + two gate fixes LANDED (2026-07-02, same day); rest PLANNED.
+Status: P5 + two gate fixes LANDED (2026-07-02, same day). Since then
+(2026-07-03, executed via grammar-spec-batching-integration.md Phases D/E):
+the **P3 KV-budget admission slice LANDED** (`--kv-budget` aggregate
+projection gate, queue-don't-OOM,
+`/stats.batch.{pending_rows,kv_bytes,kv_budget_bytes}`,
+tests/batch-kv-budget.test.ts) and the **P0 bench harness LANDED** as
+`scripts/bench-feature-matrix.ts` (the six-cell composition matrix;
+`scripts/bench-serving-load.ts` stays a client-only stack-vs-stack tool).
+Also 2026-07-03: **Tier-0 universal models (plain full-attention, e.g.
+Llama) now batch** — per-row RoPE ported (`UniversalRope.applyDynamic`)
+and gated token-exact vs mlx-lm B=2
+(tests/batched-decode-parity.test.ts "Llama 3B Tier-0"); maskArray +
+sliding universal archs stay serial (unvalidated cells). Rest PLANNED —
+P0 COMPLETED 2026-07-04: extend-join (full-attention; own mlx-lm extend oracle) + vectorized all-greedy sampling (BIT-equal A/B).
 
 **Headline result (M1 Max 32 GB, loaded machine, 4 concurrent × 128 tok,
 median of 3, wall-clock aggregate; oMLX 0.4.5.dev1, same OptiQ snapshots):**
@@ -56,7 +69,9 @@ the 2026-07-02 oMLX report (Claude artifact) and `~/.omlx` bench scripts.
 Design stance: the batch lane grows the **same parity-tier DAG as the
 serial lane** (docs/design/parity-tier-dag.md). `--l1 --batch N` stays
 bit-exact mlx-lm-B=N compat mode forever (the oracle checkpoint); the
-default/`--l3` batch path becomes the optimized one. `--batch` remains a
+default batch path becomes the optimized one. (`--l3` was deleted
+2026-07-05 — docs/design/unified-engine-frontier-plan.md §6; output-changing
+experiments are now env-flag Lab items.) `--batch` remains a
 mode switch, not a load fallback.
 
 ## 1. Current state
@@ -69,7 +84,9 @@ decode, per-row eviction (`filterKVRows` / `BatchedRotatingCache.filter`),
 per-row failure containment, drain-on-serial-waiter, `clearCache` every
 256 steps. Bit-parity with mlx-lm B=N verified at B=2 across all 4 model
 paths (`tests/batched-decode-parity.test.ts`, `tests/batched-rotating.test.ts`,
-`tests/batch-scheduler.test.ts`).
+`tests/batch-scheduler.test.ts`) — plus, since 2026-07-03, Tier-0
+universal plain-full-attention archs (Llama 3B, token-exact vs mlx-lm B=2;
+see header).
 
 `src/serve/generation-gateway.ts` (`GenerationGateway.willBatch`) routes;
 one `AsyncMutex` is the GPU/`loraState` exclusion domain. The batch lane
@@ -81,22 +98,21 @@ is deliberately compat mode: it calls `model.forwardHidden` /
 | Lever | Serial mechanism | Why batch can't use it today |
 | --- | --- | --- |
 | Quantized KV (`kv_config`) | `maybeQuantizeKv` (`src/generate.ts:101`) | `QuantizedKVCache` (`src/model/gemma4-base.ts:419`) has no dynamic-B ops (no merge/filter/extend/temporalView over the packed/scales/biases triples); `willBatch` hard-gates `kvQuant` → serial |
-| Perf-kernel fused decode SDPA | dispatched from `quantizedSdpa` (`gemma4-base.ts:1560`) | kernel M=1-specialized: `fused-decode-kernel.ts:262` rejects `B !== 1 \|\| L !== 1`, and requires empty mask (batched decode needs per-row left-pad validity) |
-| Fused N-tiled SDPA (`--fused-decode`) | `quantizedSdpaTiled` | needs 2-D mask; batched masks are 4-D `[B,1,1,S]` (`buildBatchedDecodeMask`); quantized-only anyway |
+| ~~Perf-kernel fused decode SDPA~~ | ~~dispatched from `quantizedSdpa`~~ | **DELETED 2026-07-05** (kernel + `--perf-kernel` removed; unified-engine-frontier-plan.md §6) — no longer a lever, batched or serial |
+| ~~Fused N-tiled SDPA (`--fused-decode`)~~ | ~~`quantizedSdpaTiled`~~ | **DELETED 2026-07-05** (measured 1.00×; unified-engine-frontier-plan.md §6) |
 | Compiled decode | `CompiledDecode.for()` (`generate.ts:518-525`) | trace is `[1,1]`; B changes every join/evict |
 | Prompt-cache prefix reuse | `PromptCache.take/put` in `runGeneration` | bypassed by design in v1; every batched row solo-prefills (`cachedTokens: 0`) |
 | Vectorized sampling | n/a | `#step` samples per row (B slice+sample+concat graphs per token) |
 
 Two more gap sources specific to the benchmark:
 
-- **Qwen3.5 never batches at all**: `SSMCache` (`src/model/qwen3-delta.ts`)
-  has no merge/filter, so the capability gate routes every request serial.
-  The Qwen3.5 concurrency loss is pure serial queueing. mlx-lm batches
-  Mamba-family state via `ArraysCache` — oracle exists.
+- **Qwen3.5 never batches at all** — **FIXED by P5 (2026-07-02, header
+  item 1)**: `SSMCache.mergeRows`/`filter` landed; Qwen3.5 batches
+  token-exact vs mlx-lm B=2.
 - Joins re-merge the whole batch (O(B·S) slice storm, `#mergeJoiner`)
   instead of mlx-lm's `extend`.
 
-## 2. Gap analysis: quantized KV + fused kernels at B>1
+## 2. Gap analysis: quantized KV at B>1
 
 - `QuantizedKVCache` already stores `[B,H,S,D']` triples; missing is
   dynamic-B surgery. New `src/model/batched-quant.ts`:
@@ -113,23 +129,30 @@ Two more gap sources specific to the benchmark:
 - Correctness fallback: `quantizedSdpaUnfused` (`gemma4-base.ts:1273`) is
   already B-generic with `mask.mode === "array"`; needs the batched mask
   broadcast as `[B,1,1,1,N]` for the 5-D GQA score reshape.
-- Perf kernel at B>1: (1) drop the `B !== 1` gate (keep `L !== 1`);
-  (2) grid gains a B axis (`[TG_THREADS, H·B, G]`), `SCORE_ROW` offsets
-  index `((b·KV + kv)·N + n)`; (3) replace the empty-mask requirement with
-  a per-row `leftPad` int32 `[B]` input — each row's n-loop starts at
-  `max(blockStart, leftPad[b])`, so pad bytes are never read.
+- ~~Perf kernel at B>1~~ — OBSOLETE: the perf kernel was deleted 2026-07-05
+  (unified-engine-frontier-plan.md §6); `quantizedSdpaUnfused` is the
+  quant-batched path, full stop.
 - RoPE per-row positions already solved (`ropeOffsetArr`, the S1b.1 fix).
 
 ## 3. Phases
 
 ### P0 — Baseline + engine quick wins (no new numerics)  [S–M, 2–4 days]
-- Write `scripts/bench-serving-load.ts` (N concurrent vs live server,
-  aggregate tok/s + TTFT p50/p95). Re-measure vs oMLX **with `--batch 4`
-  actually on** to size the true engine gap vs the serial-queueing gap.
+- Bench harness — **DONE 2026-07-03**: landed as
+  `scripts/bench-feature-matrix.ts` (six-cell composition matrix over live
+  SSE, TTFT p50/p95 + aggregate tok/s; integration-plan Phase E).
+  `scripts/bench-serving-load.ts` exists as the client-only stack-vs-stack
+  tool and stays one. Clean-machine run remains Josh-gated.
 - Land `extend` join (full-attention twin + `BatchedRotatingCache.extend`);
   oracle: extend `scripts/gen-batched-dynamic-golden.py` (mlx-lm `extend`).
+  **LANDED 2026-07-04** for full-attention layers (`extendKVRows`,
+  mlx-lm `BatchKVCache.extend` semantics; own oracle
+  `scripts/gen-batched-extend-golden.py`, token-for-token CPM + Llama;
+  `MLX_BUN_BATCH_EXTEND=0` = re-merge). `BatchedRotatingCache.extend`
+  (sliding layers on Gemma joins) remains open.
 - Vectorize the homogeneous-sampler fast path in `#step` (all-greedy →
-  one `ops.argmaxAxis` on `[B,V]`); mixed rows keep the loop.
+  one `ops.argmaxAxis` on `[B,V]`); mixed rows keep the loop. **LANDED
+  2026-07-04** (BIT-equal A/B, `MLX_BUN_BATCH_VEC_SAMPLE=0`;
+  tests/batch-vec-sample.test.ts).
 - Exit: gated suites unchanged; B=4 baseline recorded; joins not O(B·S).
 
 ### P1 — Quantized KV at B>1  [M–L, 1–1.5 wk]
@@ -144,19 +167,20 @@ Two more gap sources specific to the benchmark:
   (3) KL vs batched-bf16 ≤ serial quant-vs-bf16 KL; 6-task eval unchanged;
   (4) `--l1 --batch N` suites untouched.
 
-### P2 — Perf kernel at B>1 + fallback matrix  [M, ~1 wk]
-- Kernel edits per §2; dispatch when every row supported, else unfused.
-- Fallback matrix documented in server-config.md.
-- Batched frozen-oracle envelope (B=2/4 `freeze-perf-oracle` variant,
-  ≥56/64 per row, `tests/perf-kernel-oracle.test.ts` pattern).
-- Exit: dispatch counter proves engagement; envelope green; measured
-  uplift at B=4 vs P1 unfused on a clean machine; if <5%, ship the
-  fallback matrix only and close honestly.
+### P2 — Perf kernel at B>1 + fallback matrix  [OBSOLETE — kernel deleted 2026-07-05]
+The perf kernel (and `--perf-kernel` / its frozen-oracle tests) was deleted
+2026-07-05 (unified-engine-frontier-plan.md §6): it never beat the L1
+baseline in a paired A/B, so there is nothing to port to B>1. The planned
+"batched frozen-oracle envelope" gate is obsolete with it. Quant-batched
+decode stays on `quantizedSdpaUnfused` (P1).
 
 ### P3 — Admission, prompt cache, adapters, defaults  [M–L, 1–1.5 wk]
-- KV-budget admission: `projectKvBytes(model, promptLen, maxTokens)`
-  beside `fit()`; `--kv-budget`, `--prompt-concurrency`;
-  `/stats.batch.{pending_rows,kv_bytes,kv_budget_bytes}`.
+- KV-budget admission — **LANDED 2026-07-03** (integration-plan Phase D):
+  `--kv-budget <GB>` caps the batch's aggregate projected KV
+  (prompt+max_tokens per row, window-capped); over-budget joiners QUEUE
+  (FIFO), never OOM; a single request over the budget alone rejects with
+  an actionable error; `/stats.batch.{pending_rows,kv_bytes,kv_budget_bytes}`;
+  gated by tests/batch-kv-budget.test.ts. `--prompt-concurrency` still open.
 - Prompt-cache reuse under batching: `take` at admit, forward uncached
   tail; on finish/evict `extract(row)` → solo caches → `put`. Quantized
   entries are non-trimmable mid-group — match on full-prefix or STEP
@@ -204,17 +228,16 @@ Two more gap sources specific to the benchmark:
   (per-row fold; Qwen3.5's default repetition penalty had been routing every
   request serial). Exit numbers in the header table.
 
-Ordering: P1→P2 strictly ordered; P4 and P5 independent (P5 first if the
-Qwen3.5 headline matters most). Total ~5–6 wk serial, ~4 wk with P4/P5
+Ordering: P4 and P5 independent (P5 first if the Qwen3.5 headline matters
+most); P2 dropped (kernel deleted). Total ~4–5 wk serial, ~3 wk with P4/P5
 parallel.
 
 ## 4. Risks
 
 - **No external oracle for quant-batched** — mitigated by the P1 parity
   ladder; `--l1 --batch N` stays the always-reachable bit-exact checkpoint.
-- **Metal kernel at B>1**: leftPad start-index never touches pad bytes;
-  no atomics (deterministic); gate = batched frozen-oracle envelope +
-  dispatch assertions + A/B vs unfused.
+- ~~Metal kernel at B>1~~ — risk retired with P2 (perf kernel deleted
+  2026-07-05).
 - **Memory**: B× quant caches + prompt cache + weights on 24 GB — P3
   admission is the backstop; mirror mlx-lm's
   `trim_to(total − activeBatchBytes)` so LRU and live batch share one

package/docs/design/batching-v2-plan.md

@@ -1,6 +1,17 @@
 # Batching v2 — state map, debt, and the gap plan
 
-Status: **PLAN** (written 2026-07-01, read-only audit — no code changed)
+Status: **PLAN, largely EXECUTED** (written 2026-07-01 as a read-only
+audit; §1's state map is that day's snapshot). Landed since: steps 1–3
+(hotfix bundle, failure/fairness, engine hygiene — 2026-07-01); step 8
+SSM batched path (2026-07-02, superseded by batching-perf-path.md P5 —
+token-exact vs mlx-lm B=2); step 5 KV-budget admission (2026-07-03,
+executed as grammar-spec-batching-integration.md Phase D:
+`--kv-budget`, queue-don't-OOM, `/stats.batch` fields,
+tests/batch-kv-budget.test.ts — the `--prompt-concurrency` half still
+open); step 7's e1+e3 (min_p/XTC + penalties/logit_bias batch via
+per-row logits processors, 2026-07-02, batching-perf-path.md header
+item 2; e2 logprobs still routes serial). Still open: step 4 (extend
+join), step 6 (prompt-cache reuse), 7's e2, 9 (rolling), 10 (deferred).
 Scope: the `--batch N` continuous-batching engine (`src/serve/batch-scheduler.ts`,
 `src/serve/generation-gateway.ts`, `src/model/batched-mask.ts`,
 `src/model/batched-rotating.ts`, the `server.ts` wiring).
@@ -178,6 +189,9 @@ endpoints silently don't participate in. Specific debt, ranked:
   (mirror mlx-lm server.py:371-374) in `willBatch` — batchable iff every
   `makeCache()` element is `KVCache | RotatingKVCache` — plus the doc
   correction. The real SSM batch path is item (h), later. *(confidence 95)*
+  **RESOLVED**: gate landed 2026-07-01; the full SSM batch path landed
+  2026-07-02 (batching-perf-path.md P5) — server-config's Qwen3.5 claim is
+  true again.
 
 - **D2 — `#admit` materializes full-prompt logits.**
   batch-scheduler.ts:173-174 runs `logitsFromHidden(h)` on the **whole**
@@ -498,14 +512,14 @@ step a gate.
 
 | # | Item | Effort | Test gate | Notes |
 | --- | --- | --- | --- | --- |
-| 1 | **Hotfix bundle**: D2 slice-before-head; D1 capability gate + server-config.md correction; D3 route curve/`/signal`(/adapter-unmount) through `gateway.runExclusive`; stale comment server.ts:960 | **S** | gated batch suites unchanged; fast suite; new unit: willBatch(Qwen)=false | no numerics change anywhere |
-| 2 | **D5/D6 failure + fairness**: per-row onToken containment (evict row, keep batch); `#admit` try/finally; drain-on-serial-waiter | **S–M** | new gated test: throwing onToken evicts one row, siblings complete; drain test (serial completes under sustained batchable load) | mlx-lm `remove` semantics |
-| 3 | **(d) engine hygiene**: chunked+interleaved admit, pipelined decode, clearCache cadence | **M** | KL suite + mlx-lm goldens UNCHANGED; bench-serving-load A/B (clean machine, Josh-gated for quotables) | closes perf-review backlog #6 |
-| 4 | **(a) `BatchedKVCache` + `extend`** (full + rotating) | **M** | extended dynamic golden (mlx-lm `extend` driven); existing suites green | new files only — no gemma4-base.ts conflict |
-| 5 | **(c) KV-budget admission + `--prompt-concurrency`** | **M** | projection unit tests; gated queueing test; `/stats` fields | depends on 3 (interleaved admit) for prompt-concurrency semantics |
-| 6 | **(b) prompt-cache reuse under batching** | **M–L** | `cached_tokens>0` gated test; warm==cold row parity; serial suite | depends on 4 (`extract`) + 5 (byte interplay) |
-| 7 | **(e) gate widening**: e1 min_p/XTC (S) → e2 logprobs (M) → e3 penalties/logit_bias (M) | **S+M+M** | per-feature: batched==serial teacher-forced/deterministic-seed tests | e1 waits for the seed fix to merge |
-| 8 | **(h) SSM batched path** (Qwen3.5) | **M–L** | mlx-lm B=N golden for a Qwen3.5 model (ArraysCache oracle) | un-does the item-1 gate for Qwen; doc row flips back |
+| 1 | ✅ **DONE 2026-07-01** — **Hotfix bundle**: D2 slice-before-head; D1 capability gate + server-config.md correction; D3 route curve/`/signal`(/adapter-unmount) through `gateway.runExclusive`; stale comment server.ts:960 | **S** | gated batch suites unchanged; fast suite; new unit: willBatch(Qwen)=false | no numerics change anywhere |
+| 2 | ✅ **DONE 2026-07-01** — **D5/D6 failure + fairness**: per-row onToken containment (evict row, keep batch); `#admit` try/finally; drain-on-serial-waiter | **S–M** | new gated test: throwing onToken evicts one row, siblings complete; drain test (serial completes under sustained batchable load) | mlx-lm `remove` semantics |
+| 3 | ✅ **DONE 2026-07-01** — **(d) engine hygiene**: chunked+interleaved admit, pipelined decode, clearCache cadence | **M** | KL suite + mlx-lm goldens UNCHANGED; bench-serving-load A/B (clean machine, Josh-gated for quotables) | closes perf-review backlog #6 |
+| 4 | **(a) `BatchedKVCache` + `extend`** (full + rotating) — **still open** | **M** | extended dynamic golden (mlx-lm `extend` driven); existing suites green | new files only — no gemma4-base.ts conflict |
+| 5 | ✅ **DONE 2026-07-03** (admission half; `--prompt-concurrency` still open) — **(c) KV-budget admission** — executed as grammar-spec-batching-integration.md Phase D | **M** | projection unit tests; gated queueing test (tests/batch-kv-budget.test.ts); `/stats` fields | depends on 3 (interleaved admit) for prompt-concurrency semantics |
+| 6 | **(b) prompt-cache reuse under batching** — **still open** | **M–L** | `cached_tokens>0` gated test; warm==cold row parity; serial suite | depends on 4 (`extract`) + 5 (byte interplay) |
+| 7 | **(e) gate widening**: ✅ e1 min_p/XTC + e3 penalties/logit_bias **DONE 2026-07-02** (per-row logits processors, batching-perf-path.md header item 2); e2 logprobs **still open** (routes serial) | **S+M+M** | per-feature: batched==serial teacher-forced/deterministic-seed tests | e1 waits for the seed fix to merge |
+| 8 | ✅ **DONE 2026-07-02** — **(h) SSM batched path** (Qwen3.5) — landed as batching-perf-path.md P5, token-exact vs mlx-lm B=2 | **M–L** | mlx-lm B=N golden for a Qwen3.5 model (ArraysCache oracle) | un-does the item-1 gate for Qwen; doc row flips back |
 | 9 | **(g) docs pass** | **S** | doc-vs-code re-audit (this file's §1 as the checklist) | rolling — update per landing, final sweep here |
 | 10 | **(f) quantized-KV batched** | **L** | KL + 6-task + density benchmark — design doc first | DEFERRED — needs a concrete density case |
 

package/docs/design/bucketing-stage.md

@@ -166,7 +166,8 @@ future-proof empty table only.
 - **Stage 2 — train `memory-bucket` LoRA, only if Stage 0/1 miss.** Data from lucien.db, one
   example per (chunk, candidate-bucket) binary decision (gold = belongs?), shortlist regime =
   serve regime. Recipe = `chunk-finetune.ts` (sft, rank 16, lr 1e-5, seq 8192,
-  `MLX_BUN_PERF_KERNEL=0 MLX_BUN_FUSED_GELU=0`, no train-flash); adapter →
+  no train-flash; no env-flag sanitization needed since the 2026-07-05 kernel
+  deletions); adapter →
   `~/.cache/mlx-bun/adapters/memory-bucket`; pick checkpoints by the F1 judge, not val loss.
   **GPU-gated.**