@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 07516fda67

[MRV2][SD] Make Dynamic SD comatible with Full Cuda Graphs (#45953)

Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Ekagra Ranjan committed Jul 3, 2026 at 23:58 UTC 07516fda67d2133e26c0fd7386c0b0c8641e2a6e
8 files changed +436 -36
docs/features/speculative_decoding/dynamic_speculative_decoding.md
+2 -5
@@ -71,8 +71,5 @@ VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \
71
72 ## Limitations
73
74 -* only tested with Eagle and Eagle-3. Other SD methods may or may not work out of the box
75 -* only usable with Model Runner V1
76 -* not compatible with full cuda graph so we force piece-wise cuda graph with this feature
77 -
78 -We are working on enabling it on MRv2 with full cuda graph support.
74 +* Tested with Eagle, Eagle-3, and DFlash. Other SD methods may or may not work out of the box
75 +* Full Cudagraph only works with Model Runner V2. MRv1 only supports piece-wise cuda graph with this feature
tests/test_config.py
+31
@@ -34,6 +34,7 @@ from vllm.config.vllm import (
34 OptimizationLevel,
35 )
36 from vllm.platforms import current_platform
37 +from vllm.v1.attention.backend import AttentionCGSupport
38
39 DEVICE_TYPE = current_platform.device_type
40
@@ -67,6 +68,36 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
68 assert envs.VLLM_USE_V2_MODEL_RUNNER is expected
69
70
71 +@pytest.mark.parametrize(
72 + ("use_v2_model_runner", "expected_capture_sizes"),
73 + [
74 + (False, [4, 8, 12, 16]),
75 + (True, list(range(1, 17))),
76 + ],
77 +)
78 +def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(
79 + use_v2_model_runner,
80 + expected_capture_sizes,
81 +):
82 + compilation_config = CompilationConfig(
83 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
84 + cudagraph_capture_sizes=list(range(1, 17)),
85 + )
86 + compilation_config.max_cudagraph_capture_size = 16
87 + compilation_config.post_init_cudagraph_sizes()
88 +
89 + cudagraph_mode = compilation_config.resolve_cudagraph_mode_and_sizes(
90 + AttentionCGSupport.ALWAYS,
91 + "FakeAttentionBackend",
92 + uniform_decode_query_len=4,
93 + use_v2_model_runner=use_v2_model_runner,
94 + tensor_parallel_size=1,
95 + )
96 +
97 + assert cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE
98 + assert compilation_config.cudagraph_capture_sizes == expected_capture_sizes
99 +
100 +
101 @pytest.mark.parametrize(
102 ("model_config", "expected"),
103 [
tests/v1/spec_decode/test_dynamic_sd_cug.py new
+328
@@ -0,0 +1,328 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +
5 +from types import SimpleNamespace
6 +from unittest.mock import MagicMock
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm.config import (
12 + CompilationConfig,
13 + CUDAGraphMode,
14 + ParallelConfig,
15 + SchedulerConfig,
16 + VllmConfig,
17 +)
18 +from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
19 +
20 +pytestmark = pytest.mark.cpu_test
21 +
22 +
23 +def _create_vllm_config_for_dsd(
24 + max_num_seqs: int,
25 + max_spec_tokens: int,
26 + *,
27 + cudagraph_mode: str = "FULL_AND_PIECEWISE",
28 + use_dynamic_sd: bool = True,
29 + num_spec_per_batch_size: list[tuple[int, int, int]] | None = None,
30 +) -> MagicMock:
31 + """Create a minimal config that exercises DSD cudagraph dispatch.
32 +
33 + The test uses an exact capture-size grid so that every valid uniform decode
34 + shape has a directly matching FULL graph candidate.
35 +
36 + ``num_spec_per_batch_size`` lets a test supply an explicit DSD schedule of
37 + ``(range_start, range_end, num_speculative_tokens)`` tuples. When omitted,
38 + a schedule covering every query length in ``[1, max_decode_query_len]`` is
39 + generated.
40 + """
41 +
42 + max_decode_query_len = max_spec_tokens + 1
43 + max_capture_tokens = max_num_seqs * max_decode_query_len
44 +
45 + compilation_config = CompilationConfig(
46 + cudagraph_mode=cudagraph_mode,
47 + cudagraph_capture_sizes=list(range(1, max_capture_tokens + 1)),
48 + )
49 + compilation_config.max_cudagraph_capture_size = max_capture_tokens
50 + compilation_config.post_init_cudagraph_sizes()
51 +
52 + vllm_config = MagicMock(spec=VllmConfig)
53 + vllm_config.compilation_config = compilation_config
54 + vllm_config.scheduler_config = SchedulerConfig.default_factory(
55 + max_num_seqs=max_num_seqs,
56 + )
57 + vllm_config.parallel_config = ParallelConfig()
58 + # num_speculative_tokens is the max K (num_speculative_steps). The manager
59 + # recovers num_new_sampled_tokens_per_step as
60 + # decode_query_len - num_speculative_tokens; with decode_query_len =
61 + # max_spec_tokens + 1 this yields the normal per-step bonus of 1.
62 + vllm_config.num_speculative_tokens = max_spec_tokens
63 +
64 + speculative_config = MagicMock()
65 + speculative_config.uses_dynamic_speculative_decoding.return_value = use_dynamic_sd
66 + if use_dynamic_sd:
67 + # DSD reads the per-batch-size schedule; a schedule entry with K
68 + # speculative tokens maps to decode query length K + 1. By default
69 + # provide every query length in [1, max_decode_query_len] (i.e. K in
70 + # [0, max_spec_tokens]) so the manager captures a FULL decode graph for
71 + # each uniform shape.
72 + if num_spec_per_batch_size is None:
73 + num_spec_per_batch_size = [
74 + (qlen, qlen, qlen - 1) for qlen in range(1, max_decode_query_len + 1)
75 + ]
76 + speculative_config.num_speculative_tokens_per_batch_size = (
77 + num_spec_per_batch_size
78 + )
79 + else:
80 + speculative_config.num_speculative_tokens_per_batch_size = None
81 + vllm_config.speculative_config = speculative_config
82 +
83 + return vllm_config
84 +
85 +
86 +def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch):
87 + """Dynamic SD should create FULL decode candidates for every k in [1, K+1].
88 +
89 + This validates the MRv2 CudaGraphManager path directly: once candidate
90 + shapes have been built, dispatch() should pick a FULL graph for every
91 + uniform decode batch shape produced by DSD up to max_num_seqs.
92 + """
93 +
94 + max_num_seqs = 512
95 + max_spec_tokens = 7
96 + max_decode_query_len = max_spec_tokens + 1
97 +
98 + # CudaGraphManager consults PP rank helpers during initialization even
99 + # though this test only exercises CPU-side candidate generation.
100 + monkeypatch.setattr(
101 + gpu_cudagraph_utils,
102 + "get_pp_group",
103 + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
104 + )
105 +
106 + vllm_config = _create_vllm_config_for_dsd(
107 + max_num_seqs=max_num_seqs,
108 + max_spec_tokens=max_spec_tokens,
109 + )
110 + manager = gpu_cudagraph_utils.CudaGraphManager(
111 + vllm_config=vllm_config,
112 + device=torch.device("cpu"),
113 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
114 + decode_query_len=max_decode_query_len,
115 + )
116 +
117 + # dispatch() only uses the precomputed candidate table after graphs are
118 + # considered captured. The actual graph objects are irrelevant here.
119 + manager._graphs_captured = True
120 +
121 + for num_reqs in range(1, max_num_seqs + 1):
122 + for max_query_len in range(1, max_decode_query_len + 1):
123 + # Uniform decode means every request contributes the same number of
124 + # tokens, so the total token count is exactly num_reqs * query_len.
125 + num_tokens = num_reqs * max_query_len
126 + uniform_tok_count = gpu_cudagraph_utils.get_uniform_token_count(
127 + num_reqs,
128 + num_tokens,
129 + max_query_len,
130 + )
131 +
132 + # The scheduler should mark every one of these shapes as a uniform
133 + # decode batch, which is what enables FULL decode graph selection.
134 + assert uniform_tok_count == max_query_len
135 +
136 + desc = manager.dispatch(
137 + num_reqs=num_reqs,
138 + num_tokens=num_tokens,
139 + uniform_token_count=uniform_tok_count,
140 + num_active_loras=0,
141 + )
142 +
143 + # With DSD enabled, MRv2 should have captured a FULL candidate for
144 + # every k in [1, K+1], so dispatch should stay on the FULL path.
145 + assert desc.cg_mode == CUDAGraphMode.FULL
146 + assert desc.uniform_token_count == max_query_len
147 + assert desc.num_tokens == num_tokens
148 + assert desc.num_reqs == num_reqs
149 + assert desc.num_active_loras == 0
150 +
151 +
152 +def test_dynamic_sd_non_uniform_batch_falls_back_to_piecewise(monkeypatch):
153 + """DSD should use PIECEWISE when the batch is not a uniform decode batch.
154 +
155 + FULL DSD graphs are captured separately for each decode query length k.
156 + When runtime tokens are not uniform, uniform_token_count is None and those
157 + FULL candidates should be skipped in favor of the mixed-batch PIECEWISE
158 + graph under FULL_AND_PIECEWISE mode.
159 + """
160 +
161 + max_spec_tokens = 4
162 +
163 + monkeypatch.setattr(
164 + gpu_cudagraph_utils,
165 + "get_pp_group",
166 + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
167 + )
168 +
169 + vllm_config = _create_vllm_config_for_dsd(
170 + max_num_seqs=512,
171 + max_spec_tokens=max_spec_tokens,
172 + cudagraph_mode="FULL_AND_PIECEWISE",
173 + use_dynamic_sd=True,
174 + )
175 + manager = gpu_cudagraph_utils.CudaGraphManager(
176 + vllm_config=vllm_config,
177 + device=torch.device("cpu"),
178 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
179 + decode_query_len=max_spec_tokens + 1,
180 + )
181 + manager._graphs_captured = True
182 +
183 + # This shape is intentionally non-uniform: 3 tokens across 2 requests
184 + # cannot correspond to a single per-request query length.
185 + desc = manager.dispatch(
186 + num_reqs=2,
187 + num_tokens=3,
188 + uniform_token_count=None,
189 + num_active_loras=0,
190 + )
191 +
192 + assert desc.cg_mode == CUDAGraphMode.PIECEWISE
193 + assert desc.uniform_token_count is None
194 + assert desc.num_reqs is None
195 + assert desc.num_tokens == 3
196 + assert desc.num_active_loras == 0
197 +
198 +
199 +def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch):
200 + """Without DSD, only the max decode query length should get FULL graphs.
201 +
202 + Basic SD captures FULL decode graphs only for decode_query_len = K + 1.
203 + Uniform batches with smaller query lengths should therefore miss the FULL
204 + path entirely when using FULL_AND_PIECEWISE.
205 + """
206 +
207 + max_num_seqs = 512
208 + max_spec_tokens = 7
209 + max_decode_query_len = max_spec_tokens + 1
210 +
211 + monkeypatch.setattr(
212 + gpu_cudagraph_utils,
213 + "get_pp_group",
214 + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
215 + )
216 +
217 + vllm_config = _create_vllm_config_for_dsd(
218 + max_num_seqs=max_num_seqs,
219 + max_spec_tokens=max_spec_tokens,
220 + cudagraph_mode="FULL_AND_PIECEWISE",
221 + use_dynamic_sd=False,
222 + )
223 + manager = gpu_cudagraph_utils.CudaGraphManager(
224 + vllm_config=vllm_config,
225 + device=torch.device("cpu"),
226 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
227 + decode_query_len=max_decode_query_len,
228 + )
229 + manager._graphs_captured = True
230 +
231 + for num_reqs in range(1, max_num_seqs + 1):
232 + for max_query_len in range(1, max_decode_query_len):
233 + # These are still uniform decode batches, but basic SD should only
234 + # have FULL graphs for query_len == max_decode_query_len.
235 + num_tokens = num_reqs * max_query_len
236 + uniform_tok_count = gpu_cudagraph_utils.get_uniform_token_count(
237 + num_reqs,
238 + num_tokens,
239 + max_query_len,
240 + )
241 + assert uniform_tok_count == max_query_len
242 +
243 + desc = manager.dispatch(
244 + num_reqs=num_reqs,
245 + num_tokens=num_tokens,
246 + uniform_token_count=uniform_tok_count,
247 + num_active_loras=0,
248 + )
249 +
250 + assert desc.cg_mode == CUDAGraphMode.PIECEWISE
251 + assert desc.uniform_token_count is None
252 + assert desc.num_tokens == num_tokens
253 + assert desc.num_reqs is None
254 + assert desc.num_active_loras == 0
255 +
256 +
257 +def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch):
258 + """DSD should only capture FULL graphs for query lengths in the schedule.
259 +
260 + With a partial schedule of ``(1, 32, 4)`` and ``(32, 128, 3)``, only the
261 + scheduled speculative-token counts (K = 4 and K = 3) become decode query
262 + lengths (K + 1 = 5 and 4). Uniform batches at those query lengths should get
263 + FULL graphs, while every other query length (e.g. the lower values 1, 2, 3)
264 + must fall back to the mixed-batch PIECEWISE graph.
265 + """
266 +
267 + max_num_seqs = 128
268 + max_spec_tokens = 7
269 + max_decode_query_len = max_spec_tokens + 1
270 +
271 + # (range_start, range_end, num_speculative_tokens): K = 4 and K = 3 are
272 + # scheduled, so FULL decode graphs should exist for query lengths K + 1,
273 + # i.e. exactly {5, 4}.
274 + num_spec_per_batch_size = [(1, 32, 4), (32, 128, 3)]
275 + scheduled_query_lens = {entry[2] + 1 for entry in num_spec_per_batch_size}
276 +
277 + monkeypatch.setattr(
278 + gpu_cudagraph_utils,
279 + "get_pp_group",
280 + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
281 + )
282 +
283 + vllm_config = _create_vllm_config_for_dsd(
284 + max_num_seqs=max_num_seqs,
285 + max_spec_tokens=max_spec_tokens,
286 + cudagraph_mode="FULL_AND_PIECEWISE",
287 + use_dynamic_sd=True,
288 + num_spec_per_batch_size=num_spec_per_batch_size,
289 + )
290 + manager = gpu_cudagraph_utils.CudaGraphManager(
291 + vllm_config=vllm_config,
292 + device=torch.device("cpu"),
293 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
294 + decode_query_len=max_decode_query_len,
295 + )
296 + manager._graphs_captured = True
297 +
298 + for num_reqs in range(1, max_num_seqs + 1):
299 + for max_query_len in range(1, max_decode_query_len + 1):
300 + num_tokens = num_reqs * max_query_len
301 + uniform_tok_count = gpu_cudagraph_utils.get_uniform_token_count(
302 + num_reqs,
303 + num_tokens,
304 + max_query_len,
305 + )
306 + assert uniform_tok_count == max_query_len
307 +
308 + desc = manager.dispatch(
309 + num_reqs=num_reqs,
310 + num_tokens=num_tokens,
311 + uniform_token_count=uniform_tok_count,
312 + num_active_loras=0,
313 + )
314 +
315 + if max_query_len in scheduled_query_lens:
316 + # Scheduled query lengths get a dedicated FULL decode graph.
317 + assert desc.cg_mode == CUDAGraphMode.FULL
318 + assert desc.uniform_token_count == max_query_len
319 + assert desc.num_tokens == num_tokens
320 + assert desc.num_reqs == num_reqs
321 + else:
322 + # Unscheduled query lengths (including the lower values 1 and 2)
323 + # have no FULL candidate and must fall back to PIECEWISE.
324 + assert desc.cg_mode == CUDAGraphMode.PIECEWISE
325 + assert desc.uniform_token_count is None
326 + assert desc.num_tokens == num_tokens
327 + assert desc.num_reqs is None
328 + assert desc.num_active_loras == 0
vllm/config/compilation.py
+6 -2
@@ -1321,6 +1321,7 @@ class CompilationConfig:
1321 min_cg_support: "AttentionCGSupport",
1322 min_cg_attn_backend: str | None,
1323 uniform_decode_query_len: int = 1,
1324 + use_v2_model_runner: bool = False,
1325 tensor_parallel_size: int = 1,
1326 kv_cache_config: "KVCacheConfig | None" = None,
1327 max_num_reqs: int | None = None,
@@ -1421,13 +1422,16 @@ class CompilationConfig:
1422 "and make sure compilation mode is VLLM_COMPILE"
1423 )
1424
1424 - # Adjust cudagraph sizes to be a multiple of uniform_decode_query_len
1425 + # MRV1 adjusts cudagraph sizes to be a multiple of uniform_decode_query_len
1426 # to avoid: https://github.com/vllm-project/vllm/issues/28207 and temp-fix:
1427 # https://github.com/vllm-project/vllm/issues/28207#issuecomment-3504004536
1428 # Will be removed in the near future when we have separate cudagraph capture
1429 # sizes for decode and mixed prefill-decode.
1430 + # MRV2 handles cudagraph capture sizing in cudagraph_utils.py
1431 + # and doesn't need below: https://github.com/vllm-project/vllm/pull/45953
1432 if (
1430 - cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
1433 + not use_v2_model_runner
1434 + and cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
1435 and uniform_decode_query_len > 1
1436 ):
1437 self.adjust_cudagraph_sizes_for_spec_decode(
vllm/config/vllm.py
+3 -4
@@ -774,13 +774,15 @@ class VllmConfig:
774 speculative_config is None
775 or not speculative_config.uses_dynamic_speculative_decoding()
776 or not self.compilation_config.cudagraph_mode.has_full_cudagraphs()
777 + or self.use_v2_model_runner
778 ):
779 return
780
781 logger.warning_once(
782 "Dynamic speculative decoding changes the target verification "
783 "length at runtime. Overriding cudagraph_mode from %s to "
783 - "PIECEWISE for reliability.",
784 + "PIECEWISE for reliability. Use VLLM_USE_V2_MODEL_RUNNER=1 "
785 + "if you want to use full CUDA graphs.",
786 self.compilation_config.cudagraph_mode.name,
787 )
788 self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
@@ -2062,9 +2064,6 @@ class VllmConfig:
2064 ):
2065 unsupported.append(f"speculative method '{speculative_config.method}'")
2066
2065 - if speculative_config.uses_dynamic_speculative_decoding():
2066 - unsupported.append("dynamic speculative decoding")
2067 -
2067 # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle).
2068 # DFlash and DSpark use parallel drafting natively in V2 via their
2069 # own speculators.
vllm/v1/worker/gpu/cudagraph_utils.py
+58 -19
@@ -27,6 +27,7 @@ from vllm.logger import init_logger
27 from vllm.model_executor.offloader.base import get_offloader
28 from vllm.platforms import current_platform
29 from vllm.sequence import IntermediateTensors
30 +from vllm.utils.math_utils import round_up
31 from vllm.v1.kv_cache_interface import KVCacheConfig
32 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
33 from vllm.v1.worker.gpu.block_table import BlockTables
@@ -141,10 +142,7 @@ class CudaGraphManager:
142
143 self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}
144 self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}
144 - # adjust the cudagraph sizes to be a multiple of the uniform decode query length
145 - self.compilation_config.adjust_cudagraph_sizes_for_spec_decode(
146 - self.decode_query_len, self.tp_size
147 - )
145 +
146 self._init_candidates()
147
148 # Breakable CUDA graph (PW CUDA graph without torch.compile)
@@ -191,31 +189,72 @@ class CudaGraphManager:
189 decode_mode = self.cudagraph_mode.decode_mode()
190 mixed_mode = self.cudagraph_mode.mixed_mode()
191 separate_decode_routine = self.cudagraph_mode.separate_routine()
192 + max_cg_capture_size = self.compilation_config.max_cudagraph_capture_size
193
194 descs_by_token_lora: dict[tuple[int, int], list[BatchExecutionDescriptor]] = (
195 defaultdict(list)
196 )
198 - descs_by_mode = defaultdict(list)
197 + descs_by_mode: defaultdict[CUDAGraphMode, list[BatchExecutionDescriptor]] = (
198 + defaultdict(list)
199 + )
200 +
201 + # When using Dynamic SD, num_speculative_tokens is the max number of
202 + # draft tokens. The scheduler might use a smaller number so we need
203 + # to capture graphs for all possible values during decode.
204 + speculative_config = self.vllm_config.speculative_config
205 + if (
206 + speculative_config
207 + and speculative_config.uses_dynamic_speculative_decoding()
208 + ):
209 + num_spec_per_batch_size = (
210 + speculative_config.num_speculative_tokens_per_batch_size
211 + )
212 + # uses_dynamic_speculative_decoding() guarantees this is set.
213 + assert num_spec_per_batch_size is not None
214 + # decode_query_len = num_speculative_steps + num_new_sampled_tokens
215 + # _per_step. Recover num_new_sampled_tokens_per_step
216 + # from the values the manager already has.
217 + num_new_sampled_tokens_per_step = (
218 + self.decode_query_len - self.vllm_config.num_speculative_tokens
219 + )
220 + # Each entry is (range_start, range_end, num_speculative_tokens).
221 + decode_query_lens = [
222 + x[2] + num_new_sampled_tokens_per_step for x in num_spec_per_batch_size
223 + ]
224 + else:
225 + decode_query_lens = [self.decode_query_len]
226
227 for num_tokens, num_active_loras in product(
228 capture_sizes, self.lora_capture_cases
229 ):
230 # Capture uniform decode specfifc graphs if required
231 # (i.e. separate decode routine)
205 - if (
206 - separate_decode_routine
207 - and decode_mode
208 - and self.decode_query_len <= num_tokens <= max_decode_tokens
209 - ):
210 - desc = BatchExecutionDescriptor(
211 - cg_mode=decode_mode,
212 - num_tokens=num_tokens,
213 - num_reqs=num_tokens // self.decode_query_len,
214 - uniform_token_count=self.decode_query_len,
215 - num_active_loras=num_active_loras,
216 - )
217 - descs_by_mode[decode_mode].append(desc)
218 - descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)
232 + if separate_decode_routine and decode_mode:
233 + for decode_query_len in decode_query_lens:
234 + rounded_num_tokens = round_up(num_tokens, decode_query_len)
235 + rounded_num_reqs = rounded_num_tokens // decode_query_len
236 +
237 + if (
238 + rounded_num_tokens > max_decode_tokens
239 + or rounded_num_tokens > max_cg_capture_size
240 + or rounded_num_reqs > self.max_num_reqs
241 + ):
242 + continue
243 +
244 + desc = BatchExecutionDescriptor(
245 + cg_mode=decode_mode,
246 + num_tokens=rounded_num_tokens,
247 + num_reqs=rounded_num_reqs,
248 + uniform_token_count=decode_query_len,
249 + num_active_loras=num_active_loras,
250 + )
251 +
252 + # avoid duplicate graphs
253 + if desc not in descs_by_mode[decode_mode]:
254 + descs_by_mode[decode_mode].append(desc)
255 + descs_by_token_lora[
256 + (rounded_num_tokens, num_active_loras)
257 + ].append(desc)
258
259 if mixed_mode:
260 # for PIECEWISE graphs there is no limit on requests when replaying
vllm/v1/worker/gpu/model_runner.py
+4 -3
@@ -454,9 +454,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
454 attn_cg_support.min_cg_support,
455 attn_cg_support.min_cg_attn_backend,
456 self.decode_query_len,
457 - self.parallel_config.tensor_parallel_size,
458 - self.kv_cache_config,
459 - self.max_num_reqs,
457 + use_v2_model_runner=True,
458 + tensor_parallel_size=self.parallel_config.tensor_parallel_size,
459 + kv_cache_config=self.kv_cache_config,
460 + max_num_reqs=self.max_num_reqs,
461 )
462 self.cudagraph_manager = ModelCudaGraphManager(
463 self.vllm_config,
vllm/v1/worker/gpu_model_runner.py
+4 -3
@@ -6930,9 +6930,10 @@ class GPUModelRunner(
6930 min_cg_support,
6931 min_cg_attn_backend,
6932 self.uniform_decode_query_len,
6933 - self.parallel_config.tensor_parallel_size,
6934 - self.kv_cache_config,
6935 - self.max_num_reqs,
6933 + use_v2_model_runner=False,
6934 + tensor_parallel_size=self.parallel_config.tensor_parallel_size,
6935 + kv_cache_config=self.kv_cache_config,
6936 + max_num_reqs=self.max_num_reqs,
6937 is_profiling=is_profiling,
6938 )
6939 # Trigger cudagraph dispatching keys initialization after