[V1][CUDA] Full cudagraph support for FlashInfer

[fhl2000]

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.

Purpose

This PR is split from origin #20059 to support full cudagraph for FlashInfer (pure decode only), which runs pure decode batches at full cudagraph, and falls back to no cudagraph at mix prefill-decode batches. Hope to land this first before #20059.

Details include:

• Using the persistent buffer trick.
• Create many decode_warpers, one for a cudagraph batch size, as this is required by the FlashInfer API.

This PR also fixes a potential bug originally from #18581, where an assertion error will be raised when capturing if max_capture_size is greater than max_num_reqs. To resolve this, a new enum type AttentionCGSupport (adapted from #20059) is introduced to distinguish how the backend supports cudagraph, so that we can overwrite the cudagraph_batch_sizes to be not greater than max_num_seqs.
NOTE: Currently, manually setting max capture size seems impossible after the introduction of Pydantic, which blocks config --cuda_graph_sizes to be int type

Limitation:

1. FlashInfer backend currently does not support spec-decode when enabling full cudagraph (pure-decode only)
2. trtllm_batch_decode_with_kv_cache of Flashinfer is not yet considered supported in this PR.
3. Some identified cpu overheads exist but that's out of scope for this PR. That should be resolved after [Attention] Optimize FlashInfer MetadataBuilder Build call #21137 and the comment

Test Plan

lm_eval, benchmark_serving

Test Result

piecewise cudagraph (main branch)

vllm ({'pretrained': '/root/models/Qwen2.5-7B-Instruct-GPTQ-Int4', 'gpu_memory_utilization': 0.9}), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 1

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8105	±	0.0108
		strict-match	5	exact_match	↑	0.7635	±	0.0117

VLLM_ATTENTION_BACKEND=FLASHINFER python -m vllm.entrypoints.openai.api_server --model /root/models/Qwen2.5-7B-Instruct-GPTQ-Int4 --gpu-memory-utilization 0.9
python /root/vllm/benchmarks/benchmark_serving.py --model /root/models/Qwen2.5-7B-Instruct-GPTQ-Int4 --dataset-name sharegpt --dataset-path /root/ShareGPT_V3_unfiltered_cleaned_split.json --num-prompts 100 --request-rate 50

============ Serving Benchmark Result ============
Successful requests: 100
Benchmark duration (s): 11.44
Total input tokens: 23260
Total generated tokens: 21668
Request throughput (req/s): 8.74
Output token throughput (tok/s): 1893.66
Total Token throughput (tok/s): 3926.45
---------------Time to First Token----------------
Mean TTFT (ms): 290.02
Median TTFT (ms): 286.17
P99 TTFT (ms): 558.93
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 35.14
Median TPOT (ms): 21.34
P99 TPOT (ms): 154.47
---------------Inter-token Latency----------------
Mean ITL (ms): 17.92
Median ITL (ms): 14.48
P99 ITL (ms): 171.14
==================================================

full cudagraph (this PR)

vllm ({'pretrained': '/root/models/Qwen2.5-7B-Instruct-GPTQ-Int4', 'gpu_memory_utilization': 0.9, 'compilation_config': {'full_cuda_graph': True}}), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 1

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8105	±	0.0108
		strict-match	5	exact_match	↑	0.7635	±	0.0117

VLLM_ATTENTION_BACKEND=FLASHINFER python -m vllm.entrypoints.openai.api_server --model /root/models/Qwen2.5-7B-Instruct-GPTQ-Int4 --gpu-memory-utilization 0.9 --compilation-config '{"full_cuda_graph": true}'
python /root/vllm/benchmarks/benchmark_serving.py --model /root/models/Qwen2.5-7B-Instruct-GPTQ-Int4 --dataset-name sharegpt --dataset-path /root/ShareGPT_V3_unfiltered_cleaned_split.json --num-prompts 100 --request-rate 50
============ Serving Benchmark Result ============
Successful requests: 100
Benchmark duration (s): 11.31
Total input tokens: 23260
Total generated tokens: 21660
Request throughput (req/s): 8.84
Output token throughput (tok/s): 1915.59
Total Token throughput (tok/s): 3972.68
---------------Time to First Token----------------
Mean TTFT (ms): 299.74
Median TTFT (ms): 286.63
P99 TTFT (ms): 579.64
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 35.81
Median TPOT (ms): 21.53
P99 TPOT (ms): 157.76
---------------Inter-token Latency----------------
Mean ITL (ms): 17.95
Median ITL (ms): 13.83
P99 ITL (ms): 172.55
==================================================

While the median ITL is decreased by 4.5%, the overall performance increase is small anyway. I believe this can be further optimized after #20059 is landed and more cpu overheads mentioned above are reduced.

(Optional) Documentation Update

[gemini-code-assist]

Code Review

This pull request introduces full CUDA graph support for the FlashInfer attention backend, which is a great performance enhancement for pure decode scenarios. The use of a new AttentionCGSupport enum to manage different levels of CUDA graph support across backends is a solid design choice that improves code clarity and maintainability.

