[Bugfix] Use PIECEWISE cudagraphs on Blackwell if max_model_len > 131072

[mgoin]

Purpose

FIX #27057

The original issue was found because Qwen3-VL models completely lost accuracy (1% vs 86% on GSM8K) on B200 GPUs when using the default FULL_AND_PIECEWISE cudagraph_mode. The issue did not occur on Hopper at all, with PIECEWISE mode only, FlashAttention backend, or when explicitly disabling TRTLLM attention.

Because TRTLLM attention is selected dynamically based on runtime conditions (num_tokens, max_seq_len, kv_cache_dtype). During FULL CG capture, the max_seq_len is used which when greater than 128K results in FlashInfer being selected, but during actual inference without using full context length, the same conditions triggered TRTLLM selection. This created a graph/runtime mismatch where captured graphs referenced FlashInfer kernels but runtime attempted to execute TRTLLM kernels, producing incorrect results. I was able to see this behavior on any model with default max_model_len>128K

By enforcing PIECEWISE mode in this PR to disable cuda graph capture of attention, we can avoid this issue of dynamism. In the future we should see if we can made TRTLLM support larger context lengths to support FULL graphs

Test Plan

Test Result

Reproduction on B200 on main:

vllm serve gradientai/Llama-3-8B-Instruct-Gradient-1048k
python tests/evals/gsm8k/gsm8k_eval.py

Running GSM8K evaluation: 1319 questions, 5-shot
Evaluating: 100%|████████████████████████████████████████| 1319/1319 [01:04<00:00, 20.56it/s]

Results:
Accuracy: 0.006
Invalid responses: 0.775
Total latency: 64.173 s
Questions per second: 20.554

vllm serve gradientai/Llama-3-8B-Instruct-Gradient-1048k --max-model-len=100K
python tests/evals/gsm8k/gsm8k_eval.py

Running GSM8K evaluation: 1319 questions, 5-shot
Evaluating: 100%|███████████████████████████████████████| 1319/1319 [00:11<00:00, 118.01it/s]

Results:
Accuracy: 0.579
Invalid responses: 0.004
Total latency: 11.190 s
Questions per second: 117.872

Running on this PR:

vllm serve gradientai/Llama-3-8B-Instruct-Gradient-1048k
(APIServer pid=2588191) WARNING 10-17 13:50:11 [vllm.py:385] NVIDIA Blackwell TRTLLM attention cannot support max_model_len >= 131072 (found 1048576), causing dynamic dispatching that breaks full cudagraphs. Overriding cudagraph_mode to PIECEWISE.
python tests/evals/gsm8k/gsm8k_eval.py

Running GSM8K evaluation: 1319 questions, 5-shot
Evaluating: 100%|███████████████████████████████████████| 1319/1319 [00:12<00:00, 107.95it/s]

Results:
Accuracy: 0.578
Invalid responses: 0.004
Total latency: 12.233 s
Questions per second: 107.825

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

The pull request introduces a bug fix that uses PIECEWISE cudagraphs on Blackwell architecture if the max_model_len exceeds 131072. The code changes modify the VllmConfig class to check for this condition and override the cudagraph_mode accordingly. The changes also include adding warning messages to the logger.

[ywang96]

Amazing - thank you!

[yewentao256]

LGTM, thanks for the work!

[xinli-sw]

Added this issue to Flashinfer to track the long term fix for FULL CG support of TRTLLM backend

flashinfer-ai/flashinfer#1968

[benchislett]

FYI, it's possible that #28479 caused a regression that breaks this fix. One path forward is #28755, another is just to patch the check added by here to also consider max model len.