Repository navigation
perf(speculative): close batched DFlash gaps left by the uniform clamp #2442
Description
Follow-up to #2353 / #2440 (see TECHNICAL_REPORTS/pr-2440-dflash-batched-uniform-accept.en.md). Three gaps in the Qwen 3.5 batched DFlash path. Items 2 and 3 are correctness/coverage and may land together in one PR before item 1, which needs measurement. Run the local gate once at the end of each PR and keep pushes to a minimum.
1. Throughput: per-row accepts instead of the uniform clamp
clamp_uniform_acceptance (src/lib/mlxcel-core/src/drafter/dflash/round_loop_batched.rs:188, called at :557) cuts every round to the shortest live row. On qwen3.5-4b-4bit with qwen3.5-4b-dflash, B = 4 commits 1.23 drafts per row-round (walk: 4.11, 55 of 57 rounds clamped) versus about 3.4 at B = 1, so batching gains little. rollback_speculative_cache asserts a uniform slice (src/models/qwen3_5.rs:2932).
Evaluate porting the layout Gemma 4 uses under MtpBatchedGenerator (src/lib/mlxcel-core/src/speculative/mtp/round_loop_batched.rs:126): per-row valid prefix and RoPE position (DivergentVerifyRows, src/models/gemma4.rs:1669; build_divergent_verify_mask, :4201; mask_stale_key_gap, src/lib/mlxcel-core/src/utils.rs:618), applied to the Qwen 3.5 attention KV and a per-row GDN conv/recurrent replay.
- Greedy B = 4 output matches four B = 1 runs (
greedy_parity_dflash_qwen35_batched_b4_matches_b1); this is a hard gate. Landed: the gate passes with per-row accepts as the default under the calibrated 1.0-logit near-tie margin of fix(speculative): batched DFlash verify forward is not byte-identical to a single-row forward #2851, and the identical-prompts control passes too. - Committed accepts per row-round approach the walk's count, and aggregate tok/s at B = 4 is measured against the clamp on a quiet box. Keep the clamp if the port does not win. Landed: committed accepts 1.036 against the clamp's 0.323 per row-round, and B = 4 aggregate throughput 116.05 tok/s against 83.06 (1.397x, three alternating runs per arm, 128 tokens per row).
MLXCEL_DFLASH_PER_ROW_ACCEPTS=0restores the clamp.
Status 2026-11-06: item 1 implemented and measured. Per-row accepts compact the Qwen 3.5 attention KV back to per-row contiguous prefixes, rotate and mask the verify block at each row's logical positions, and replay each GDN layer over that row's accepted prefix. Method and numbers: docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md. The drafter's context is still 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.
Status 2026-11-03: still open. Not implemented because neither qwen3.5-4b-4bit nor qwen3.5-4b-dflash is present on the validation host, so the hard parity gate cannot run. Shipping an unvalidated per-row accepts port would forfeit the temperature-0 contract.
2. Routing: queued equal-length requests skip the burst
execute_prefill_action (src/server/batch/scheduler/request_memory.rs:23-33) sends any queue of 2+ to execute_batched_prefill when max_batch_prefill > 1 (default 4, src/server/config.rs:1241). Only execute_prefill calls try_speculative_burst (src/server/batch/scheduler/prefill.rs:227), so those rows get no speculative decoding at all, B = 1 or batched. --max-batch-prefill 1 was needed to reach the burst.
- Decide and document in the
execute_prefill_actiondoc comment. If unintended, route a burst-eligible head (should_burst_for_sequence) throughexecute_prefill. Landed in fix(speculative): route queued DFlash heads to the burst #2447: a burst-eligible head now takesexecute_prefillthroughqueued_head_takes_speculative_burst, and the decision is documented onexecute_prefill_action. - Scheduler test: DFlash dispatch, default
max_batch_prefill, two equal-length queued requests reachtry_run_burst_batchedas one window; fails on current code. Landed in fix(speculative): route queued DFlash heads to the burst #2447:queued_equal_length_dflash_requests_reach_the_batched_burst_as_one_window.
3. Batched burst skips the exactness probe
run_dflash_burst gates on exactness_allows(bs) before drafter load (src/server/batch/speculative_burst.rs:1724-1760). run_dflash_burst_batched (:2788) checks only supports_batched(). The doc at :1671 still calls Qwen 3.5 permissive, though dflash_exactness_allows (src/models/qwen3_5.rs:1610) now probes.
- Apply the same probe and
DeclineToClassicwarning beforeensure_loaded; fix the stale doc. Landed in fix(speculative): route queued DFlash heads to the burst #2447:run_dflash_burst_batchedapplies the shareddflash_target_exactness_allowsprobe beforeensure_loaded, and the stale "permissive on Qwen 3.5" doc is corrected. - Test that a failing verdict declines the window without loading the drafter. Landed in fix(speculative): route queued DFlash heads to the burst #2447:
qwen35_batched_dflash_declines_a_failing_exactness_verdict_before_drafter_io.
Verification
cargo clippy --features cuda --all-targets -- -D warnings
cargo test --profile test-fast --features cuda --lib server::batch -- --test-threads=1
MLXCEL_SDPA_DETERMINISTIC=1 gpu-lock run --tag dflash-parity -- cargo test --release --features cuda --test speculative_parity qwen35_dflash_batched -- --ignored --test-threads=1Server check: mlxcel-server --model <target> --model-draft <drafter> --draft-kind dflash with default flags, two concurrent temperature-0 requests run as one batched burst and match two sequential runs.
Activity
fix(speculative): keep batched prefill while classic rows decode
inureyes commentedon Oct 18, 2026
Items 2 and 3 landed in #2447; item 1 remains open here.
--max-batch-prefill. The batched DFlash window is now opt-in (MLXCEL_ENABLE_DFLASH_BATCH), so by default each queued DFlash request runs the exact B = 1 burst.run_dflash_burst_batchedapplies the B = 1 exactness probe before loading the drafter.Finding relevant to item 1: on GB10 (qwen3.5-4b-4bit with qwen3.5-4b-dflash, temperature 0,
MLXCEL_SDPA_DETERMINISTIC=1) a B = 2 batched DFlash burst on unchanged main (--max-batch-prefill 1) diverged from the same two requests served one at a time, at block 4 and 16, with and without the prompt cache. A window of two identical prompts, where no round clamps (clamped_rounds=0), diverged as well, so the batched[B, block]verify forward is not byte-identical to the single-row one on this host, independent of the uniform clamp. The B = 1 burst and classic batched decode matched. The parity gate in item 1 needs to cover this before the opt-in can become the default. Details and the full run matrix are in the #2447 description.fix(speculative): keep batched DFlash windows on by default
inureyes commentedon Oct 19, 2026
Correction to the earlier comment: PR #2447 no longer adds
MLXCEL_ENABLE_DFLASH_BATCH. Batched DFlash windows stay on by default, and with the routing fix equal-length queued DFlash requests now reach the batched burst as one window under default flags. The batched burst also runs the exactness probe the B = 1 path runs.Still open here, item 1: on GB10 the batched verify forward is not byte-identical to the B = 1 forward. Two concurrent temperature-0 requests through the batched burst can differ from the same requests run one at a time, including two identical prompts with no clamped round. Main already showed this with
--max-batch-prefill 1. Per-row accepts and a passing B = 4 versus B = 1 parity gate belong with that work.fix(speculative): route queued DFlash heads to the burst (#2447)
10 remaining items
inureyes commentedon Nov 4, 2026
Item 1's prerequisite is measured and closed in #2853.
The batched verify forward is not bitwise identical to the B = 1 forward, and the accept policy is not the cause. Four identical prompts in one B = 4 burst (
clamped_rounds=0, no ragged round) part from B = 1 at the same token on every row, so the divergence is inside the batched target forward. At that position a one-shot B = 1 arbiter ranks the batched token first and the B = 1 token 0.25 logits behind, which is the evaluation-path near-tie class (kernel dispatch by leading dimension), not a structural break.So the gate
greedy_parity_dflash_qwen35_batched_b4_matches_b1is re-expressed as a calibrated 1.0-logit near-tie margin, with a committed margin-removal control that rejects every observed mismatch at margin 0. It passes on the M3 Ultra host at block 16, with and withoutMLXCEL_SDPA_DETERMINISTIC=1. Byte identity was not taken: it would require per-row evaluation of the target verify forward against per-row KV/GDN cache views (the Gemma 4 per-row architecture), which the clamp's 1.07x batching win does not justify.Item 1 is therefore no longer blocked on a byte-identity prerequisite. Per-row accepts can now be evaluated under the margin-gated parity, with its own throughput measurement against the clamp arms (B = 1 77.7, B = 4 clamp 83.0 tok/s aggregate). The runtime
dflash_exactness_allowsprobe still measures only the B = 1 block-versus-chain pair; a B>1 probe is not part of this result.test(speculative): gate batched DFlash parity on measured jitter (#2853)
inureyes commentedon Nov 6, 2026
Item 1 is implemented in #2865 and measured on the M3 Ultra (
qwen3.5-4b-4bitwithqwen3.5-4b-dflash, block 16, 128 tokens per row, three alternating runs per arm): B = 4 aggregate 116.05 tok/s against the clamp's 83.06 (1.397x), committed accepts 1.036 against 0.323 per row-round, 67 rounds against 96. Per-row accepts are the default;MLXCEL_DFLASH_PER_ROW_ACCEPTS=0restores the clamp, and a quantized or paged KV cache selects it automatically. The hard parity gate and its identical-prompts control pass under the #2851 near-tie margin. Caveat: the drafter context is still 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 follow-up. Method and numbers:docs/benchmark_results/dflash-per-row-accepts-m3ultra-2026-11-06.md.perf(speculative): commit per-row accepts in batched DFlash