diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b4c470b47..6ca70d5ea 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index de2585ea7..644f29560 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3884ca164..2824e562f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 5e5e57cac..76f5cc13a 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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, *, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 0fd05931f..b802dda43 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/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()) diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index ba545ccf7..67de78596 100644 --- a/python/sglang/srt/managers/utils.py +++ b/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