The PR also includes an important bug fix to prevent graph capture for unsupported batch sizes, which is crucial for stability. I've identified one critical issue where a data structure for padding batch sizes is not updated after filtering the capture sizes, which could lead to runtime errors. I've provided a suggestion to fix this. Overall, this is a valuable contribution.

[gemini-code-assist]

This change correctly filters self.cudagraph_batch_sizes to prevent capturing graphs for sizes larger than max_num_seqs for PURE_DECODE_ONLY backends. However, the pad_for_cudagraph method, which is used at runtime to determine the padded graph size, relies on a mapping (bs_to_padded_graph_size) that was initialized with the original, unfiltered cudagraph_batch_sizes.

This discrepancy can lead to a KeyError at runtime. For example, if a batch with num_decodes is processed, pad_for_cudagraph might return a padded size that was filtered out and for which no CUDA graph was captured. This will cause a lookup failure in _decode_wrappers_cudagraph.

To fix this, you should re-initialize the padding map after filtering self.cudagraph_batch_sizes.

Suggested change

	self.cudagraph_batch_sizes = [
	size for size in self.cudagraph_batch_sizes
	if size <= max_num_seqs]
	self.cudagraph_batch_sizes = [
	size for size in self.cudagraph_batch_sizes
	if size <= max_num_seqs
	]
	self.vllm_config.compilation_config.init_with_cudagraph_sizes(
	self.cudagraph_batch_sizes)

[fhl2000]

Good catch for pad_for_cudagraph , though I think it would not affect the final correctness.

[tdoublep]

I find that the following hangs:

 VLLM_ATTENTION_BACKEND=FLASHINFER  vllm serve models/Llama-3.1-8B-Instruct --no-enable-prefix-caching --compilation-config='{"full_cuda_graph": true}' --max-num-seqs 2

at this point:

Capturing CUDA graph shapes: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  2.42it/s]
INFO 07-22 10:32:54 [gpu_model_runner.py:2404] Graph capturing finished in 1 secs, took 0.42 GiB

It just hangs here - could this be related?

[tdoublep]

if I remove the --max-num-seqs then it works fine, so I think it is indeed related.

[fhl2000]

I find overriding gpu model runner's cudagraph_batch_sizes would be enough. And vllm_config.compilation_config.init_with_cudagraph_sizes method does not actually override cudagraph_batch_sizes of compilation config after its first call.

[tdoublep]

Hmm unfortunately it doesn't happen on main + FlashAttention. The hang is 100% reproducible using the code from this PR.

[fhl2000]

I will try my best to figure it out.

[fhl2000]

I have tested --max-num-seqs being one of [2,4,8,16,24,32, 40] leads to hangs, while [1,48,56,...] work normally. The stuck occurs in a final dummy_run after all capturing in gpu_worker.py around lines 285~292, which runs into cudagraph replay (nums_tokens = max_num_seqs) without creating attn_metadata. I guess something weird happened in Flashinfer.

[fhl2000]

@tdoublep Take the new fix! It should be fine now. Could you please also test if it works for you?

[tdoublep]

Yes, seems to work now. Thanks!

[fhl2000]

cc @ProExpertProg @LucasWilkinson @SageMoore

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[tdoublep]

This PR will be super-helpful for enabling full (decode-only) CUDA graphs for hybrid (mamba/attention) models in V1, where right now we need to use FlashInfer. I am testing these changes with my branch now.

[tdoublep]

A few questions but otherwise looks good. I'm keen to see this PR merged because it should push us close to the point when we can deprecate V0 for hybrid models (which require FlashInfer currently).

[tdoublep]

Could you explain a bit more why this change is needed?

[fhl2000]

Sure! The current dummy_run is not aware of whether it is in the status of warming up or needs to trigger cudagraph. It doesn't set up skip_cuda_graph to forward_ctx, so it blindly expects the model executions to go through cudagraph. However, even after the full cudagraphs are captured (buffer address is solid), if attn_metadata is not built correctly in dummy_run (nothing has been going through Flashinfer's plan function), it may access incorrect values in these buffers and may potentially fall into an infinite loop. I think always activating this part is also not bad for piecewise cudagraph.

[fhl2000]

Hey, just found that always activating it here would cause CI failures for FlexAttention. So I have to enable this only when full_cuda_graph and not enforce eager.

[tdoublep]

need to use piecewise cudagraph or no cudagraph for mixed prefill-decode batches

In the mixed prefill-decode case: when will piecewise cudagraph be used and when will no cuda graph be used?

[fhl2000]

I am sorry, this comment is directly from #20059, I didn't realize there are no piecewise cudagraph exists in this PR when enabling full_cuda_graph. That is only allowed after that PR, which introduced a new cudagraph_mode for config that enables decouple cuda graph logic from vllm compilation. In that PR, when cudagraph_mode is FULL, it will run full cudagraph for pure decode if attention support is PURE_DECODE_ONLY, and fall back to piecewise cudagraph for other situations that are incompatible with full cudagraph when vllm compilation is on. However, if vllm compilation is disabled, then just turn to no cudagraph.

[fhl2000]

It's fixed now