@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b2dc864bb6

[Bugfix][Spec Decode] Keep default CUDA graph sizes memory-safe (#54418)

Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com>

Kevin H. Luu committed Aug 30, 2026 at 11:02 UTC b2dc864bb668da328aee8a8b0b72a0ad13c82252

3 files changed +73 -67

tests/compile/test_config.py

+64 -56

@@ -542,34 +542,36 @@ def _mock_config_for_cudagraph_sizes(

542

543

544 @pytest.mark.parametrize(

545 - ("max_num_seqs", "num_speculative_tokens"),

545 + ("max_num_seqs", "num_speculative_tokens", "widest_is_captured"),

546 [

547 # No speculation: the 2x headroom under the platform ceiling, unchanged.

548 - (8, 0),

549 - (32, 0),

548 + (8, 0, True),

549 + (32, 0, True),

550 # Speculating, but the widest decode batch still fits under the ceiling.

551 - (64, 7),

552 - (256, 1),

553 - # Widest decode batch above the ceiling, which must not cut below it.

554 - (32, 16),

555 - (64, 16),

551 + (64, 7, True),

552 + (256, 1, True),

553 + # Wider decode batches must not raise the memory-safety ceiling.

554 + (32, 16, False),

555 + (64, 16, False),

556 + # Exact H200 regressions: DFlash (2176) and suffix decoding (6400).

557 + (128, 16, False),

558 + (256, 24, False),

559 # Widest decode batch off the capture stride, above the ceiling.

557 - (33, 16),

560 + (33, 16, False),

561 # ... and off the stride while below it, where the ceiling stands but

562 # the generated sizes would otherwise stop at 400.

560 - (24, 16),

563 + (24, 16, True),

564 ],

565 )

563 -def test_default_cudagraph_capture_size_covers_widest_uniform_decode(

564 - max_num_seqs, num_speculative_tokens

566 +def test_default_cudagraph_capture_size_respects_platform_ceiling(

567 + max_num_seqs, num_speculative_tokens, widest_is_captured

568 ):

566 - """The widest uniform decode batch has to be inside the captured range.

569 + """Uniform decode coverage must not raise the platform default ceiling.

570

571 A decode step presents up to `max_num_seqs * (1 + num_speculative_tokens)`

569 - tokens. If that size is not captured the decode path silently falls back to

570 - eager while the resolved mode still reports full decode graphs, so neither

571 - the platform's token-denominated ceiling nor the capture stride may stop

572 - short of it.

572 + tokens. Off-stride sizes within the platform ceiling are captured exactly.

573 + Larger batches fall back to eager execution instead of expanding the

574 + default capture range and risking OOM during engine initialization.

575 """

576 compilation_config = CompilationConfig(

577 cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE

@@ -592,20 +594,21 @@ def test_default_cudagraph_capture_size_covers_widest_uniform_decode(

594 default_max_graph_size,

595 )

596 widest_uniform_decode = max_num_seqs * decode_query_len

595 - expected_max_size = max(token_grid_max, widest_uniform_decode)

596 - assert compilation_config.max_cudagraph_capture_size == expected_max_size

597 + assert compilation_config.max_cudagraph_capture_size == token_grid_max

598 assert (

599 compilation_config.max_cudagraph_capture_size

600 == compilation_config.cudagraph_capture_sizes[-1]

601 )

601 - assert widest_uniform_decode in compilation_config.cudagraph_capture_sizes

602 + assert (

603 + widest_uniform_decode in compilation_config.cudagraph_capture_sizes

604 + ) is widest_is_captured

605

606

604 -def test_default_cudagraph_capture_sizes_keep_token_grid_bounded():

605 - """Keep the token grid bounded; only request counts may exceed the default.

607 +def test_default_cudagraph_capture_sizes_keep_all_sizes_bounded():

608 + """Keep both token and uniform-decode grids under the platform default.

609

607 - This guards the measured 581 versus 100 capture-size regression at

608 - max_num_seqs=512 with 16 draft tokens.

610 + This guards both the measured 581 versus 100 capture-count regression and

611 + the graph-memory regression from capturing shapes up to 8704 tokens.

612 """

613 max_num_seqs = 512

614 decode_query_len = 17

@@ -640,12 +643,19 @@ def test_default_cudagraph_capture_sizes_keep_token_grid_bounded():

643 request_counts += list(range(256, max_request_count + 1, 16))

644 request_counts.append(max_request_count)

645 expected_sizes = sorted(

643 - set(token_grid + [count * decode_query_len for count in request_counts])

646 + set(

647 + token_grid

648 + + [

649 + count * decode_query_len

650 + for count in request_counts

651 + if count * decode_query_len <= default_max_graph_size

652 + ]

653 + )

654 )

655

656 assert compilation_config.cudagraph_capture_sizes == expected_sizes

657 assert all(

648 - size <= default_max_graph_size or size % decode_query_len == 0

658 + size <= default_max_graph_size

659 for size in compilation_config.cudagraph_capture_sizes

660 )

661

@@ -675,12 +685,13 @@ def test_cudagraph_capture_sizes_respect_sequence_parallelism():

685 VllmConfig._set_cudagraph_sizes(config)

686

687 assert all(size % 2 == 0 for size in compilation_config.cudagraph_capture_sizes)

678 - assert 544 in compilation_config.cudagraph_capture_sizes

688 + assert 544 not in compilation_config.cudagraph_capture_sizes

689 + assert compilation_config.max_cudagraph_capture_size == 512

690

691

692 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")

682 -def test_real_vllm_config_captures_widest_ngram_decode_batch():

683 - """Real config post-init preserves the widest static ngram decode batch."""

693 +def test_real_vllm_config_caps_widest_ngram_decode_batch():

694 + """Real config post-init preserves the default capture ceiling."""

695 model_config = ModelConfig(model="facebook/opt-125m")

696 speculative_config = SpeculativeConfig(

697 prompt_lookup_min=1,

@@ -708,11 +719,11 @@ def test_real_vllm_config_captures_widest_ngram_decode_batch():

719 compilation_config=compilation_config,

720 )

721

711 - assert 544 in config.compilation_config.cudagraph_capture_sizes

722 + assert 544 not in config.compilation_config.cudagraph_capture_sizes

723 assert (

724 config.compilation_config.max_cudagraph_capture_size

725 == config.compilation_config.cudagraph_capture_sizes[-1]

715 - == 544

726 + == 512

727 )

728

729

@@ -750,13 +761,12 @@ def test_default_cudagraph_capture_size_unchanged_without_speculation(max_num_se

761 )

762

763

753 -def test_default_cudagraph_capture_size_covers_a_single_speculative_token():

754 - """One speculative token is already enough to lose the widest batch.

764 +def test_single_speculative_token_does_not_raise_default_capture_size():

765 + """One speculative token must not raise the default capture ceiling.

766

767 MTP at depth 1 gives a query length of 2, so 300 requests is a 600-token

757 - decode batch against the platform's token ceiling. Gating the unit conversion on a

758 - deeper speculative width would leave this common configuration dispatching

759 - its largest decode steps eager.

768 + decode batch against the platform's token ceiling. The widest batch falls

769 + back to eager execution rather than expanding the capture range.

770 """

771 compilation_config = CompilationConfig(

772 cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE

@@ -774,21 +784,19 @@ def test_default_cudagraph_capture_size_covers_a_single_speculative_token():

784 1024 if current_platform.is_device_capability_family(100) else 512

785 )

786 token_grid_max = min(300 * 2 * 2, default_max_graph_size)

777 - widest_uniform_decode = 300 * 2

778 - expected_max_size = max(token_grid_max, widest_uniform_decode)

779 - assert compilation_config.max_cudagraph_capture_size == expected_max_size

787 + assert compilation_config.max_cudagraph_capture_size == token_grid_max

788 assert (

789 compilation_config.max_cudagraph_capture_size

790 == compilation_config.cudagraph_capture_sizes[-1]

791 )

784 - assert 600 in compilation_config.cudagraph_capture_sizes

792 + assert 600 not in compilation_config.cudagraph_capture_sizes

793

794

787 -def test_default_cudagraph_capture_size_caps_requests_not_tokens():

788 - """The ceiling admits at most the platform default number of requests.

795 +def test_default_cudagraph_capture_size_caps_tokens():

796 + """The ceiling remains a platform-bounded token count.

797

790 - `max_num_seqs` far above the platform default must not push the capture

791 - range up without bound just because each request is now several tokens wide.

798 + `max_num_seqs` and speculative width must not push the default capture range

799 + above the platform's memory-safety guard.

800 """

801 compilation_config = CompilationConfig(

802 cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE

@@ -805,16 +813,15 @@ def test_default_cudagraph_capture_size_caps_requests_not_tokens():

813 default_max_graph_size = (

814 1024 if current_platform.is_device_capability_family(100) else 512

815 )

808 - decode_query_len = 17

809 - token_grid_max = min(1024 * decode_query_len * 2, default_max_graph_size)

810 - widest_covered_decode = min(1024, default_max_graph_size) * decode_query_len

811 - expected_max_size = max(token_grid_max, widest_covered_decode)

812 - assert compilation_config.max_cudagraph_capture_size == expected_max_size

816 + assert compilation_config.max_cudagraph_capture_size == default_max_graph_size

817 assert (

818 compilation_config.max_cudagraph_capture_size

819 == compilation_config.cudagraph_capture_sizes[-1]

820 )

817 - assert expected_max_size == default_max_graph_size * 17

821 + assert all(

822 + size <= default_max_graph_size

823 + for size in compilation_config.cudagraph_capture_sizes

824 + )

825

826

827 def _widest_covered_request_count(capture_sizes, query_len, max_num_seqs):

@@ -844,20 +851,21 @@ def test_default_cudagraph_capture_sizes_cover_every_dynamic_decode_width():

851 cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE

852 )

853 config = _mock_config_for_cudagraph_sizes(

847 - max_num_seqs=256,

854 + max_num_seqs=128,

855 num_speculative_tokens=16,

856 max_num_batched_tokens=32768,

857 compilation_config=compilation_config,

851 - # 16 draft tokens up to a batch of 16, then 2 out to 256.

852 - num_speculative_tokens_per_batch_size=[(1, 16, 16), (17, 256, 2)],

858 + # 16 draft tokens up to a batch of 16, then 2 out to 128. Both

859 + # tier maxima fit under the 512-token default capture ceiling.

860 + num_speculative_tokens_per_batch_size=[(1, 16, 16), (17, 128, 2)],

861 )

862

863 VllmConfig._set_cudagraph_sizes(config)

864

865 sizes = compilation_config.cudagraph_capture_sizes

858 - # The wide tier only ever runs to a batch of 16, the narrow one to 256.

859 - assert _widest_covered_request_count(sizes, 17, 256) >= 16

860 - assert _widest_covered_request_count(sizes, 3, 256) == 256

866 + # The wide tier only ever runs to a batch of 16, the narrow one to 128.

867 + assert _widest_covered_request_count(sizes, 17, 128) >= 16

868 + assert _widest_covered_request_count(sizes, 3, 128) == 128

869

870

871 def test_dynamic_decode_capture_clamps_configured_width():

vllm/config/compilation.py

+2 -3

@@ -703,9 +703,8 @@ class CompilationConfig:

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. When the uniform decode query length

707 - exceeds one, uniform decode sizes are appended and can raise the final value

708 - above this default.

706 + increase startup time and memory usage. Uniform decode sizes are appended

707 + only within this default ceiling.

708 """

709

710 dynamic_shapes_config: DynamicShapesConfig = field(

vllm/config/vllm.py

+7 -8

@@ -1953,10 +1953,9 @@ class VllmConfig:

1953

1954 `max_num_batched_tokens` is also appended to the list if it fits

1955 within `max_cudagraph_capture_size`, so the max batch size is captured

1956 - even when off-stride. Likewise, when the uniform decode query length

1957 - exceeds one, the widest uniform decode batch (`max_num_seqs *

1958 - decode_query_len`) is appended when it fits, since it need not land on

1959 - an 8- or 16-token stride.

1956 + even when off-stride. Uniform decode sizes are appended when they fit

1957 + within the platform's default capture ceiling, since they need not land

1958 + on an 8- or 16-token stride.

1959

1960 In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`

1961 will be the final sizes to capture cudagraph (in ascending order).

@@ -2085,6 +2084,7 @@ class VllmConfig:

2084 n * query_len

2085 for query_len, tier_max_reqs in decode_tiers

2086 for n in request_counts(tier_max_reqs)

2087 + if n * query_len <= max_cudagraph_capture_size

2088 }

2089 )

2090 max_num_tokens = self.scheduler_config.max_num_batched_tokens

@@ -2134,10 +2134,9 @@ class VllmConfig:

2134 and max_num_tokens not in cudagraph_capture_sizes

2135 ):

2136 cudagraph_capture_sizes.append(max_num_tokens)

2137 - # These extend past the token-strided ceiling, which counts one

2138 - # token per request. valid_max_size below raises the final

2139 - # max_cudagraph_capture_size to the widest of them. They remain

2140 - # filtered by max_num_tokens.

2137 + # Preserve the platform's default capture ceiling. Larger

2138 + # uniform decode batches fall back to eager execution unless

2139 + # users explicitly configure wider capture sizes.

2140 cudagraph_capture_sizes += [

2141 size for size in uniform_decode_sizes if size <= max_num_tokens

2142 ]