[Bugfix][Model Runner V2] Order uniform decodes first so spec decodes aren't misclassified as prefills (#47381)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Woosuk Kwon committed
Jul 9, 2026 at 06:26 UTC
85b3a7264b6c4e4e89a1c45a2c4ccfd1b8c342dc
3 files changed
+86
-3
tests/v1/worker/test_gpu_batch_ordering.py
new
+70
@@ -0,0 +1,70 @@
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).
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
+
11
+import torch
12
+
13
+from vllm.v1.attention.backend import CommonAttentionMetadata
14
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
15
+from vllm.v1.worker.gpu.model_runner import sort_batch_req_ids
16
+
17
+
18
+def _make_common_attn_metadata(query_lens: list[int]) -> CommonAttentionMetadata:
19
+ num_reqs = len(query_lens)
20
+ num_tokens = sum(query_lens)
21
+ query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32)
22
+ torch.cumsum(
23
+ torch.tensor(query_lens, dtype=torch.int32), 0, out=query_start_loc[1:]
24
+ )
25
+ seq_lens = torch.tensor([1000 + q for q in query_lens], dtype=torch.int32)
26
+ return CommonAttentionMetadata(
27
+ query_start_loc=query_start_loc,
28
+ query_start_loc_cpu=query_start_loc,
29
+ seq_lens=seq_lens,
30
+ seq_lens_cpu_upper_bound=seq_lens,
31
+ max_seq_len=int(seq_lens.max()),
32
+ num_reqs=num_reqs,
33
+ num_actual_tokens=num_tokens,
34
+ max_query_len=max(query_lens),
35
+ block_table_tensor=torch.zeros(num_reqs, 1, dtype=torch.int32),
36
+ slot_mapping=torch.zeros(num_tokens, dtype=torch.int64),
37
+ )
38
+
39
+
40
+def test_sort_batch_req_ids_no_spec():
41
+ # decode_query_len == 1: plain ascending order (decodes first).
42
+ num_tokens_per_req = {"p1": 100, "d1": 1, "p2": 7, "d2": 1}
43
+ assert sort_batch_req_ids(num_tokens_per_req, 1) == ["d1", "d2", "p2", "p1"]
44
+
45
+
46
+def test_sort_batch_req_ids_spec_decode():
47
+ # decode_query_len == 2 (MTP k=1): uniform decodes lead, then the 1-token
48
+ # chunked-prefill tail, then longer prefills.
49
+ num_tokens_per_req = {"tail": 1, "d1": 2, "p1": 100, "d2": 2}
50
+ assert sort_batch_req_ids(num_tokens_per_req, 2) == ["d1", "d2", "tail", "p1"]
51
+
52
+
53
+def test_spec_decodes_lead_short_prefill_tail():
54
+ # With the fixed ordering, split_decodes_and_prefills classifies the
55
+ # uniform 2-token decodes as decodes even when a 1-token prefill tail is
56
+ # in the batch (indexer-style: require_uniform, threshold=1+k).
57
+ num_tokens_per_req = {"tail": 1, **{f"d{i}": 2 for i in range(8)}}
58
+ req_ids = sort_batch_req_ids(num_tokens_per_req, 2)
59
+ query_lens = [num_tokens_per_req[r] for r in req_ids]
60
+ assert query_lens == [2] * 8 + [1]
61
+
62
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
63
+ split_decodes_and_prefills(
64
+ _make_common_attn_metadata(query_lens),
65
+ decode_threshold=2,
66
+ require_uniform=True,
67
+ )
68
+ )
69
+ assert (num_decodes, num_prefills) == (8, 1)
70
+ assert (num_decode_tokens, num_prefill_tokens) == (16, 1)
vllm/model_executor/models/deepseek_v2.py
+6
-1
@@ -1156,7 +1156,12 @@ class DeepseekV2MLAAttention(nn.Module):
1156
cache_config,
1157
quant_config,
1158
prefix,
1159
- skip_topk=_skip_topk,
1159
+ # MTP layers must never start with skip_topk=True: their indexer
1160
+ # computes indices at draft step 0, and the runtime toggle
1161
+ # (set_skip_topk, index_share_for_mtp_iteration) only exists in
1162
+ # the V1 proposer. A frozen True would leave the draft reading a
1163
+ # never-written topk buffer.
1164
+ skip_topk=_skip_topk and not is_mtp_layer,
1165
)
1166
1167
def forward(
vllm/v1/worker/gpu/model_runner.py
+10
-2
@@ -854,9 +854,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
854
num_tokens_per_req = scheduler_output.num_scheduled_tokens
855
num_reqs = len(num_tokens_per_req)
856
857
- # Decode first, then prefill.
857
# batch_idx -> req_id
859
- req_ids = sorted(num_tokens_per_req, key=num_tokens_per_req.get) # type: ignore[arg-type]
858
+ req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)
859
numtoks_iter = map(num_tokens_per_req.get, req_ids)
860
num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
861
@@ -1607,3 +1606,12 @@ class ExecuteModelState(NamedTuple):
1606
hidden_states: torch.Tensor | None
1607
aux_hidden_states: list[torch.Tensor] | None
1608
finished_req_ids: set[str]
1609
+
1610
+
1611
+def sort_batch_req_ids(
1612
+ num_tokens_per_req: dict[str, int], decode_query_len: int
1613
+) -> list[str]:
1614
+ # Order decode -> short_extend -> prefill; split_decodes_and_prefills
1615
+ # relies on uniform decodes (query_len == decode_query_len) leading.
1616
+ key = lambda r: ((num := num_tokens_per_req[r]) != decode_query_len, num)
1617
+ return sorted(num_tokens_per_req, key=key)