Repository navigation
fix(cuda): attribute and remove the fixed cost of multi-row speculative verify - #1795
Merged
Merged
Conversation
…step-wise The fixed device cost of a multi-row speculative verify on GB10 (issue #1782) is not graph capture. nsys kernel tables per verify round for the Qwen 3.5 DFlash pairing put the row-independent term in two places: the drafter, whose bf16 checkpoint was converted to f16 at load while the target stays bf16 on Ampere and later CUDA, so MLX promoted every drafter matmul to float32 and re-upcast all 540M drafter weights on every round (72 copy_v<__half, float> launches plus f32 cutlass GEMMs, about 240 launches and 30 to 38 ms per round at every block width); and the 24 gated-delta linear-attention layers, which took the 64-row chunked scan for every verify block, paying a fixed 63-step Horner loop per layer (about 4400 launches per 2-row round, 1.65x that at block 8 and 16 because the rollback replay runs the same scan). The byte and op graph budgets (MLX_MAX_MB_PER_BUFFER, MLX_MAX_OPS_PER_BUFFER) do nothing at 2 rows and are worth 5 to 10% at block 8; graphs off is faster than graphs on at block 8; the materialized-mask attention on the head_dim-256 layers is under 2 ms per round. The per-row term is the qmv to qmm_sm80 switch at M >= 8, not attention. Drafters loaded through mlxcel-core (DFlash, Muse Glimmer, Inkling MTP, Qwen 3.5 MTP) now follow the target loaders' load-time dtype rule (apply_drafter_load_dtype_policy): f16 on Apple Silicon and pre-Ampere CUDA, bf16 on Ampere and later CUDA, with MLXCEL_KEEP_BF16 and MLXCEL_CUDA_F16_NORMALIZE honored the same way, so the drafter lands in the dtype of the residual stream it reads. The gated-delta ops fallback now honors chain parity for an unmasked scalar-gated block of at most CHAIN_PARITY_SEQUENTIAL_MAX_T (32) rows by running the sequential loop of fused single steps. That loop is gated_delta_step applied T times, the same arithmetic the classic decode chain runs, so the block is bit-identical to T single-token steps (the contract the Metal chain-parity kernel keeps and the ops fallback used to forfeit), and it launches T fused steps per layer instead of the fixed scan. Wider blocks and prefill keep the chunked scan. MLXCEL_GDN_CHAIN_PARITY=0 restores the scan and MLXCEL_CUDA_F16_NORMALIZE=1 restores the f16 drafter, so the same binary A/Bs both changes. Reaches every speculative path in the tree (four drafter families plus the GDN fallback shared by Qwen3Next, Qwen 3.5, Kimi Linear and Kimi K3), not only Qwen 3.5; only the Qwen 3.5 DFlash pairing was measured. Before numbers, controls and the kernel tables are in docs/benchmark_results/dflash-verify-fixed-cost-gb10-2026-09-11.md; the after table follows in a later commit. Tests: a pure policy unit test and a bit-exactness test of the chain-parity block against T single steps on the ops path. Refs #1782
After the two fixes, on the same binary with n = 3 per arm: Qwen 3.5 DFlash block 2 goes from 0.44x to 1.24x classic (72.57 tok/s, range 72.21 to 72.99 against a classic maximum of 59.38), block 3 to 1.31x and block 4 to 1.32x; the 2-row verify round's device work drops from 42.6 to 18.6 ms, 1.09x a classic step instead of 2.5x. Greedy text is byte-identical to classic on every arm, including block 16, which diverged before the change. The same-binary kill-switch arms reproduce the before table within 1% and show each fix alone landing near 36 tok/s, both together at 72.6, because the round is a serial chain of host draft build, host verify build and device sync. Widths 6 and above do not win (1.01x at 6, 0.92x at 7, 0.84x at 8, 0.83x at 16): the multirow qmv kernel dispatches accumulator widths of 2, 4 or 8 and the dispatcher switches to qmm_sm80 at M >= 8, so the checkpoint's default width of 16 is on the losing side. That width policy, and MLX_MAX_OPS_PER_BUFFER=100 giving +3% on the classic arm, are recorded as follow-ups rather than changed here. Refs #1782
…regime The test's trailing sanity check on the chunked arm compared bf16 outputs through `rms_rel`, whose in-dtype reduction returns a near-zero denominator for bf16 and reported a drift of 1e8 while the two arms agreed to 1e-7 relative (rms 0.054839782 against 0.054839786). Compare in f32, keep q and k at unit-or-smaller norm the way the model's RMS normalization does, and guard against a degenerate near-zero reference so the bit-exactness asserts cannot pass vacuously. The bit-exactness asserts themselves passed on every run. The drift bound stays at the original 1e-2. Measured on GB10 (CUDA, test-fast profile) under the final input regime the chunked arm's drift against the sequential reference is rms_rel = 2.72e-3 on three identical runs, 3.7x under the bound, so the bound keeps catching a scan that stops tracking the recurrence. An intermediate 5e-2 was set while the 1e8 misreading was still unexplained and is withdrawn; the test now prints the measured drift under --nocapture so the derivation is reproducible. Refs #1782
inureyes
force-pushed
the
fix/issue-1782-verify-graph-capture
branch
from
September 11, 2026 08:20
302577f to
206ce98
Compare
This was referenced Sep 11, 2026
inureyes
added a commit
that referenced
this pull request
Sep 11, 2026
Merge origin/main at 0c21aed (#1795 removed the fixed cost of multi-row speculative verify, #1796 and CI chores) and drop the Laguna loader's private copy of the drafter dtype rule in favor of `apply_drafter_load_dtype_policy`, so the Laguna drafter and the other DFlash-family drafters follow one rule (f16 where the target loaders convert, bf16 on Ampere-and-later CUDA, `MLXCEL_KEEP_BF16` and `MLXCEL_CUDA_F16_NORMALIZE` honored) and cannot drift from the target's residual-stream dtype. Refs #1351
inureyes
added a commit
that referenced
this pull request
Sep 11, 2026
Re-measured after merging main at 0c21aed, which carries #1795 (drafter dtype policy and gated-delta step chaining). Neither fix reaches this pairing: the Laguna drafter was already bf16 and Laguna has no gated-delta layers, so the exactness probe still declines the GB10 (107246 of 200704 logit bytes at block 16, 107466 at block 8) and the throughput curve did not move. The sweep now covers block sizes 2 to 16 continuously (n=3 per width, n=5 off, same binary, GPU lock held): best 1.00x at block 8 with its range inside the off arm's, 0.88x at the default 16, no sub-8 cliff on the NVFP4 path and no crossover past 8. The docs row, the PR body and both reports carry the table, the host state (GPU exclusive, CPU shared with another session's rustc), the per-round attribution, and #1797 for any served-width change; the default stays at 16. Refs #1351
Closed
8 tasks done
inureyes
added a commit
that referenced
this pull request
Sep 11, 2026
## Summary Adds the Poolside Laguna DFlash drafter (`poolside/Laguna-*-DFlash`) as `mlxcel_core::drafter::laguna_dflash`, implements `SpeculativeTarget` on the Laguna target, routes `model_type: laguna` drafters through `load_drafter`, integrates Laguna into the server's `DFlashTargetModel` design (#1751, #1762) as its own drafter family, and wires the pairing into both `mlxcel-server` and offline `mlxcel generate --draft-kind dflash` (the offline DFlash arm did not exist for any target before; it is Laguna-only here). Like the LFM2 and Muse Glimmer arms, the pairing runs behind the measured block-versus-chain exactness gate (`LagunaModel::dflash_exactness_allows`, `MLXCEL_MTP_ALLOW_INEXACT` override) on both the server and the offline arm. ## Where the landed code and the issue text disagree - The forward follows vLLM's `laguna_dflash.py` (vllm-project/vllm#46853), which is what the published checkpoint is served with: every drafter layer passes the projected context through its own `input_layernorm` before the K/V projection (the issue has `ctx_qkv = qkv_proj(ctx_i)`), and the sliding window is a per-query limit over `[context | block]` rather than a fixed 511-row context view. With the checkpoint's norm weights all exactly 1.0 the first difference is numerically a no-op today; the second only matters past 512 context tokens. - The drafter's initial context is every captured prompt row (vLLM's DFlash proposer; `FirstHiddenRows::EveryPromptRow`), not the last row the Qwen 3.5 port uses. - `src/server/speculative_dispatch.rs` needed no Laguna entry: dispatch is by drafter kind; the target arm lives in `dflash_target.rs` and `speculative_burst.rs`. - The sliding target caches are armed with `enable_speculative_buffer(block_size)` so rollback is a tail trim; the drafter's context cache is a temporal buffer capped at `window - 1`, not a ring. ## Measurements (GB10 sm_121, NVFP4 target, post-rebase on main `0c21aed2` with #1795, `MLXCEL_MTP_ALLOW_INEXACT=1`) **Exactness probe, re-run after #1795:** still declines this host at block 16 (107246 of 200704 logit bytes differ at the first verify position) and at block 8 (107466); `qmv_wide` off does not change it. Neither #1795 fix applies to this pairing: the Laguna loader already kept the drafter in bf16 (it now goes through the shared `apply_drafter_load_dtype_policy`), and Laguna has no gated-delta layers. Without the override the offline arm exits with the verdict and the server burst declines to classic decode, so the feature ships inert here by default. **Throughput A/B, same binary, feature off (no `--draft-model`) versus on per block size.** Raw code prompts, no chat template, 200 tokens, GPU lock held for the whole sweep (GPU exclusive, 0 percent at start); the host was not otherwise idle: another session's `cargo build` / `rustc` ran at full CPU throughout (load 2.6 to 3.9), which is the same for both arms. Decode tok/s as the CLI reports it; n runs, min to max; per-round device time is the synchronized verify (drafter plus target) and the drafter's host-side graph build is listed separately. | Configuration | n | tok/s mean (min to max) | Mean accepted | Device ms/round | Drafter build ms/round | vs off | |---|---|---|---|---|---|---| | code 0 (`retry_with_backoff`), off | 5 | 31.18 (30.41 to 31.98) | | | | | | code 0, block 2 | 3 | 16.19 (16.12 to 16.25) | 0.84 | 80 | 30 | 0.52x | | code 0, block 3 | 3 | 21.61 (21.52 to 21.77) | 1.52 | 83 | 30 | 0.69x | | code 0, block 4 | 3 | 27.05 (26.61 to 27.53) | 2.21 | 85 | 31 | 0.87x | | code 0, block 5 | 3 | 28.33 (26.21 to 29.51) | 2.55 | 89 | 34 | 0.91x | | code 0, block 6 | 3 | 30.46 (30.05 to 31.26) | 2.90 | 92 | 33 | 0.98x | | code 0, block 7 | 3 | 30.24 (30.12 to 30.47) | 3.06 | 97 | 34 | 0.97x | | code 0, block 8 | 3 | 31.26 (31.07 to 31.36) | 3.33 | 101 | 34 | 1.00x | | code 0, block 9 | 3 | 30.58 (30.37 to 30.75) | 3.33 | 106 | 33 | 0.98x | | code 0, block 10 | 3 | 30.02 (29.68 to 30.19) | 3.42 | 109 | 35 | 0.96x | | code 0, block 12 | 3 | 28.94 (28.66 to 29.41) | 3.52 | 116 | 37 | 0.93x | | code 0, block 16 (default) | 3 | 27.41 (26.90 to 27.90) | 3.55 | 127 | 35 | 0.88x | | code 1 (`lru_get`), off | 3 | 31.80 (31.50 to 32.38) | | | | | | code 1, block 4 | 3 | 26.10 (25.74 to 26.47) | 2.08 | 83 | 31 | 0.82x | | code 1, block 6 | 3 | 30.42 (30.11 to 30.64) | 2.92 | 91 | 34 | 0.96x | | code 1, block 8 | 3 | 29.95 (29.60 to 30.56) | 3.17 | 102 | 34 | 0.94x | | code 1, block 10 | 3 | 30.16 (29.88 to 30.58) | 3.44 | 109 | 35 | 0.95x | **No configuration is a net win on this host, before or after #1795.** The best width is block 8 at 1.00x on code 0 (31.26 versus 31.18, its whole range inside the off arm's 30.41 to 31.98) and 0.94x on code 1; the default block 16 is 0.88x. The pre-rebase sweep on an idle host gave the same picture (block 8 at 0.97x, block 16 at 0.86x). There is no sub-8 cliff on this NVFP4 path (5, 6, 7 sit on a smooth curve) and no crossover past 8. The default stays at 16 (`DEFAULT_DFLASH_BLOCK_SIZE`); if a served width is ever retuned for Laguna, block 8 is the least-bad on GB10 and that change belongs to #1797, not this PR. **Why the ceiling is where it is.** The device time of a verify round is a fixed 77 ms plus about 3.1 ms per row (80 ms at 2 rows, 127 ms at 16) against a 31 to 32 ms classic token, 2.5x a single-token step even at 2 rows, plus about 30 to 35 ms of host-side drafter graph construction per round, one classic token's worth. The classic step is graph-replayed while the multi-row verify runs eagerly and launch-bound; that is a runtime property of the CUDA backend on this host, unchanged by #1795 for this pairing, and out of scope for a drafter port. Even perfect acceptance at block 8 would reach about 17 ms per token; at the measured 3.3 to 3.6 accepted, 0.88x to 1.00x. Thinking-channel prompts (1.1 to 1.8 accepted) are about 0.4x. Every greedy-id divergence from classic decode was checked with the block-versus-chain probe in `laguna_real_checkpoint_probe`: the chain arm's top-2 margin is 0.0 or 0.125 (one bf16 ulp) at each, so it is the `M >= 2` versus `M = 1` quantized-kernel property the exactness gate measures, not cache or mask logic; the tiny f32 model is token-exact with forced accept lengths 0, 1, 2 and full across the window wrap. The drafter's per-position accuracy along the reference path is 0.88, 1.00, 0.75, 0.62 for d_0 to d_3 on code and 0.88, 0.50, 0.25 inside `<think>`. Server path: same counters as offline (`rounds=60 proposed=835 accepted=67` on the issue's chat prompt). ## Acceptance criteria - Drafter directory loads through `load_drafter` and is refused as a standalone model: verified (detection test plus real runs). - Greedy speculative output token-identical to non-speculative: **not met on this host**, unchanged after #1795: the probe declines by default (107246 of 200704 bytes differ at block 16, 107466 at block 8); with the override, identical on 1 of 6 prompts and differing at one or two bf16 tie positions per 128 tokens on the other 5. Documented in `docs/supported-models.md` and the report. - Rollback trims full and sliding target caches and the drafter context caches: verified by `greedy_invariant_with_forced_rejections` (oracle drafter forcing accept lengths 0, 1, 2, full over a window of 6) and `greedy_invariant_with_real_drafter`. - Mean accepted length above 3 on the code prompt: 1.12 on the issue's chat prompt (the model answers inside `<think>`); 3.55 at block 16 and 3.33 at block 8 on the raw code body (n=3, post-rebase). - Higher decode throughput: **not met on this host** at any block size 2 to 16, before or after #1795 (best 1.00x at block 8, n=3, range inside the off arm's; 0.88x at the default block 16), for the launch-bound reason above; the feature is gated off here by the exactness probe regardless. - `docs/supported-models.md` and the detection test: done. - `cargo test --workspace --profile test-fast --features metal,accelerate`: not runnable on this Linux/CUDA host; unrun. Ran instead on the merged tree: core `drafter::laguna_dflash`, `drafter::dflash`, `drafter::tests` (123 passed); binary `models::laguna_dflash_tests`, `models::detection_tests`, `models::laguna_tests`, `server::batch::dflash_target`, `server::speculative_dispatch_tests` (27 passed, 1 ignored real-checkpoint probe). - `cargo clippy --profile test-fast --features cuda --lib --bins --tests -- -D warnings` (both crates) and `cargo fmt --all -- --check`: clean. ## Changes during review - Merged `origin/main` (nine PRs, including the `DFlashTargetModel` server design from #1751 and #1762 and the drafter-refusal wording from #1770); the branch's own burst trait was dropped in favor of main's and Laguna became a fourth drafter family with the shared measured exactness gate. Merged again at `0c21aed2` to pick up #1795; the Laguna loader now uses the shared `apply_drafter_load_dtype_policy`. - `from_weights` validates the `qkv_proj`, `o_proj`, `g_proj` and `fc` row counts against the config so a mismatched checkpoint is a load error rather than an MLX reshape throw inside the first draft, and the per-round `copy` of the captured hidden input was removed. - Implementation review: the offline arm now refuses a non-Laguna DFlash drafter by name (the pairing gate the server already ran) and refuses sampling requests (the round loop is greedy-only); the drafter prefers its one trained block width so a wider `--draft-block-size` cannot engage the adaptive width controller at an unprobed width; context rows older than the window are dropped before `aux_hidden_norms` and `fc` instead of after; `concatenate_many` replaces the pairwise fold; the rotating caches take the Muse Glimmer slack rule and warn instead of panicking; a short rollback trim logs an error in release builds; `load_drafter` routes on the same Laguna predicate detection uses; `attention_bias: true` is rejected and `intermediate_size` is checked at load; `MLXCEL_PRINT_TOKEN_IDS` keeps its single stderr contract with a prompt-ids line added; two tests pin context-offset continuity and the fused q/k/v row order. - Security and performance review: `greedy_only` is true so the server burst declines sampling requests to classic decode instead of silently returning greedy output; `dflash_config.block_size` (2..=512) and the sliding windows (2..=2^20) are bounded at config load because the drafter's `config.json` is untrusted and sizes allocations and `i32` mask arithmetic; the sanitizer checks split q/k/v shapes before fusing and refuses prefixed-plus-unprefixed duplicates; projection input widths and norm lengths are checked at load; the context K/V comes from a row-sliced K/V-only projection on plain checkpoints, the context cache exposes its buffer instead of returning a copy, and `combine_hidden` builds one `concatenate_many` node; the offline arm refuses a token bias. The regression check after each fix commit reproduced the same 128 ids and `accepted=98` on the code prompt. Closes #1351
9 tasks
This was referenced Sep 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The fixed device cost of a multi-row speculative verify on GB10 is not graph capture. nsys kernel tables per verify round on the Qwen 3.5 DFlash pairing (
qwen3.5-4b-4bit+qwen3.5-4b-dflash, the pairing onmain) put the row-independent term in two places: the drafter ran in float32, becauseDFlashDrafter::loadconverted its bf16 checkpoint to f16 unconditionally while the target keeps bf16 on Ampere and later CUDA, so MLX promoted every drafter matmul to f32 and re-upcast all 540M drafter weights each round (30 to 38 ms per round at every width); and the 24 gated-delta layers took the 64-row chunked scan for every verify block, a fixed 63-step Horner loop per layer (about 4400 launches per 2-row round, 1.65x that at block 8 and 16 through the rollback replay). The issue's three graph-side candidates are secondary (the two budget knobs move nothing at 2 rows and 5 to 10% at block 8; graphs off is faster than graphs on at block 8; replay is healthy), the materialized-mask attention is under 2 ms per round, and the per-row term is theqmvtoqmm_sm80switch atM >= 8. Full attribution, controls and kernel tables:docs/benchmark_results/dflash-verify-fixed-cost-gb10-2026-09-11.md.What changed
mlxcel-core(DFlash, Muse Glimmer, Inkling MTP, Qwen 3.5 MTP) follow the target loaders' load-time dtype rule (apply_drafter_load_dtype_policy): f16 on Apple Silicon and pre-Ampere CUDA, bf16 on Ampere and later CUDA,MLXCEL_KEEP_BF16andMLXCEL_CUDA_F16_NORMALIZEhonored the same way.MLXCEL_CUDA_F16_NORMALIZE=1restores the f16 drafter on this host.CHAIN_PARITY_SEQUENTIAL_MAX_T) by runninggated_delta_stepTtimes, the classic chain's own arithmetic, so the block is bit-identical toTsingle-token steps (a test pins that) and launchesTfused steps per layer instead of the fixed scan. Wider blocks and prefill keep the scan.MLXCEL_GDN_CHAIN_PARITY=0restores it.docs/environment-variables.mdrows for the three variables, a dated GB10 record underdocs/benchmark_results/, and the bilingual technical report (force-added).Measurements (GB10 sm_121, same binary, n = 3, idle host, GPU lock held, min to max)
Same-binary A/B: both kill switches set reproduces the before table within 1% (25.98 at block 2, 28.83 at block 8); each fix alone lands near 36 tok/s at block 2, both together at 72.6, because the round is a serial chain of host draft build, host verify build and device sync and each fix removes a different link. The 2-row verify is now 1.09x a classic step instead of 2.5x. Block 16 diverged from classic greedy text before the change and matches after it.
Blast radius, and what was not measured
The dtype policy reaches every drafter family loaded through
mlxcel-core; the gated-delta change reaches the chain-parity call sites (Qwen 3.5 verify and rollback replay) on every non-Metal backend and no othergated_delta_opscaller. Nothing in the MLX pin or the C++ bridge changed. Only the Qwen 3.5 DFlash pairing was measured; Muse Glimmer, Inkling MTP, Qwen 3.5 MTP and Laguna (#1771, not onmain) share the code path without a number yet.Not fixed: widths of 6 and above still lose here (the multirow
qmvkernel's 8-wide accumulator from 5 rows,qmm_sm80from 8), and the checkpoint default of 16 is on that side;--draft-block-size 4is the measured setting for this pairing, and a hardware-gated default width is left as a follow-up. Separate finding for its own issue:MLX_MAX_OPS_PER_BUFFER=100gave +3% on the classic arm of one model (60.20 vs 58.42, ranges disjoint), which belongs to every CUDA decode path; not changed here.Test plan
cargo fmt --all -- --checkdrafter::dflash::drafterunit tests on this host under--profile test-fast --features cuda(policy and env-flag parsing, 9 passed)models::gated_delta::testson this host under the same profile (8 passed, including the new bit-exactness test; its drift sanity check compares in f32 becauserms_relreduces in the input dtype and misreads bf16, and measuresrms_rel = 2.72e-3on three identical runs against the original 1e-2 bound)cargo clippy --profile test-fast --features cuda --lib --tests -p mlxcel-core -p mlxcel -- -D warningsmake verify-test-cuda(whole-workspace, cannot run under the 600 s watchdog on this host; an unscopedcargo test --libalso aborts here with acudaStreamEndCaptureC++ abort unrelated to this change)cargo test --workspace --profile test-fast --features metal,accelerate: not runnable on this Linux/CUDA host, unrunCloses #1782