[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>

tests/v1/worker/test_gpu_batch_ordering.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Batch ordering in Model Runner V2 (vllm.v1.worker.gpu).
+
+split_decodes_and_prefills assumes decode -> short_extend -> prefill request
+ordering. With spec decode (decode_query_len > 1), a shorter chunked-prefill
+tail sorted in front of the uniform decodes misclassifies every decode as a
+prefill.
+"""
+
+import torch
+
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+from vllm.v1.worker.gpu.model_runner import sort_batch_req_ids
+
+
+def _make_common_attn_metadata(query_lens: list[int]) -> CommonAttentionMetadata:
+    num_reqs = len(query_lens)
+    num_tokens = sum(query_lens)
+    query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32)
+    torch.cumsum(
+        torch.tensor(query_lens, dtype=torch.int32), 0, out=query_start_loc[1:]
+    )
+    seq_lens = torch.tensor([1000 + q for q in query_lens], dtype=torch.int32)
+    return CommonAttentionMetadata(
+        query_start_loc=query_start_loc,
+        query_start_loc_cpu=query_start_loc,
+        seq_lens=seq_lens,
+        seq_lens_cpu_upper_bound=seq_lens,
+        max_seq_len=int(seq_lens.max()),
+        num_reqs=num_reqs,
+        num_actual_tokens=num_tokens,
+        max_query_len=max(query_lens),
+        block_table_tensor=torch.zeros(num_reqs, 1, dtype=torch.int32),
+        slot_mapping=torch.zeros(num_tokens, dtype=torch.int64),
+    )
+
+
+def test_sort_batch_req_ids_no_spec():
+    # decode_query_len == 1: plain ascending order (decodes first).
+    num_tokens_per_req = {"p1": 100, "d1": 1, "p2": 7, "d2": 1}
+    assert sort_batch_req_ids(num_tokens_per_req, 1) == ["d1", "d2", "p2", "p1"]
+
+
+def test_sort_batch_req_ids_spec_decode():
+    # decode_query_len == 2 (MTP k=1): uniform decodes lead, then the 1-token
+    # chunked-prefill tail, then longer prefills.
+    num_tokens_per_req = {"tail": 1, "d1": 2, "p1": 100, "d2": 2}
+    assert sort_batch_req_ids(num_tokens_per_req, 2) == ["d1", "d2", "tail", "p1"]
+
+
+def test_spec_decodes_lead_short_prefill_tail():
+    # With the fixed ordering, split_decodes_and_prefills classifies the
+    # uniform 2-token decodes as decodes even when a 1-token prefill tail is
+    # in the batch (indexer-style: require_uniform, threshold=1+k).
+    num_tokens_per_req = {"tail": 1, **{f"d{i}": 2 for i in range(8)}}
+    req_ids = sort_batch_req_ids(num_tokens_per_req, 2)
+    query_lens = [num_tokens_per_req[r] for r in req_ids]
+    assert query_lens == [2] * 8 + [1]
+
+    num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
+        split_decodes_and_prefills(
+            _make_common_attn_metadata(query_lens),
+            decode_threshold=2,
+            require_uniform=True,
+        )
+    )
+    assert (num_decodes, num_prefills) == (8, 1)
+    assert (num_decode_tokens, num_prefill_tokens) == (16, 1)

vllm/model_executor/models/deepseek_v2.py

@@ -1156,7 +1156,12 @@ class DeepseekV2MLAAttention(nn.Module):
             cache_config,
             quant_config,
             prefix,
-            skip_topk=_skip_topk,
+            # MTP layers must never start with skip_topk=True: their indexer
+            # computes indices at draft step 0, and the runtime toggle
+            # (set_skip_topk, index_share_for_mtp_iteration) only exists in
+            # the V1 proposer. A frozen True would leave the draft reading a
+            # never-written topk buffer.
+            skip_topk=_skip_topk and not is_mtp_layer,
         )
 
     def forward(

vllm/v1/worker/gpu/model_runner.py

@@ -854,9 +854,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         num_tokens_per_req = scheduler_output.num_scheduled_tokens
         num_reqs = len(num_tokens_per_req)
 
-        # Decode first, then prefill.
         # batch_idx -> req_id
-        req_ids = sorted(num_tokens_per_req, key=num_tokens_per_req.get)  # type: ignore[arg-type]
+        req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)
         numtoks_iter = map(num_tokens_per_req.get, req_ids)
         num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
 
@@ -1607,3 +1606,12 @@ class ExecuteModelState(NamedTuple):
     hidden_states: torch.Tensor | None
     aux_hidden_states: list[torch.Tensor] | None
     finished_req_ids: set[str]
+
+
+def sort_batch_req_ids(
+    num_tokens_per_req: dict[str, int], decode_query_len: int
+) -> list[str]:
+    # Order decode -> short_extend -> prefill; split_decodes_and_prefills
+    # relies on uniform decodes (query_len == decode_query_len) leading.
+    key = lambda r: ((num := num_tokens_per_req[r]) != decode_query_len, num)
+    return sorted(num_tokens_per_req, key=key)