Support scheduler_recv_interval (recv skipper) under DP-attention (#30457)

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>

python/sglang/srt/managers/schedule_batch.py

@@ -1896,6 +1896,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     can_run_dp_cuda_graph: bool = False
     can_run_dp_breakable_cuda_graph: bool = False
     tbo_split_seq_index: Optional[int] = None
+    # Rank-consistent forward mode for the recv skipper, derived from the MLP
+    # sync all-gather (the TBO-only `global_forward_mode` is None without TBO).
+    recv_skipper_forward_mode: Optional[ForwardMode] = None
     spec_verify_tier_num_tokens: int = -1
 
     # For processing logprobs

python/sglang/srt/managers/scheduler.py

@@ -1726,9 +1726,7 @@ class Scheduler(
             model_config=self.model_config,
             max_recv_per_poll=self.max_recv_per_poll,
             stream_output=lambda *a, **kw: self.output_streamer.stream_output(*a, **kw),
-            get_last_forward_mode=lambda: (
-                self.last_batch.forward_mode if self.last_batch is not None else None
-            ),
+            get_last_batch=lambda: self.last_batch,
             scripted_scheduler_hook=self.scripted_scheduler_hook,
         )
 

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

@@ -11,6 +11,7 @@ from sglang.srt.distributed.parallel_state import get_tp_group
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.scheduler_recv_skipper import SchedulerRecvSkipper
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@@ -257,6 +258,15 @@ def prepare_mlp_sync_batch_raw(
             batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather
         )
 
+    # Set on `local_batch`, not `batch_to_gather`: for PREBUILT batches the
+    # scheduler's `last_batch` is the prebuilt batch, not its inner idle batch.
+    if local_batch is not None and not skip_all_gather:
+        local_batch.recv_skipper_forward_mode = (
+            SchedulerRecvSkipper.derive_forward_mode(
+                mlp_sync_info.tp0_info[:, 5].tolist()
+            )
+        )
+
     if _ENABLE_METRICS_DP_ATTENTION and local_batch is not None:
         local_batch.dp_cooperation_info = mlp_sync_info.dp_cooperation_info
 

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

@@ -61,7 +61,7 @@ class SchedulerRequestReceiver:
     model_config: ModelConfig
     max_recv_per_poll: int
     stream_output: Callable[..., None]
-    get_last_forward_mode: Callable[[], Any]
+    get_last_batch: Callable[[], Any]
     scripted_scheduler_hook: Optional[ScriptedSchedulerHook] = None
 
     def recv_limit_reached(self, num_recv_reqs: int) -> bool:
@@ -79,7 +79,7 @@ class SchedulerRequestReceiver:
             self.scripted_scheduler_hook.step()
 
         if self.recv_skipper is not None:
-            if not self.recv_skipper.handle(self.get_last_forward_mode()):
+            if not self.recv_skipper.handle(self.get_last_batch()):
                 return []
 
         recv_reqs = self._pull_raw_reqs()

python/sglang/srt/managers/scheduler_recv_skipper.py

@@ -1,7 +1,14 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
+
 from sglang.srt.environ import envs
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.server_args import ServerArgs
 
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+
 
 class SchedulerRecvSkipper:
     @staticmethod
@@ -10,9 +17,24 @@ class SchedulerRecvSkipper:
             return None
         return SchedulerRecvSkipper(server_args)
 
+    @staticmethod
+    def derive_forward_mode(gathered_modes: List[int]) -> Optional[ForwardMode]:
+        """Collapse the gathered per-DP-rank forward modes into one weight-table
+        bucket; the input is rank-identical, so the recv decision is too."""
+        active = set(gathered_modes) - {
+            ForwardMode.IDLE.value,
+            ForwardMode.PREBUILT.value,
+        }
+        if not active:
+            return None  # globally idle: same bucket as "no last batch"
+        if active - {ForwardMode.DECODE.value, ForwardMode.TARGET_VERIFY.value}:
+            return ForwardMode.EXTEND  # any extend-like rank: prompt recv
+        if active == {ForwardMode.TARGET_VERIFY.value}:
+            return ForwardMode.TARGET_VERIFY
+        return ForwardMode.DECODE
+
     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._use_synced_mode = server_args.enable_dp_attention
         self._counter = 0
         self._threshold = server_args.scheduler_recv_interval
         # All can be tuned if needed
@@ -23,11 +45,21 @@ class SchedulerRecvSkipper:
             None: envs.SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_NONE.get(),
         }
 
-    def handle(self, last_forward_mode: ForwardMode):
+    def _pick_mode(self, last_batch: Optional[ScheduleBatch]) -> Optional[ForwardMode]:
+        # The recv decision must be identical on every rank in the request
+        # broadcast. Local modes differ across DP ranks (IDLE vs DECODE), so
+        # use the rank-consistent mode derived from the MLP sync all-gather.
+        if last_batch is None:
+            return None
+        if self._use_synced_mode:
+            return last_batch.recv_skipper_forward_mode
+        return last_batch.forward_mode
+
+    def handle(self, last_batch: Optional[ScheduleBatch]) -> bool:
         should_recv = False
 
         last_weight = self._weight_of_forward_mode.get(
-            last_forward_mode, self._default_weight
+            self._pick_mode(last_batch), self._default_weight
         )
         self._counter += last_weight