[PD] Fix PP dynamic chunking with DP attention (#17339)

Signed-off-by: Shangming Cai <csmthu@gmail.com>

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, model_config.hidden_size),
-                            current_seq_len,
+                        dtype=model_config.dtype,
-                            self.tp_worker.model_runner.model_config.hidden_size,
-                        ),
-                        dtype=self.tp_worker.model_runner.model_config.dtype,
                         device="cuda",
                     ),
                     "residual": torch.zeros(
-                        (
+                        (current_seq_len, model_config.hidden_size),
-                            current_seq_len,
+                        dtype=model_config.dtype,
-                            self.tp_worker.model_runner.model_config.hidden_size,
-                        ),
-                        dtype=self.tp_worker.model_runner.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(
+                forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
-                    model_worker_batch, self.tp_worker.model_runner
+                _ = model_runner.forward(
-                )
-                _ = self.tp_worker.model_runner.forward(
                     forward_batch=forward_batch, pp_proxy_tensors=pp_proxy
                 )