mlx-bun 0.2.0 → 0.4.0

package/docs/design/parallel-slots.md

@@ -1,567 +0,0 @@
-# Parallel slots — configurable batched serving
-
-Status: **S0 + S1b + S2 SHIPPED** (S2's continuous inject/evict shipped
-with the scheduler; its byte-budget admission half landed 2026-07-03 as
-`--kv-budget`; S3+ deferred)
-Owner: serving layer
-Default: **off** (`--batch 1` = today's serialized path, untouched)
-
-> **Flag rename note:** This document was drafted with `--slots N` as the design name.
-> The shipped flag is `--batch N` (`--decode-concurrency` accepted as an mlx_lm.server alias).
-> Occurrences of `--slots` in pre-ship text below refer to today's `--batch N`.
-
-## Goal
-
-Let the server process more than one generation request concurrently,
-gated entirely behind a `--batch N` flag (originally designed as `--slots N`). `N=1` runs exactly what we
-have today (the serialized promise chain). `N>1` loads a batch
-scheduler that runs multiple sequences through one forward pass.
-
-This is a server-path feature — serving speed is the user metric.
-Direct/embedding mode stays batch=1.
-
-> **v2 execution plan:** this document remains the PRIMITIVE-design record
-> (masks, caches, lanes). The living gap/debt plan and sequencing for the
-> batching engine is [batching-v2-plan.md](batching-v2-plan.md). Landed from
-> it (2026-07-01, steps 1–3): cache-capability gate (hybrid/SSM models route
-> serial), slice-before-head admit, ONE lock domain (curve `/generate`,
-> `/signal`, adapter mount/unmount via `gateway.runExclusive`), per-row
-> onToken failure containment, drain-on-serial-waiter (no starvation), and
-> the engine-hygiene transplant (pipelined decode via asyncEval, clearCache
-> 1/256 steps, chunked admit prefill interleaved with decode steps).
-> Landed since: **SSM batched path** (Qwen3.5 batches, token-exact vs
-> mlx-lm B=2 — 2026-07-02, batching-perf-path.md P5); **per-row logits
-> processors** (repetition/presence/frequency penalties + logit_bias +
-> min_p/XTC batch — 2026-07-02, same doc); **grammar under batching** (B1
-> per-row matchers — structured-output.md); **KV-budget admission**
-> (`--kv-budget`, queue-don't-OOM, `/stats.batch` fields — 2026-07-03,
-> grammar-spec-batching-integration.md Phase D); **Tier-0 universal
-> coverage** (plain full-attention archs, e.g. Llama, batch via per-row
-> RoPE, token-exact vs mlx-lm B=2 — 2026-07-03, same doc's execution log;
-> maskArray + sliding universal archs stay serial).
-
-## STATUS (2026-06-14) — LIVE for full-attention models
-
-**`--batch N` now serves B>1** for full-attention models (CPM). The numerics
-were proven first (the batched **forward primitive** `BatchedDecodeMaskCache` +
-per-layer fixes = bit-parity with mlx-lm B=N across all 4 models; static-B AND
-dynamic-B `mergeKVRows`/`filterKVRows` vs mlx-lm `BatchKVCache.merge`/`.extract`/
-`.filter`, CPM L1), then the **scheduler** was built and wired in:
-- `BatchScheduler` (`src/serve/batch-scheduler.ts`) — the detached async driver:
-  admission queue → running batch → continuous inject (`mergeKVRows`) / evict
-  (`filterKVRows`) → per-row sample + accounting. Gated teacher-forced (KL).
-- `GenerationGateway` (`src/serve/generation-gateway.ts`) — lane picker +
-  `AsyncMutex` (serial↔batched GPU/`loraState` exclusivity) + the `_is_batchable`
-  predicate. Wired into both `handleChat` call sites → per-row SSE fan-out.
-- End-to-end: `tests/batch-serving.test.ts` (live CPM `--batch 2` server).
-
-**Sliding-window (Gemma) dynamic-B is now served too** (2026-06-14):
-`BatchedRotatingCache` (`src/model/batched-rotating.ts`, port of mlx-lm
-`BatchRotatingKVCache` incl. the ring-wrap rolled mask) is bit-exact vs mlx-lm
-model-free (`tests/batched-rotating.test.ts`); the scheduler assembles each
-layer's cache by type and Gemma 12B scheduled greedy == the mlx-lm B=2 golden
-with eviction. So `--batch N` batches BOTH CPM and Gemma.
-
-**This meets the mlx-lm-parity bar.** mlx-lm's batched path is bf16 (its
-quantized batching is NYI — `BatchRotatingKVCache.to_quantized` raises), so bf16
-continuous batching IS the drop-in. `--batch N` is therefore a bf16 MODE: KV
-quant unset → bf16 (so the batch path engages out of the box, Option B); serial
-default stays mixed-precision (optiq parity); an explicit `--kv-quant` under
-`--batch N` routes those requests to serial (warned).
-
-**Batched + mixed-precision KV quant is NOT a parity gap** — no ancestor does it
-(mlx-lm NYI, optiq no batching), so there's no bit-exact oracle. It's an OPTIONAL
-novel extension: a memory-density win (batching × 4-bit KV compound), gated by
-KL + 6-task quality (per the tree framing), deferred. Other follow-ups (all
-optional polish): the `extend` join op (today joins re-merge), KV-budget
-admission (**LANDED 2026-07-03** — `--kv-budget`), prompt-cache reuse under
-batching, and **paged KV** (zero-padding-waste allocation).
-
-## Why it helps (and when it doesn't)
-
-The win is **throughput via memory bandwidth**. Decode on Apple Silicon
-is bandwidth-bound (~273 GB/s on the M4 Pro): every token streams the
-entire weight set through the GPU to produce one token. Batch B
-sequences and you stream the weights **once** to produce **B** tokens.
-Aggregate tok/s scales ~linearly until the compute roofline.
-
-The crucial design point — **continuous batching, not static
-batching**:
-
-- *Static batching*: form a batch, run to completion, then accept the
-  next. A request arriving mid-flight waits for the whole batch to
-  finish. This is the wrong design; do not build it.
-- *Continuous batching* (iteration-level scheduling): the scheduler
-  works per **decode step**. After each token it can inject a
-  newly-arrived request (after a short prefill) and retire finished
-  ones. A late request does **not** wait for the in-flight one to
-  finish — it joins at the next step.
-
-So the benefit window is "any time two requests are alive in the same
-wall-clock window" — not "a backlog has formed." Because generations
-last seconds, requests overlap even at modest rates (Little's law:
-concurrent ≈ arrival_rate × mean_generation_time).
-
-When it does **not** help: traffic so light that requests never overlap
-in time. Then there's only ever one live sequence and batching is idle
-— which is exactly why `--batch 1` is and stays the default.
-
-## Current code: what's batch=1 and what isn't
-
-Good news: the hot path is already written over a batch dim `B`.
-
-- `Attention.forward` reads `const [B, L] = x.shape` and threads `B`
-  through every reshape/transpose/SDPA — `src/model/gemma4.ts:152`.
-- `KVCache` stores `[B, H, S, D]` and grows/slices generically —
-  `src/model/gemma4-base.ts:204` (STEP=256 growth).
-- The byte-capped LRU prompt cache already accounts bytes, not count —
-  `src/prompt-cache.ts` (the precedent for a byte-budgeted KV pool).
-
-`B=1` is baked into a finite, findable set of places:
-
-| Location | Hardcoded |
-| --- | --- |
-| `src/generate.ts:295,315,378` | prompt/decode tensors built `[1, L]` / `[1, 1]` |
-| `src/generate.ts:320,235` | `hLast` slice `[1, …]`; `sampleStep` reshape `[1, V]` |
-| `src/generate.ts:247` | sampler history `[1]`, one sampled token/step |
-| `src/model/gemma4.ts:647-653` | per-layer-input reshapes `[1, L, …]` *(pre-S1b inventory — this entry was fixed in S1b; e4b per-layer-input is now B-generic)* |
-| `src/model/gemma4-base.ts:210` | **one scalar `offset` per cache, shared across rows** |
-| `src/model/gemma4.ts:174,206` | **scalar RoPE offset** for the whole batch |
-| `src/model/gemma4-base.ts:256` | `makeMask` from that single offset (returns empty mask at N=1) |
-
-## The one genuinely hard problem: per-sequence position
-
-Batched sequences are different lengths and at different decode
-positions, but the cache has a single `offset`, RoPE takes a scalar
-offset, and the mask is built from one offset.
-
-Solution (mlx-lm's, and ours): **left-padding + per-row mask + per-row
-RoPE offsets.**
-
-- Right-align every sequence so "the current token" is the same column
-  for all rows → one global advancing offset still works.
-- Pad on the left; mask the padding with a real `[B, 1, 1, S]` boolean
-  mask — **even at the N=1 decode step**, which today returns an empty
-  mask (`gemma4-base.ts:257`). That shortcut is the first thing to change.
-- RoPE positions become a per-row vector. **We already have the
-  primitive**: `ops.ropeDynamic` + the `cache.ropeOffsetArr` array-offset
-  path (`gemma4.ts:179,206`), currently used only by compiled decode.
-
-Highest-risk sub-item: Gemma3 alternates sliding-window
-(`RotatingKVCache`) and full-attention layers. Per-row offsets through
-the ring cache's wrap semantics is the trickiest mask math in the
-project. Validate the full-attention mask path first; bring sliding
-layers online separately under teacher-forced parity.
-
-### Reuse: batched PREFILL already exists (big de-risk, found 2026-06-14)
-
-The training path already solves batched padded forward and proves it:
-- `src/train/forward.ts` — `buildBatchedPadMask(B, L, validLengths, window)`
-  builds the `[B,1,L,L]` mask (causal AND per-row key-validity, per distinct
-  window) and `BatchedMaskCache` routes it through the stock attention by
-  wrapping a real cache and overriding only `makeMask`. Comment confirms
-  "KVCache/RotatingKVCache are already shape-generic over B."
-- `tests/train-batch-e2e.test.ts` proves a B=2 padded forward's per-row loss
-  equals two independent B=1 forwards within bf16 tolerance — i.e. the mask
-  makes a real row's logits independent of the other (padded) row. That is
-  exactly the teacher-forced gate S1 needs, already passing.
-- `tests/train-batch.test.ts` unit-tests the mask (fast, no model).
-
-So batched **prefill** is ~done and reusable (offset-0, right-padded). The
-genuinely new work is batched **decode**: a growing per-row KV with per-row
-offsets, per-row RoPE (array-offset path), and a per-row N=1 key-validity
-mask. First decode brick landed 2026-06-14: `src/model/batched-mask.ts`
-`buildBatchedDecodeMask` (left-padded, nonzero-offset counterpart) +
-`tests/batched-decode-mask.test.ts` (fast). The two mask builders consolidate
-when prefill is wired into serving (S1a).
-
-### S1b decode: design traps & sequencing (found 2026-06-14)
-
-Reading the attention path surfaced two real subtleties — the kind that
-produce plausible-but-wrong logits, so flagged before writing the cache:
-
-1. **Prefill/decode padding conventions differ.** Training's batched prefill
-   RIGHT-pads (real tokens left-aligned, positions 0..len-1 → scalar RoPE
-   offset 0 is correct). Batched decode wants LEFT-padding (rows right-aligned
-   so every row's next write lands in the same column → one advancing offset).
-   Left-padding then forces **per-row RoPE** (real position = column −
-   leftPad[b]) via the array-offset path, and means the prefill KV must be
-   re-laid-out (right-aligned) before decode. First cut: prefill each request
-   solo, copy its KV into the right-aligned row of a shared `[B,H,Smax,D]`
-   decode buffer, then batch the decode.
-2. **RoPE offset-timing trap (FIXED 2026-06-14).** `Attention.forward` read
-   `cache.ropeOffsetArr` twice — for K *before* `updateAndFetch` and for Q
-   *after* it (gemma4.ts ~179 vs ~195/201). Since `updateAndFetch` advances
-   `offset`, a per-row position array derived from the live offset would give
-   K and Q different positions. Fixed by capturing `offsetArr` once alongside
-   `offset` (line ~174). Behavior-preserving (real caches leave it unset;
-   compiled-decode passes a constant) — verified bit-exact against
-   `tests/compiled-decode.test.ts` (12B, 7/7).
-
-**Sequencing (parity gate first, on the riskiest piece):**
-- **S1b.1** — gated teacher-forced *decode* parity harness: manually assemble a
-  left-padded 2-row KV from two solo prefills, run K batched decode steps with
-  a `BatchedDecodeMaskCache` (delegates KV, returns `buildBatchedDecodeMask`,
-  exposes per-row `ropeOffsetArr` pinned to the pre-write offset), and assert
-  each row's logits match its solo greedy decode within bf16 tolerance. Gated
-  behind an opt-in env + local weights (like `train-batch-e2e`). Isolates
-  decode-batching numerics from the prefill-assembly plumbing.
-  **DONE 2026-06-14** — `BatchedDecodeMaskCache` (src/model/batched-mask.ts) +
-  `tests/batched-decode-parity.test.ts` (`MLX_BUN_TEST_BATCH_DECODE=1`,
-  MiniCPM5-1B). Passes: the unpadded row is **bit-exact** vs its solo decode
-  (batching isolates rows); the left-padded row matches within bounded bf16
-  reduction-order noise (≤0.23, stable over 8 steps — the shared buffer shifts
-  its keys by leftPad columns, so the attention sum accumulates in a different
-  order). **The gate earned its keep**: it caught `LlamaAttention`
-  (src/model/minicpm5.ts) roping Q/K with the scalar `cache.offset` and ignoring
-  `ropeOffsetArr` → left-padded rows mis-positioned (logit diff 8.7). Fixed to
-  use the dynamic array-offset path when present; MiniCPM5 solo parity stays
-  bit-exact vs the oracle. (Gemma4's `Attention` already had the array path; its
-  capture-once fix is verified bit-exact vs compiled-decode.) Lesson: each model
-  family's attention needs the per-row offset path wired separately.
-- **S1b.2** — production KV assembly (solo prefill → left-aligned decode buffer).
-- **S1b.3** — scheduler + B-token/step generate loop + per-row stream fan-out.
-
-### The validation matrix — 3 parity layers × every model path
-
-**Current state (2026-06-14).** Roster: **CPM (MiniCPM5-1B) + 3 Gemmas
-(e4b / 12B / 26B)**. Every NEW model family (e.g. **Qwen**, planned) adds its
-OWN handler path and its own copy of the three layers below — nothing here is
-generic. This matrix is a snapshot; re-derive it when the roster changes.
-
-**Framing (refined 2026-06-14 — it's a TREE, not strict layers).** mlx-lm is the
-base library; **mlx-optiq and mlx-bun are both optimizations on top of it**
-(optiq = mixed-precision per-layer KV quant, "drop-in for mlx_lm.server", no
-batching; mlx-bun = per-model fused kernels + batched serving). Bit-parity with
-an ancestor is just an **easy correctness test where we overlap one** — not a
-porting goal. The "L1/L2" labels below are **parity anchors / correctness
-checkpoints**, not a strict stack: L1 = "matches mlx-lm (bf16)", L2 = "matches
-optiq (mixed-precision quant)"; beyond them is the no-oracle **Lab** class ("our
-own optimizations, no ancestor does them → no oracle" — formerly labeled "L3";
-the `--l3` product tier was deleted 2026-07-05,
-docs/design/unified-engine-frontier-plan.md §6). A combo no ancestor does (e.g. **batched mixed-precision
-serving** — mlx-lm can't do mixed-precision, optiq can't batch) is genuinely new
-territory: gate it by KL + the 6-task quality + benchmarks, never bit-exactness.
-The anchors are still useful — we can always dial back to a bit-exact-vs-ancestor
-config as a correctness checkpoint and build outward.
-
-- **L1 — mlx-lm parity (standard KV).** bf16 KV, **bit-exact vs mlx-lm**. The
-  foundation oracle (and, since 2026-07-05, the naked default). Path: monolith
-  `Gemma4Model` / `MiniCPM5Model`, `ops.sdpa`.
-- **L2 — mlx-optiq parity (mixed-precision KV quant).** quant/mixed KV,
-  **bit-exact vs optiq**. Path: `QuantizedKVCache` + `quantizedSdpa` (fused-sdpa
-  follows `--kv-quant`: on for `config` — the optiq golden composition; off →
-  *unfused* for uniform/bf16).
-- **L3 — DELETED as a product mode 2026-07-05**
-  (docs/design/unified-engine-frontier-plan.md §6): no output-changing kernel
-  (perf-kernel, fused-decode, fused-gelu) beat the L1 baseline in a paired A/B,
-  so the kernels were deleted and `--l3` hard-errors. Its no-oracle territory is
-  now the **Lab**: output-changing experiments are env-flag items with a bench +
-  expiry, gated by **KL vs the compatible upstream + the 6 eval-task mean
-  (`src/eval/tasks/`) + perf numbers** — never bit-exactness — until one earns
-  promotion.
-
-Degradation lever: bf16 KV drops L2→L1. The surviving kill switches
-(`--compiled-decode`, `--compiled-activations`, `--fused-sdpa`) are per-kernel
-escapes. Each layer must stand alone and stay reachable (don't-delete-
-optionality as a fallback ladder).
-
-So a feature like batched decode (`--batch`) is "done" only when green at
-**every (model × layer)** cell — and each cell may need its own fix:
-
-| Model · path | L1 (bf16 / mlx-lm) | L2 (quant / optiq) | ~~L3~~ (tier deleted 2026-07-05) |
-| --- | --- | --- | --- |
-| CPM (MiniCPM5) | ✅ **oracle-verified** (= mlx-lm B=2) | ☐ quant + unfused-mask | — |
-| Gemma 12B | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + sliding-window | — |
-| Gemma e4b | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + same | — |
-| Gemma 26B | ✅ **oracle-verified** (= mlx-lm B=2, short-ctx) | ☐ generated + MoE | — |
-| *(Qwen, future)* | *new path* | *new path* | — |
-
-**🎯 L1 batched decode COMPLETE (2026-06-14d): the entire L1 row is green** —
-all 4 model paths bit-parity with mlx-lm B=2 (CPM full-attn, Gemma 12B
-sliding+full, e4b per-layer-input+KV-share, 26B MoE). Only one per-path fix was
-needed (e4b's `[1,L,…]` per-layer-input hardcode); sliding (RotatingKVCache),
-KV-sharing, and MoE were all already B-generic. Caveat: short-context (pre-wrap)
-— ring-wrap (>window) is the one remaining L1 follow-up. Next row: **L2**
-(quantized KV → `quantizedSdpaUnfused` + 4-D mask, vs optiq). (The former
-"L3: perf kernels under batching" row is void — those kernels were deleted
-2026-07-05, docs/design/unified-engine-frontier-plan.md §6.)
-
-Harness: `tests/batched-decode-parity.test.ts` has two paths — (a)
-`runBatchedDecodeParity` = the internal-consistency KL check (B=N vs our B=1,
-NOT the gate); (b) **`realBatchedGreedy` = the REAL gate**: left-pad → one
-`BatchedDecodeMaskCache` (handles prefill at offset 0 AND per-step decode) →
-batch-prefill → greedy batch-decode, trajectory compared bit-for-bit to the
-mlx-lm B=N oracle golden (`goldens/batched-golden-*.json`, machine-layered via
-`tests/goldens.ts` — per-GPU overrides in `goldens/<machine>/` — generated by
-`scripts/gen-batched-golden.py` in the oracle venv).
-
-**CPM L1 oracle-verified 2026-06-14d:** mlx-bun's real batched greedy trajectory
-== mlx-lm B=2 exactly (both rows incl. the left-padded one, 8 steps). Confirms
-`BatchedDecodeMaskCache` == mlx-lm `BatchKVCache` (left-pad + per-row offset RoPE
-+ `j>=leftPad` mask + fused bool sdpa), and that the wrapper handles batched
-PREFILL too (offset 0 → left-pad prefill mask + `ropeOffsetArr = -leftPad`). The
-earlier "KL 0.26" was measured vs the wrong oracle (our own B=1).
-
-**Gate philosophy (revised twice, 2026-06-14):** batched decode is NOT bit-exact
-vs single-stream (batching changes the attention kernel's reduction order). The
-RIGHT oracle is **mlx-lm's OWN batch mode at the same B** — mlx-lm B=N ≡ mlx-bun
-B=N, **bit-exact**. Comparing our B=N to our own B=1 (the KL harness) measures
-the wrong thing and can't distinguish a real bug from inherent batched-kernel
-behavior. mlx-lm's batch path (oracle venv, confirmed 2026-06-14):
-`BatchKVCache`/`BatchRotatingKVCache(left_padding)` — left-padded, a **per-row
-`offset` array** (`-left_padding`, `+=1`/step → real per-row position, feeding
-RoPE), mask `j >= left_padding[b]` (`create_causal_mask(left_padding=…)`), and
-for bf16 the **same fused bool-mask `mx.fast.scaled_dot_product_attention`** we
-use. So our approach matches in principle and the **additive-mask idea is wrong**
-(would deviate). Oracle generator: `scripts/gen-batched-golden.py` (drives
-mlx-lm's real batch-prefill + greedy batch-decode, dumps per-row trajectories +
-last-position logits). CPM B=2 oracle captured 2026-06-14.
-
-**Remaining gap:** the KL harness solo-prefills each row then *assembles* a
-batched cache, and compares to *solo* decode. mlx-lm batch-PREFILLS (B=N) — so
-even prefill KV differs (B=1 vs B=N kernel). To gate faithfully we need a real
-mlx-bun **batched prefill + batched decode** path, compared bit-exact to the
-mlx-lm B=N golden. (The KL B2-vs-B1 harness stays as a cheap internal-consistency
-check, not the gate.)
-
-**Gemma 12B L1 — RESOLVED 2026-06-14d: it was the wrong oracle, not a bug.**
-The real `realBatchedGreedy` path (batch-prefill + decode) matches the mlx-lm
-B=2 oracle **exactly** for Gemma 12B (both rows incl. left-padded, sliding
-layers via RotatingKVCache→BatchRotatingKVCache, short-context). So the "padded
-row KL ~0.26" was entirely an artifact of the KL harness comparing our B=2 to
-our own B=1 (different prefill + reduction order) — mlx-lm B=2 diverges from its
-own B=1 the same way. No additive-mask fix was needed; the bool+fused path is
-correct (it's what mlx-lm uses). The KL B2-vs-B1 harness stays only as an
-internal-consistency check, never the gate.
-
-**Caveat:** verified at short context (Lmax ≪ 1024 window → pre-wrap, sliding
-window inactive). **Ring-wrap (context > window) is a separate validation** —
-needs a long-context golden (mlx-lm BatchRotatingKVCache wrapped) + mlx-bun
-RotatingKVCache wrapped past the window.
-
-Batched-decode landmines found by the modes analysis (2026-06-14):
-- **e4b**: `computePerLayerInputs` hardcodes `[1,L,…]` (gemma4.ts:647) → B>1
-  breaks; plus KV-sharing (sharer `SharedKv`).
-- **Gemmas**: `RotatingKVCache` sliding-window ring — per-row decode mask + ring
-  write (the hard one). 26B adds MoE (`GatherQMM`/`SwitchGLU`) batched routing.
-- **Generated handlers (all 3 Gemmas)** repeat the K/Q rope timing trap —
-  they read `cache.ropeOffsetArr` twice around `updateAndFetch`
-  (generated/gemma4-12b.ts:63 & :74). The fix goes in the GENERATOR
-  `scripts/gen-model.ts`, then regenerate (the files are DO-NOT-EDIT).
-- **Quant + batched**: quant batched decode goes through `quantizedSdpaUnfused`
-  — correct. (The old "perf kernel bypassed under batching" debt is resolved by
-  deletion: the perf/fused-decode kernels were deleted 2026-07-05,
-  docs/design/unified-engine-frontier-plan.md §6 — no batched-aware fused decode
-  kernel is planned.) `quantizedSdpaUnfused` must be validated with a 4-D bool mask.
-- **All cells**: compiled-decode forced off under `--batch>1`; LoRA batches only
-  within one adapter group (else drain).
-
-## KV memory: dynamic allocation, not static partition
-
-The naive model — reserve `budget / slots` per row, fixed — is rejected:
-it wastes memory and imposes an arbitrary per-slot context cap. We
-allocate **by need**. Three rungs, in increasing order of "who needs it
-most" fidelity and cost:
-
-1. **Static partition** — *not built.* Fixed per-slot ceiling. Wasteful.
-2. **Dynamic shared-budget, contiguous caches (padded batch)** — the
-   shippable path. Admission is governed by *total projected bytes*
-   against one KV budget (mirrors the byte-capped `PromptCache`); each
-   sequence grows its own contiguous cache on demand (KVCache already
-   grows in 256-token steps). Short prompts use little; long ones take
-   more. **Caveat**: a padded batch shares one `[B, H, S, D]` buffer, so
-   the buffer width tracks the *longest live sequence* — short rows still
-   carry padding to that width. Partial "who needs it most."
-3. **Paged KV (the true goal)** — fixed-size KV blocks in a shared pool,
-   gathered by a custom paged-attention Metal kernel. No padding waste;
-   the neediest sequence grabs more blocks, others fewer; maximizes
-   concurrent rows per GB. This is what "dynamically allocate based on
-   who needs it most" actually means. Cost: a custom attention kernel +
-   a block manager. Feasible — the project already ships custom Metal
-   kernel infrastructure (`src/mlx/metal-kernel.ts`).
-
-Plan: ship rung 2 (S1–S2) for the throughput win + dynamic admission
-with no new kernel; build rung 3 (S3+) as the memory-density upgrade.
-Budget accounting is total-bytes from day one.
-
-**Current state (2026-06-14):** the batched *primitive* is **rung 2** —
-one shared `[B, H, S, D]` buffer per layer (`KVCache`/`RotatingKVCache`),
-growing in 256-token steps, width = longest live sequence, left-pad
-"waste" on short rows. Rung-2 allocation IS wired into serving via
-`BatchScheduler` + `GenerationGateway` (S1b SHIPPED). But the
-**budget/admission half of rung 2 is NOT built** (no `B × S_max`
-projection, no KV budget enforcement) — that's scheduler work (S2, NOT
-STARTED). *(Since LANDED 2026-07-03: `--kv-budget` aggregate projection
-gate in the scheduler — grammar-spec-batching-integration.md Phase D.)*
-And **rung 3 (paged): the v1 CACHE CLASS shipped** —
-`PagedKVCache`/`BlockPool` behind default-off `--paged-kv`
-(docs/design/paged-kv-cache.md): serial batch=1, Gemma4, bf16,
-gather-then-stock-sdpa (no custom kernel yet), gated bit-exact vs the
-plain path. The rung-3 PAYOFF (batched block allocation replacing the
-pad+concat merge/extend, block-level CoW prefix sharing, the fused paged
-kernel) is still the deferred follow-up — the shipped slice is the block
-manager + correctness proof, not the density win.
-
-KV quantization is a force multiplier: 4-bit KV (`generate.ts:67`,
-already supported) is ~4× smaller, so the same budget buys ~4× the slots
-or ~4× context per slot. Batching + KV-quant compound.
-
-Note: the existing `2e9` (`server.ts:494`) is the **prompt (prefix-reuse)
-cache**, not live decode KV. Live caches are made fresh per generation
-(`server.ts:517`). Slots introduce a *new* budget: the pool of live
-decode caches.
-
-## The scheduler (port of mlx-lm `BatchGenerator` / `ResponseGenerator`)
-
-When `slots > 1`, the single promise chain (`server.ts:487`) is replaced
-by a scheduler owning:
-
-- a **pending queue** of admitted requests;
-- a **running batch** of ≤ `decodeConcurrency` rows;
-- **admission control** against the KV budget;
-- a **batchable gate** (mlx-lm `_is_batchable`): same model, compatible
-  sampler params, **no fixed seed**, no vision, compatible adapter set;
-- **prefill phase** (≤ `promptConcurrency` rows) vs **decode phase**;
-- **draining**: an incompatible request waits for the batch to clear,
-  then runs in its own batch.
-
-Streaming fan-out is the main plumbing rework: today `generate()` yields
-one token to one SSE stream (`server.ts:964,1037`). Batched, each step
-produces B tokens that must route to B independent response streams,
-each with its own `StopSequencer`/tool-call router. Touches
-`server.ts`, `responses.ts`, `anthropic.ts`.
-
-### Concrete plan (the `--batch N` engine, build order — 2026-06-14)
-
-Started: `--batch N` flag landed (renamed from `--slots`; `--decode-concurrency`
-accepted as an mlx_lm.server alias). N=1 → today's serial path; N>1 → the engine
-below. The engine is a **mode** (whole server), not a load-fallback (keeps
-determinism + the drop-in promise).
-
-**Dynamic-B cache ops** (port of mlx-lm's cache methods — the new piece our
-fixed-B verified forward doesn't have):
-- **`merge(perRowCaches)`** — stack N single-row caches into one `[B,H,S,D]`
-  batch, left-pad-aligned (generalizes `realBatchedGreedy`'s assemble).
-- **`filter(keep)`** — slice the B axis to drop finished rows (eviction).
-- **`extend(other)`** — append a freshly-prefilled batch's rows (insertion).
-- mlx-lm's flow: a joining request is **right-padded + prefilled** in a prompt
-  batch, then `finalize()` (dynamic-roll) → left-pad, then `merge`/`extend` into
-  the running generation batch. Our verified attention/mask/RoPE drop straight in.
-
-**Engine loop (Bun-async, NOT Python threads):** an async loop owns the running
-batch. Each iteration: (1) admit waiting+batchable requests up to N (prefill +
-`extend`); (2) one batched decode step (the verified forward); (3) sample
-per-row, push each row's token to its own SSE `ReadableStream` (per-row
-`StopSequencer` + tool parser + sampler); (4) `filter` out rows that hit
-EOS/stop/max-tokens, resolve their streams. HTTP handlers `await` their row's
-stream; no thread, just the event loop + `mx.async_eval` pipelining.
-
-**Build sequence:** (1) **DONE 2026-06-14** — dynamic-B cache `merge`/`filter` +
-a pure unit test, then a gated parity test vs an mlx-lm *dynamic* golden (rows
-join/leave mid-stream). `scripts/gen-batched-dynamic-golden.py` drives mlx-lm's
-real `BatchKVCache.merge`/`.extract`/`.filter` through {A,B}→join C→evict A→{B,C};
-`realDynamicBatchedGreedy` (`tests/batched-decode-parity.test.ts`) runs the same
-scenario through `mergeKVRows`/`filterKVRows` and matches all 3 per-row greedy
-trajectories token-for-token (CPM L1). Join = re-merge of extracted advanced-offset
-rows + a fresh solo prefill (proves `mergeKVRows` on non-fresh rows); `extend`
-(keep-the-running-batch optimization, mlx-lm's actual join op) is deferred into
-the scheduler. Added `BatchedDecodeMaskCache.releaseRopeArr()`. (Originally
-full-attention only; sliding-window `BatchedRotatingCache` landed 2026-06-14 —
-see the STATUS block above.)
-(2) **DONE 2026-06-14** — the async loop + per-row SSE wired into `createServer`
-behind `--batch N`. `BatchScheduler` (`src/serve/batch-scheduler.ts`) is the
-detached driver; `GenerationGateway` (`src/serve/generation-gateway.ts`) picks
-the lane and an `AsyncMutex` keeps serial and batched off the GPU/`loraState`
-together. Both `handleChat` call sites route through `gateway.run`; per-row
-onToken closures give per-row SSE fan-out for free. Gated: `tests/batch-scheduler.test.ts`
-(teacher-forced KL, evict + join) + `tests/batch-serving.test.ts` (live CPM
-server, concurrent + streaming + drain). No serial regression (server 17/17,
-tools 13/13).
-(3) **DONE 2026-06-14** — `_is_batchable` gate is the `GenerationGateway.willBatch`
-predicate (full-attention model + no vision/adapter/repetition-penalty/user-seed
-→ batch; else serial). *(The predicate has since widened: repetition
-penalty + logits extras now batch via per-row processors, 2026-07-02.)*
-`B×S_max` memory admission is still TODO — **DONE 2026-07-03**
-(`--kv-budget`, see the v2-plan pointer above).
-
-## Fallbacks (must degrade, never break)
-
-- **Vision** → always solo. Image embeddings + bidirectional mask need
-  offset-0 single-sequence prefill (`gemma4.ts:690` throws otherwise).
-- **LoRA** → batch only requests sharing the same adapter set, else
-  drain to solo. `loraState.active` is a single per-generation field
-  (`generate.ts:187`); cache namespaces are adapter-keyed
-  (`prompt-cache.ts:18`). Per-row adapters is a much harder problem;
-  group-by-adapter + drain is accepted (confirmed fine for our use).
-- **Fixed seed / unusual sampler** → drain to solo (matches `_is_batchable`).
-- **Compiled decode** (`generate.ts:348`) → start **disabled** under
-  batching (single-sequence shapeless replay); extend to `[B,1]` later.
-- **Prompt cache** → per-row take/put still works, but the
-  "single queue ⇒ take/reinsert is race-free" invariant
-  (`prompt-cache.ts:7`, `lora.ts:7`) now depends on the **scheduler**
-  being the sole owner of cache mutations. Keep all take/put inside the
-  scheduler loop.
-
-## Config surface
-
-- `ServerOptions.batch` (default 1); `kvBudgetBytes` (shipped 2026-07-03);
-  later `promptConcurrency`.
-- CLI: `--batch N` (shipped; `--slots N` was the original design name, `--decode-concurrency` accepted as mlx_lm.server alias); `--kv-budget <GB>` (shipped 2026-07-03).
-- `batch === 1` → today's serialized promise chain, untouched.
-  `batch > 1` → scheduler.
-
-## Phasing (incremental, parity-gated)
-
-- **S0 — config seam (DONE).** `--slots N` (now shipped as `--batch N`) / `ServerOptions.slots`
-  plumbed, validated, surfaced (ready card + `/stats`). `N>1` warns that
-  batched execution lands in S1 and runs serially. No behavior change.
-- **S1 — static 2-wide, base model only. (DONE — S1b LIVE as `--batch N`)** Split by the reuse finding:
-  - **S1a (prefill)** — reuse the training machinery (`buildBatchedPadMask`
-    / `BatchedMaskCache`); wire the serving path to prefill B prompts in one
-    forward. Mostly generalization + consolidating the two mask builders.
-  - **S1b (decode) — SHIPPED 2026-06-14.** `BatchScheduler`
-    (`src/serve/batch-scheduler.ts`) + `GenerationGateway`
-    (`src/serve/generation-gateway.ts`) wired into `createServer` behind
-    `--batch N`. L1 oracle-verified for CPM + all 3 Gemmas (full-attn +
-    sliding-window via `BatchedRotatingCache`). Per-row SSE fan-out live.
-    See STATUS block above for details.
-  **Teacher-forced gate**: a 2-row batch must produce per-row logits matching
-  two solo runs within bf16 tolerance (the train-batch-e2e methodology, now
-  applied to decode).
-- **S2 — N-wide + continuous injection/eviction. (DONE)** Rows retire on EOS;
-  queued requests prefill into freed rows (shipped with the scheduler).
-  Dynamic byte-budget admission landed 2026-07-03 as `--kv-budget`
-  (aggregate projection gate, queue-don't-OOM —
-  grammar-spec-batching-integration.md Phase D).
-- **S3+ — paged KV (v1 cache class SHIPPED; density payoff DEFERRED)** —
-  the block manager (`src/model/paged-kv.ts`) shipped behind default-off
-  `--paged-kv` at serial scope, bit-exact-gated (paged-kv-cache.md); the
-  batched block allocation, block-level CoW, and the custom
-  paged-attention kernel remain the follow-ups. KV-quant under batch
-  landed separately (Phase 3.1 + milestone 2); LoRA-group batching open.
-
-Every phase ships default-off behind `batch=1`; the serialized path is
-never removed (only bypassed).
-
-## Open questions / to measure
-
-- Sliding-window per-row masking (above) — retire this correctness risk
-  early.
-- `async_eval` decode pipelining (`generate.ts:386`) builds step n+1
-  before reading n; one row stopping mid-pipeline while others continue
-  needs care.
-- Throughput ceiling: the B where we cross bandwidth→compute bound, and
-  aggregate tok/s vs solo — measured via `benchmark.sh` (preflight-gated;
-  loaded-machine numbers don't count).
-
-## References
-
-- mlx-lm continuous batching: `BatchGenerator` / `ResponseGenerator`,
-  flags `--decode-concurrency` (default 32), `--prompt-concurrency`
-  (default 8); `_is_batchable` gate; batch draining.
-  https://deepwiki.com/ml-explore/mlx-lm/3.3-http-server ·
-  https://github.com/ml-explore/mlx-lm/issues/499
-- optiq: **no** batching — single-process serialized, advises external
-  max-concurrency; each in-flight request keeps a resident KV cache.
-  https://mlx-optiq.com/docs/serve