diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index 5331fc033..9733cd3e1 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/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: diff --git a/python/sglang/srt/managers/scheduler_recv_skipper.py b/python/sglang/srt/managers/scheduler_recv_skipper.py index 69c3e19a5..1a31c84c0 100644 --- a/python/sglang/srt/managers/scheduler_recv_skipper.py +++ b/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