[Bugfix][PP] Never drop a decoding request from the sampled-token broadcast (#54436)
Signed-off-by: ArcheyChen <45146080+ArcheyChen@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
Ausar committed
Sep 1, 2026 at 07:49 UTC
6bafc049aae6c26e210162630914ee9177a4b586
6 files changed
+87
-22
tests/v1/worker/test_gpu_batch_shard.py
-1
@@ -96,7 +96,6 @@ def _make_batch(
96
num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
97
is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
98
has_prefill=False,
99
- max_seq_len_np=None,
99
input_ids=torch.zeros(num_tokens, dtype=torch.int32, device=DEVICE),
100
positions=torch.arange(num_tokens, dtype=torch.int64, device=DEVICE),
101
is_padding=torch.zeros(num_tokens, dtype=torch.bool, device=DEVICE),
tests/v1/worker/test_pp_utils.py
new
+83
@@ -0,0 +1,83 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Which rows the PP sampled-token broadcast must carry."""
4
+
5
+from unittest.mock import Mock
6
+
7
+import numpy as np
8
+
9
+from vllm.v1.worker.gpu import pp_utils
10
+
11
+
12
+def _batch(num_computed, prefill_len, num_scheduled):
13
+ return Mock(
14
+ num_reqs=len(num_computed),
15
+ num_computed_tokens_np=np.array(num_computed, dtype=np.int32),
16
+ prefill_len_np=np.array(prefill_len, dtype=np.int32),
17
+ num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32),
18
+ )
19
+
20
+
21
+def test_excludes_non_final_prefill_chunks():
22
+ """Unchanged behaviour: a chunk that does not finish its prefill is skipped."""
23
+ # Row 0 is a middle prefill chunk and produces no sample; row 1 finishes its
24
+ # prefill this step and therefore does.
25
+ batch = _batch(
26
+ num_computed=[512, 1000],
27
+ prefill_len=[4096, 1004],
28
+ num_scheduled=[448, 4],
29
+ )
30
+
31
+ mask = pp_utils.compute_need_sampled_mask(batch)
32
+
33
+ assert mask is not None
34
+ assert mask.tolist() == [False, True]
35
+
36
+
37
+def test_none_when_no_row_samples():
38
+ """Unchanged behaviour: an all-prefill batch needs no broadcast at all."""
39
+ batch = _batch(
40
+ num_computed=[0, 512],
41
+ prefill_len=[4096, 4096],
42
+ num_scheduled=[448, 448],
43
+ )
44
+
45
+ assert pp_utils.compute_need_sampled_mask(batch) is None
46
+
47
+
48
+def test_keeps_decoding_request_past_its_length_cap():
49
+ """A decoding request must never be dropped from the broadcast.
50
+
51
+ Speculative decoding advances `num_computed_tokens` several tokens per step,
52
+ so it can overrun `prompt_len + max_tokens` while the scheduler is still
53
+ running the request. Predicting "this one is finishing" and skipping its
54
+ broadcast freezes the earlier pipeline stages' `last_sampled_tokens` and
55
+ `draft_tokens` while the last rank keeps advancing its own, and the stages
56
+ then diverge permanently.
57
+ """
58
+ batch = _batch(
59
+ # 14176 computed tokens is already past this request's own
60
+ # prompt_len + max_tokens; the scheduler is still running it.
61
+ num_computed=[14176],
62
+ prefill_len=[12175],
63
+ num_scheduled=[8],
64
+ )
65
+
66
+ mask = pp_utils.compute_need_sampled_mask(batch)
67
+
68
+ assert mask is not None
69
+ assert mask.tolist() == [True]
70
+
71
+
72
+def test_decode_row_ahead_of_a_prefill_chunk():
73
+ """Row order does not matter: only whether the row finishes its prefill."""
74
+ batch = _batch(
75
+ num_computed=[10, 512],
76
+ prefill_len=[8, 4096],
77
+ num_scheduled=[1, 448],
78
+ )
79
+
80
+ mask = pp_utils.compute_need_sampled_mask(batch)
81
+
82
+ assert mask is not None
83
+ assert mask.tolist() == [True, False]
vllm/v1/worker/gpu/input_batch.py
-4
@@ -85,9 +85,6 @@ class InputBatch:
85
# == np.any(is_prefilling_np)
86
has_prefill: bool
87
88
- # [num_reqs] only populated when pipeline parallelism is enabled.
89
- max_seq_len_np: np.ndarray | None
90
-
88
# [num_tokens_after_padding]
89
input_ids: torch.Tensor
90
# [num_tokens_after_padding]
@@ -194,7 +191,6 @@ class InputBatch:
191
num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
192
is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
193
has_prefill=False,
197
- max_seq_len_np=None,
194
input_ids=input_ids,
195
positions=positions,
196
is_padding=is_padding,
vllm/v1/worker/gpu/model_runner.py
-6
@@ -1311,11 +1311,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1311
)
1312
seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)
1313
1314
- max_seq_len_np = None
1315
- if self.use_pp:
1316
- # max_seq_len is only consumed by the PP `compute_need_sampled_mask`
1317
- max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]
1318
-
1314
prompt_lens = None
1315
if self.model_config.rswa_window is not None:
1316
# prompt_lens is only used in R-SWA case.
@@ -1344,7 +1339,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1339
num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,
1340
is_prefilling_np=batch_req_state.is_prefilling_np,
1341
has_prefill=batch_req_state.has_prefill,
1347
- max_seq_len_np=max_seq_len_np,
1342
input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
1343
positions=self.input_buffers.positions[:num_tokens_after_padding],
1344
is_padding=self.input_buffers.is_padding[:num_tokens_after_padding],
vllm/v1/worker/gpu/pcp_manager.py
-3
@@ -575,9 +575,6 @@ class PCPManager:
575
num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,
576
is_prefilling_np=local_is_prefilling_np,
577
has_prefill=bool(local_is_prefilling_np.any()),
578
- max_seq_len_np=global_batch.max_seq_len_np[local_to_global_batch_req_idx_np]
579
- if global_batch.max_seq_len_np is not None
580
- else None,
578
input_ids=input_buffers.input_ids[:num_local_tokens_padded],
579
positions=input_buffers.positions[:num_local_tokens_padded],
580
is_padding=is_padding,
vllm/v1/worker/gpu/pp_utils.py
+4
-8
@@ -34,19 +34,15 @@ class PendingRecv:
34
35
def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
36
"""Return a bool array of shape `[input_batch.num_reqs]` marking requests
37
- with outputs that might be needed in a subsequent (decode) step.
38
- Returns None if no sampled outputs are needed in the requests' next step."""
37
+ that produce a sampled token this step, and therefore must have that token
38
+ (and the draft block proposed from it) propagated to the earlier PP stages.
39
+ Returns None if no request in the batch produces a sample."""
40
41
old_computed = input_batch.num_computed_tokens_np
42
prefill_len = input_batch.prefill_len_np
42
- max_seq_len = input_batch.max_seq_len_np
43
- assert max_seq_len is not None # always populated under PP
43
# Exclude non-final prefill chunks (they don't produce a sample).
44
produces_sample = old_computed + input_batch.num_scheduled_tokens >= prefill_len
46
- # Exclude requests that we know are finished.
47
- not_finishing = np.maximum(old_computed, prefill_len) + 1 < max_seq_len
48
- need_sampled_mask = produces_sample & not_finishing
49
- return need_sampled_mask if need_sampled_mask.any() else None
45
+ return produces_sample if produces_sample.any() else None
46
47
48
class PPHandler: