From 7314360c9eebb3b4b460ec49f6606ef0a08fcae2 Mon Sep 17 00:00:00 2001 From: LostFox11 <140548160+LostFox11@users.noreply.github.com> Date: Thu, 1 Oct 2026 10:18:55 +0800 Subject: [PATCH] [Perf][PP] Skip sampled-token broadcasts whose requests leave the engine (#58542) Signed-off-by: LostFox11 Signed-off-by: Nick Hill Co-authored-by: LostFox11 Co-authored-by: Nick Hill --- tests/v1/worker/test_gpu_batch_ordering.py | 2 + tests/v1/worker/test_pp_utils.py | 87 +++++++++++++++++++++- vllm/v1/worker/gpu/input_batch.py | 4 + vllm/v1/worker/gpu/model_runner.py | 7 ++ vllm/v1/worker/gpu/pcp_manager.py | 1 + vllm/v1/worker/gpu/pp_utils.py | 15 ++-- vllm/v1/worker/gpu/ubatch_utils.py | 1 + 7 files changed, 111 insertions(+), 6 deletions(-) diff --git a/tests/v1/worker/test_gpu_batch_ordering.py b/tests/v1/worker/test_gpu_batch_ordering.py index db9e080e4e..498ff7a728 100644 --- a/tests/v1/worker/test_gpu_batch_ordering.py +++ b/tests/v1/worker/test_gpu_batch_ordering.py @@ -57,6 +57,7 @@ def _make_runner( runner.decode_query_len = decode_query_len runner.adaptive_verification = None runner.pcp_manager = None + runner.use_pp = False runner.req_states = SimpleNamespace( req_id_to_index={req_id: i for i, req_id in enumerate(req_states)}, # 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(): runner.decode_query_len = 8 runner.adaptive_verification = None runner.pcp_manager = None + runner.use_pp = False # State arrays in state-index order: a prefilling request, then two decodes. runner.req_states = SimpleNamespace( req_id_to_index={"prefilling": 0, "decode_a": 1, "decode_b": 2}, diff --git a/tests/v1/worker/test_pp_utils.py b/tests/v1/worker/test_pp_utils.py index cfc92479f5..a6209abb97 100644 --- a/tests/v1/worker/test_pp_utils.py +++ b/tests/v1/worker/test_pp_utils.py @@ -9,12 +9,26 @@ import numpy as np from vllm.v1.worker.gpu import pp_utils -def _batch(num_computed, prefill_len, num_scheduled): +def _batch( + num_computed, + prefill_len, + num_scheduled, + is_prefilling=None, + max_seq_len=None, +): return Mock( num_reqs=len(num_computed), num_computed_tokens_np=np.array(num_computed, dtype=np.int32), prefill_len_np=np.array(prefill_len, dtype=np.int32), num_scheduled_tokens=np.array(num_scheduled, dtype=np.int32), + is_prefilling_np=( + np.array(is_prefilling, dtype=np.bool_) + if is_prefilling is not None + else None + ), + max_seq_len_np=( + np.array(max_seq_len, dtype=np.int32) if max_seq_len is not None else None + ), ) @@ -81,3 +95,74 @@ def test_decode_row_ahead_of_a_prefill_chunk(): assert mask is not None assert mask.tolist() == [True, False] + + +def test_excludes_final_prefill_chunk_reaching_its_length_cap(): + """A max_tokens=1 final chunk hands its token off and leaves the engine. + + Disaggregated prefill submits max_tokens=1, so the request finishes with + the final chunk's single sample and no follow-up step on this engine + reads the broadcast payload. + """ + batch = _batch( + num_computed=[1000], + prefill_len=[1004], + num_scheduled=[4], + is_prefilling=[True], + max_seq_len=[1005], + ) + + assert pp_utils.compute_need_sampled_mask(batch) is None + + +def test_keeps_final_prefill_chunk_with_generation_budget_left(): + """The same final chunk, but the request keeps decoding here.""" + batch = _batch( + num_computed=[1000], + prefill_len=[1004], + num_scheduled=[4], + is_prefilling=[True], + max_seq_len=[1004 + 256], + ) + + mask = pp_utils.compute_need_sampled_mask(batch) + + assert mask is not None + assert mask.tolist() == [True] + + +def test_length_cap_overrun_decode_row_is_kept(): + """Decoding rows are never excluded, even past their length cap. + + Speculative decoding can transiently overrun prompt_len + max_tokens + while the scheduler still runs the request; dropping the row would + desynchronize the PP stages. + """ + batch = _batch( + num_computed=[14176], + prefill_len=[12175], + num_scheduled=[8], + is_prefilling=[False], + max_seq_len=[12176], + ) + + mask = pp_utils.compute_need_sampled_mask(batch) + + assert mask is not None + assert mask.tolist() == [True] + + +def test_finishing_prefill_row_does_not_drop_a_sibling_row(): + """Only finishing rows are excluded; a decoding sibling keeps its row.""" + batch = _batch( + num_computed=[1000, 10], + prefill_len=[1004, 8], + num_scheduled=[4, 1], + is_prefilling=[True, False], + max_seq_len=[1005, 1024], + ) + + mask = pp_utils.compute_need_sampled_mask(batch) + + assert mask is not None + assert mask.tolist() == [False, True] diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 1fc3737837..e2f68ad598 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -103,6 +103,9 @@ class InputBatch: # [num_reqs] per-request prompt length, only populated for R-SWA. prompt_lens: torch.Tensor | None + # [num_reqs] only populated for non-dummy pipeline-parallel batches. + max_seq_len_np: np.ndarray | None = None + # Longest query the batch may contain. Set when a cudagraph descriptor promises # a query length this batch's own split does not reach, so attention metadata # stays valid for every replay the graph serves. @@ -203,6 +206,7 @@ class InputBatch: num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32), is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_), has_prefill=False, + max_seq_len_np=None, decode_graph_eligible=True, input_ids=input_ids, positions=positions, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ebc7a5058d..14ff11f448 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1274,6 +1274,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): ] is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np has_prefill = bool(is_prefilling_np.any()) + max_seq_len_np = None + if self.use_pp and has_prefill: + # consumed by the PP `compute_need_sampled_mask` for batches with prefill. + max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np] num_draft_tokens_np = None if draft_tokens: @@ -1312,6 +1316,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): has_prefill=has_prefill, prefill_runs_as_decode_np=prefill_runs_as_decode_np, decode_graph_eligible=decode_graph_eligible, + max_seq_len_np=max_seq_len_np, ) return batch_state, get_uniform_decode_token_count( num_reqs, num_toks, max_query_len, decode_graph_eligible @@ -1499,6 +1504,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): has_prefill=batch_req_state.has_prefill, prefill_runs_as_decode_np=batch_req_state.prefill_runs_as_decode_np, decode_graph_eligible=batch_req_state.decode_graph_eligible, + max_seq_len_np=batch_req_state.max_seq_len_np, input_ids=self.input_buffers.input_ids[:num_tokens_after_padding], positions=self.input_buffers.positions[:num_tokens_after_padding], is_padding=is_padding, @@ -2392,6 +2398,7 @@ class BatchReqState(NamedTuple): prefill_len_np: np.ndarray # [num_reqs] num_computed_prefill_tokens_np: np.ndarray # [num_reqs] is_prefilling_np: np.ndarray # [num_reqs] + max_seq_len_np: np.ndarray | None # [num_reqs], None in non-pp case has_prefill: bool prefill_runs_as_decode_np: np.ndarray | None # [num_reqs] decode_graph_eligible: bool diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index f740e4caf3..89e6b1064b 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -649,6 +649,7 @@ class PCPManager: prefill_len_np=local_prefill_len_np, num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np, is_prefilling_np=local_is_prefilling_np, + max_seq_len_np=None, has_prefill=local_has_prefill, decode_graph_eligible=not local_has_prefill, prefill_runs_as_decode_np=None, diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 0a246c98ce..818d80759a 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -37,14 +37,19 @@ class PendingRecv: def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: """Return a bool array of shape `[input_batch.num_reqs]` marking requests - that produce a sampled token this step, and therefore must have that token - (and the draft block proposed from it) propagated to the earlier PP stages. - Returns None if no request in the batch produces a sample.""" + with outputs that might be needed in a subsequent (decode) step. + Returns None if no sampled outputs are needed in the requests' next step.""" old_computed = input_batch.num_computed_tokens_np prefill_len = input_batch.prefill_len_np # Exclude non-final prefill chunks (they don't produce a sample). - produces_sample = old_computed + input_batch.num_scheduled_tokens >= prefill_len - return produces_sample if produces_sample.any() else None + need_sampled_mask = old_computed + input_batch.num_scheduled_tokens >= prefill_len + if input_batch.max_seq_len_np is not None: + # Also exclude final prefill chunks whose single sampled token reaches + # the request's length cap. + finished_prefill = prefill_len + 1 >= input_batch.max_seq_len_np + finished_prefill &= input_batch.is_prefilling_np + need_sampled_mask &= ~finished_prefill + return need_sampled_mask if need_sampled_mask.any() else None class PPHandler: diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 6889103ed1..3fa9783bf5 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -194,6 +194,7 @@ def _slice_input_batch( if input_batch.prompt_lens is None else input_batch.prompt_lens[req_start:req_stop] ), + max_seq_len_np=None, fast_prefill=None, )