@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0a94d85a66

[Bugfix][MRV2] Require all requests to be decoding for uniform-decode dispatch (#51865)

Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Co-authored-by: Janelle Cai <janelle.cai@modal.com>

Nick Hill committed Aug 11, 2026 at 17:15 UTC 0a94d85a66499cad8297ead86a470967de5c0212
9 files changed +341 -54
tests/v1/spec_decode/test_dynamic_sd_cug.py
+72
@@ -16,6 +16,7 @@ from vllm.config import (
16 VllmConfig,
17 )
18 from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
19 +from vllm.v1.worker.utils import get_uniform_decode_token_count
20
21 pytestmark = pytest.mark.cpu_test
22
@@ -196,6 +197,77 @@ def test_dynamic_sd_non_uniform_batch_falls_back_to_piecewise(monkeypatch):
197 assert desc.num_active_loras == 0
198
199
200 +def test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph(monkeypatch):
201 + """A batch holding K+1-token prompt chunks must not replay a decode graph.
202 +
203 + Prompt chunks of exactly K + 1 tokens give the batch a spec-decode shape
204 + by coincidence; replaying the FULL graph over them corrupts every request
205 + in the batch. See https://github.com/vllm-project/vllm/issues/49918.
206 + """
207 +
208 + max_spec_tokens = 7
209 + 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=64,
219 + max_spec_tokens=max_spec_tokens,
220 + use_dynamic_sd=False,
221 + )
222 + manager = gpu_cudagraph_utils.CudaGraphManager(
223 + vllm_config=vllm_config,
224 + device=torch.device("cpu"),
225 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
226 + decode_query_len=decode_query_len,
227 + )
228 + manager._graphs_captured = True
229 +
230 + num_decodes, num_chunks = 6, 2
231 + num_reqs = num_decodes + num_chunks
232 + num_tokens = num_reqs * decode_query_len
233 +
234 + # The batch shape is indistinguishable from a full batch of spec decodes.
235 + assert (
236 + gpu_cudagraph_utils.get_uniform_token_count(
237 + num_reqs, num_tokens, decode_query_len
238 + )
239 + == decode_query_len
240 + )
241 +
242 + # Two of the requests are decode_query_len tokens into a longer prompt.
243 + uniform_tok_count = get_uniform_decode_token_count(
244 + num_reqs, num_tokens, decode_query_len, has_prefill=True
245 + )
246 + assert uniform_tok_count is None
247 + desc = manager.dispatch(
248 + num_reqs=num_reqs,
249 + num_tokens=num_tokens,
250 + uniform_token_count=uniform_tok_count,
251 + num_active_loras=0,
252 + )
253 + assert desc.cg_mode == CUDAGraphMode.PIECEWISE
254 + assert desc.uniform_token_count is None
255 +
256 + # The same shape with every request decoding still gets its FULL graph.
257 + uniform_tok_count = get_uniform_decode_token_count(
258 + num_reqs, num_tokens, decode_query_len, has_prefill=False
259 + )
260 + assert uniform_tok_count == decode_query_len
261 + desc = manager.dispatch(
262 + num_reqs=num_reqs,
263 + num_tokens=num_tokens,
264 + uniform_token_count=uniform_tok_count,
265 + num_active_loras=0,
266 + )
267 + assert desc.cg_mode == CUDAGraphMode.FULL
268 + assert desc.uniform_token_count == decode_query_len
269 +
270 +
271 def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch):
272 """Without DSD, only the max decode query length should get FULL graphs.
273
tests/v1/worker/test_gpu_batch_ordering.py
+149 -2
@@ -1,18 +1,28 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Batch ordering in Model Runner V2 (vllm.v1.worker.gpu).
3 +"""Batch ordering and classification in Model Runner V2 (vllm.v1.worker.gpu).
4
5 split_decodes_and_prefills assumes decode -> short_extend -> prefill request
6 ordering. With spec decode (decode_query_len > 1), a shorter chunked-prefill
7 tail sorted in front of the uniform decodes misclassifies every decode as a
8 prefill.
9 +
10 +Conversely, a prompt chunk of exactly decode_query_len tokens has a decode
11 +batch's shape, and must not be classified as a uniform decode batch.
12 """
13
14 +import ast
15 +from pathlib import Path
16 +from types import SimpleNamespace
17 +from typing import Any
18 +
19 +import numpy as np
20 import torch
21
22 from vllm.v1.attention.backend import CommonAttentionMetadata
23 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
15 -from vllm.v1.worker.gpu.model_runner import sort_batch_req_ids
24 +from vllm.v1.worker.gpu.model_runner import GPUModelRunner, sort_batch_req_ids
25 +from vllm.v1.worker.utils import get_uniform_decode_token_count
26
27
28 def _make_common_attn_metadata(query_lens: list[int]) -> CommonAttentionMetadata:
@@ -37,6 +47,72 @@ def _make_common_attn_metadata(query_lens: list[int]) -> CommonAttentionMetadata
47 )
48
49
50 +def _make_runner(
51 + req_states: dict[str, tuple[int, int]], decode_query_len: int = 8
52 +) -> Any:
53 + """Build a runner stub from {req_id: (num_computed_tokens, prefill_len)}."""
54 + prefill_lens = np.array([s[1] for s in req_states.values()], dtype=np.int32)
55 + num_computed = np.array([s[0] for s in req_states.values()], dtype=np.int32)
56 + runner: Any = GPUModelRunner.__new__(GPUModelRunner)
57 + runner.decode_query_len = decode_query_len
58 + runner.req_states = SimpleNamespace(
59 + req_id_to_index={req_id: i for i, req_id in enumerate(req_states)},
60 + # The runner keeps this as min(num_computed_tokens, prefill_len).
61 + num_computed_prefill_tokens=np.minimum(num_computed, prefill_lens),
62 + prefill_len=SimpleNamespace(np=prefill_lens),
63 + )
64 + return runner
65 +
66 +
67 +def _uniform_token_count(
68 + req_states: dict[str, tuple[int, int]],
69 + query_len: int,
70 + dummy_run: bool = False,
71 +) -> int | None:
72 + """Classify via the runner's pre-dispatch gather, as execute_model does."""
73 + scheduler_output = SimpleNamespace(
74 + num_scheduled_tokens={req_id: query_len for req_id in req_states},
75 + total_num_scheduled_tokens=query_len * len(req_states),
76 + )
77 + runner = _make_runner(req_states, decode_query_len=query_len)
78 + _, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, dummy_run)
79 + return uniform_tok_count
80 +
81 +
82 +def test_spec_decode_batch_is_uniform_decode():
83 + # decode_query_len == 8 (K=7): every request past its prefill.
84 + decodes = {f"d{i}": (16, 16) for i in range(6)}
85 + assert _uniform_token_count(decodes, 8) == 8
86 +
87 +
88 +def test_prompt_chunk_of_decode_query_len_is_not_uniform_decode():
89 + # Two prompt chunks whose query length coincides with the K+1 spec-decode
90 + # query length must reject the batch (issue #49918).
91 + batch = {f"d{i}": (16, 16) for i in range(6)}
92 + batch.update({f"p{i}": (8, 40) for i in range(2)})
93 + assert _uniform_token_count(batch, 8) is None
94 +
95 + # Same collision without spec decoding: 1-token prompt vs 1-token decode.
96 + assert _uniform_token_count({"p0": (0, 1)}, 1) is None
97 +
98 + # A fresh prompt of exactly decode_query_len tokens.
99 + assert _uniform_token_count({"p0": (0, 8)}, 8) is None
100 +
101 +
102 +def test_dummy_batches_stay_uniform_decode():
103 + # Dummy batches have no request state to consult and must stay classified
104 + # by shape alone; rejecting them would make capture and replay disagree.
105 + scheduler_output = SimpleNamespace(
106 + num_scheduled_tokens={f"_dummy_req_{i}": 8 for i in range(4)},
107 + total_num_scheduled_tokens=32,
108 + )
109 + state, uniform_tok_count = _make_runner({}).gather_batch_req_state(
110 + scheduler_output, True
111 + )
112 + assert state is None
113 + assert uniform_tok_count == 8
114 +
115 +
116 def test_sort_batch_req_ids_no_spec():
117 # decode_query_len == 1: plain ascending order (decodes first).
118 num_tokens_per_req = {"p1": 100, "d1": 1, "p2": 7, "d2": 1}
@@ -68,3 +144,74 @@ def test_spec_decodes_lead_short_prefill_tail():
144 )
145 assert (num_decodes, num_prefills) == (8, 1)
146 assert (num_decode_tokens, num_prefill_tokens) == (16, 1)
147 +
148 +
149 +def test_uniform_decode_uses_state_index_not_batch_position():
150 + """The gather must read each request's own state, not its batch slot.
151 +
152 + Batch position and state index diverge as requests come and go; here the
153 + only prefilling request sits at state index 0, so a gather keyed on batch
154 + position would read its row and wrongly reject the two decodes.
155 + """
156 + runner: Any = GPUModelRunner.__new__(GPUModelRunner)
157 + runner.decode_query_len = 8
158 + # State arrays in state-index order: a prefilling request, then two decodes.
159 + runner.req_states = SimpleNamespace(
160 + req_id_to_index={"prefilling": 0, "decode_a": 1, "decode_b": 2},
161 + num_computed_prefill_tokens=np.array([8, 16, 16], dtype=np.int32),
162 + prefill_len=SimpleNamespace(np=np.array([40, 16, 16], dtype=np.int32)),
163 + )
164 + scheduler_output = SimpleNamespace(
165 + num_scheduled_tokens={"decode_a": 8, "decode_b": 8},
166 + total_num_scheduled_tokens=16,
167 + )
168 +
169 + state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False)
170 + assert not state.has_prefill
171 + assert uniform_tok_count == 8
172 +
173 + # Scheduling the prefilling request alongside them must reject the batch.
174 + scheduler_output.num_scheduled_tokens = {
175 + "decode_a": 8,
176 + "prefilling": 8,
177 + "decode_b": 8,
178 + }
179 + scheduler_output.total_num_scheduled_tokens = 24
180 + state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False)
181 + assert state.has_prefill
182 + assert uniform_tok_count is None
183 +
184 +
185 +def test_uniform_decode_predicate():
186 + # Shape and prefill state must both pass.
187 + assert get_uniform_decode_token_count(2, 16, 8, False) == 8
188 + assert get_uniform_decode_token_count(2, 16, 8, True) is None
189 + # 12 tokens over 2 requests is no shared query length.
190 + assert get_uniform_decode_token_count(2, 12, 8, False) is None
191 +
192 +
193 +def test_no_speculator_dispatches_on_query_length_alone():
194 + """No speculator may pick its cudagraph from a shape-only test.
195 +
196 + Written over the whole package because speculators are added by copying an
197 + existing one: the shape-only call already reappeared verbatim in a
198 + speculator added after the original call sites were fixed.
199 + """
200 + import vllm.v1.worker.gpu.spec_decode as spec_decode
201 +
202 + shape_only = {"get_uniform_token_count", "is_uniform_query_len"}
203 + root = Path(spec_decode.__file__).parent
204 + offenders = []
205 + for path in sorted(root.rglob("speculator.py")):
206 + for node in ast.walk(ast.parse(path.read_text())):
207 + if not isinstance(node, ast.Call):
208 + continue
209 + func = node.func
210 + name = func.id if isinstance(func, ast.Name) else getattr(func, "attr", "")
211 + if name in shape_only:
212 + offenders.append(f"{path.relative_to(root)}:{node.lineno} {name}")
213 +
214 + assert not offenders, (
215 + "these speculators classify a decode batch by query length alone; use "
216 + f"get_uniform_decode_token_count instead: {offenders}"
217 + )
vllm/v1/worker/gpu/cudagraph_utils.py
+7 -7
@@ -35,7 +35,7 @@ from vllm.v1.worker.gpu.block_table import BlockTables
35 from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
36 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
37 from vllm.v1.worker.gpu.model_states.interface import ModelState
38 -from vllm.v1.worker.utils import AttentionGroup
38 +from vllm.v1.worker.utils import AttentionGroup, is_uniform_query_len
39
40 logger = init_logger(__name__)
41
@@ -90,17 +90,17 @@ def _is_compatible(
90
91
92 def get_uniform_token_count(
93 - num_reqs: int,
94 - num_tokens: int,
95 - max_query_len: int,
93 + num_reqs: int, num_tokens: int, max_query_len: int
94 ) -> int | None:
95 """
96 Return the uniform token count if batch is uniform, else None.
97 A batch is uniform if all requests have the same number of tokens.
98 +
99 + Shape test only, valid for batches known to be decoding by construction
100 + (e.g. dummy runs). Scheduled batches must use
101 + `get_uniform_decode_token_count`.
102 """
101 - if (max_query_len == num_tokens // num_reqs) and (
102 - num_tokens == max_query_len * num_reqs
103 - ):
103 + if is_uniform_query_len(num_reqs, num_tokens, max_query_len):
104 return max_query_len
105 return None
106
vllm/v1/worker/gpu/input_batch.py
+4 -4
@@ -76,6 +76,8 @@ class InputBatch:
76 num_computed_prefill_tokens_np: np.ndarray
77 # [num_reqs] CPU bool array == (num_computed_prefill_tokens_np < prefill_len_np).
78 is_prefilling_np: np.ndarray
79 + # == np.any(is_prefilling_np)
80 + has_prefill: bool
81
82 # [num_reqs] only populated when pipeline parallelism is enabled.
83 max_seq_len_np: np.ndarray | None
@@ -101,10 +103,7 @@ class InputBatch:
103
104 @classmethod
105 def make_dummy(
104 - cls,
105 - num_reqs: int,
106 - num_tokens: int,
107 - input_buffers: InputBuffers,
106 + cls, num_reqs: int, num_tokens: int, input_buffers: InputBuffers
107 ) -> "InputBatch":
108 assert 0 < num_reqs <= num_tokens
109 device = input_buffers.device
@@ -175,6 +174,7 @@ class InputBatch:
174 prefill_len_np=np.zeros(num_reqs, dtype=np.int32),
175 num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
176 is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
177 + has_prefill=False,
178 max_seq_len_np=None,
179 input_ids=input_ids,
180 positions=positions,
vllm/v1/worker/gpu/model_runner.py
+83 -28
@@ -133,7 +133,11 @@ from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
133 from vllm.v1.worker.gpu.states import RequestState
134 from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
135 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
136 -from vllm.v1.worker.utils import KVBlockZeroer, copy_kv_cache_blocks_inplace
136 +from vllm.v1.worker.utils import (
137 + KVBlockZeroer,
138 + copy_kv_cache_blocks_inplace,
139 + get_uniform_decode_token_count,
140 +)
141
142 logger = init_logger(__name__)
143
@@ -974,21 +978,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
978 scheduler_output.kv_cache_block_copies,
979 )
980
977 - def prepare_inputs(
978 - self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor
979 - ) -> InputBatch:
980 - num_tokens = scheduler_output.total_num_scheduled_tokens
981 - num_tokens_after_padding = batch_desc.num_tokens
982 - assert num_tokens > 0
983 - if envs.VLLM_MOE_SKIP_PADDING:
984 - # Mark trailing cudagraph-padding rows so kernels can skip work for
985 - # them when supported.
986 - self.input_buffers.is_padding[:num_tokens].fill_(False)
987 - self.input_buffers.is_padding[num_tokens:num_tokens_after_padding].fill_(
988 - True
989 - )
981 + def gather_batch_req_state(
982 + self, scheduler_output: SchedulerOutput, dummy_run: bool
983 + ) -> tuple["BatchReqState | None", int | None]:
984 + """Gather CPU request state for the scheduled batch, in batch order.
985 + Returns (batch_state, uniform_decode_token_count)
986 + """
987 num_tokens_per_req = scheduler_output.num_scheduled_tokens
988 num_reqs = len(num_tokens_per_req)
989 + num_toks = scheduler_output.total_num_scheduled_tokens
990 + max_query_len = max(scheduler_output.num_scheduled_tokens.values())
991 +
992 + if dummy_run:
993 + # Dummy batches are uniform by construction.
994 + return None, get_uniform_token_count(num_reqs, num_toks, max_query_len)
995
996 # batch_idx -> req_id
997 req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)
@@ -997,7 +1000,46 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1000
1001 idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids)
1002 idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs)
1003 + prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np]
1004 + num_computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens[
1005 + idx_mapping_np
1006 + ]
1007 + is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
1008 +
1009 + batch_state = BatchReqState(
1010 + req_ids=req_ids,
1011 + num_scheduled_tokens=num_scheduled_tokens,
1012 + idx_mapping_np=idx_mapping_np,
1013 + prefill_len_np=prefill_len_np,
1014 + num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
1015 + is_prefilling_np=is_prefilling_np,
1016 + has_prefill=bool(is_prefilling_np.any()),
1017 + )
1018 + return batch_state, get_uniform_decode_token_count(
1019 + num_reqs, num_toks, max_query_len, batch_state.has_prefill
1020 + )
1021 +
1022 + def prepare_inputs(
1023 + self,
1024 + scheduler_output: SchedulerOutput,
1025 + batch_req_state: "BatchReqState",
1026 + batch_desc: BatchExecutionDescriptor,
1027 + ) -> InputBatch:
1028 + num_tokens = scheduler_output.total_num_scheduled_tokens
1029 + num_tokens_after_padding = batch_desc.num_tokens
1030 + assert num_tokens > 0
1031 + if envs.VLLM_MOE_SKIP_PADDING:
1032 + # Mark trailing cudagraph-padding rows so kernels can skip work for
1033 + # them when supported.
1034 + is_padding = self.input_buffers.is_padding
1035 + is_padding[:num_tokens].fill_(False)
1036 + is_padding[num_tokens:num_tokens_after_padding].fill_(True)
1037 +
1038 + req_ids = batch_req_state.req_ids
1039 + num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
1040 + idx_mapping_np = batch_req_state.idx_mapping_np
1041 idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
1042 + num_reqs = len(req_ids)
1043
1044 # Get the number of draft tokens for each request.
1045 draft_tokens = scheduler_output.scheduled_spec_decode_tokens
@@ -1039,20 +1081,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1081 num_reqs_padded = batch_desc.num_reqs or num_reqs
1082 query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
1083 query_start_loc_np[0] = 0
1042 - np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1 : num_reqs + 1])
1084 + np.cumsum(num_scheduled_tokens_np, out=query_start_loc_np[1 : num_reqs + 1])
1085 # Pad for full CUDA graph mode.
1086 # Some attention backends like FA3 require query_start_loc to be non-decreasing.
1087 query_start_loc_np[num_reqs + 1 :] = num_tokens
1088 async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
1089 query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
1090 query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1]
1049 - prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np]
1050 - computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens
1051 - num_computed_prefill_tokens_np = computed_prefill_tokens_np[idx_mapping_np]
1052 - is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
1091
1092 # Get prefill tokens if any.
1055 - if np.any(is_prefilling_np):
1093 + if batch_req_state.has_prefill:
1094 prepare_prefill_inputs(
1095 self.input_buffers.input_ids,
1096 self.req_states.next_prefill_tokens,
@@ -1106,7 +1144,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1144 seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
1145 np.add(
1146 num_computed_tokens_np,
1109 - num_scheduled_tokens,
1147 + num_scheduled_tokens_np,
1148 out=seq_lens_cpu_upper_bound_np[:num_reqs],
1149 )
1150 seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)
@@ -1129,7 +1167,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1167 idx_mapping_np=idx_mapping_np,
1168 expanded_idx_mapping=expanded_idx_mapping,
1169 expanded_local_pos=expanded_local_pos,
1132 - num_scheduled_tokens=num_scheduled_tokens,
1170 + num_scheduled_tokens=num_scheduled_tokens_np,
1171 num_tokens=num_tokens,
1172 num_tokens_after_padding=num_tokens_after_padding,
1173 num_draft_tokens=total_num_draft_tokens,
@@ -1140,9 +1178,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1178 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
1179 dcp_local_seq_lens=dcp_local_seq_lens,
1180 num_computed_tokens_np=num_computed_tokens_np,
1143 - prefill_len_np=prefill_len_np,
1144 - num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
1145 - is_prefilling_np=is_prefilling_np,
1181 + prefill_len_np=batch_req_state.prefill_len_np,
1182 + num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,
1183 + is_prefilling_np=batch_req_state.is_prefilling_np,
1184 + has_prefill=batch_req_state.has_prefill,
1185 max_seq_len_np=max_seq_len_np,
1186 input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
1187 positions=self.input_buffers.positions[:num_tokens_after_padding],
@@ -1275,8 +1314,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1314 # Get batch descriptor and sync across DP ranks.
1315 num_reqs = len(scheduler_output.num_scheduled_tokens)
1316 num_toks = scheduler_output.total_num_scheduled_tokens
1278 - max_query_len = max(scheduler_output.num_scheduled_tokens.values())
1279 - uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len)
1317 + batch_req_state, uniform_tok_count = self.gather_batch_req_state(
1318 + scheduler_output, dummy_run
1319 + )
1320
1321 num_active_loras = 0
1322 if self.lora_config:
@@ -1311,7 +1351,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1351 if not dummy_run:
1352 # Common case.
1353 # Prepare all the inputs and copy to the input buffers.
1314 - input_batch = self.prepare_inputs(scheduler_output, batch_desc)
1354 + assert batch_req_state is not None
1355 + input_batch = self.prepare_inputs(
1356 + scheduler_output, batch_req_state, batch_desc
1357 + )
1358 block_tables, slot_mappings = self.prepare_attn(input_batch)
1359 # Mamba "align" pre-copy: migrate recurrent state across block
1360 # boundaries before the forward. Runs only on real batches, and
@@ -1790,6 +1833,18 @@ class ExecuteModelState(NamedTuple):
1833 routed_experts: RoutedExpertsTensors | None
1834
1835
1836 +class BatchReqState(NamedTuple):
1837 + """CPU request state for a scheduled batch, in batch (sorted) order."""
1838 +
1839 + req_ids: list[str]
1840 + num_scheduled_tokens: np.ndarray # [num_reqs]
1841 + idx_mapping_np: np.ndarray # [num_reqs]
1842 + prefill_len_np: np.ndarray # [num_reqs]
1843 + num_computed_prefill_tokens_np: np.ndarray # [num_reqs]
1844 + is_prefilling_np: np.ndarray # [num_reqs]
1845 + has_prefill: bool
1846 +
1847 +
1848 def sort_batch_req_ids(
1849 num_tokens_per_req: dict[str, int], decode_query_len: int
1850 ) -> list[str]:
vllm/v1/worker/gpu/pcp_manager.py
+1
@@ -529,6 +529,7 @@ class PCPManager:
529 prefill_len_np=local_prefill_len_np,
530 num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,
531 is_prefilling_np=local_is_prefilling_np,
532 + has_prefill=bool(local_is_prefilling_np.any()),
533 max_seq_len_np=global_batch.max_seq_len_np[local_to_global_batch_req_idx_np]
534 if global_batch.max_seq_len_np is not None
535 else None,
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+4 -6
@@ -13,10 +13,7 @@ from vllm.triton_utils import tl, triton
13 from vllm.v1.kv_cache_interface import KVCacheConfig
14 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
15 from vllm.v1.worker.gpu.block_table import BlockTables
16 -from vllm.v1.worker.gpu.cudagraph_utils import (
17 - BatchExecutionDescriptor,
18 - get_uniform_token_count,
19 -)
16 +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
17 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
18 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
19 from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -24,7 +21,7 @@ from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
21 SpeculatorCudaGraphManager,
22 )
23 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
27 -from vllm.v1.worker.utils import AttentionGroup
24 +from vllm.v1.worker.utils import AttentionGroup, get_uniform_decode_token_count
25
26 logger = init_logger(__name__)
27
@@ -274,12 +271,13 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
271
272 # When all requests are decoding (no true prefills), each has
273 # num_speculative_steps + 1 tokens, enabling FULL graph replay.
277 - uniform_token_count = get_uniform_token_count(
274 + uniform_token_count = get_uniform_decode_token_count(
275 num_reqs,
276 # Use the actual number of tokens without padding added by
277 # the target model during FULL cudagraph.
278 num_tokens,
279 max_query_len,
280 + input_batch.has_prefill,
281 )
282 prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
283 self.prefill_cudagraph_manager,
vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
+3 -7
@@ -12,9 +12,6 @@ from vllm.logger import init_logger
12 from vllm.triton_utils import tl, triton
13 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
14 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
15 -from vllm.v1.worker.gpu.cudagraph_utils import (
16 - get_uniform_token_count,
17 -)
15 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
16 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
17 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
@@ -22,6 +19,7 @@ from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
19 )
20 from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
21 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
22 +from vllm.v1.worker.utils import get_uniform_decode_token_count
23
24 logger = init_logger(__name__)
25
@@ -190,10 +188,8 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
188
189 # When all requests are decoding (no true prefills), each has
190 # num_speculative_steps + 1 tokens, enabling FULL graph replay.
193 - uniform_token_count = get_uniform_token_count(
194 - num_reqs,
195 - num_tokens,
196 - max_query_len,
191 + uniform_token_count = get_uniform_decode_token_count(
192 + num_reqs, num_tokens, max_query_len, input_batch.has_prefill
193 )
194 batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
195 self.cudagraph_manager,
vllm/v1/worker/utils.py
+18
@@ -611,6 +611,24 @@ def copy_kv_cache_blocks_inplace(
611 blocks[dst_indices] = blocks[src_indices]
612
613
614 +def is_uniform_query_len(num_reqs: int, num_tokens: int, max_query_len: int) -> bool:
615 + """Whether every request in the batch has the same query length.
616 +
617 + Shape test only; use ``get_uniform_decode_token_count`` to classify a
618 + scheduled batch, since a prompt chunk can have a decode batch's shape.
619 + """
620 + return num_reqs > 0 and num_tokens == max_query_len * num_reqs
621 +
622 +
623 +def get_uniform_decode_token_count(
624 + num_reqs: int, num_tokens: int, max_query_len: int, has_prefill: bool
625 +) -> int | None:
626 + """Per-request token count of a uniform decode batch, or None."""
627 + if not has_prefill and is_uniform_query_len(num_reqs, num_tokens, max_query_len):
628 + return max_query_len
629 + return None
630 +
631 +
632 def is_residual_scattered_for_sp(
633 vllm_config: VllmConfig, num_input_tokens: int
634 ) -> bool: