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

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},

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]

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,

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

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,

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:

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,
     )