[Perf][PP] Skip sampled-token broadcasts whose requests leave the engine (#58542)
Signed-off-by: LostFox11 <wangziyue17@huawei.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: LostFox11 <wangziyue17@huawei.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
LostFox11 committed
Oct 1, 2026 at 10:18 UTC
7314360c9eebb3b4b460ec49f6606ef0a08fcae2
7 files changed
+111
-6
tests/v1/worker/test_gpu_batch_ordering.py
+2
@@ -57,6 +57,7 @@ def _make_runner(
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).
@@ -246,6 +247,7 @@ def test_uniform_decode_uses_state_index_not_batch_position():
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},
tests/v1/worker/test_pp_utils.py
+86
-1
@@ -9,12 +9,26 @@ import numpy as np
9
from vllm.v1.worker.gpu import pp_utils
10
11
12
-def _batch(num_computed, prefill_len, num_scheduled):
12
+def _batch(
13
+ num_computed,
14
+ prefill_len,
15
+ num_scheduled,
16
+ is_prefilling=None,
17
+ max_seq_len=None,
18
+):
19
return Mock(
20
num_reqs=len(num_computed),
21
num_computed_tokens_np=np.array(num_computed, dtype=np.int32),
22
prefill_len_np=np.array(prefill_len, dtype=np.int32),
23
num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32),
24
+ is_prefilling_np=(
25
+ np.array(is_prefilling, dtype=np.bool_)
26
+ if is_prefilling is not None
27
+ else None
28
+ ),
29
+ max_seq_len_np=(
30
+ np.array(max_seq_len, dtype=np.int32) if max_seq_len is not None else None
31
+ ),
32
)
33
34
@@ -81,3 +95,74 @@ def test_decode_row_ahead_of_a_prefill_chunk():
95
96
assert mask is not None
97
assert mask.tolist() == [True, False]
98
+
99
+
100
+def test_excludes_final_prefill_chunk_reaching_its_length_cap():
101
+ """A max_tokens=1 final chunk hands its token off and leaves the engine.
102
+
103
+ Disaggregated prefill submits max_tokens=1, so the request finishes with
104
+ the final chunk's single sample and no follow-up step on this engine
105
+ reads the broadcast payload.
106
+ """
107
+ batch = _batch(
108
+ num_computed=[1000],
109
+ prefill_len=[1004],
110
+ num_scheduled=[4],
111
+ is_prefilling=[True],
112
+ max_seq_len=[1005],
113
+ )
114
+
115
+ assert pp_utils.compute_need_sampled_mask(batch) is None
116
+
117
+
118
+def test_keeps_final_prefill_chunk_with_generation_budget_left():
119
+ """The same final chunk, but the request keeps decoding here."""
120
+ batch = _batch(
121
+ num_computed=[1000],
122
+ prefill_len=[1004],
123
+ num_scheduled=[4],
124
+ is_prefilling=[True],
125
+ max_seq_len=[1004 + 256],
126
+ )
127
+
128
+ mask = pp_utils.compute_need_sampled_mask(batch)
129
+
130
+ assert mask is not None
131
+ assert mask.tolist() == [True]
132
+
133
+
134
+def test_length_cap_overrun_decode_row_is_kept():
135
+ """Decoding rows are never excluded, even past their length cap.
136
+
137
+ Speculative decoding can transiently overrun prompt_len + max_tokens
138
+ while the scheduler still runs the request; dropping the row would
139
+ desynchronize the PP stages.
140
+ """
141
+ batch = _batch(
142
+ num_computed=[14176],
143
+ prefill_len=[12175],
144
+ num_scheduled=[8],
145
+ is_prefilling=[False],
146
+ max_seq_len=[12176],
147
+ )
148
+
149
+ mask = pp_utils.compute_need_sampled_mask(batch)
150
+
151
+ assert mask is not None
152
+ assert mask.tolist() == [True]
153
+
154
+
155
+def test_finishing_prefill_row_does_not_drop_a_sibling_row():
156
+ """Only finishing rows are excluded; a decoding sibling keeps its row."""
157
+ batch = _batch(
158
+ num_computed=[1000, 10],
159
+ prefill_len=[1004, 8],
160
+ num_scheduled=[4, 1],
161
+ is_prefilling=[True, False],
162
+ max_seq_len=[1005, 1024],
163
+ )
164
+
165
+ mask = pp_utils.compute_need_sampled_mask(batch)
166
+
167
+ assert mask is not None
168
+ assert mask.tolist() == [False, True]
vllm/v1/worker/gpu/input_batch.py
+4
@@ -103,6 +103,9 @@ class InputBatch:
103
# [num_reqs] per-request prompt length, only populated for R-SWA.
104
prompt_lens: torch.Tensor | None
105
106
+ # [num_reqs] only populated for non-dummy pipeline-parallel batches.
107
+ max_seq_len_np: np.ndarray | None = None
108
+
109
# Longest query the batch may contain. Set when a cudagraph descriptor promises
110
# a query length this batch's own split does not reach, so attention metadata
111
# stays valid for every replay the graph serves.
@@ -203,6 +206,7 @@ class InputBatch:
206
num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
207
is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
208
has_prefill=False,
209
+ max_seq_len_np=None,
210
decode_graph_eligible=True,
211
input_ids=input_ids,
212
positions=positions,
vllm/v1/worker/gpu/model_runner.py
+7
@@ -1274,6 +1274,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1274
]
1275
is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
1276
has_prefill = bool(is_prefilling_np.any())
1277
+ max_seq_len_np = None
1278
+ if self.use_pp and has_prefill:
1279
+ # consumed by the PP `compute_need_sampled_mask` for batches with prefill.
1280
+ max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]
1281
1282
num_draft_tokens_np = None
1283
if draft_tokens:
@@ -1312,6 +1316,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1316
has_prefill=has_prefill,
1317
prefill_runs_as_decode_np=prefill_runs_as_decode_np,
1318
decode_graph_eligible=decode_graph_eligible,
1319
+ max_seq_len_np=max_seq_len_np,
1320
)
1321
return batch_state, get_uniform_decode_token_count(
1322
num_reqs, num_toks, max_query_len, decode_graph_eligible
@@ -1499,6 +1504,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1504
has_prefill=batch_req_state.has_prefill,
1505
prefill_runs_as_decode_np=batch_req_state.prefill_runs_as_decode_np,
1506
decode_graph_eligible=batch_req_state.decode_graph_eligible,
1507
+ max_seq_len_np=batch_req_state.max_seq_len_np,
1508
input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
1509
positions=self.input_buffers.positions[:num_tokens_after_padding],
1510
is_padding=is_padding,
@@ -2392,6 +2398,7 @@ class BatchReqState(NamedTuple):
2398
prefill_len_np: np.ndarray # [num_reqs]
2399
num_computed_prefill_tokens_np: np.ndarray # [num_reqs]
2400
is_prefilling_np: np.ndarray # [num_reqs]
2401
+ max_seq_len_np: np.ndarray | None # [num_reqs], None in non-pp case
2402
has_prefill: bool
2403
prefill_runs_as_decode_np: np.ndarray | None # [num_reqs]
2404
decode_graph_eligible: bool
vllm/v1/worker/gpu/pcp_manager.py
+1
@@ -649,6 +649,7 @@ class PCPManager:
649
prefill_len_np=local_prefill_len_np,
650
num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,
651
is_prefilling_np=local_is_prefilling_np,
652
+ max_seq_len_np=None,
653
has_prefill=local_has_prefill,
654
decode_graph_eligible=not local_has_prefill,
655
prefill_runs_as_decode_np=None,
vllm/v1/worker/gpu/pp_utils.py
+10
-5
@@ -37,14 +37,19 @@ class PendingRecv:
37
38
def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
39
"""Return a bool array of shape `[input_batch.num_reqs]` marking requests
40
- that produce a sampled token this step, and therefore must have that token
41
- (and the draft block proposed from it) propagated to the earlier PP stages.
42
- Returns None if no request in the batch produces a sample."""
40
+ with outputs that might be needed in a subsequent (decode) step.
41
+ Returns None if no sampled outputs are needed in the requests' next step."""
42
old_computed = input_batch.num_computed_tokens_np
43
prefill_len = input_batch.prefill_len_np
44
# Exclude non-final prefill chunks (they don't produce a sample).
46
- produces_sample = old_computed + input_batch.num_scheduled_tokens >= prefill_len
47
- return produces_sample if produces_sample.any() else None
45
+ need_sampled_mask = old_computed + input_batch.num_scheduled_tokens >= prefill_len
46
+ if input_batch.max_seq_len_np is not None:
47
+ # Also exclude final prefill chunks whose single sampled token reaches
48
+ # the request's length cap.
49
+ finished_prefill = prefill_len + 1 >= input_batch.max_seq_len_np
50
+ finished_prefill &= input_batch.is_prefilling_np
51
+ need_sampled_mask &= ~finished_prefill
52
+ return need_sampled_mask if need_sampled_mask.any() else None
53
54
55
class PPHandler:
vllm/v1/worker/gpu/ubatch_utils.py
+1
@@ -194,6 +194,7 @@ def _slice_input_batch(
194
if input_batch.prompt_lens is None
195
else input_batch.prompt_lens[req_start:req_stop]
196
),
197
+ max_seq_len_np=None,
198
fast_prefill=None,
199
)
200