[PP] Skip PP output communication for pure chunked prefill batches (#26148)

python/sglang/srt/environ.py

@@ -288,6 +288,9 @@ class Envs:
     # Scheduler: others:
     SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1)  # in seconds. Set if you observe high memory accumulation over a long serving period.
     SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP = EnvBool(False)
+    # PP: skip output send/recv when the entire batch consists of non-final chunked prefill requests,
+    # since process_batch_result_prefill discards next_token_ids for those anyway.
+    SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM = EnvBool(False)
     SGLANG_SCHEDULER_MAX_RECV_PER_POLL = EnvInt(-1)
     SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False)
     SGLANG_RADIX_FORCE_MISS = EnvBool(False)

python/sglang/srt/managers/schedule_batch.py

@@ -1488,6 +1488,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
     # For chunked prefill in PP
     chunked_req: Optional[Req] = None
+    contains_last_prefill_chunk: bool = True
 
     # Sampling info
     sampling_info: SamplingBatchInfo = None

python/sglang/srt/managers/scheduler.py

@@ -2674,6 +2674,11 @@ class Scheduler(
             self.spec_algorithm,
             chunked_req=self.chunked_req,
         )
+
+        new_batch.contains_last_prefill_chunk = (
+            self.chunked_req is None or len(can_run_list) != 1
+        )
+
         self.max_prefill_bs = max(self.max_prefill_bs, len(can_run_list))
         if self.enable_hierarchical_cache:
             # todo (zhiqiang): disable cuda graph execution if hicache loading triggered

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -208,6 +208,8 @@ class SchedulerBatchResultProcessor:
             next_token_ids = next_token_ids.tolist()
             self._move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
 
+            self._validate_pp_skip_output_comm(batch, result)
+
             hidden_state_offset = 0
 
             # Check finish conditions
@@ -415,6 +417,47 @@ class SchedulerBatchResultProcessor:
         logprob_pt += num_input_logprobs
         return logprob_pt
 
+    @staticmethod
+    def _validate_pp_skip_output_comm(
+        batch: ScheduleBatch,
+        result: Union[GenerationBatchResult, EmbeddingBatchResult],
+    ):
+        """Validate PP skip output comm correctness.
+
+        - When skip=True: all reqs must be middle chunks (inflight_middle_chunks > 0)
+          so placeholder zeros are never consumed via req.output_ids.append().
+        - When skip=False: at least one req should consume next_token_ids
+          (inflight_middle_chunks <= 0), otherwise warn.
+        """
+        if not envs.SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM.get():
+            return
+
+        if not getattr(result, "skipped_output_comm", False):
+            if batch.forward_mode.is_extend() and not batch.forward_mode.is_prebuilt():
+                has_consumed_output = any(
+                    req.inflight_middle_chunks <= 0
+                    for req in batch.reqs
+                    if not req.finished() and not req.is_retracted
+                )
+                if not has_consumed_output and len(batch.reqs) > 0:
+                    chunks = list([r.inflight_middle_chunks for r in batch.reqs])
+                    logger.warning(
+                        f"PP non-skip output comm: no req consumed next_token_ids. "
+                        f"contains_last_prefill_chunk={batch.contains_last_prefill_chunk}, "
+                        f"num_reqs={len(batch.reqs)}, all inflight_middle_chunks={chunks}"
+                    )
+            return
+
+        for req in batch.reqs:
+            if not req.finished() and not req.is_retracted:
+                assert req.inflight_middle_chunks > 0, (
+                    f"PP skip output comm invariant violated: req {req.rid} "
+                    f"has inflight_middle_chunks={req.inflight_middle_chunks} "
+                    f"but output was skipped (contains_last_prefill_chunk="
+                    f"{batch.contains_last_prefill_chunk}). "
+                    f"Placeholder zeros would be appended to output_ids."
+                )
+
     def _append_prefill_hidden_states(
         self,
         *,

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -29,7 +29,11 @@ from sglang.srt.managers.utils import (
     get_logprob_dict_from_result,
     get_logprob_from_pp_outputs,
 )
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.forward_batch_info import (
+    ForwardBatch,
+    ForwardMode,
+    PPProxyTensors,
+)
 from sglang.srt.observability.req_time_stats import set_time_batch
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.utils import DynamicGradMode, broadcast_pyobj, point_to_point_pyobj
@@ -41,6 +45,18 @@ if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import Scheduler
 
 
+def _pp_can_skip_output_comm(batch: ScheduleBatch) -> bool:
+    """Check if output send/recv can be skipped for this batch."""
+    return (
+        envs.SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM.get()
+        and batch is not None
+        and batch.forward_mode == ForwardMode.EXTEND
+        and len(batch.reqs) == 1
+        and not batch.contains_last_prefill_chunk
+        and not batch.return_logprob
+    )
+
+
 @dataclass
 class PPBatchMetadata:
     can_run_cuda_graph: bool
@@ -1032,6 +1048,30 @@ class SchedulerPPMixin:
             ),
         )
 
