From 85b3a7264b6c4e4e89a1c45a2c4ccfd1b8c342dc Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 9 Jul 2026 06:26:27 -0700 Subject: [PATCH] [Bugfix][Model Runner V2] Order uniform decodes first so spec decodes aren't misclassified as prefills (#47381) Signed-off-by: Woosuk Kwon Signed-off-by: Nick Hill Co-authored-by: Claude Fable 5 Co-authored-by: Nick Hill --- tests/v1/worker/test_gpu_batch_ordering.py | 70 ++++++++++++++++++++++ vllm/model_executor/models/deepseek_v2.py | 7 ++- vllm/v1/worker/gpu/model_runner.py | 12 +++- 3 files changed, 86 insertions(+), 3 deletions(-) create mode 100644 tests/v1/worker/test_gpu_batch_ordering.py diff --git a/tests/v1/worker/test_gpu_batch_ordering.py b/tests/v1/worker/test_gpu_batch_ordering.py new file mode 100644 index 0000000000..2c2115dc81 --- /dev/null +++ b/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) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 4290609cbd..ac187151a0 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/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( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f5f51a8263..1f5dd4d2fb 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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)