Skip to content

fix(cuda): attribute and remove the fixed cost of multi-row speculative verify - #1795

Merged
inureyes merged 3 commits into
mainfrom
fix/issue-1782-verify-graph-capture
Sep 11, 2026
Merged

inureyes merged 3 commits into
mainfrom
fix/issue-1782-verify-graph-capture

Conversation

@inureyes

@inureyes inureyes commented Sep 11, 2026 •

Copy link
Copy Markdown
Member

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 on main) put the row-independent term in two places: the drafter ran in float32, because DFlashDrafter::load converted 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 the qmv to qmm_sm80 switch at M >= 8. Full attribution, controls and kernel tables: docs/benchmark_results/dflash-verify-fixed-cost-gb10-2026-09-11.md.

What changed

  • Drafters loaded through 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_BF16 and MLXCEL_CUDA_F16_NORMALIZE honored the same way. MLXCEL_CUDA_F16_NORMALIZE=1 restores the f16 drafter on this host.
  • The gated-delta ops fallback honors chain parity for an unmasked scalar-gated block of at most 32 rows (CHAIN_PARITY_SEQUENTIAL_MAX_T) by running gated_delta_step T times, the classic chain's own arithmetic, so the block is bit-identical to T single-token steps (a test pins that) and launches T fused steps per layer instead of the fixed scan. Wider blocks and prefill keep the scan. MLXCEL_GDN_CHAIN_PARITY=0 restores it.
  • docs/environment-variables.md rows for the three variables, a dated GB10 record under docs/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)

Configuration before tok/s after tok/s vs classic after round device sync before / after greedy text == classic
classic 58.42 (58.14 to 58.62) 58.33 (57.27 to 59.38) yes
block 2 25.82 (25.76 to 25.88) 72.57 (72.21 to 72.99) 1.24x 42.6 / 18.6 ms yes
block 3 76.42 (76.36 to 76.51) 1.31x / 24.4 ms yes
block 4 31.17 (30.93 to 31.38) 76.88 (76.62 to 77.28) 1.32x 58.5 / 29.9 ms yes
block 6 59.06 (58.84 to 59.19) 1.01x / 47.4 ms yes
block 8 28.55 (28.47 to 28.64) 48.84 (48.48 to 49.32) 0.84x 81.6 / 59.2 ms yes
block 16 (checkpoint default) 31.04 (30.57 to 31.45) 48.63 (48.48 to 48.83) 0.83x 78.7 / 64.1 ms yes (was NO)

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 other gated_delta_ops caller. 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 on main) share the code path without a number yet.

Not fixed: widths of 6 and above still lose here (the multirow qmv kernel's 8-wide accumulator from 5 rows, qmm_sm80 from 8), and the checkpoint default of 16 is on that side; --draft-block-size 4 is 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=100 gave +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 -- --check
  • drafter::dflash::drafter unit tests on this host under --profile test-fast --features cuda (policy and env-flag parsing, 9 passed)
  • models::gated_delta::tests on this host under the same profile (8 passed, including the new bit-exactness test; its drift sanity check compares in f32 because rms_rel reduces in the input dtype and misreads bf16, and measures rms_rel = 2.72e-3 on three identical runs against the original 1e-2 bound)
  • cargo clippy --profile test-fast --features cuda --lib --tests -p mlxcel-core -p mlxcel -- -D warnings
  • make verify-test-cuda (whole-workspace, cannot run under the 600 s watchdog on this host; an unscoped cargo test --lib also aborts here with a cudaStreamEndCapture C++ abort unrelated to this change)
  • cargo test --workspace --profile test-fast --features metal,accelerate: not runnable on this Linux/CUDA host, unrun

Closes #1782

…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
@inureyes inureyes added status:review Under review type:performance Performance improvements priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:inference Generation, sampling, decoding (incl. speculative, DRY) platform:linux Linux (CUDA / packaging) specific status:done Completed and removed status:review Under review labels Sep 11, 2026
…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
inureyes force-pushed the fix/issue-1782-verify-graph-capture branch from 302577f to 206ce98 Compare September 11, 2026 08:20
@inureyes
inureyes merged commit e391ae9 into main Sep 11, 2026
14 checks passed
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
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
@inureyes
inureyes deleted the fix/issue-1782-verify-graph-capture branch September 14, 2026 11:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:inference Generation, sampling, decoding (incl. speculative, DRY) platform:linux Linux (CUDA / packaging) specific priority:medium Medium priority status:done Completed type:performance Performance improvements

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf(cuda): the multi-row speculative verify forward is launch-bound on GB10

1 participant