[EXL3] allow CUDA-graph decode for dense checkpoints by priming exl3_gemm autotune - #314
Conversation
…gemm exl3_gemm autotunes with timing launches, so _require_enforce_eager refuses non-eager execution for every dense tensor_storage checkpoint. Decode capture sizes are a known finite set, so prime every serialized shard over them during process_weights_after_loading -- the same pattern _warm_decode_shapes already uses for the online-K6 path -- and permit decode-only capture behind VLLM_EXL3_GRAPH_DECODE. Measured on Qwen3.8-27B EXL3 K4 (RTX PRO 6000, SM120, TP1): decode throughput +92%/+84%/+98% at C1/C4/C8, distribution parity vs eager exactly 0.000000 KLD over 32 contexts.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (1)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Tested this PR on AIBoss: RTX 5090 (SM120), driver 610.57.04, TP1, GG r34 base image The derivative image contains the exact Confirmed fail-closed behavior first: non-eager Decode resultsGreedy, fixed 512 output tokens,
The eager run used 131072 max context / 0.85 utilization; graph runs used native 262144 / 0.98, so the graph-vs-eager percentages are preliminary rather than a controlled performance claim. MTP1 versus graph MTP0 used the same graph/native-context configuration and improved C1/C4/C8 by 57.1% / 60.7% / 51.0%. MTP1 accepted 3536/4145 draft tokens (85.31%). Capture and operational checks
Correctness caveatA separate 32-prompt free-running greedy corpus was 29/32 token-sequence exact between the preliminary eager run and graph MTP0; the three differences appeared after autoregressive branch points. Because those two runs did not have matched memory/context settings, I am not treating this as a regression or as parity proof. A matched rerun was interrupted by an unrelated NAS outage. This test also does not reproduce the PR's full-distribution KLD harness, so it makes no model-fidelity claim. Bottom line: graph startup/capture, fail-closed gating, MTP1 capture-row priming, tools, vision, and near-native context all work on SM120/RTX 5090, with a large decode gain. Matched eager/graph numerical parity remains the one item I would close before calling this production-qualified. |
|
Follow-up with the matched eager/graph run that was missing from my earlier qualification. Both servers used the same RTX 5090, image/model revision, Results:
Graph gain: +148.1%, +135.0%, +133.6%. Exact cross-mode token/text parity was 29/32, not 32/32. Repeating the corpus without restarting gave 32/32 identical sequences within eager and 32/32 within graph, so the three differences are reproducible mode differences rather than run-to-run nondeterminism. All three first divergences are near-tie decisions:
Across 967 identical pre-divergence chosen tokens, chosen-token logprob mean absolute delta was 0.00994 and maximum was 0.62842. The generated answers remain semantically valid, but exact eager/graph top-1 parity does not hold on this SM120/RTX 5090 test. Receipt SHA-256 values:
This closes the configuration mismatch from my previous comment and narrows the residual discrepancy to graph/eager numerical behavior at tied or near-tied decode decisions. |
|
Final AIBoss/RTX 5090 qualification disposition for this head (
No NVRM/Xid faults occurred. The host has been returned to its healthy NVFP4 production service. Full receipts and source provenance are preserved with SHA-256 manifest in |
|
Thank you for the qualification - it found a real defect in my evidence, and the follow-up changes what the HOLD should be measured against. Three parts: a retraction, a refuted hypothesis, and a control that I think resolves the open question. 1. My parity receipt was measuring the wrong path. Retracted.I reported 2. I reproduced your finding on a second device, and it is slightly larger hereProper harness (real decode steps through the OpenAI endpoint, 32 prompts x 32 tokens, 1x RTX PRO 6000 Blackwell (SM120), driver 595.58.03, dense
Same shape as yours - internally deterministic, cross-mode divergence at near-ties, e.g. 3. The zero-priming hypothesis: tested, refutedMy first suspicion was this patch's fault: Primed instead on realistic activations (
Unchanged to five digits. Not the cause. I am keeping the randomised priming anyway - probing a timing-based autotuner with zeros is indefensible on its own terms - but it is not a fix and I will not present it as one. 4. The control that I think resolves the HOLDIf the drift is not from priming, is it from EXL3 at all? Same harness, same prompts, same seed, on the unquantised BF16
The BF16 model drifts by the same amount as the EXL3 model. So the reproducible near-tie drift you measured is a property of this build's CUDA-graph decode path in general, not of this patch or of EXL3's kernels. This patch does not introduce it; it makes graphs reachable for dense EXL3, and graphs then behave exactly as they do for an unquantised model on the same stack. That suggests the gate should be restated: not "graph == eager bit-exact" - no graph path here satisfies that, BF16 included - but "graph drift within the same envelope as an unquantised model on the same stack", which is measured above and is satisfied. If you would rather hold on the stronger criterion, that is a decision about CUDA graphs in this build, and it should apply to BF16 too. Honest limit on all of this, mine and yours: 32 x 32 greedy tokens is a numerical probe, not a task evaluation. Neither of us has shown whether flipped near-ties matter downstream, and that remains open - for graphs generally, not for EXL3. Changed in this PR since your test
Receipts and method: |
exl3_gemm's autotuner picks a kernel configuration by measured time, so an all-zero probe can select a different winner than a real activation does. Measured on SM120: this does NOT change eager-vs-graph near-tie drift (25/32 vs 24/32 exact sequences, identical logprob deltas), and BF16 with no EXL3 code drifts the same amount (24/32), so the drift is ambient to the graph path. Kept because probing a timing-based autotuner with zeros is indefensible on its own terms.
|
Data point: this PR's doctrine (initialise every autotune/plan bucket eagerly, before any capture) turned out to generalise to B12X. Our b12x warm ran once per device; bits=4/5 mixed-Trellis plans therefore initialised lazily — potentially inside decode-graph capture — which is the surviving hypothesis for a plausible-garbage-under-replay corruption we isolated (the bits=5 GEMM itself proves bit-clean eagerly, cos 1.000000 on real checkpoint tensors). Warm is now keyed on |
1 participant
Companion to #312. Makes CUDA graphs usable for dense EXL3 checkpoints.
Problem
Exl3Config._require_enforce_eager()raises for every checkpoint withoutrank_sliced_metadata, becauseexl3_gemmautotunes with timing launches that fault under capture, and the comment notes that "m-bucketing means a warmup pass cannot reliably cover every bucket". So a dense checkpoint is eager-only, which on a 27B dense model costs 46-50 % of decode throughput (measured below).Why the m-bucketing objection does not hold for decode
gemm_autotune_hash(exllamav3_ext/quant/exl3_gemm.cu) mixesMIN(roundup_pow2(MAX(size_m,2)),16),size_k,size_n,K,c_fp32, device, cc, SM count and the codebook selector — shape and codebook only, no weight pointers. The effective bucket set is therefore{2,4,8,16}, and vLLM's decode capture sizes are a known finite list. Priming the exact capture sizes is a strict superset of bucket coverage. The first call per device also does a lazycudaMallocinDevCtx::get_locks, which likewise must happen before capture.Change
_graph_decode_capture_rows()builds the row plan:cudagraph_capture_sizes, plus the supersetadjust_cudagraph_sizes_for_spec_decode()can produce after weight loading, plussize // qper-request rows, clamped tomax_cudagraph_capture_size.Exl3LinearMethod._prime_graph_decode_shapes()runs one zero-filled fp16 GEMM per plan row for every shard that routes through_exl3_gemm, memoised in_EXL3_GEMM_PRIMED_SIGNATURESkeyed(device, m, k, n, bits, codebook)— mirroring the existing_EXL3_ONLINE_WARMED_SIGNATURES. K6+mcgshards are skipped (native B12X, already prepared and warmed). One arena per geometry. A priming failure is converted into the enforce-eagerValueError, so the relaxation stays fail-closed._require_enforce_eagerpermits non-eager only whenVLLM_EXL3_GRAPH_DECODE=1and the mode captures decode only (mixed_mode() == NONE) and no DBO/ubatching and the capture list is non-empty; otherwise it raises the original error with the failing condition named. Default is unchanged: off._require_eager_moe_experts()keeps refusing for non-rank-sliced EXL3 MoE, whose per-expert row counts only the router knows. R7 rank-sliced and online-K6 paths are untouched.VLLM_EXL3_GRAPH_DECODEregistration forvllm/envs.pyis attached in the sibling patch file in the linked repo; happy to fold it into this PR if you prefer one commit.Verification
r34 image, 1x RTX PRO 6000 Blackwell (SM120), driver 595.58.03, TP1, dense EXL3 Qwen3.8-27B (193 EXL3 matrices: MLP K4/
mcg,lm_headK6/mul1, attention BF16 for the b6 overlay).Startup, with
--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'and no--enforce-eager:Exactly three priming lines, one per shard geometry, and no mixed prefill-decode capture bar.
Throughput (
--max-num-seqs 8, greedy,ignore_eos, 256 output tokens, warmup discarded):For scale, graphs are worth only +7.7/+9.0/+11.4 % to the BF16 model on the same box; eager EXL3 pays per-call dispatch on 193 quantized matmuls, which is why the win is ~9x larger.
Correctness: greedy text and a multimodal request both answer correctly, and distribution parity against eager is exact — full-vocabulary
KL(eager || graphs)over 32 contexts x 2047 positions is 0.000000 with top-1 agreement 1.000000.Negative controls behave as intended: without the env var, and with the default mixed-capturing mode, startup still raises with the failing condition named;
--enforce-eageris unaffected.Not covered, deliberately
Mixed/prefill capture (refused, token counts not enumerable), non-rank-sliced MoE experts (refused), DBO/ubatching (refused), and the pre-existing rank-sliced path (untouched). Spec-decode row plans are a superset computed at weight-load time; if a future capture-size transform appears the failure mode is a fault at capture, not silent corruption.