A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 15:15 UTC
[Fix] Misc Fixes in ViT CUDA Graph (#38040)
Signed-off-by: Baorun Mu <bmu@nvidia.com>
Baorun (Lauren) Mu committed
May 14, 2026 at 11:49 UTC
a7737cb4f339aa3fe9b232b8986750102afcbf65
4 files changed
+242
-21
tests/v1/cudagraph/test_encoder_cudagraph.py
+172
@@ -32,6 +32,68 @@ from vllm.v1.worker.encoder_cudagraph_defs import (
32
# ---------------------------------------------------------------------------
33
34
35
+class _MockCompilationConfig:
36
+ """Minimal mock for VllmConfig.compilation_config."""
37
+
38
+ def __init__(
39
+ self,
40
+ token_budgets: list[int] | None = None,
41
+ max_mm_items: int = 0,
42
+ ):
43
+ self.encoder_cudagraph_token_budgets = token_budgets or []
44
+ self.encoder_cudagraph_max_vision_items_per_batch = max_mm_items
45
+ self.encoder_cudagraph_max_frames_per_batch = None
46
+
47
+
48
+class _MockMultimodalConfig:
49
+ mm_encoder_tp_mode = "replicate"
50
+
51
+ def get_limit_per_prompt(self, modality: str) -> int:
52
+ # Image-only mocks — return 0 for "video" to short-circuit the
53
+ # max_frames_per_batch branch, so tests don't need a video-frame mock.
54
+ return 0
55
+
56
+
57
+class _MockModelConfig:
58
+ multimodal_config = _MockMultimodalConfig()
59
+
60
+
61
+class _MockParallelConfig:
62
+ tensor_parallel_size = 1
63
+
64
+
65
+class _MockVllmConfig:
66
+ """Minimal mock for VllmConfig used in __init__ tests."""
67
+
68
+ def __init__(
69
+ self,
70
+ token_budgets: list[int] | None = None,
71
+ max_mm_items: int = 0,
72
+ ):
73
+ self.compilation_config = _MockCompilationConfig(token_budgets, max_mm_items)
74
+ self.model_config = _MockModelConfig()
75
+ self.parallel_config = _MockParallelConfig()
76
+
77
+
78
+class _MockModel:
79
+ """Minimal mock implementing SupportsEncoderCudaGraph for __init__."""
80
+
81
+ def __init__(self, min_budget: int = 4, max_budget: int = 128):
82
+ self._min_budget = min_budget
83
+ self._max_budget = max_budget
84
+
85
+ def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
86
+ return EncoderCudaGraphConfig(
87
+ modalities=["image"],
88
+ input_key_by_modality={"image": "pixel_values"},
89
+ buffer_keys=["dummy_buf"],
90
+ out_hidden_size=32,
91
+ )
92
+
93
+ def get_encoder_cudagraph_budget_range(self, vllm_config):
94
+ return (self._min_budget, self._max_budget)
95
+
96
+
97
def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
98
"""Create a minimal EncoderCudaGraphManager with only token_budgets set.
99
@@ -760,3 +822,113 @@ class TestEncoderCudaGraphVideoReplay:
822
assert len(vid_result) == 1
823
assert img_result[0].shape == (4, _HIDDEN)
824
assert vid_result[0].shape == (8, _HIDDEN)
825
+
826
+
827
+# ---------------------------------------------------------------------------
828
+# __init__ invariant validation tests (no GPU required)
829
+# ---------------------------------------------------------------------------
830
+
831
+
832
+class TestInitInvariantValidation:
833
+ """Ensure max_batch_size <= min(token_budgets) for all config paths."""
834
+
835
+ def _make_mgr(
836
+ self,
837
+ token_budgets=None,
838
+ max_mm_items=0,
839
+ min_budget=4,
840
+ max_budget=128,
841
+ ):
842
+ vllm_config = _MockVllmConfig(token_budgets, max_mm_items)
843
+ model = _MockModel(min_budget, max_budget)
844
+ return EncoderCudaGraphManager(
845
+ vllm_config=vllm_config,
846
+ device=torch.device("cpu"),
847
+ dtype=torch.float32,
848
+ model=model,
849
+ )
850
+
851
+ # --- Finding 1: fully auto-inferred ---
852
+
853
+ def test_auto_inferred_invariant_holds(self):
854
+ mgr = self._make_mgr(min_budget=64, max_budget=16384)
855
+ assert mgr.max_batch_size <= min(mgr.token_budgets)
856
+
857
+ def test_auto_inferred_small_range(self):
858
+ mgr = self._make_mgr(min_budget=4, max_budget=128)
859
+ assert mgr.max_batch_size <= min(mgr.token_budgets)
860
+
861
+ # --- Finding 2: fully user-specified, bad combo ---
862
+
863
+ def test_user_specified_bad_combo_raises(self):
864
+ with pytest.raises(ValueError, match="must be <= smallest token budget"):
865
+ self._make_mgr(token_budgets=[64], max_mm_items=256)
866
+
867
+ def test_user_specified_valid_combo(self):
868
+ mgr = self._make_mgr(token_budgets=[64, 128], max_mm_items=32)
869
+ assert mgr.max_batch_size == 32
870
+ assert mgr.token_budgets == [64, 128]
871
+
872
+ def test_user_specified_exact_boundary(self):
873
+ # max_mm_items == min(budgets) is OK (per_image_output = 1)
874
+ mgr = self._make_mgr(token_budgets=[64, 128], max_mm_items=64)
875
+ assert mgr.max_batch_size == 64
876
+
877
+ # --- Finding 3: user provides only max_mm_items ---
878
+
879
+ def test_user_max_mm_items_only_adjusts_budgets(self):
880
+ # model min_budget=64, user max_mm_items=128 → budgets start at 128
881
+ mgr = self._make_mgr(max_mm_items=128, min_budget=64, max_budget=16384)
882
+ assert mgr.max_batch_size == 128
883
+ assert min(mgr.token_budgets) >= 128
884
+
885
+ def test_user_max_mm_items_smaller_than_min_budget(self):
886
+ # max_mm_items=2, model min=4 → budgets start at 4 (>= 2), OK
887
+ mgr = self._make_mgr(max_mm_items=2, min_budget=4, max_budget=128)
888
+ assert mgr.max_batch_size == 2
889
+ assert min(mgr.token_budgets) >= 2
890
+
891
+ # --- Finding 4: user provides only budgets ---
892
+
893
+ def test_user_budgets_only_caps_max_batch_size(self):
894
+ # user budgets start at 32, model min_budget=64
895
+ # without fix: max_batch_size = min(128//64, 64) = 2 → OK
896
+ # but if user budgets=[16, 64]:
897
+ # without fix: max_batch_size = min(128//4, 4) = 4 > 16? No.
898
+ # Let's use a case that triggers it:
899
+ # model min=64, max=16384 → max_budget//min_budget = 256
900
+ # user budgets=[32, 64] → min = 32
901
+ # without fix: max_batch_size = min(256, 64) = 64 > 32 → BUG
902
+ # with fix: max_batch_size = min(256, 32) = 32 → OK
903
+ mgr = self._make_mgr(token_budgets=[32, 64], min_budget=64, max_budget=16384)
904
+ assert mgr.max_batch_size <= min(mgr.token_budgets)
905
+ assert mgr.max_batch_size == 32
906
+
907
+ # --- Finding 5/6: bad model budget range ---
908
+
909
+ def test_zero_min_budget_raises(self):
910
+ with pytest.raises(ValueError, match="Both must be positive"):
911
+ self._make_mgr(min_budget=0, max_budget=128)
912
+
913
+ def test_negative_max_budget_raises(self):
914
+ with pytest.raises(ValueError, match="Both must be positive"):
915
+ self._make_mgr(min_budget=4, max_budget=-1)
916
+
917
+ def test_min_greater_than_max_raises(self):
918
+ with pytest.raises(ValueError, match="min_budget=200 > max_budget=100"):
919
+ self._make_mgr(min_budget=200, max_budget=100)
920
+
921
+ # --- Finding 7: user-provided budgets with non-positive values ---
922
+
923
+ def test_user_budgets_zero_raises(self):
924
+ """Non-positive budgets should be caught at config validation."""
925
+ from vllm.config.compilation import CompilationConfig
926
+
927
+ with pytest.raises(ValueError, match="must be positive"):
928
+ CompilationConfig(encoder_cudagraph_token_budgets=[0, 128])
929
+
930
+ def test_user_budgets_negative_raises(self):
931
+ from vllm.config.compilation import CompilationConfig
932
+
933
+ with pytest.raises(ValueError, match="must be positive"):
934
+ CompilationConfig(encoder_cudagraph_token_budgets=[-1, 64])
vllm/config/compilation.py
+8
@@ -1014,6 +1014,14 @@ class CompilationConfig:
1014
"non-negative (None = auto-infer)"
1015
)
1016
1017
+ if self.encoder_cudagraph_token_budgets and any(
1018
+ b <= 0 for b in self.encoder_cudagraph_token_budgets
1019
+ ):
1020
+ raise ValueError(
1021
+ f"All encoder_cudagraph_token_budgets must be positive, "
1022
+ f"got {self.encoder_cudagraph_token_budgets}"
1023
+ )
1024
+
1025
if self.backend == "":
1026
self.backend = current_platform.get_compile_backend()
1027
vllm/model_executor/models/qwen3_vl.py
+10
-9
@@ -1768,14 +1768,12 @@ class Qwen3VLForConditionalGeneration(
1768
EncoderCudaGraphConfig,
1769
)
1770
1771
- modalities = ["image"]
1772
- # NOTE: When EVS (Efficient Video Sampling) pruning is enabled, the number
1773
- # of tokens becomes data-dependent (i.e., the retained tokens are
1774
- # dynamically selected based on inter-frame differences) and therefore
1775
- # cannot be captured by CUDA Graphs. As a result, video CUDA Graphs are
1776
- # only enabled when EVS is disabled.
1777
- if not self.is_multimodal_pruning_enabled:
1778
- modalities.append("video")
1771
+ # When EVS pruning is enabled, embed_multimodal post-processes both
1772
+ # image and video embeddings (mrope positions are appended for image,
1773
+ # prune+append for video). The encoder CUDA graph path bypasses that
1774
+ # post-process, producing inconsistent embedding formats vs eager. So
1775
+ # disable CUDA graph for all modalities when pruning is on.
1776
+ modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"]
1777
1778
return EncoderCudaGraphConfig(
1779
modalities=modalities,
@@ -1923,7 +1921,10 @@ class Qwen3VLForConditionalGeneration(
1921
)
1922
1923
spatial_merge_size = self.visual.spatial_merge_size
1926
- per_mm_item_output = token_budget // max_batch_size
1924
+ # Ceil so the buffer fits the worst case of one item using the full
1925
+ # budget. Floor under-allocates when budget is not a multiple of
1926
+ # max_batch_size.
1927
+ per_mm_item_output = (token_budget + max_batch_size - 1) // max_batch_size
1928
1929
frames_per_item = max_frames_per_batch // max_batch_size
1930
if frames_per_item > 1:
vllm/v1/worker/encoder_cudagraph.py
+52
-12
@@ -72,25 +72,65 @@ class EncoderCudaGraphManager:
72
73
multimodal_config = vllm_config.model_config.multimodal_config
74
75
+ # Invariant: max_batch_size <= min_token_budget.
76
+ # This ensures per_image_output = budget // max_batch_size >= 1
77
+ # for every captured budget, preventing reshape crashes on empty
78
+ # tensors during CUDA graph capture. Validated/enforced below for
79
+ # each configuration path.
80
if user_budgets and user_max_vision_items > 0:
76
- # Fully user-specified
81
+ # Fully user-specified: validate the invariant.
82
self.token_budgets = sorted(user_budgets)
83
self.max_batch_size = user_max_vision_items
84
+ min_tok = min(self.token_budgets)
85
+ if self.max_batch_size > min_tok:
86
+ raise ValueError(
87
+ f"encoder_cudagraph_max_vision_items_per_batch "
88
+ f"({self.max_batch_size}) must be <= smallest token "
89
+ f"budget ({min_tok}). With budgets="
90
+ f"{self.token_budgets}, per_image_output = "
91
+ f"{min_tok} // {self.max_batch_size} = "
92
+ f"{min_tok // self.max_batch_size}, which would cause "
93
+ f"a capture failure. Either increase the smallest "
94
+ f"budget or decrease max_vision_items_per_batch."
95
+ )
96
else:
80
- # Auto-infer missing values from model
97
+ # Auto-infer missing values from model.
98
min_budget, max_budget = model.get_encoder_cudagraph_budget_range(
99
vllm_config
100
)
84
- self.token_budgets = (
85
- sorted(user_budgets)
86
- if user_budgets
87
- else self._generate_budgets(min_budget, max_budget)
88
- )
89
- self.max_batch_size = (
90
- user_max_vision_items
91
- if user_max_vision_items > 0
92
- else max_budget // min_budget
93
- )
101
+ if min_budget <= 0 or max_budget <= 0:
102
+ raise ValueError(
103
+ f"Invalid encoder cudagraph budget range: "
104
+ f"min_budget={min_budget}, max_budget={max_budget}. "
105
+ f"Both must be positive."
106
+ )
107
+ if min_budget > max_budget:
108
+ raise ValueError(
109
+ f"Invalid encoder cudagraph budget range: "
110
+ f"min_budget={min_budget} > max_budget={max_budget}."
111
+ )
112
+
113
+ if user_max_vision_items > 0:
114
+ # User provided max_vision_items only; adjust auto-inferred
115
+ # budgets so min(budgets) >= max_batch_size.
116
+ self.max_batch_size = user_max_vision_items
117
+ effective_min = max(min_budget, user_max_vision_items)
118
+ self.token_budgets = self._generate_budgets(effective_min, max_budget)
119
+ elif user_budgets:
120
+ # User provided budgets only; cap auto-inferred
121
+ # max_batch_size to min(user_budgets).
122
+ self.token_budgets = sorted(user_budgets)
123
+ self.max_batch_size = min(
124
+ max_budget // min_budget,
125
+ min(self.token_budgets),
126
+ )
127
+ else:
128
+ # Fully auto-inferred.
129
+ self.token_budgets = self._generate_budgets(min_budget, max_budget)
130
+ self.max_batch_size = min(
131
+ max_budget // min_budget,
132
+ min(self.token_budgets),
133
+ )
134
135
assert multimodal_config is not None
136
if multimodal_config.get_limit_per_prompt("video") == 0: