From 71cb9d0302e90bbadee274d908b0aca4e21f814e Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Tue, 20 Jan 2026 11:27:13 +0800 Subject: [PATCH] [PD] Fix PP dynamic chunking with DP attention (#17339) Signed-off-by: Shangming Cai --- .../sglang/srt/managers/scheduler_pp_mixin.py | 42 +++++++++++-------- 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index b785f8073..e793350ad 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -16,13 +16,18 @@ from sglang.srt.disaggregation.base.conn import KVPoll from sglang.srt.disaggregation.utils import DisaggregationMode, poll_and_all_reduce from sglang.srt.distributed.parallel_state import P2PWork from sglang.srt.environ import envs +from sglang.srt.layers.dp_attention import ( + get_attention_dp_rank, + get_attention_dp_size, + is_dp_attention_enabled, +) from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.utils import ( GenerationBatchResult, get_logprob_dict_from_result, get_logprob_from_pp_outputs, ) -from sglang.srt.model_executor.forward_batch_info import PPProxyTensors +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import DynamicGradMode, broadcast_pyobj, point_to_point_pyobj @@ -536,6 +541,8 @@ class SchedulerPPMixin: latencies: List[float] = [] if self.pp_group.is_first_rank: + model_runner = self.tp_worker.model_runner + model_config = model_runner.model_config input_ids_list = [] for i in range(128): chunk_size = int( @@ -582,25 +589,29 @@ class SchedulerPPMixin: ) current_seq_len = len(req.fill_ids) + + if is_dp_attention_enabled(): + # For profiling, we only have one request on PP0 + # Set global_num_tokens to indicate this rank has tokens, others have 0 + dp_size = get_attention_dp_size() + global_num_tokens = [0] * dp_size + dp_rank = get_attention_dp_rank() + global_num_tokens[dp_rank] = current_seq_len + batch.global_num_tokens = global_num_tokens + batch.global_num_tokens_for_logprob = global_num_tokens + proxy_tensors = { "hidden_states": torch.zeros( - ( - current_seq_len, - self.tp_worker.model_runner.model_config.hidden_size, - ), - dtype=self.tp_worker.model_runner.model_config.dtype, + (current_seq_len, model_config.hidden_size), + dtype=model_config.dtype, device="cuda", ), "residual": torch.zeros( - ( - current_seq_len, - self.tp_worker.model_runner.model_config.hidden_size, - ), - dtype=self.tp_worker.model_runner.model_config.dtype, + (current_seq_len, model_config.hidden_size), + dtype=model_config.dtype, device="cuda", ), } - from sglang.srt.managers.scheduler_pp_mixin import PPProxyTensors pp_proxy = PPProxyTensors(proxy_tensors) @@ -612,12 +623,9 @@ class SchedulerPPMixin: start = time.perf_counter() batch.prepare_for_extend() model_worker_batch = batch.get_model_worker_batch() - from sglang.srt.model_executor.forward_batch_info import ForwardBatch - forward_batch = ForwardBatch.init_new( - model_worker_batch, self.tp_worker.model_runner - ) - _ = self.tp_worker.model_runner.forward( + forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) + _ = model_runner.forward( forward_batch=forward_batch, pp_proxy_tensors=pp_proxy )