[Bugfix][MRV2] Dispatch uniform decode to a padded FULL cudagraph (#53407)
Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: xiaohuguo2023 <149615094+xiaohuguo2023@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
xiaohuguo2023 committed
Aug 25, 2026 at 07:59 UTC
c74dd2d208389305e3a6ced992ee0632f939c0a0
2 files changed
+247
-24
tests/v1/cudagraph/test_cudagraph_manager.py
+230
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections import defaultdict
5
from contextlib import contextmanager
6
from types import SimpleNamespace
7
from unittest.mock import MagicMock, patch
@@ -109,3 +110,232 @@ def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
110
manager.capture(create_forward_fn)
111
112
mock_cuda_graph.assert_called_once()
113
+
114
+
115
+_DECODE_QUERY_LEN = 3
116
+
117
+
118
+def _create_decode_vllm_config(
119
+ capture_sizes: list[int],
120
+ num_speculative_tokens: int = 0,
121
+ dynamic_spec_num_tokens: list[int] | None = None,
122
+) -> MagicMock:
123
+ compilation_config = CompilationConfig(
124
+ cudagraph_mode="FULL_AND_PIECEWISE",
125
+ cudagraph_capture_sizes=capture_sizes,
126
+ )
127
+ compilation_config.max_cudagraph_capture_size = capture_sizes[-1]
128
+ compilation_config.post_init_cudagraph_sizes()
129
+
130
+ vllm_config = MagicMock(spec=VllmConfig)
131
+ vllm_config.compilation_config = compilation_config
132
+ vllm_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=8)
133
+ vllm_config.parallel_config = ParallelConfig()
134
+ vllm_config.num_speculative_tokens = num_speculative_tokens
135
+ if dynamic_spec_num_tokens is None:
136
+ vllm_config.speculative_config = None
137
+ else:
138
+ speculative_config = MagicMock()
139
+ speculative_config.uses_dynamic_speculative_decoding.return_value = True
140
+ # Each entry is (range_start, range_end, num_speculative_tokens); only
141
+ # the third element is read by the manager.
142
+ speculative_config.num_speculative_tokens_per_batch_size = [
143
+ (0, 0, n) for n in dynamic_spec_num_tokens
144
+ ]
145
+ vllm_config.speculative_config = speculative_config
146
+ return vllm_config
147
+
148
+
149
+def _make_spec_decode_manager(
150
+ monkeypatch,
151
+ decode_query_len: int = _DECODE_QUERY_LEN,
152
+ capture_sizes: list[int] | None = None,
153
+ num_speculative_tokens: int = 0,
154
+ dynamic_spec_num_tokens: list[int] | None = None,
155
+) -> gpu_cudagraph_utils.CudaGraphManager:
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: object(),
165
+ )
166
+ manager = gpu_cudagraph_utils.CudaGraphManager(
167
+ vllm_config=_create_decode_vllm_config(
168
+ capture_sizes or [1, 2, 4, 8, 16, 24],
169
+ num_speculative_tokens=num_speculative_tokens,
170
+ dynamic_spec_num_tokens=dynamic_spec_num_tokens,
171
+ ),
172
+ device=torch.device("cpu"),
173
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
174
+ decode_query_len=decode_query_len,
175
+ )
176
+ manager._graphs_captured = True
177
+ return manager
178
+
179
+
180
+def test_uniform_decode_pads_up_to_full_graph(monkeypatch):
181
+ manager = _make_spec_decode_manager(monkeypatch)
182
+ assert [
183
+ (desc.cg_mode, desc.num_tokens) for desc in manager._candidates[(12, 0)]
184
+ ] == [
185
+ (CUDAGraphMode.FULL, 18),
186
+ (CUDAGraphMode.PIECEWISE, 16),
187
+ ]
188
+
189
+ desc = manager.dispatch(
190
+ num_reqs=4,
191
+ num_tokens=12,
192
+ uniform_token_count=_DECODE_QUERY_LEN,
193
+ num_active_loras=0,
194
+ )
195
+
196
+ assert desc.cg_mode == CUDAGraphMode.FULL
197
+ assert desc.uniform_token_count == _DECODE_QUERY_LEN
198
+ assert desc.num_tokens == 18
199
+ assert desc.num_reqs == 6
200
+
201
+
202
+def test_uniform_decode_exact_match_is_not_over_padded(monkeypatch):
203
+ manager = _make_spec_decode_manager(monkeypatch)
204
+
205
+ desc = manager.dispatch(
206
+ num_reqs=3,
207
+ num_tokens=9,
208
+ uniform_token_count=_DECODE_QUERY_LEN,
209
+ num_active_loras=0,
210
+ )
211
+
212
+ assert desc.cg_mode == CUDAGraphMode.FULL
213
+ assert desc.num_tokens == 9
214
+ assert desc.num_reqs == 3
215
+
216
+
217
+def test_mixed_batch_never_selects_a_uniform_decode_graph(monkeypatch):
218
+ manager = _make_spec_decode_manager(monkeypatch)
219
+
220
+ desc = manager.dispatch(
221
+ num_reqs=2,
222
+ num_tokens=12,
223
+ uniform_token_count=None,
224
+ num_active_loras=0,
225
+ )
226
+
227
+ assert desc.cg_mode == CUDAGraphMode.PIECEWISE
228
+ assert desc.uniform_token_count is None
229
+ assert desc.num_tokens == 16
230
+
231
+
232
+def test_mixed_batch_at_decode_only_token_count_still_gets_a_graph(monkeypatch):
233
+ """A mixed batch must not fall to eager where only decode graphs are staged.
234
+
235
+ ``round_up(size, decode_query_len)`` lands FULL decode graphs on token counts
236
+ the PIECEWISE ladder never uses -- with capture sizes [1, 2, 4, 8, 16, 24]
237
+ and query length 3, FULL covers 3, 6, 9 and 18 while PIECEWISE has only
238
+ 1, 2, 4, 8, 16, 24. Building each mode's candidate ranges independently keeps
239
+ a PIECEWISE graph reachable there. Deriving the ranges from the staged sizes
240
+ instead offers a mixed batch nothing but decode descriptors, whose
241
+ ``uniform_token_count`` it can never match, so it runs fully eager.
242
+ """
243
+ manager = _make_spec_decode_manager(monkeypatch)
244
+
245
+ decode_only_token_counts = sorted(
246
+ {desc.num_tokens for desc in manager._capture_descs[CUDAGraphMode.FULL]}
247
+ - {desc.num_tokens for desc in manager._capture_descs[CUDAGraphMode.PIECEWISE]}
248
+ )
249
+ # Guard the premise: with no such token counts this test would cover nothing.
250
+ assert decode_only_token_counts
251
+
252
+ for num_tokens in decode_only_token_counts:
253
+ desc = manager.dispatch(
254
+ num_reqs=1,
255
+ num_tokens=num_tokens,
256
+ uniform_token_count=None,
257
+ num_active_loras=0,
258
+ )
259
+ assert desc.cg_mode == CUDAGraphMode.PIECEWISE, num_tokens
260
+ assert desc.uniform_token_count is None, num_tokens
261
+ assert desc.num_tokens >= num_tokens, num_tokens
262
+
263
+
264
+def test_uniform_decode_beyond_capture_ladder_falls_back(monkeypatch):
265
+ manager = _make_spec_decode_manager(monkeypatch)
266
+
267
+ desc = manager.dispatch(
268
+ num_reqs=9,
269
+ num_tokens=27,
270
+ uniform_token_count=_DECODE_QUERY_LEN,
271
+ num_active_loras=0,
272
+ )
273
+
274
+ assert desc.cg_mode == CUDAGraphMode.NONE
275
+
276
+
277
+def test_dynamic_spec_decode_shared_token_count_stays_reachable(monkeypatch):
278
+ """Every captured decode graph must remain reachable from dispatch().
279
+
280
+ Under dynamic speculative decoding ``decode_query_lens`` is a list, and the
281
+ staging loop rounds each capture size up to *every* query length. Several
282
+ decode graphs therefore land on the same ``num_tokens`` -- e.g. capture size
283
+ 2 stages both ``round_up(2, 1) == 2`` and ``round_up(2, 2) == 2``. Building
284
+ the candidate ranges one descriptor at a time would give all but the first
285
+ of those an empty range, so they would never enter a candidate list and
286
+ their batches would silently fall through to PIECEWISE.
287
+ """
288
+ manager = _make_spec_decode_manager(
289
+ monkeypatch,
290
+ decode_query_len=4,
291
+ capture_sizes=[2, 4, 6, 8],
292
+ num_speculative_tokens=3,
293
+ dynamic_spec_num_tokens=[0, 1, 2, 3], # -> decode_query_lens [1, 2, 3, 4]
294
+ )
295
+
296
+ full_descs = manager._capture_descs[CUDAGraphMode.FULL]
297
+ by_num_tokens: dict[int, list] = defaultdict(list)
298
+ for desc in full_descs:
299
+ by_num_tokens[desc.num_tokens].append(desc)
300
+ # Guard the premise: without collisions this test would not cover the bug.
301
+ assert any(len(descs) > 1 for descs in by_num_tokens.values())
302
+
303
+ for desc in full_descs:
304
+ assert desc in manager._candidates[(desc.num_tokens, 0)], desc
305
+ assert (
306
+ manager.dispatch(
307
+ num_reqs=desc.num_reqs,
308
+ num_tokens=desc.num_tokens,
309
+ uniform_token_count=desc.uniform_token_count,
310
+ num_active_loras=0,
311
+ )
312
+ == desc
313
+ ), desc
314
+
315
+
316
+@pytest.mark.parametrize(
317
+ "decode_query_len,capture_sizes",
318
+ [(1, [1, 2, 4, 8]), (2, [2, 4, 8, 16]), (8, [8, 16, 32, 64])],
319
+)
320
+def test_divisor_query_len_dispatch_is_unchanged(
321
+ monkeypatch, decode_query_len, capture_sizes
322
+):
323
+ manager = _make_spec_decode_manager(
324
+ monkeypatch,
325
+ decode_query_len=decode_query_len,
326
+ capture_sizes=capture_sizes,
327
+ )
328
+
329
+ max_reqs = capture_sizes[-1] // decode_query_len
330
+ for num_reqs in range(1, max_reqs + 1):
331
+ num_tokens = num_reqs * decode_query_len
332
+ desc = manager.dispatch(
333
+ num_reqs=num_reqs,
334
+ num_tokens=num_tokens,
335
+ uniform_token_count=decode_query_len,
336
+ num_active_loras=0,
337
+ )
338
+ expected = min(s for s in capture_sizes if s >= num_tokens)
339
+ assert desc.cg_mode == CUDAGraphMode.FULL, num_tokens
340
+ assert desc.num_tokens == expected, num_tokens
341
+ assert desc.num_reqs == expected // decode_query_len, num_tokens
vllm/v1/worker/gpu/cudagraph_utils.py
+17
-24
@@ -4,7 +4,7 @@ import gc
4
from collections import defaultdict
5
from collections.abc import Callable
6
from dataclasses import dataclass
7
-from itertools import product
7
+from itertools import groupby, product
8
from typing import TYPE_CHECKING, Any, NamedTuple, Protocol
9
10
import torch
@@ -194,9 +194,6 @@ class CudaGraphManager:
194
separate_decode_routine = self.cudagraph_mode.separate_routine()
195
max_cg_capture_size = self.compilation_config.max_cudagraph_capture_size
196
197
- descs_by_token_lora: dict[tuple[int, int], list[BatchExecutionDescriptor]] = (
198
- defaultdict(list)
199
- )
197
descs_by_mode: defaultdict[CUDAGraphMode, list[BatchExecutionDescriptor]] = (
198
defaultdict(list)
199
)
@@ -244,7 +241,6 @@ class CudaGraphManager:
241
num_active_loras=num_active_loras,
242
)
243
descs_by_mode[decode_mode].append(desc)
247
- descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)
244
# Capture uniform decode specfifc graphs if required
245
# (i.e. separate decode routine)
246
elif separate_decode_routine and decode_mode and not self.varlen_decode:
@@ -270,9 +266,6 @@ class CudaGraphManager:
266
# avoid duplicate graphs
267
if desc not in descs_by_mode[decode_mode]:
268
descs_by_mode[decode_mode].append(desc)
273
- descs_by_token_lora[
274
- (rounded_num_tokens, num_active_loras)
275
- ].append(desc)
269
270
if mixed_mode:
271
# for PIECEWISE graphs there is no limit on requests when replaying
@@ -290,27 +283,27 @@ class CudaGraphManager:
283
num_active_loras=num_active_loras,
284
)
285
descs_by_mode[mixed_mode].append(desc)
293
- descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)
294
-
295
- if not descs_by_token_lora:
296
- return
297
-
298
- all_token_counts = sorted({k[0] for k in descs_by_token_lora})
299
- current_range_start = 0
300
- for token_cg_size in all_token_counts:
301
- for i in range(current_range_start, token_cg_size + 1):
302
- for num_active_loras in self.lora_capture_cases:
303
- staging_key = (token_cg_size, num_active_loras)
304
- if staging_key in descs_by_token_lora:
305
- self._candidates[(i, num_active_loras)] = descs_by_token_lora[
306
- staging_key
307
- ]
308
- current_range_start = token_cg_size + 1
286
287
for mode, descs in descs_by_mode.items():
288
descs.sort(key=lambda d: d.num_tokens, reverse=True)
289
self._capture_descs[mode] = descs
290
291
+ for mode in (CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE):
292
+ mode_descs = tuple(reversed(descs_by_mode.get(mode, [])))
293
+ for num_active_loras in self.lora_capture_cases:
294
+ lora_descs = [
295
+ d for d in mode_descs if d.num_active_loras == num_active_loras
296
+ ]
297
+ current_range_start = 0
298
+ # Dynamic speculative decoding can produce multiple graphs with the same
299
+ # num_tokens. Group them so each graph covers the same candidate range.
300
+ for num_tokens, group in groupby(lora_descs, lambda d: d.num_tokens):
301
+ matching = list(group)
302
+ for i in range(current_range_start, num_tokens + 1):
303
+ key = (i, num_active_loras)
304
+ self._candidates.setdefault(key, []).extend(matching)
305
+ current_range_start = num_tokens + 1
306
+
307
def needs_capture(self) -> bool:
308
return len(self._capture_descs) > 0
309