[Perf] Raise Blackwell CUDA graph capture default to 1024 (#49390)
Lucas Wilkinson committed
Aug 7, 2026 at 15:18 UTC
021b7d985b54264393b4c339f2823d783080942a
3 files changed
+48
-8
tests/compile/test_config.py
+32
@@ -492,6 +492,38 @@ def test_cudagraph_sizes_post_init(
492
)
493
494
495
+@pytest.mark.skipif(
496
+ not current_platform.support_static_graph_mode(),
497
+ reason="Skip if static graph mode is not supported",
498
+)
499
+@pytest.mark.parametrize(
500
+ ("is_blackwell", "expected_max_size"), [(False, 512), (True, 1024)]
501
+)
502
+def test_blackwell_cudagraph_default(is_blackwell, expected_max_size):
503
+ vllm_config = VllmConfig()
504
+ vllm_config.model_config = MagicMock(enforce_eager=False)
505
+ vllm_config.scheduler_config = SchedulerConfig(
506
+ max_num_seqs=512,
507
+ max_num_batched_tokens=2048,
508
+ max_model_len=2048,
509
+ is_encoder_decoder=False,
510
+ )
511
+ vllm_config.compilation_config = CompilationConfig(
512
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
513
+ )
514
+
515
+ with patch.object(
516
+ current_platform,
517
+ "is_device_capability_family",
518
+ return_value=is_blackwell,
519
+ ):
520
+ vllm_config._set_cudagraph_sizes()
521
+
522
+ assert (
523
+ vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size
524
+ )
525
+
526
+
527
@pytest.mark.skipif(
528
not current_platform.support_static_graph_mode(),
529
reason="Skip if not cudagraph mode supported",
vllm/config/compilation.py
+4
-4
@@ -700,10 +700,10 @@ class CompilationConfig:
700
[1, 2, 4] + list(range(8, 256, 8)) + list(
701
range(256, max_cudagraph_capture_size + 1, 16))
702
703
- If not specified, max_cudagraph_capture_size is set to min(max_num_seqs*2,
704
- 512) by default. This voids OOM in tight memory scenarios with small
705
- max_num_seqs, and prevents capture of many large graphs (>512) that would
706
- greatly increase startup time with limited performance benefit.
703
+ If not specified, max_cudagraph_capture_size is capped at 512 by default,
704
+ or 1024 on data center Blackwell GPUs. This avoids OOM in tight memory
705
+ scenarios with small max_num_seqs, and limits capture of large graphs that
706
+ increase startup time and memory usage.
707
"""
708
709
dynamic_shapes_config: DynamicShapesConfig = field(
vllm/config/vllm.py
+12
-4
@@ -1808,7 +1808,9 @@ class VllmConfig:
1808
capture as:
1809
1810
```python
1811
- max_graph_size = min(max_num_seqs * 2, 512)
1811
+ default_max_graph_size = 1024 if is_data_center_blackwell else 512
1812
+ max_graph_size = min(max_num_seqs * decode_query_len * 2,
1813
+ default_max_graph_size)
1814
# 1, 2, 4, then multiples of 8 up to 256 and then multiples of 16
1815
# up to max_graph_size
1816
cudagraph_capture_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list(
@@ -1832,8 +1834,8 @@ class VllmConfig:
1834
Example:
1835
With `max_num_batched_tokens = 8192`, and typical sequences
1836
averaging ~32 tokens, most practical batch sizes fall below 256.
1835
- However, the system will still allow capture sizes up to 512 if
1836
- shape and memory permit.
1837
+ However, the system will still allow capture sizes up to the
1838
+ platform default if shape and memory permit.
1839
1840
Note:
1841
If users explicitly specify cudagraph capture sizes in the
@@ -1856,9 +1858,15 @@ class VllmConfig:
1858
self.compilation_config.max_cudagraph_capture_size
1859
)
1860
if max_cudagraph_capture_size is None:
1861
+ from vllm.platforms import current_platform
1862
+
1863
decode_query_len = 1 + self.num_speculative_tokens
1864
+ default_max_graph_size = (
1865
+ 1024 if current_platform.is_device_capability_family(100) else 512
1866
+ )
1867
max_cudagraph_capture_size = min(
1861
- self.scheduler_config.max_num_seqs * decode_query_len * 2, 512
1868
+ self.scheduler_config.max_num_seqs * decode_query_len * 2,
1869
+ default_max_graph_size,
1870
)
1871
max_num_tokens = self.scheduler_config.max_num_batched_tokens
1872
max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size)