perf(speculative): commit per-row accepts in batched DFlash

docs/CONTINUOUS_BATCHING.md

@@ -708,7 +708,16 @@ kernels run, so a row can pick a different token where two candidates are within
 about a quarter logit. The parity gate admits a first mismatch only inside a
 calibrated 1.0-logit near-tie margin, and a four-identical-prompts control
 shows the divergence is in the batched verify forward rather than the accept
-policy (#2851). On the Gemma 4
+policy (#2851). The round then commits each row's own accept count (#2442):
+the attention KV caches are compacted to per-row contiguous valid prefixes,
+queries and keys rotate at each row's logical positions with the row's stale
+gap masked, and each GDN layer replays its state over that row's accepted
+prefix. `MLXCEL_DFLASH_PER_ROW_ACCEPTS=0` restores the uniform clamp, which a
+quantized or paged KV cache also selects automatically. On an M3 Ultra the
+per-row arm measured 1.397x the clamp's B = 4 aggregate throughput (committed
+accepts 1.036 against 0.323 per row-round); the drafter's context is still a
+shared rectangular prefix, so the walk itself is lower than the clamp arm's,
+and per-row drafter context is a follow-up. On the Gemma 4
 geometries whose MTP verify runs row by row (the 31B everywhere, the 12B on
 CUDA) the batched burst is decode-exact since issue #2190: each row is
 prefilled and verified through the B=1 calls and only the cache write is

docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md

@@ -0,0 +1,74 @@
+# Batched DFlash per-row accepts on M3 Ultra, 2026-11-06
+
+Issue #2442 item 1 follow-up to #2353 / #2440 and #2851. The batched Qwen 3.5
+DFlash round loop used to clamp every round to the shortest live row's accept
+because the attention KV cache is rectangular. That left the committed accept
+count far below the speculative walk, so batching bought little. This measures
+the per-row accept port that replaces the clamp.
+
+## What changed
+
+- `DFlashBatchedGenerator` commits each row's own walk count. The target gets
+  each row's pre-round logical valid end, compacts the attention KV caches back
+  to per-row contiguous prefixes, and replays each GDN layer's conv and
+  recurrent state over that row's accepted prefix
+  (`Qwen35Model::rollback_speculative_cache_rows`).
+- The verify forward rotates queries and keys at each row's logical positions
+  and masks the row's stale gap (`build_divergent_row_mask`,
+  `attend_per_position_rows`), the layout Gemma 4's divergent MTP verify uses.
+- `MLXCEL_DFLASH_PER_ROW_ACCEPTS=0` restores the uniform clamp. A quantized,
+  paged or head-trimmed KV cache selects the clamp automatically.
+
+## Method
+
+- Host: Apple M3 Ultra (80-core, 512 GB), macOS 26.6.2, Metal + Accelerate.
+- Checkpoints: `models/qwen3.5-4b-4bit` target, `models/qwen3.5-4b-dflash`
+  drafter, block size 16, temperature 0.
+- Prompts: the four prompts of the batched DFlash parity harness, truncated to
+  the shortest tokenization so the `[B, L]` prefill is rectangular.
+- 128 tokens per row. Throughput is the round-loop wall clock from
+  `DFlashBatchedRunOutput::stats` (prefill excluded), aggregate over rows.
+- Both arms run in one process behind
+  `throughput_dflash_qwen35_batched_clamp_vs_per_row`, two discarded warm-up
+  runs, three measured runs per arm alternating clamp, per-row, clamp,
+  per-row, ... under `caffeinate -i` and `scripts/with_indexers_paused.sh`.
+  A separate worktree's Rust build was compiling during the run (a few of the
+  80 cores); both arms alternated under that same load, and the clamp and
+  B = 1 cells reproduce the quiet-box #2442 record (83.0 against 82.96-83.00,
+  78.09 against 77.70-77.83), so the comparison is not load-shifted.
+
+## Throughput at B = 4
+
+| run | clamp tok/s | per-row tok/s |
+|---:|---:|---:|
+| 0 | 83.04 | 115.82 |
+| 1 | 83.22 | 116.24 |
+| 2 | 83.06 | 116.05 |
+| median | 83.06 | 116.05 |
+
+Per-row is **1.397x** the clamp. The clamp's three runs span 0.2%, the per-row
+arm's 0.36%. For context, B = 1 measured 78.09 tok/s, so the per-row arm is
+1.49x over the B = 1 round loop and the clamp is 1.06x.
+
+## Accepts
+
+| arm | rounds | clamped rounds | walk per row-round | committed per row-round |
+|---|---:|---:|---:|---:|
+| clamp | 96 | 93 | 2.938 | 0.323 |
+| per-row | 67 | 0 | 1.036 | 1.036 |
+
+The clamp numbers reproduce the #2442 re-scope measurement exactly (walk
+2.938, committed 0.323, 93 of 96 rounds clamped). Per-row commits 3.2x the
+clamp's accept count and finishes in 30% fewer rounds.
+
+The walk itself falls from 2.938 to 1.036: the drafter's context is still a
+shared rectangular prefix (the shortest live row's committed prefix), while the
+target now advances each row by its own accept, so the drafter's context lags
+the rows that accept more. Giving the drafter per-row context is a separate
+change; this port was measured and decided on its own.
+
+## Correctness
+
+`greedy_parity_dflash_qwen35_batched_b4_matches_b1` and the
+`identical_prompts_dflash_qwen35_b4_matches_b1` control pass with the port
+enabled under the calibrated 1.0-logit evaluation-path margin of #2851.

docs/environment-variables.md

@@ -621,6 +621,7 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1
 | `MLXCEL_ENABLE_MTP_B1` | `0`/`false`/`no`/`off` to disable, any other value to force on | adaptive (per hardware) | Manual override for the singleton (B=1) MTP burst, in both directions. When set it pins the decision and disables adaptive profiling (issue #333). When unset, the adaptive policy decides (see `MLXCEL_MTP_ADAPTIVE`); with `MLXCEL_MTP_ADAPTIVE=0` the decision is the static per-hardware default (issue #165): non-batchable targets (`gemma4_unified` 12B pairs, whose only decode path is B=1) default **on** everywhere (measured across three prompts: 1.90x to 3.14x on M5 Max, 1.74x to 2.61x on M3 Ultra, 0.95x to 1.48x on M1 Ultra, where the prose end is a 5% loss rather than a gain); batch-capable targets (the 31B + bf16 assistant) default **on from Apple GPU generation 15** (M3, M4, M5), and fall back to classic decode on generation 13 (M1, M2). This was `has_neural_accelerator` (M5 only) until issue #1217, on the strength of ~1.2 to 1.4× on M5 Max against a ~0.75 to 0.96× regression on M1 Ultra, both measured before #1194/#1199/#1203/#1208/#1215 and neither re-measured since. M3 Ultra, which the old predicate lumped with M1 Ultra, was never measured on this pairing at all until #1217 did it: 1.95× (prose), 2.41× (source code) and 2.65× (enumeration) on 2026-08-20 under the #1215 protocol, at a verify-round cost of 1.51 classic steps against 2.71 on M1 Ultra. The discriminator is the `use_qmv_wide` split (an affine-quantized projection at `M >= 2` runs as one wide pass from generation 15 and as `K` narrow passes below it), not the Neural Accelerator. M4 is grouped by that shared dispatch rather than measured. Generation 13 keeps declining and the width sweep supports that: round cost fits `0.83 + 0.170 K` classic steps for this pairing on M3 Ultra against `1.14 + 0.090 K` for the 12B pairing there, so the bf16 drafter costs about 1.9x as much per block position, and carrying that ratio onto generation 13's `1.35 + 0.346 K` puts a block-4 round near 3.6 classic steps, which 2.96 to 3.99 emitted tokens would only just cover. See `docs/benchmark_results/mtp-b1-gate-m3ultra-2026-08-20.md`. The 31B singleton exactness correction in #1983 restores the automatic narrow retry on M5 Max; its attention cost means these historical pre-gate speedups are not current guarantees. See the [matched QAT/non-QAT report](benchmark_results/gemma4-31b-mtp-exactness-2026-09-26.md). |
 | `MLXCEL_ENABLE_MTP_BATCH` | truthy value | off | **Advanced.** Forces the batched Gemma 4 MTP burst path for parity/debug testing. Not governed by the adaptive policy (issue #333), which scopes to the validated, byte-identical B=1 path. On the row-wise geometries (Gemma 4 31B, and the 12B on CUDA) the batched burst is decode-exact (issue #2190) and runs on CUDA for windows of up to 4 rows whose prompts fit the sliding window and whose `max_prompt_len + 4 * (max_tokens + 1)` stays within the window plus the 32-token rollback buffer; other windows decline to classic decode. A window forms only from requests that are queued together, share a seed and carry no token bias (`--ignore-eos` adds one). With this flag set, a queued burst-eligible head skips the classic batched prefill so the window can form under the default `--max-batch-prefill` (issue #2442); before that, it formed only with `--max-batch-prefill 1`. See `docs/benchmark_results/gemma4-mtp-batched-gb10-2026-10-07.md`. |
 | `MLXCEL_ENABLE_MTP_BATCH_RAGGED` | exactly `1`, `true`, `TRUE`, `yes` or `on` (case-sensitive) | off | **Advanced, experimental.** Only matters with `MLXCEL_ENABLE_MTP_BATCH`. Drops the equal-prompt-length requirement for a batched Gemma 4 MTP window: rows are left-padded to the longest prompt and keep per-row positions. A window whose longest prompt exceeds the target's sliding window declines and returns to classic decode. DFlash windows stay equal-length. |
+| `MLXCEL_DFLASH_PER_ROW_ACCEPTS` | `0`/`false`/`no`/`off` to disable; unset or any other value to enable | on | Per-row accept commits in the batched Qwen 3.5 DFlash round loop (issue #2442). When on, each row commits its own speculative-walk accept count instead of clamping the round to the shortest live row (the uniform clamp of issue #2353): the attention KV caches are compacted back to a per-row contiguous valid prefix, the block's queries and keys rotate at each row's own logical positions with that row's stale gap masked out, and each linear-attention (GDN) layer replays its conv and recurrent state over the row's own accepted prefix. The clamp stays reachable with `0` for A/B attribution; it is also taken automatically when the target's attention caches are quantized, paged or head-trimmed, because those cannot be compacted row by row. Measured on an M3 Ultra (`qwen3.5-4b-4bit` plus `qwen3.5-4b-dflash`, block 16, 128 tokens per row, temperature 0): B = 4 aggregate decode 116.05 tok/s against the clamp's 83.06 (1.397x) and B = 1's 78.09, with committed accepts per row-round 1.036 against the clamp's 0.323. Greedy parity still passes the calibrated `greedy_parity_dflash_qwen35_batched_b4_matches_b1` near-tie gate. The drafter's context remains a shared rectangular prefix, so the walk itself falls from the clamp arm's 2.938 to 1.036; per-row drafter context is a separate follow-up. See `docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md`. |
 | `MLXCEL_METAL4_ATTENTION` | `0`/`false`/`no`/`off` to disable; unset or any other value to enable | on (where the hardware has it) | **Advanced, diagnostic kill switch.** Forces the M5 neural-accelerator fused attention route off on hardware that has it, so `layers::metal4_causal_attention` is skipped and the ordinary SDPA path runs instead. Off-switch only: it cannot turn the route on where `has_neural_accelerator && macos_supports_na` is false, so setting it on an M1 is inert. Use it to A/B the route when an M5 disagrees numerically with an earlier Apple GPU generation; before this switch the only way was to patch `should_use_metal4_attention` and rebuild, which is what the #1182 M5 investigation had to do. It does not touch f32 GEMM precision: the f32 numeric-agreement failures first filed as an M5 problem (issue #1065) were caused by `MLX_ENABLE_TF32`, reproduce on CUDA too, and did not move with this switch. Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on non-Metal builds. |
 | `MLXCEL_NATIVE_WINDOW_SDPA` | `0`/`off` for the legacy mask route everywhere; `1`/`on` for the native route wherever it is statically eligible; `auto` for the native route only where MLX reports a fused kernel for the shape | `auto` | **Advanced, A/B selector (issue #2621).** Chooses how a causal multi-query prefill over a sliding window reaches MLX's SDPA on Metal once the returned keys outgrow the window: through upstream MLX's native `window_size` argument with `mask_mode = "causal"` (native), or through the explicit full-width windowed-causal mask the array-mask SDPA always took (legacy). The decision lives in `mlxcel_core::window_sdpa`; static eligibility is Metal, a positive window, no softcap, no attention sink, no extra mask array (padding, segments, per-head, cache validity), a query block longer than 8 rows and no longer than the keys, keys that outgrow the window, matching query/value head dims in {64, 72, 80, 96, 128, 192, 256}, and float16/bfloat16/float32. Everything else, decode and small verify blocks included, keeps its existing route whatever this variable says, and CPU/CUDA/ROCm never see a positive window (there it selects MLX's graph fallback). `auto` additionally asks MLX's own dispatch (`ScaledDotProductAttention::use_fallback`) whether the native arm runs a fused kernel; where it would run the graph fallback (on a pre-M5 GPU: head dim 256 below a fresh 2048-row half-precision prefill, every head-dim-256 continuation chunk) the legacy mask route runs unchanged. `1` forces the native entry through the graph fallback as well, which is what a kernel A/B needs. Gemma 3 and Gemma 4 build no sliding mask when the route is planned native (`RotatingKVCache::attend` carries the window). The model-level plan reads the K/V dtype the sliding cache returns (`RotatingKVCache::fetched_kv_dtype`), so a `--kv-cache-mode` whose ring returns FP16 K/V under bf16 activations (`fp16+turbo4`) plans legacy and builds the mask once per forward as before, and so does a mix of sliding caches that return different dtypes (`--turbo-boundary-v` keeps the outer layers FP16); should a layer still decline a native plan, it builds the same mask itself, which costs a mask build per layer but never changes semantics. Read once per process. Measured on an M1 Ultra in `docs/benchmark_results/native-window-sdpa-m1ultra-2026-10-26.md`. |
 | `MLXCEL_GDN_CHAIN_PARITY` | `0` to disable, any other value (or unset) to enable | on | **Advanced, diagnostic escape hatch.** Gates the chain-parity gated-delta Metal kernel used by Qwen 3.5 MTP's speculative verify and rollback-replay paths (issue #1165). The standard gated-delta kernel carries float32 recurrent state across a `T = K` verify block and rounds it to the storage dtype only once at the end, while the classic single-token decode chain rounds after every token; a `T = K` verify block is therefore NOT bit-identical to `K` consecutive single-token decode steps unless the state is rounded per in-block step. The chain-parity kernel (`gated_delta_step_seqpar`) does that rounding, which is what makes Qwen 3.5 MTP's temperature-0 output byte-identical to classic decode. **Setting this to `0` forfeits that exactness contract**, restoring the pre-#1165 block numerics for A/B attribution of the parity kernel's own cost and acceptance effect; do not set it to `0` in a deployment that needs byte-identical speculative output. On the non-Metal ops fallback (CUDA, CPU) the flag selects, for an unmasked scalar-gated verify block of at most 32 rows (`CHAIN_PARITY_SEQUENTIAL_MAX_T`), the sequential loop of fused single steps instead of the 64-row chunked scan (issue #1782). That loop is the same arithmetic as the classic decode chain, so it is bit-identical to `T` single-token steps, and it launches far fewer kernels: the chunked scan pays a fixed 63-step Horner loop per linear-attention layer whatever the row count (about 4400 launches and a fifth of the device time per 2-row verify round on GB10 with `qwen3.5-4b-4bit`, twice that when a partial accept replays through it). `0` restores the chunked scan there for A/B attribution. **The kernel is necessary but not sufficient**: byte-identity also requires every quantized projection to dispatch to the same MLX kernel at `M = block_size` as at `M = 1`, which is not true on every GPU generation or at every block width, so the runtime probe behind `MLXCEL_MTP_ALLOW_INEXACT` is what actually decides whether MTP engages. See `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` for the measured kernel cost (inside the dispatch-noise band). |