[Bugfix][Spec Decode] Capture the widest uniform decode batch by default (#50488)
Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Signed-off-by: Rahul Chalamala <rahul@modal.com> Co-authored-by: Janelle Cai <janelle.cai@modal.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Devin <devin@cognition.ai>
Rahul Chalamala committed
Aug 28, 2026 at 22:38 UTC
d3d79ffc1e036a3263ba8497a3596a4767caf799
6 files changed
+621
-30
tests/compile/test_config.py
+441
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import copy
4
from contextlib import nullcontext
5
+from types import SimpleNamespace
6
from unittest.mock import MagicMock, patch
7
8
import pytest
@@ -15,13 +16,16 @@ from vllm.compilation.passes.utility.fix_functionalization import (
16
from vllm.config import (
17
CompilationConfig,
18
CUDAGraphMode,
19
+ ModelConfig,
20
ParallelConfig,
21
SchedulerConfig,
22
+ SpeculativeConfig,
23
VllmConfig,
24
)
25
from vllm.config.compilation import CompilationMode, PassConfig
26
from vllm.engine.arg_utils import EngineArgs
27
from vllm.platforms import current_platform
28
+from vllm.utils.math_utils import cdiv
29
from vllm.utils.torch_utils import (
30
_is_torch_equal_or_newer,
31
is_torch_equal,
@@ -497,6 +501,443 @@ def test_cudagraph_sizes_post_init(
501
)
502
503
504
+def _mock_config_for_cudagraph_sizes(
505
+ max_num_seqs: int,
506
+ num_speculative_tokens: int,
507
+ max_num_batched_tokens: int,
508
+ compilation_config: CompilationConfig,
509
+ num_speculative_tokens_per_batch_size: list[tuple[int, int, int]] | None = None,
510
+) -> MagicMock:
511
+ """Mock VllmConfig wired up enough to run `_set_cudagraph_sizes`.
512
+
513
+ `num_speculative_tokens` and `uniform_decode_query_len` are filled in by
514
+ calling the real property functions, so the tests below cover the
515
+ derivation from `speculative_config` and not just the arithmetic
516
+ downstream of it.
517
+ """
518
+ config = MagicMock(spec=VllmConfig)
519
+ config.compilation_config = compilation_config
520
+ config.scheduler_config = SchedulerConfig.default_factory(
521
+ max_num_seqs=max_num_seqs,
522
+ max_num_batched_tokens=max_num_batched_tokens,
523
+ )
524
+ config.parallel_config = ParallelConfig()
525
+ config.model_config = MagicMock()
526
+ config.model_config.enforce_eager = False
527
+ config.performance_mode = None
528
+ config.diffusion_config = None
529
+ schedule = num_speculative_tokens_per_batch_size
530
+ config.speculative_config = (
531
+ SimpleNamespace(
532
+ num_speculative_tokens=num_speculative_tokens,
533
+ num_speculative_tokens_per_batch_size=schedule,
534
+ uses_dynamic_speculative_decoding=lambda: schedule is not None,
535
+ )
536
+ if num_speculative_tokens
537
+ else None
538
+ )
539
+ config.num_speculative_tokens = VllmConfig.num_speculative_tokens.fget(config)
540
+ config.uniform_decode_query_len = VllmConfig.uniform_decode_query_len.fget(config)
541
+ return config
542
+
543
+
544
+@pytest.mark.parametrize(
545
+ ("max_num_seqs", "num_speculative_tokens"),
546
+ [
547
+ # No speculation: the 2x headroom under the platform ceiling, unchanged.
548
+ (8, 0),
549
+ (32, 0),
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),
556
+ # Widest decode batch off the capture stride, above the ceiling.
557
+ (33, 16),
558
+ # ... and off the stride while below it, where the ceiling stands but
559
+ # the generated sizes would otherwise stop at 400.
560
+ (24, 16),
561
+ ],
562
+)
563
+def test_default_cudagraph_capture_size_covers_widest_uniform_decode(
564
+ max_num_seqs, num_speculative_tokens
565
+):
566
+ """The widest uniform decode batch has to be inside the captured range.
567
+
568
+ 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.
573
+ """
574
+ compilation_config = CompilationConfig(
575
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
576
+ )
577
+ config = _mock_config_for_cudagraph_sizes(
578
+ max_num_seqs=max_num_seqs,
579
+ num_speculative_tokens=num_speculative_tokens,
580
+ max_num_batched_tokens=32768,
581
+ compilation_config=compilation_config,
582
+ )
583
+
584
+ VllmConfig._set_cudagraph_sizes(config)
585
+
586
+ decode_query_len = 1 + num_speculative_tokens
587
+ default_max_graph_size = (
588
+ 1024 if current_platform.is_device_capability_family(100) else 512
589
+ )
590
+ token_grid_max = min(
591
+ max_num_seqs * decode_query_len * 2,
592
+ default_max_graph_size,
593
+ )
594
+ 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 (
598
+ compilation_config.max_cudagraph_capture_size
599
+ == compilation_config.cudagraph_capture_sizes[-1]
600
+ )
601
+ assert widest_uniform_decode in compilation_config.cudagraph_capture_sizes
602
+
603
+
604
+def test_default_cudagraph_capture_sizes_keep_token_grid_bounded():
605
+ """Keep the token grid bounded; only request counts may exceed the default.
606
+
607
+ This guards the measured 581 versus 100 capture-size regression at
608
+ max_num_seqs=512 with 16 draft tokens.
609
+ """
610
+ max_num_seqs = 512
611
+ decode_query_len = 17
612
+ compilation_config = CompilationConfig(
613
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
614
+ )
615
+ config = _mock_config_for_cudagraph_sizes(
616
+ max_num_seqs=max_num_seqs,
617
+ num_speculative_tokens=16,
618
+ max_num_batched_tokens=32768,
619
+ compilation_config=compilation_config,
620
+ )
621
+
622
+ with patch.object(
623
+ current_platform,
624
+ "is_device_capability_family",
625
+ return_value=False,
626
+ ):
627
+ default_max_graph_size = (
628
+ 1024 if current_platform.is_device_capability_family(100) else 512
629
+ )
630
+ VllmConfig._set_cudagraph_sizes(config)
631
+
632
+ token_grid_max = min(max_num_seqs * decode_query_len * 2, default_max_graph_size)
633
+ token_grid = [size for size in [1, 2, 4] if size <= token_grid_max]
634
+ token_grid += list(range(8, min(token_grid_max + 1, 256), 8))
635
+ token_grid += list(range(256, token_grid_max + 1, 16))
636
+
637
+ max_request_count = min(max_num_seqs, default_max_graph_size)
638
+ request_counts = [count for count in [1, 2, 4] if count <= max_request_count]
639
+ request_counts += list(range(8, min(max_request_count + 1, 256), 8))
640
+ request_counts += list(range(256, max_request_count + 1, 16))
641
+ request_counts.append(max_request_count)
642
+ expected_sizes = sorted(
643
+ set(token_grid + [count * decode_query_len for count in request_counts])
644
+ )
645
+
646
+ assert compilation_config.cudagraph_capture_sizes == expected_sizes
647
+ assert all(
648
+ size <= default_max_graph_size or size % decode_query_len == 0
649
+ for size in compilation_config.cudagraph_capture_sizes
650
+ )
651
+
652
+
653
+def test_cudagraph_capture_sizes_respect_sequence_parallelism():
654
+ """Sequence-parallel capture sizes stay divisible by tensor parallel size."""
655
+ compilation_config = CompilationConfig(
656
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
657
+ )
658
+ compilation_config.pass_config.enable_sp = True
659
+ config = _mock_config_for_cudagraph_sizes(
660
+ max_num_seqs=32,
661
+ num_speculative_tokens=16,
662
+ max_num_batched_tokens=32768,
663
+ compilation_config=compilation_config,
664
+ )
665
+ config.parallel_config = SimpleNamespace(tensor_parallel_size=2)
666
+ config.update_sizes_for_sequence_parallelism = lambda sizes: (
667
+ VllmConfig.update_sizes_for_sequence_parallelism(config, sizes)
668
+ )
669
+
670
+ with patch.object(
671
+ current_platform,
672
+ "is_device_capability_family",
673
+ return_value=False,
674
+ ):
675
+ VllmConfig._set_cudagraph_sizes(config)
676
+
677
+ assert all(size % 2 == 0 for size in compilation_config.cudagraph_capture_sizes)
678
+ assert 544 in compilation_config.cudagraph_capture_sizes
679
+
680
+
681
+@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."""
684
+ model_config = ModelConfig(model="facebook/opt-125m")
685
+ speculative_config = SpeculativeConfig(
686
+ prompt_lookup_min=1,
687
+ prompt_lookup_max=1,
688
+ num_speculative_tokens=16,
689
+ method="ngram",
690
+ )
691
+ scheduler_config = SchedulerConfig.default_factory(
692
+ max_num_seqs=32,
693
+ max_num_batched_tokens=32768,
694
+ )
695
+ compilation_config = CompilationConfig(
696
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
697
+ )
698
+
699
+ with patch.object(
700
+ current_platform,
701
+ "is_device_capability_family",
702
+ return_value=False,
703
+ ):
704
+ config = VllmConfig(
705
+ model_config=model_config,
706
+ speculative_config=speculative_config,
707
+ scheduler_config=scheduler_config,
708
+ compilation_config=compilation_config,
709
+ )
710
+
711
+ assert 544 in config.compilation_config.cudagraph_capture_sizes
712
+ assert (
713
+ config.compilation_config.max_cudagraph_capture_size
714
+ == config.compilation_config.cudagraph_capture_sizes[-1]
715
+ == 544
716
+ )
717
+
718
+
719
+@pytest.mark.parametrize("max_num_seqs", [8, 32, 256, 300, 512, 600, 1024, 2048])
720
+def test_default_cudagraph_capture_size_unchanged_without_speculation(max_num_seqs):
721
+ """Without speculation the default must reproduce the historical formula.
722
+
723
+ The platform ceiling bounds a request count, and without speculation a
724
+ request is one token, so `min(max_num_seqs, ceiling) * 1` can never lift it.
725
+ The result must match `min(max_num_seqs * 2, ceiling)` bit for bit.
726
+ """
727
+ compilation_config = CompilationConfig(
728
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
729
+ )
730
+ config = _mock_config_for_cudagraph_sizes(
731
+ max_num_seqs=max_num_seqs,
732
+ num_speculative_tokens=0,
733
+ max_num_batched_tokens=1_000_000,
734
+ compilation_config=compilation_config,
735
+ )
736
+
737
+ with patch.object(
738
+ current_platform,
739
+ "is_device_capability_family",
740
+ return_value=False,
741
+ ):
742
+ VllmConfig._set_cudagraph_sizes(config)
743
+
744
+ default_max_graph_size = (
745
+ 1024 if current_platform.is_device_capability_family(100) else 512
746
+ )
747
+ assert compilation_config.max_cudagraph_capture_size == min(
748
+ max_num_seqs * 2,
749
+ default_max_graph_size,
750
+ )
751
+
752
+
753
+def test_default_cudagraph_capture_size_covers_a_single_speculative_token():
754
+ """One speculative token is already enough to lose the widest batch.
755
+
756
+ 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.
760
+ """
761
+ compilation_config = CompilationConfig(
762
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
763
+ )
764
+ config = _mock_config_for_cudagraph_sizes(
765
+ max_num_seqs=300,
766
+ num_speculative_tokens=1,
767
+ max_num_batched_tokens=1_000_000,
768
+ compilation_config=compilation_config,
769
+ )
770
+
771
+ VllmConfig._set_cudagraph_sizes(config)
772
+
773
+ default_max_graph_size = (
774
+ 1024 if current_platform.is_device_capability_family(100) else 512
775
+ )
776
+ 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
780
+ assert (
781
+ compilation_config.max_cudagraph_capture_size
782
+ == compilation_config.cudagraph_capture_sizes[-1]
783
+ )
784
+ assert 600 in compilation_config.cudagraph_capture_sizes
785
+
786
+
787
+def test_default_cudagraph_capture_size_caps_requests_not_tokens():
788
+ """The ceiling admits at most the platform default number of requests.
789
+
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.
792
+ """
793
+ compilation_config = CompilationConfig(
794
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
795
+ )
796
+ config = _mock_config_for_cudagraph_sizes(
797
+ max_num_seqs=1024,
798
+ num_speculative_tokens=16,
799
+ max_num_batched_tokens=1_000_000,
800
+ compilation_config=compilation_config,
801
+ )
802
+
803
+ VllmConfig._set_cudagraph_sizes(config)
804
+
805
+ default_max_graph_size = (
806
+ 1024 if current_platform.is_device_capability_family(100) else 512
807
+ )
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
813
+ assert (
814
+ compilation_config.max_cudagraph_capture_size
815
+ == compilation_config.cudagraph_capture_sizes[-1]
816
+ )
817
+ assert expected_max_size == default_max_graph_size * 17
818
+
819
+
820
+def _widest_covered_request_count(capture_sizes, query_len, max_num_seqs):
821
+ """Widest request count `CudaGraphManager` can build a decode graph for.
822
+
823
+ Mirrors `_init_candidates`: a capture size is rounded up to a multiple of
824
+ the tier's query length, then dropped once the implied request count runs
825
+ past `max_num_seqs`.
826
+ """
827
+ reachable = [
828
+ cdiv(size, query_len)
829
+ for size in capture_sizes
830
+ if cdiv(size, query_len) <= max_num_seqs
831
+ ]
832
+ return max(reachable, default=0)
833
+
834
+
835
+def test_default_cudagraph_capture_sizes_cover_every_dynamic_decode_width():
836
+ """Each scheduled draft width needs sizes over the range it applies to.
837
+
838
+ Dynamic speculative decoding picks the width from the batch size, so
839
+ scaling by the widest one alone leaves the narrower tiers short: sizes
840
+ built from query length 17 round up to multiples of 3 that imply more
841
+ requests than the scheduler can run, and coverage stops at 227 of 256.
842
+ """
843
+ compilation_config = CompilationConfig(
844
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
845
+ )
846
+ config = _mock_config_for_cudagraph_sizes(
847
+ max_num_seqs=256,
848
+ num_speculative_tokens=16,
849
+ max_num_batched_tokens=32768,
850
+ 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)],
853
+ )
854
+
855
+ VllmConfig._set_cudagraph_sizes(config)
856
+
857
+ 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
861
+
862
+
863
+def test_dynamic_decode_capture_clamps_configured_width():
864
+ compilation_config = CompilationConfig(
865
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
866
+ )
867
+ config = _mock_config_for_cudagraph_sizes(
868
+ max_num_seqs=16,
869
+ num_speculative_tokens=3,
870
+ max_num_batched_tokens=32768,
871
+ compilation_config=compilation_config,
872
+ num_speculative_tokens_per_batch_size=[(1, 16, 5)],
873
+ )
874
+
875
+ VllmConfig._set_cudagraph_sizes(config)
876
+
877
+ sizes = compilation_config.cudagraph_capture_sizes
878
+ assert _widest_covered_request_count(sizes, 4, 16) == 16
879
+ assert 6 not in sizes
880
+
881
+
882
+def test_dynamic_decode_capture_covers_schedule_gap_and_tail():
883
+ compilation_config = CompilationConfig(
884
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
885
+ )
886
+ config = _mock_config_for_cudagraph_sizes(
887
+ max_num_seqs=17,
888
+ num_speculative_tokens=4,
889
+ max_num_batched_tokens=32768,
890
+ compilation_config=compilation_config,
891
+ num_speculative_tokens_per_batch_size=[(1, 2, 4), (5, 5, 1)],
892
+ )
893
+
894
+ VllmConfig._set_cudagraph_sizes(config)
895
+
896
+ sizes = compilation_config.cudagraph_capture_sizes
897
+ assert 20 in sizes
898
+ assert 34 in sizes
899
+
900
+
901
+def test_default_cudagraph_capture_size_still_clamped_by_token_budget():
902
+ """Decode coverage does not override the `max_num_batched_tokens` clamp.
903
+
904
+ A batch wider than the token budget cannot be scheduled in the first place,
905
+ so there is no decode step of that size to capture a graph for.
906
+ """
907
+ compilation_config = CompilationConfig(
908
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
909
+ )
910
+ config = _mock_config_for_cudagraph_sizes(
911
+ max_num_seqs=32,
912
+ num_speculative_tokens=16,
913
+ max_num_batched_tokens=512,
914
+ compilation_config=compilation_config,
915
+ )
916
+
917
+ VllmConfig._set_cudagraph_sizes(config)
918
+
919
+ # The token budget still clamps the final capture size.
920
+ default_max_graph_size = (
921
+ 1024 if current_platform.is_device_capability_family(100) else 512
922
+ )
923
+ decode_query_len = 17
924
+ token_grid_max = min(
925
+ 32 * decode_query_len * 2,
926
+ default_max_graph_size,
927
+ )
928
+ widest_capturable_decode = min(
929
+ 32 * decode_query_len,
930
+ 512,
931
+ )
932
+ expected_max_size = max(token_grid_max, widest_capturable_decode)
933
+ assert compilation_config.max_cudagraph_capture_size == expected_max_size
934
+ assert (
935
+ compilation_config.max_cudagraph_capture_size
936
+ == compilation_config.cudagraph_capture_sizes[-1]
937
+ )
938
+ assert 544 not in compilation_config.cudagraph_capture_sizes
939
+
940
+
941
@pytest.mark.skipif(
942
not current_platform.support_static_graph_mode(),
943
reason="Skip if static graph mode is not supported",
tests/v1/cudagraph/test_cudagraph_manager.py
+8
-10
@@ -171,7 +171,7 @@ _DECODE_QUERY_LEN = 3
171
def _create_decode_vllm_config(
172
capture_sizes: list[int],
173
num_speculative_tokens: int = 0,
174
- dynamic_spec_num_tokens: list[int] | None = None,
174
+ dynamic_spec_schedule: list[tuple[int, int, int]] | None = None,
175
) -> MagicMock:
176
compilation_config = CompilationConfig(
177
cudagraph_mode="FULL_AND_PIECEWISE",
@@ -185,16 +185,12 @@ def _create_decode_vllm_config(
185
vllm_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=8)
186
vllm_config.parallel_config = ParallelConfig()
187
vllm_config.num_speculative_tokens = num_speculative_tokens
188
- if dynamic_spec_num_tokens is None:
188
+ if dynamic_spec_schedule is None:
189
vllm_config.speculative_config = None
190
else:
191
speculative_config = MagicMock()
192
speculative_config.uses_dynamic_speculative_decoding.return_value = True
193
- # Each entry is (range_start, range_end, num_speculative_tokens); only
194
- # the third element is read by the manager.
195
- speculative_config.num_speculative_tokens_per_batch_size = [
196
- (0, 0, n) for n in dynamic_spec_num_tokens
197
- ]
193
+ speculative_config.num_speculative_tokens_per_batch_size = dynamic_spec_schedule
194
vllm_config.speculative_config = speculative_config
195
return vllm_config
196
@@ -204,7 +200,7 @@ def _make_spec_decode_manager(
200
decode_query_len: int = _DECODE_QUERY_LEN,
201
capture_sizes: list[int] | None = None,
202
num_speculative_tokens: int = 0,
207
- dynamic_spec_num_tokens: list[int] | None = None,
203
+ dynamic_spec_schedule: list[tuple[int, int, int]] | None = None,
204
) -> gpu_cudagraph_utils.CudaGraphManager:
205
monkeypatch.setattr(
206
gpu_cudagraph_utils,
@@ -220,7 +216,7 @@ def _make_spec_decode_manager(
216
vllm_config=_create_decode_vllm_config(
217
capture_sizes or [1, 2, 4, 8, 16, 24],
218
num_speculative_tokens=num_speculative_tokens,
223
- dynamic_spec_num_tokens=dynamic_spec_num_tokens,
219
+ dynamic_spec_schedule=dynamic_spec_schedule,
220
),
221
device=torch.device("cpu"),
222
cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
@@ -343,7 +339,9 @@ def test_dynamic_spec_decode_shared_token_count_stays_reachable(monkeypatch):
339
decode_query_len=4,
340
capture_sizes=[2, 4, 6, 8],
341
num_speculative_tokens=3,
346
- dynamic_spec_num_tokens=[0, 1, 2, 3], # -> decode_query_lens [1, 2, 3, 4]
342
+ # max_num_seqs is 8, and K narrows as the batch grows, so the schedule
343
+ # covers K 3, 2, 1 and 0, giving decode_query_lens [1, 2, 3, 4].
344
+ dynamic_spec_schedule=[(1, 2, 3), (3, 4, 2), (5, 6, 1), (7, 8, 0)],
345
)
346
347
full_descs = manager._capture_descs[CUDAGraphMode.FULL]
tests/v1/spec_decode/test_dynamic_sd_cug.py
+35
-4
@@ -151,6 +151,37 @@ def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch)
151
assert desc.num_active_loras == 0
152
153
154
+def test_dynamic_sd_cudagraphs_use_clamped_query_length(monkeypatch):
155
+ """Clamp configured K=5 to K=3, yielding query length 4 rather than 6."""
156
+ monkeypatch.setattr(
157
+ gpu_cudagraph_utils,
158
+ "get_pp_group",
159
+ lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
160
+ )
161
+ monkeypatch.setattr(
162
+ gpu_cudagraph_utils.current_platform,
163
+ "get_global_graph_pool",
164
+ lambda: None,
165
+ )
166
+
167
+ vllm_config = _create_vllm_config_for_dsd(
168
+ max_num_seqs=16,
169
+ max_spec_tokens=3,
170
+ num_spec_per_batch_size=[(1, 16, 5)],
171
+ )
172
+ manager = gpu_cudagraph_utils.CudaGraphManager(
173
+ vllm_config=vllm_config,
174
+ device=torch.device("cpu"),
175
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
176
+ decode_query_len=4,
177
+ )
178
+
179
+ full_query_lens = {
180
+ desc.uniform_token_count for desc in manager._capture_descs[CUDAGraphMode.FULL]
181
+ }
182
+ assert full_query_lens == {4}
183
+
184
+
185
def test_dynamic_sd_non_uniform_batch_falls_back_to_piecewise(monkeypatch):
186
"""DSD should use PIECEWISE when the batch is not a uniform decode batch.
187
@@ -331,7 +362,7 @@ def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch):
362
def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch):
363
"""DSD should only capture FULL graphs for query lengths in the schedule.
364
334
- With a partial schedule of ``(1, 32, 4)`` and ``(32, 128, 3)``, only the
365
+ With a partial schedule of ``(1, 31, 4)`` and ``(32, 128, 3)``, only the
366
scheduled speculative-token counts (K = 4 and K = 3) become decode query
367
lengths (K + 1 = 5 and 4). Uniform batches at those query lengths should get
368
FULL graphs, while every other query length (e.g. the lower values 1, 2, 3)
@@ -343,9 +374,9 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch):
374
max_decode_query_len = max_spec_tokens + 1
375
376
# (range_start, range_end, num_speculative_tokens): K = 4 and K = 3 are
346
- # scheduled, so FULL decode graphs should exist for query lengths K + 1,
347
- # i.e. exactly {5, 4}.
348
- num_spec_per_batch_size = [(1, 32, 4), (32, 128, 3)]
377
+ # scheduled in non-overlapping tiers, so FULL decode graphs should exist
378
+ # for query lengths K + 1, i.e. exactly {5, 4}.
379
+ num_spec_per_batch_size = [(1, 31, 4), (32, 128, 3)]
380
scheduled_query_lens = {entry[2] + 1 for entry in num_spec_per_batch_size}
381
382
monkeypatch.setattr(
vllm/config/compilation.py
+3
-1
@@ -703,7 +703,9 @@ 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.
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.
709
"""
710
711
dynamic_shapes_config: DynamicShapesConfig = field(
vllm/config/vllm.py
+122
-6
@@ -618,6 +618,27 @@ class VllmConfig:
618
return self.num_speculative_tokens
619
return 0
620
621
+ @property
622
+ def uniform_decode_query_len(self) -> int:
623
+ """Query length of every request in a uniform decode batch.
624
+
625
+ A decode step submits one query for the newly sampled token plus one
626
+ for each draft token, so the widest uniform decode batch the scheduler
627
+ can build is `max_num_seqs * uniform_decode_query_len` tokens. Anything
628
+ that has to cover a decode batch reads this, so the sizing rule cannot
629
+ drift between the places that apply it.
630
+
631
+ This deliberately does not derive from the KV slots a drafter reserves
632
+ past the target's query range, which is a *reservation* contract rather
633
+ than a query-length one. The two do not differ by a constant: DFlash
634
+ reserves `num_speculative_tokens + 1` slots yet still verifies `1 +
635
+ num_speculative_tokens` queries, while EAGLE reserves
636
+ `num_speculative_tokens` and verifies the same `1 + n`. Deriving one
637
+ from the other would under-size EAGLE by a full request width, which is
638
+ the failure this property exists to prevent.
639
+ """
640
+ return 1 + self.num_speculative_tokens
641
+
642
@property
643
def use_v2_model_runner(self) -> bool:
644
use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
@@ -1897,8 +1918,10 @@ class VllmConfig:
1918
1919
```python
1920
default_max_graph_size = 1024 if is_data_center_blackwell else 512
1900
- max_graph_size = min(max_num_seqs * decode_query_len * 2,
1901
- default_max_graph_size)
1921
+ decode_query_len = self.uniform_decode_query_len
1922
+ max_graph_size = min(
1923
+ max_num_seqs * decode_query_len * 2, default_max_graph_size
1924
+ )
1925
# 1, 2, 4, then multiples of 8 up to 256 and then multiples of 16
1926
# up to max_graph_size
1927
cudagraph_capture_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list(
@@ -1906,7 +1929,10 @@ class VllmConfig:
1929
1930
`max_num_batched_tokens` is also appended to the list if it fits
1931
within `max_cudagraph_capture_size`, so the max batch size is captured
1909
- even when off-stride.
1932
+ even when off-stride. Likewise, when the uniform decode query length
1933
+ exceeds one, the widest uniform decode batch (`max_num_seqs *
1934
+ decode_query_len`) is appended when it fits, since it need not land on
1935
+ an 8- or 16-token stride.
1936
1937
In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`
1938
will be the final sizes to capture cudagraph (in ascending order).
@@ -1945,17 +1971,98 @@ class VllmConfig:
1971
max_cudagraph_capture_size = (
1972
self.compilation_config.max_cudagraph_capture_size
1973
)
1974
+ # Decode sizes to cover, in tokens. Populated only when a request
1975
+ # is more than one token wide and only when the default is computed
1976
+ # here, so an explicit capture range is left exactly as configured.
1977
+ uniform_decode_sizes: list[int] = []
1978
if max_cudagraph_capture_size is None:
1979
from vllm.platforms import current_platform
1980
1951
- decode_query_len = 1 + self.num_speculative_tokens
1981
default_max_graph_size = (
1982
1024 if current_platform.is_device_capability_family(100) else 512
1983
)
1984
+ decode_query_len = self.uniform_decode_query_len
1985
+ max_num_seqs = self.scheduler_config.max_num_seqs
1986
max_cudagraph_capture_size = min(
1956
- self.scheduler_config.max_num_seqs * decode_query_len * 2,
1957
- default_max_graph_size,
1987
+ max_num_seqs * decode_query_len * 2, default_max_graph_size
1988
)
1989
+ if decode_query_len > 1:
1990
+ # A uniform decode batch is decode_query_len tokens per
1991
+ # request, so the widest one is far outside this ceiling.
1992
+ # Coverage comes from appending the decode sizes rather than
1993
+ # extending the token-strided grid. Extending that grid to
1994
+ # the widest decode size produces 581 sizes at
1995
+ # max_num_seqs=512 and 16 draft tokens, versus 100 with the
1996
+ # request-count grid.
1997
+ #
1998
+ # The grid would not buy decode coverage anyway. Dispatch
1999
+ # requires an exact multiple of decode_query_len, so a
2000
+ # token-strided entry is only usable when it happens to be
2001
+ # one; at query length 17 a captured 560 rounds to 561 and
2002
+ # is rejected. Scaling a request-count grid keeps every
2003
+ # entry usable and the count comparable to the non-
2004
+ # speculative case.
2005
+ def request_counts(max_reqs: int) -> list[int]:
2006
+ # At most the platform default number of requests,
2007
+ # mirroring the one-token-per-request decode ceiling.
2008
+ max_reqs = min(max_reqs, default_max_graph_size)
2009
+ counts = [n for n in (1, 2, 4) if n <= max_reqs]
2010
+ counts += list(range(8, min(max_reqs + 1, 256), 8))
2011
+ counts += list(range(256, max_reqs + 1, 16))
2012
+ return sorted(set(counts + [max_reqs]))
2013
+
2014
+ # Dynamic speculative decoding picks the draft width from
2015
+ # the batch size, so a decode step is only uniform within a
2016
+ # tier and each tier needs its own sizes. Scaling by the
2017
+ # widest one alone leaves the narrower tiers short: the
2018
+ # manager rounds a capture size up to a multiple of the
2019
+ # tier's query length and drops it once the implied request
2020
+ # count exceeds max_num_seqs, so at query length 3 sizes
2021
+ # built from 17 stop covering at 227 of 256 requests.
2022
+ decode_tiers = [(decode_query_len, max_num_seqs)]
2023
+ speculative_config = self.speculative_config
2024
+ if (
2025
+ speculative_config is not None
2026
+ and speculative_config.uses_dynamic_speculative_decoding()
2027
+ ):
2028
+ from vllm.v1.spec_decode.dynamic.utils import (
2029
+ build_dynamic_sd_schedule_lookup,
2030
+ )
2031
+
2032
+ schedule = (
2033
+ speculative_config.num_speculative_tokens_per_batch_size
2034
+ )
2035
+ assert schedule is not None
2036
+ # Read the tiers off the dense lookup the scheduler
2037
+ # runs on, so the clamp against num_speculative_tokens
2038
+ # and the carry-forward through gaps and the tail
2039
+ # cannot drift from it. Validation lives elsewhere; an
2040
+ # invalid schedule keeps the single-tier default.
2041
+ try:
2042
+ dense_schedule = build_dynamic_sd_schedule_lookup(
2043
+ schedule,
2044
+ vllm_max_batch_size=max_num_seqs,
2045
+ vllm_num_speculative_tokens=self.num_speculative_tokens,
2046
+ )
2047
+ except ValueError:
2048
+ pass
2049
+ else:
2050
+ # Ascending batch size, so the last write per
2051
+ # query length is the widest batch running at it.
2052
+ widest_batch: dict[int, int] = {}
2053
+ for batch_size, num_spec in enumerate(
2054
+ dense_schedule[1:], start=1
2055
+ ):
2056
+ widest_batch[num_spec + 1] = batch_size
2057
+ decode_tiers = list(widest_batch.items())
2058
+
2059
+ uniform_decode_sizes = sorted(
2060
+ {
2061
+ n * query_len
2062
+ for query_len, tier_max_reqs in decode_tiers
2063
+ for n in request_counts(tier_max_reqs)
2064
+ }
2065
+ )
2066
max_num_tokens = self.scheduler_config.max_num_batched_tokens
2067
max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size)
2068
@@ -2003,6 +2110,13 @@ class VllmConfig:
2110
and max_num_tokens not in cudagraph_capture_sizes
2111
):
2112
cudagraph_capture_sizes.append(max_num_tokens)
2113
+ # These extend past the token-strided ceiling, which counts one
2114
+ # token per request. valid_max_size below raises the final
2115
+ # max_cudagraph_capture_size to the widest of them. They remain
2116
+ # filtered by max_num_tokens.
2117
+ cudagraph_capture_sizes += [
2118
+ size for size in uniform_decode_sizes if size <= max_num_tokens
2119
+ ]
2120
# de-duplicate and sort the sizes
2121
cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes))
2122
@@ -2010,6 +2124,8 @@ class VllmConfig:
2124
self.parallel_config.tensor_parallel_size > 1
2125
and self.compilation_config.pass_config.enable_sp
2126
):
2127
+ # Sequence parallelism only captures TP-divisible sizes, so a
2128
+ # wider non-divisible decode batch cannot be captured under SP.
2129
cudagraph_capture_sizes = self.update_sizes_for_sequence_parallelism(
2130
cudagraph_capture_sizes
2131
)
vllm/v1/worker/gpu/cudagraph_utils.py
+12
-9
@@ -34,6 +34,7 @@ from vllm.sequence import IntermediateTensors
34
from vllm.utils.math_utils import round_up
35
from vllm.utils.torch_utils import current_stream
36
from vllm.v1.kv_cache_interface import KVCacheConfig
37
+from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
38
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
39
from vllm.v1.worker.gpu.block_table import BlockTables
40
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
@@ -209,21 +210,23 @@ class CudaGraphManager:
210
speculative_config
211
and speculative_config.uses_dynamic_speculative_decoding()
212
):
212
- num_spec_per_batch_size = (
213
- speculative_config.num_speculative_tokens_per_batch_size
214
- )
215
- # uses_dynamic_speculative_decoding() guarantees this is set.
216
- assert num_spec_per_batch_size is not None
213
# decode_query_len = num_speculative_steps + num_new_sampled_tokens
214
# _per_step. Recover num_new_sampled_tokens_per_step
215
# from the values the manager already has.
216
num_new_sampled_tokens_per_step = (
217
self.decode_query_len - self.vllm_config.num_speculative_tokens
218
)
223
- # Each entry is (range_start, range_end, num_speculative_tokens).
224
- decode_query_lens = [
225
- x[2] + num_new_sampled_tokens_per_step for x in num_spec_per_batch_size
226
- ]
219
+ dense_schedule = build_dynamic_sd_schedule_lookup(
220
+ speculative_config.num_speculative_tokens_per_batch_size,
221
+ vllm_max_batch_size=self.max_num_reqs,
222
+ vllm_num_speculative_tokens=self.vllm_config.num_speculative_tokens,
223
+ )
224
+ decode_query_lens = sorted(
225
+ {
226
+ num_spec + num_new_sampled_tokens_per_step
227
+ for num_spec in dense_schedule[1:]
228
+ }
229
+ )
230
else:
231
decode_query_lens = [self.decode_query_len]
232