feat(scheduler): support scheduler recv skip in dp

python/sglang/srt/managers/scheduler_dp_attn_mixin.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable, Optional
 
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
 import torch
 
 from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer
@@ -127,6 +128,19 @@ def _update_gather_batch(
     # Check forward mode for cuda graph
     batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
 
+def _get_global_forward_mode(forward_modes):
+    forward_modes_excluding_idle = [
+        x for x in forward_modes if x != ForwardMode.IDLE.value
+    ]
+
+    if not forward_modes_excluding_idle:
+        return ForwardMode.IDLE
+
+    forward_mode_agree = all(value == forward_modes_excluding_idle[0] for value in forward_modes_excluding_idle)
+    global_forward_mode = (
+        ForwardMode(forward_modes_excluding_idle[0]) if forward_mode_agree else None
+    )
+    return global_forward_mode
 
 def prepare_mlp_sync_batch_raw(
     local_batch: ScheduleBatch,
@@ -207,6 +221,11 @@ def prepare_mlp_sync_batch_raw(
             )
         )
 
+        if mlp_sync_info.global_forward_mode is None:
+            forward_modes = mlp_sync_info.tp0_info[:, 5].tolist()
+            global_forward_mode = _get_global_forward_mode(forward_modes)
+            mlp_sync_info.global_forward_mode = global_forward_mode
+
     need_idle_batch = skip_all_gather or max(mlp_sync_info.global_num_tokens) > 0
     if need_idle_batch:
         batch_to_gather = local_batch
@@ -255,6 +274,11 @@ class SchedulerDPAttnMixin:
         """
         if need_sync if need_sync is not None else self.require_mlp_sync:
             batch = self.prepare_mlp_sync_batch(batch)
+        # if log_stats:
+        #     self.log_prefill_stats_late(batch)
+        if batch is not None and hasattr(batch, "global_forward_mode"):
+            if self.recv_skipper is not None:
+                self.recv_skipper.update_counter_dp(batch.global_forward_mode)
         return batch
 
     def get_idle_batch(self: Scheduler) -> ScheduleBatch:

python/sglang/srt/managers/scheduler_recv_skipper.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from sglang.srt.environ import envs
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.server_args import ServerArgs
@@ -11,9 +13,8 @@ class SchedulerRecvSkipper:
         return SchedulerRecvSkipper(server_args)
 
     def __init__(self, server_args: ServerArgs):
-        # Can be supported if needed, but may need e.g. `global_forward_mode`
-        assert not server_args.enable_dp_attention
-        self._counter = 0
+        self._enable_dp_attention = server_args.enable_dp_attention
+        self._counter = -1
         self._threshold = server_args.scheduler_recv_interval
         # All can be tuned if needed
         self._default_weight = envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DEFAULT.get()
@@ -23,16 +24,54 @@ class SchedulerRecvSkipper:
             None: envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_NONE.get(),
         }
 
-    def handle(self, last_forward_mode: ForwardMode):
-        should_recv = False
+    def update_counter_dp(self, global_forward_mode: Optional[ForwardMode]) -> None:
+        """
+        Update counter based on global forward mode (for dp_attention mode).
+        This should be called AFTER the sync point (prepare_mlp_sync_batch) to ensure
+        all DP ranks have the same counter value.
 
-        last_weight = self._weight_of_forward_mode.get(
-            last_forward_mode, self._default_weight
-        )
-        self._counter += last_weight
+        Args:
+            global_forward_mode: The global forward mode synchronized across DP ranks.
+        """
+        self._counter += self._weight_of_forward_mode.get(global_forward_mode, self._default_weight)
 
-        if self._counter >= self._threshold:
-            self._counter = 0
-            should_recv = True
+    def handle(
+        self,
+        last_forward_mode: ForwardMode,
+    ):
+        """
+        Determine whether to receive requests based on forward mode.
 
-        return should_recv
+        For dp_attention mode:
+            - This method only CHECKS if counter >= threshold (does not update counter)
+            - Counter is updated separately via update_counter_dp() after sync point
+            - This ensures all DP ranks make the same skip/recv decision
+
+        For non-dp_attention mode:
+            - check and update counter in one call
+
+        Args:
+            last_forward_mode: The local forward mode from the last batch.
+        Returns:
+            bool: True if should receive requests, False otherwise.
+        """
+        if self._enable_dp_attention:
+            if self._counter == -1:
+                self._counter = 0
+                return True
+
+            if self._counter >= self._threshold:
+                self._counter = 0
+                return True
+            return False
+        else:
+            last_weight = self._weight_of_forward_mode.get(
+                last_forward_mode, self._default_weight
+            )
+            self._counter += last_weight
+
+            if self._counter >= self._threshold:
+                self._counter = 0
+                return True
+
+            return False