| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 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 |
| 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: |
| 29 | num_reqs = len(query_lens) |
| 30 | num_tokens = sum(query_lens) |
| 31 | query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32) |
| 32 | torch.cumsum( |
| 33 | torch.tensor(query_lens, dtype=torch.int32), 0, out=query_start_loc[1:] |
| 34 | ) |
| 35 | seq_lens = torch.tensor([1000 + q for q in query_lens], dtype=torch.int32) |
| 36 | return CommonAttentionMetadata( |
| 37 | query_start_loc=query_start_loc, |
| 38 | query_start_loc_cpu=query_start_loc, |
| 39 | seq_lens=seq_lens, |
| 40 | seq_lens_cpu_upper_bound=seq_lens, |
| 41 | max_seq_len=int(seq_lens.max()), |
| 42 | num_reqs=num_reqs, |
| 43 | num_actual_tokens=num_tokens, |
| 44 | max_query_len=max(query_lens), |
| 45 | block_table_tensor=torch.zeros(num_reqs, 1, dtype=torch.int32), |
| 46 | slot_mapping=torch.zeros(num_tokens, dtype=torch.int64), |
| 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.adaptive_verification = None |
| 59 | runner.pcp_manager = None |
| 60 | runner.use_pp = False |
| 61 | runner.req_states = SimpleNamespace( |
| 62 | req_id_to_index={req_id: i for i, req_id in enumerate(req_states)}, |
| 63 | # The runner keeps this as min(num_computed_tokens, prefill_len). |
| 64 | num_computed_prefill_tokens=np.minimum(num_computed, prefill_lens), |
| 65 | prefill_len=SimpleNamespace(np=prefill_lens), |
| 66 | ) |
| 67 | return runner |
| 68 | |
| 69 | |
| 70 | def _uniform_token_count( |
| 71 | req_states: dict[str, tuple[int, int]], |
| 72 | query_len: int, |
| 73 | dummy_run: bool = False, |
| 74 | ) -> int | None: |
| 75 | """Classify via the runner's pre-dispatch gather, as execute_model does.""" |
| 76 | scheduler_output = SimpleNamespace( |
| 77 | num_scheduled_tokens={req_id: query_len for req_id in req_states}, |
| 78 | total_num_scheduled_tokens=query_len * len(req_states), |
| 79 | scheduled_spec_decode_tokens={}, |
| 80 | ) |
| 81 | runner = _make_runner(req_states, decode_query_len=query_len) |
| 82 | _, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, dummy_run) |
| 83 | return uniform_tok_count |
| 84 | |
| 85 | |
| 86 | def test_spec_decode_batch_is_uniform_decode(): |
| 87 | # decode_query_len == 8 (K=7): every request past its prefill. |
| 88 | decodes = {f"d{i}": (16, 16) for i in range(6)} |
| 89 | assert _uniform_token_count(decodes, 8) == 8 |
| 90 | |
| 91 | |
| 92 | def test_adaptive_verification_sizes_only_batches_with_drafts(): |
| 93 | decodes = {"d0": (16, 16), "d1": (16, 16)} |
| 94 | runner = _make_runner(decodes, decode_query_len=8) |
| 95 | manager = SimpleNamespace( |
| 96 | get_num_tokens=lambda _num_tokens_per_req, _draft_tokens: 12 |
| 97 | ) |
| 98 | runner.adaptive_verification = manager |
| 99 | scheduler_output = SimpleNamespace( |
| 100 | num_scheduled_tokens={req_id: 8 for req_id in decodes}, |
| 101 | total_num_scheduled_tokens=16, |
| 102 | scheduled_spec_decode_tokens={}, |
| 103 | ) |
| 104 | |
| 105 | state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) |
| 106 | assert state is not None |
| 107 | assert state.num_tokens == 16 |
| 108 | assert uniform_tok_count == 8 |
| 109 | |
| 110 | scheduler_output.scheduled_spec_decode_tokens = { |
| 111 | req_id: [1, 2] for req_id in decodes |
| 112 | } |
| 113 | state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) |
| 114 | assert state is not None |
| 115 | assert state.num_tokens == 12 |
| 116 | assert uniform_tok_count is None |
| 117 | |
| 118 | |
| 119 | def test_prompt_chunk_of_decode_query_len_is_not_uniform_decode(): |
| 120 | # Two prompt chunks whose query length coincides with the K+1 spec-decode |
| 121 | # query length must reject the batch (issue #49918). |
| 122 | batch = {f"d{i}": (16, 16) for i in range(6)} |
| 123 | batch.update({f"p{i}": (8, 40) for i in range(2)}) |
| 124 | assert _uniform_token_count(batch, 8) is None |
| 125 | |
| 126 | # Same collision without spec decoding: 1-token prompt vs 1-token decode. |
| 127 | assert _uniform_token_count({"p0": (0, 1)}, 1) is None |
| 128 | |
| 129 | # A fresh prompt of exactly decode_query_len tokens. |
| 130 | assert _uniform_token_count({"p0": (0, 8)}, 8) is None |
| 131 | |
| 132 | |
| 133 | def test_one_new_prompt_token_over_context_runs_as_decode(): |
| 134 | # A prompt tail padded with placeholder drafts and a one-token mid-prompt |
| 135 | # chunk both compute like decodes. A fresh one-token prompt and a real |
| 136 | # multi-token chunk don't. |
| 137 | batch = {"d0": (16, 16), "tail": (128, 129), "mid": (64, 129)} |
| 138 | batch.update({"fresh": (0, 1), "chunk": (120, 129)}) |
| 139 | runner = _make_runner(batch, decode_query_len=4) |
| 140 | scheduler_output = SimpleNamespace( |
| 141 | num_scheduled_tokens={"d0": 4, "tail": 4, "mid": 1, "fresh": 1, "chunk": 4}, |
| 142 | total_num_scheduled_tokens=14, |
| 143 | scheduled_spec_decode_tokens={"d0": [1, 2, 3], "tail": [-1] * 3}, |
| 144 | ) |
| 145 | state, _ = runner.gather_batch_req_state(scheduler_output, False) |
| 146 | assert state is not None |
| 147 | runs_as_decode = dict(zip(state.req_ids, state.prefill_runs_as_decode_np)) |
| 148 | assert runs_as_decode == { |
| 149 | "d0": False, |
| 150 | "tail": True, |
| 151 | "mid": True, |
| 152 | "fresh": False, |
| 153 | "chunk": False, |
| 154 | } |
| 155 | |
| 156 | |
| 157 | def test_one_token_prompt_tail_with_prior_context_is_uniform_decode(): |
| 158 | # Replaying the last prompt token over existing context (full prefix-cache |
| 159 | # hit or P/D KV import) computes exactly like a decode. |
| 160 | batch = {"d0": (16, 16), "d1": (16, 16), "p0": (128, 129)} |
| 161 | assert _uniform_token_count(batch, 1) == 1 |
| 162 | |
| 163 | # A fresh 1-token prompt alongside it has no prior state. |
| 164 | batch["p1"] = (0, 1) |
| 165 | assert _uniform_token_count(batch, 1) is None |
| 166 | |
| 167 | # Only one-token tails qualify, not K+1 chunks with prior context. |
| 168 | assert _uniform_token_count({"d0": (16, 16), "p0": (120, 129)}, 8) is None |
| 169 | |
| 170 | |
| 171 | def test_padded_prompt_tail_with_prior_context_is_uniform_decode(): |
| 172 | # The scheduler pads a one-token prompt tail with K placeholder drafts to |
| 173 | # keep the K+1 spec-decode shape; it verifies like any other decode. |
| 174 | batch = {"d0": (16, 16), "d1": (16, 16), "p0": (128, 129)} |
| 175 | runner = _make_runner(batch, decode_query_len=4) |
| 176 | scheduler_output = SimpleNamespace( |
| 177 | num_scheduled_tokens={req_id: 4 for req_id in batch}, |
| 178 | total_num_scheduled_tokens=12, |
| 179 | scheduled_spec_decode_tokens={req_id: [-1] * 3 for req_id in batch}, |
| 180 | ) |
| 181 | state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) |
| 182 | assert state.has_prefill and state.decode_graph_eligible |
| 183 | assert uniform_tok_count == 4 |
| 184 | |
| 185 | # Without the placeholder drafts, it is a real 4-token prompt chunk. |
| 186 | del scheduler_output.scheduled_spec_decode_tokens["p0"] |
| 187 | state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) |
| 188 | assert not state.decode_graph_eligible |
| 189 | assert uniform_tok_count is None |
| 190 | |
| 191 | |
| 192 | def test_dummy_batches_stay_uniform_decode(): |
| 193 | # Dummy batches have no request state to consult and must stay classified |
| 194 | # by shape alone; rejecting them would make capture and replay disagree. |
| 195 | scheduler_output = SimpleNamespace( |
| 196 | num_scheduled_tokens={f"_dummy_req_{i}": 8 for i in range(4)}, |
| 197 | total_num_scheduled_tokens=32, |
| 198 | ) |
| 199 | state, uniform_tok_count = _make_runner({}).gather_batch_req_state( |
| 200 | scheduler_output, True |
| 201 | ) |
| 202 | assert state is None |
| 203 | assert uniform_tok_count == 8 |
| 204 | |
| 205 | |
| 206 | def test_sort_batch_req_ids_no_spec(): |
| 207 | # decode_query_len == 1: plain ascending order (decodes first). |
| 208 | num_tokens_per_req = {"p1": 100, "d1": 1, "p2": 7, "d2": 1} |
| 209 | assert sort_batch_req_ids(num_tokens_per_req, {}, 1) == ["d1", "d2", "p2", "p1"] |
| 210 | |
| 211 | |
| 212 | def test_sort_batch_req_ids_spec_decode(): |
| 213 | # decode_query_len == 2 (MTP k=1): uniform decodes lead, then the 1-token |
| 214 | # chunked-prefill tail, then longer prefills. |
| 215 | num_tokens_per_req = {"tail": 1, "d1": 2, "p1": 100, "d2": 2} |
| 216 | assert sort_batch_req_ids(num_tokens_per_req, {}, 2) == ["d1", "d2", "tail", "p1"] |
| 217 | |
| 218 | |
| 219 | def test_spec_decodes_lead_short_prefill_tail(): |
| 220 | # With the fixed ordering, split_decodes_and_prefills classifies the |
| 221 | # uniform 2-token decodes as decodes even when a 1-token prefill tail is |
| 222 | # in the batch (indexer-style: require_uniform, threshold=1+k). |
| 223 | num_tokens_per_req = {"tail": 1, **{f"d{i}": 2 for i in range(8)}} |
| 224 | req_ids = sort_batch_req_ids(num_tokens_per_req, {}, 2) |
| 225 | query_lens = [num_tokens_per_req[r] for r in req_ids] |
| 226 | assert query_lens == [2] * 8 + [1] |
| 227 | |
| 228 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 229 | split_decodes_and_prefills( |
| 230 | _make_common_attn_metadata(query_lens), |
| 231 | decode_threshold=2, |
| 232 | require_uniform=True, |
| 233 | ) |
| 234 | ) |
| 235 | assert (num_decodes, num_prefills) == (8, 1) |
| 236 | assert (num_decode_tokens, num_prefill_tokens) == (16, 1) |
| 237 | |
| 238 | |
| 239 | def test_uniform_decode_uses_state_index_not_batch_position(): |
| 240 | """The gather must read each request's own state, not its batch slot. |
| 241 | |
| 242 | Batch position and state index diverge as requests come and go; here the |
| 243 | only prefilling request sits at state index 0, so a gather keyed on batch |
| 244 | position would read its row and wrongly reject the two decodes. |
| 245 | """ |
| 246 | runner: Any = GPUModelRunner.__new__(GPUModelRunner) |
| 247 | runner.decode_query_len = 8 |
| 248 | runner.adaptive_verification = None |
| 249 | runner.pcp_manager = None |
| 250 | runner.use_pp = False |
| 251 | # State arrays in state-index order: a prefilling request, then two decodes. |
| 252 | runner.req_states = SimpleNamespace( |
| 253 | req_id_to_index={"prefilling": 0, "decode_a": 1, "decode_b": 2}, |
| 254 | num_computed_prefill_tokens=np.array([8, 16, 16], dtype=np.int32), |
| 255 | prefill_len=SimpleNamespace(np=np.array([40, 16, 16], dtype=np.int32)), |
| 256 | ) |
| 257 | scheduler_output = SimpleNamespace( |
| 258 | num_scheduled_tokens={"decode_a": 8, "decode_b": 8}, |
| 259 | total_num_scheduled_tokens=16, |
| 260 | scheduled_spec_decode_tokens={}, |
| 261 | ) |
| 262 | |
| 263 | state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) |
| 264 | assert not state.has_prefill |
| 265 | assert uniform_tok_count == 8 |
| 266 | |
| 267 | # Scheduling the prefilling request alongside them must reject the batch. |
| 268 | scheduler_output.num_scheduled_tokens = { |
| 269 | "decode_a": 8, |
| 270 | "prefilling": 8, |
| 271 | "decode_b": 8, |
| 272 | } |
| 273 | scheduler_output.total_num_scheduled_tokens = 24 |
| 274 | state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False) |
| 275 | assert state.has_prefill |
| 276 | assert uniform_tok_count is None |
| 277 | |
| 278 | |
| 279 | def test_uniform_decode_predicate(): |
| 280 | # Shape and decode-graph eligibility must both pass. |
| 281 | assert get_uniform_decode_token_count(2, 16, 8, True) == 8 |
| 282 | assert get_uniform_decode_token_count(2, 16, 8, False) is None |
| 283 | # 12 tokens over 2 requests is no shared query length. |
| 284 | assert get_uniform_decode_token_count(2, 12, 8, True) is None |
| 285 | |
| 286 | |
| 287 | def test_no_speculator_dispatches_on_query_length_alone(): |
| 288 | """No speculator may pick its cudagraph from a shape-only test. |
| 289 | |
| 290 | Written over the whole package because speculators are added by copying an |
| 291 | existing one: the shape-only call already reappeared verbatim in a |
| 292 | speculator added after the original call sites were fixed. |
| 293 | """ |
| 294 | import vllm.v1.worker.gpu.spec_decode as spec_decode |
| 295 | |
| 296 | shape_only = {"get_uniform_token_count", "is_uniform_query_len"} |
| 297 | root = Path(spec_decode.__file__).parent |
| 298 | offenders = [] |
| 299 | for path in sorted(root.rglob("speculator.py")): |
| 300 | for node in ast.walk(ast.parse(path.read_text())): |
| 301 | if not isinstance(node, ast.Call): |
| 302 | continue |
| 303 | func = node.func |
| 304 | name = func.id if isinstance(func, ast.Name) else getattr(func, "attr", "") |
| 305 | if name in shape_only: |
| 306 | offenders.append(f"{path.relative_to(root)}:{node.lineno} {name}") |
| 307 | |
| 308 | assert not offenders, ( |
| 309 | "these speculators classify a decode batch by query length alone; use " |
| 310 | f"get_uniform_decode_token_count instead: {offenders}" |
| 311 | ) |