mlx-bun 0.0.9 → 0.0.11

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

@@ -0,0 +1,255 @@
+# Batching perf path — promote continuous batching out of compat mode
+
+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):**
+
+| Model | oMLX | mlx-bun `--batch 4` | mean TTFT (theirs/ours) |
+| --- | --- | --- | --- |
+| MiniCPM5-1B | 339.0 | **345.4** | 397 / **119 ms** |
+| gemma-4-e4b | **89.8** | 87.1 (−3%) | 815 / **397 ms** |
+| Qwen3.5-4B | **100.9** | 99.5 (−1%) | 848 / **368 ms** |
+
+What landed (all bit-gated, `tsc` 0, suites green):
+1. **P5 — SSM batched path** (`SSMCache.mergeRows`/`filter` in
+   `src/model/qwen3-delta.ts`; "ssm" kind in `batch-scheduler.ts` merge/wrap/
+   filter; per-row RoPE via `ropeOffsetArr` in `qwen3_5.ts`; gate relaxed in
+   `generation-gateway.ts`, kill switch `MLX_BUN_BATCH_SSM=0`). Oracle:
+   `tests/qwen35-batched-parity.test.ts` — solo-prefill+merge batched greedy
+   is **token-for-token identical** to mlx-lm B=2 left-padded batch
+   (`goldens/apple-m1-max/batched-golden-qwen35.json`). No ssm_mask needed:
+   rows solo-prefill unpadded, decode feeds one real token per row.
+2. **Logits processors batch** (`generation-gateway.ts`): the per-row sampler
+   folds `makeLogitsProcessors` over a per-row device-side token history
+   (generate()'s pushHistory), so repetition/presence/frequency penalties +
+   logit_bias + min_p/XTC no longer route serial. Load-bearing discovery:
+   **Qwen3.5 ships a default repetition penalty in generation_config.json**,
+   which routed EVERY request serial under `--batch N` — the SSM port alone
+   changed nothing over HTTP until this landed.
+3. Diagnostic breadcrumb: `MLX_BUN_LANE_DEBUG=1` logs per-request lane
+   decisions + shape (server.ts) — how (2) was found.
+
+Bench-method traps found while measuring — **both FIXED 2026-07-02**:
+- `serve`/`bench` accepted `--model` but IGNORED it (positional-only fuzzy
+  query; auto-pick silently loaded the wrong model). Now: `--model
+  <path-or-query>` is an explicit override — a directory with config.json
+  loads via `scanSnapshot` (no registry dependency), anything else resolves
+  like the positional. Precedence: `--model` > positional > `--query` >
+  auto-pick.
+- Event-loop starvation during generation was the SERIAL lane, not batch
+  (batch's drive loop yields per step — `/stats` answered in ~8 ms mid-
+  generation once measured properly): serial decode is an unbroken microtask
+  chain, so `/stats`//health/new accepts stalled for the whole generation
+  (measured 2.5 s on a 512-token run). Fixed with a ≥25 ms rate-limited
+  macrotask hop wrapped around onToken in `GenerationGateway.run`'s serial
+  branch — covers all four API surfaces, streaming and not. After: `/stats`
+  10–44 ms mid-generation, decode unchanged (cpm5 serial 193 tok/s). Motivated by the oMLX head-to-head (M1 Max,
+2026-07-02): at 4 concurrent requests oMLX's continuous batching gets
+1.46–1.89× our aggregate throughput (Qwen3.5-4B 97.3 vs 66.5 tok/s;
+MiniCPM5 328 vs 173; e4b 85 vs 52) and our mean TTFT queue-inflates to
+1.1–3.8 s. Single-stream we match or beat them (e4b 55.8 vs 53.7,
+Qwen3.5 66.3 vs 63.9, TTFT 3–7× better everywhere); the gap is entirely
+the serial lane going flat under parallel load. Full comparison + numbers:
+the 2026-07-02 oMLX report (Claude artifact) and `~/.omlx` bench scripts.
+
+Design stance: the batch lane grows the **same parity-tier DAG as the
+serial lane** (docs/design/parity-tier-dag.md). `--l1 --batch N` stays
+bit-exact mlx-lm-B=N compat mode forever (the oracle checkpoint); the
+default 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
+
+`src/serve/batch-scheduler.ts` (`BatchScheduler`) is a real
+continuous-batching engine: iteration-level scheduling, pipelined decode
+(`#step` via `ops.asyncEvalAll`, mlx-lm `GenerationBatch._step`
+semantics), chunked joiner prefill (2048-token chunks) interleaved with
+decode, per-row eviction (`filterKVRows` / `BatchedRotatingCache.filter`),
+per-row failure containment, drain-on-serial-waiter, `clearCache` every
+256 steps. Bit-parity with mlx-lm B=N verified at B=2 across all 4 model
+paths (`tests/batched-decode-parity.test.ts`, `tests/batched-rotating.test.ts`,
+`tests/batch-scheduler.test.ts`) — 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
+is deliberately compat mode: it calls `model.forwardHidden` /
+`logitsFromHidden` directly, bypassing `generate()` and every perf lever.
+
+### Levers the batch lane lacks, and why
+
+| Lever | Serial mechanism | Why batch can't use it today |
+| --- | --- | --- |
+| Quantized KV (`kv_config`) | `maybeQuantizeKv` (`src/generate.ts:101`) | `QuantizedKVCache` (`src/model/gemma4-base.ts:419`) has no dynamic-B ops (no merge/filter/extend/temporalView over the packed/scales/biases triples); `willBatch` hard-gates `kvQuant` → serial |
+| ~~Perf-kernel fused decode SDPA~~ | ~~dispatched from `quantizedSdpa`~~ | **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** — **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 at B>1
+
+- `QuantizedKVCache` already stores `[B,H,S,D']` triples; missing is
+  dynamic-B surgery. New `src/model/batched-quant.ts`:
+  `BatchedQuantizedKVCache` = quantized twin of `BatchedRotatingCache`
+  (merge/extend/filter/extract/temporalView = the `mergeKVRows`/`filterKVRows`
+  B-axis ops from `src/model/batched-mask.ts:171,196`, applied to all three
+  triples). Left-pad must align to quantization group boundaries — pad in
+  whole 256-token STEP granules (both caches already grow in STEP=256).
+- Conversion point mirrors `maybeQuantizeKv`: row solo-prefills bf16,
+  `toQuantized` per layer per `kv_config` before merge.
+- **Scope cut**: shipped `kv_config.json` quantize full-attention layers
+  only ⇒ P1 = quantized full layers + bf16 `BatchedRotatingCache` sliding
+  layers (the serial `config` composition). Batched rotating-quant deferred.
+- Correctness fallback: `quantizedSdpaUnfused` (`gemma4-base.ts:1273`) is
+  already B-generic with `mask.mode === "array"`; needs the batched mask
+  broadcast as `[B,1,1,1,N]` for the 5-D GQA score reshape.
+- ~~Perf kernel at B>1~~ — 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]
+- 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. **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]
+- `src/model/batched-quant.ts` + scheduler wiring (layer kind
+  `"full" | "rot" | "qfull"`); solo-prefill bf16 → `toQuantized` → extend.
+- Batched mask 5-D broadcast; validate `quantizedSdpaUnfused`.
+- `willBatch`: `kvQuant` stops being a serial gate when quant layers are
+  full-attention-only.
+- Exit (parity ladder; no external oracle — mlx-lm quant batching NYI):
+  (1) B=1-row batched-quant **bit-exact vs serial `--kv-quant config`**;
+  (2) row-isolation within the calibrated bf16-reduction envelope;
+  (3) KL vs batched-bf16 ≤ serial quant-vs-bf16 KL; 6-task eval unchanged;
+  (4) `--l1 --batch N` suites untouched.
+
+### P2 — Perf kernel at B>1 + fallback matrix  [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 — **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
+  boundaries. Real `prefillMs/decodeMs`; `BatchStats.cachedTokens` real.
+- Single-adapter batches: adapter-spec = batch compatibility key
+  (join if match, drain if not; `loraState.active` set once per epoch).
+- Default review: if batch lane at 1 live row is within ~2% of serial,
+  propose flipping the shipped default to `--batch 4` (Josh decision,
+  numerics change vs serial default); else README recommends `--batch 4`
+  for agent workloads.
+- Exit: TTFT p50 <500 ms / p95 <1.5 s at 4 concurrent short prompts;
+  `cached_tokens>0` on repeats under `--batch 2`; oversized-request
+  queueing test (no OOM); B=4 aggregate within 15% of oMLX on cpm5 + e4b.
+
+### P4 — Burst decode  [REVISED after reading + porting oMLX's source]
+- **oMLX-style event-loop bursting: BUILT AND REFUTED (2026-07-02).** Read
+  their actual implementation (`omlx/engine_core.py _step_burst` — full
+  source ships in `/Applications/oMLX.app/Contents/Resources/omlx/`): it is
+  NOT a GPU trick — it runs K `scheduler.step()` calls per
+  `run_in_executor` hand-off with adaptive time budgets (100 ms single /
+  30 ms concurrent, 64-step cap) because each per-token hand-off ping-pongs
+  the GIL with asyncio/uvicorn (~1 ms/token; their comment: 80 vs 74 tok/s).
+  Ported faithfully to `BatchScheduler.#drive` and measured: **regression**
+  — cpm5 B=4 345→289, batch-lane B=1 149→121, TTFT +~100 ms (SSE flushes
+  wait out the budget). Bun's `setImmediate` hop costs microseconds — no
+  GIL to dodge — so the port kept the cost and lost the win. Reverted;
+  numbers restored (cpm5 B=4 349.1). Breadcrumb comment left at the drive
+  loop. This also explains their ~21-token SSE bursts (measurement trap in
+  the comparison report).
+- **Still open (different mechanism, not refuted): device-side step
+  chaining** — generalize `generateInner`'s pipeline to depth k (k chained
+  step graphs fed device-side, one host sync per k tokens). That attacks
+  OUR real per-token cost (FFI graph-build + sync readback), which is not
+  what oMLX's burst attacks. Exit unchanged: bit-exact vs depth-1, cpm5
+  single-stream gap vs oMLX (178 vs ~214) is the target; no regression at
+  k=1.
+
+### P5 — SSM batched path (Qwen3.5)  [DONE 2026-07-02 — see header]
+- Landed in one day, much cheaper than the ~1 wk estimate: `ssm_mask` under
+  left-padding turned out unnecessary (our join model solo-prefills unpadded,
+  so SSM state never sees a pad token; mlx-lm's masked-pad prefill is
+  equivalent — proven by the token-exact oracle match). Batch ops are plain
+  B-axis concat/take on the two state slots.
+- Bonus fix that the exit benchmark forced out: logits processors now batch
+  (per-row fold; Qwen3.5's default repetition penalty had been routing every
+  request serial). Exit numbers in the header table.
+
+Ordering: 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~~ — 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
+  ceiling. GPU OOM is uncatchable — queueing test before default changes.
+- **Ring-wrap at B>1 short-context-verified only**: land a long-context
+  bf16 batched golden before any P3 default flip. Quantized rotating out
+  of scope.
+- **Default flip changes numerics** vs today's serial default — explicit
+  Josh decision with doc trail.
+- **Benchmark hygiene**: quotable numbers via `bench-serving-load.ts` on a
+  clean machine; TTFT at the SSE first-token event, not server-side
+  (oMLX's own server log over-reports by excluding ~180 ms/req cache
+  bookkeeping — measure them wall-clock only).
+- **Regression fence**: all existing gated batch suites + the serial suite
+  green unchanged at every phase.