+    def _pp_make_skip_output_result(
+        self: Scheduler,
+        batch: ScheduleBatch,
+        mb_metadata: Optional[PPBatchMetadata],
+    ):
+        bs = len(batch.reqs)
+        placeholder = torch.zeros(bs, dtype=torch.int64, device=self.device)
+        # next_pp_outputs = None so non-last ranks skip forwarding
+        # (pp_outputs is None gate). Placeholder carried in
+        # batch_result.next_token_ids for process_batch_result_prefill.
+        batch.output_ids = placeholder
+        batch_result = GenerationBatchResult(
+            logits_output=None,
+            pp_hidden_states_proxy_tensors=None,
+            next_token_ids=placeholder,
+            can_run_cuda_graph=(
+                mb_metadata.can_run_cuda_graph if mb_metadata else False
+            ),
+            skipped_output_comm=True,
+        )
+        d2h_event = self.device_module.Event()
+        d2h_event.record(self.device_module.current_stream())
+        return None, batch_result, d2h_event
+
     def _pp_prep_batch_result(
         self: Scheduler,
         batch: ScheduleBatch,
@@ -1076,9 +1116,13 @@ class SchedulerPPMixin:
         send_output_work = []
         if self.pp_group.is_last_rank:
             # send ready PP output to rank 0
-            if mbs[next_first_rank_mb_id] is not None:
+            target = mbs[next_first_rank_mb_id]
+            if target is not None:
                 q_event, pp_outputs_to_send = last_rank_comm_queue.popleft()
-                if not mbs[next_first_rank_mb_id].forward_mode.is_prebuilt():
+                if (
+                    not target.forward_mode.is_prebuilt()
+                    and not _pp_can_skip_output_comm(target)
+                ):
                     self.device_module.current_stream().wait_event(q_event)
                     with torch.profiler.record_function("send_res_dict_to_next_stage"):
                         send_output_work = self._pp_send_dict_to_next_stage(
@@ -1139,14 +1183,20 @@ class SchedulerPPMixin:
 
         def _do_recv():
             nonlocal next_pp_outputs, batch_result, d2h_event
-            if mbs[next_mb_id] is None or mbs[next_mb_id].forward_mode.is_prebuilt():
+            target = mbs[next_mb_id]
+            if target is None or target.forward_mode.is_prebuilt():
+                return
+            if _pp_can_skip_output_comm(target):
+                next_pp_outputs, batch_result, d2h_event = (
+                    self._pp_make_skip_output_result(target, mb_metadata[next_mb_id])
+                )
                 return
             with torch.profiler.record_function("recv_res_dict_from_prev_stage"):
                 next_pp_outputs = PPProxyTensors(self._pp_recv_dict_from_prev_stage())
             with self.copy_stream_ctx:
                 self.copy_stream.wait_stream(self.schedule_stream)
                 batch_result = self._pp_prep_batch_result(
-                    mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs
+                    target, mb_metadata[next_mb_id], next_pp_outputs
                 )
                 d2h_event = self.device_module.Event()
                 d2h_event.record(self.device_module.current_stream())

python/sglang/srt/managers/utils.py

@@ -32,6 +32,11 @@ class GenerationBatchResult:
     num_correct_drafts_per_req_cpu: Optional[List[int]] = None
     can_run_cuda_graph: bool = False
 
+    # PP skip output comm: True when output send/recv was skipped and
+    # next_token_ids are placeholder zeros. Used by process_batch_result_prefill
+    # to validate that skipped output is never consumed.
+    skipped_output_comm: bool = False
+
     # For output processing
     extend_input_len_per_req: Optional[List[int]] = None
     extend_logprob_start_len_per_req: Optional[List[int]] = None