Introduce SchedulerDPAttnAdapter to own DP-attention state (#25611)

python/sglang/srt/disaggregation/decode.py

@@ -1670,7 +1670,7 @@ class SchedulerDisaggregationDecodeMixin:
             self.running_batch = self.update_running_batch(self.running_batch)
             ret = self.running_batch if not self.running_batch.is_empty() else None
 
-        ret = self.maybe_prepare_mlp_sync_batch(ret)
+        ret = self.maybe_prepare_mlp_sync_batch(self.dp_attn_adapter, ret)
         if ret:
             set_schedule_time_batch(ret)
         return ret

python/sglang/srt/disaggregation/prefill.py

@@ -381,7 +381,7 @@ class SchedulerDisaggregationPrefillMixin:
         self.process_prefill_chunk()
 
         batch = self.get_new_batch_prefill()
-        batch = self.maybe_prepare_mlp_sync_batch(batch)
+        batch = self.maybe_prepare_mlp_sync_batch(self.dp_attn_adapter, batch)
 
         if batch:
             set_schedule_time_batch(batch)

python/sglang/srt/managers/scheduler.py

@@ -164,6 +164,9 @@ from sglang.srt.managers.schedule_policy import (
     PrefillAdder,
     SchedulePolicy,
 )
+from sglang.srt.managers.scheduler_components.dp_attn import (
+    SchedulerDPAttnAdapter,
+)
 from sglang.srt.managers.scheduler_components.request_receiver import (
     SchedulerRequestReceiver,
 )
@@ -583,6 +586,20 @@ class Scheduler(
             ),
         )
 
+        self.dp_attn_adapter = SchedulerDPAttnAdapter(
+            tp_group=self.tp_group,
+            req_to_token_pool=self.req_to_token_pool,
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            tree_cache=self.tree_cache,
+            offload_tags=self.offload_tags,
+            ps=self.ps,
+            server_args=self.server_args,
+            model_config=self.model_config,
+            enable_overlap=self.enable_overlap,
+            spec_algorithm=self.spec_algorithm,
+            get_require_mlp_sync=lambda: self.require_mlp_sync,
+        )
+
         self.is_initializing = False
 
     def init_zbal_on_npu(self):
@@ -2226,7 +2243,9 @@ class Scheduler(
             # Before merging the new batch into running batch:
             # 1. All new batches are none -> need_mlp_sync remains true (sync is needed for decode batch).
             # 2. All new batches are some (prefill / idle) -> we do not need prepare mlp sync one more time.
-            new_batch = self.maybe_prepare_mlp_sync_batch(new_batch)
+            new_batch = self.maybe_prepare_mlp_sync_batch(
+                self.dp_attn_adapter, new_batch
+            )
             need_mlp_sync = new_batch is None
 
         if new_batch is not None:
@@ -2244,7 +2263,9 @@ class Scheduler(
                 ret = None
 
         # Handle DP attention and log stats
-        ret = self.maybe_prepare_mlp_sync_batch(ret, need_sync=need_mlp_sync)
+        ret = self.maybe_prepare_mlp_sync_batch(
+            self.dp_attn_adapter, ret, need_sync=need_mlp_sync
+        )
 
         # Handle ngram embedding
         ret = self._maybe_prepare_ngram_embedding(ret)

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

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable
+
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+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
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerDPAttnAdapter:
+    tp_group: "GroupCoordinator"
+    req_to_token_pool: ReqToTokenPool
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    tree_cache: BasePrefixCache
+    offload_tags: set[str]
+    ps: ParallelState
+    server_args: ServerArgs
+    model_config: ModelConfig
+    enable_overlap: bool
+    spec_algorithm: SpeculativeAlgorithm
+    get_require_mlp_sync: Callable[[], bool]

python/sglang/srt/managers/scheduler_dp_attn_mixin.py

@@ -15,7 +15,7 @@ from sglang.srt.utils.common import require_mlp_tp_gather
 
 if TYPE_CHECKING:
     from sglang.srt.distributed.parallel_state import GroupCoordinator
-    from sglang.srt.managers.scheduler import Scheduler
+    from sglang.srt.managers.scheduler_components.dp_attn import SchedulerDPAttnAdapter
 
 
 _ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get()
@@ -226,22 +226,26 @@ def prepare_mlp_sync_batch_raw(
 
 
 class SchedulerDPAttnMixin:
-    def prepare_mlp_sync_batch(self: Scheduler, local_batch: ScheduleBatch):
+    @staticmethod
+    def prepare_mlp_sync_batch(
+        self: "SchedulerDPAttnAdapter", local_batch: ScheduleBatch
+    ):
         return prepare_mlp_sync_batch_raw(
             local_batch,
             dp_size=self.server_args.dp_size,
             attn_tp_size=self.ps.attn_tp_size,
             attn_cp_size=self.ps.attn_cp_size,
             tp_group=self.tp_group,
-            get_idle_batch=self.get_idle_batch,
+            get_idle_batch=lambda: SchedulerDPAttnMixin.get_idle_batch(self),
             disable_cuda_graph=self.server_args.disable_cuda_graph,
             require_mlp_tp_gather=require_mlp_tp_gather(self.server_args),
             disable_overlap_schedule=self.server_args.disable_overlap_schedule,
             offload_tags=self.offload_tags,
         )
 
+    @staticmethod
     def maybe_prepare_mlp_sync_batch(
-        self: Scheduler,
+        self: "SchedulerDPAttnAdapter",
         batch: Optional[ScheduleBatch],
         need_sync: Optional[bool] = None,
     ) -> Optional[ScheduleBatch]:
@@ -251,13 +255,14 @@ class SchedulerDPAttnMixin:
 
         Args:
             batch: The batch to process
-            need_sync: If specified, overrides self.require_mlp_sync for prepare_mlp_sync_batch decision
+            need_sync: If specified, overrides self.get_require_mlp_sync() for prepare_mlp_sync_batch decision
         """
-        if need_sync if need_sync is not None else self.require_mlp_sync:
+        if need_sync if need_sync is not None else self.get_require_mlp_sync():
-            batch = self.prepare_mlp_sync_batch(batch)
+            batch = SchedulerDPAttnMixin.prepare_mlp_sync_batch(self, batch)
         return batch
 
-    def get_idle_batch(self: Scheduler) -> ScheduleBatch:
+    @staticmethod
+    def get_idle_batch(self: "SchedulerDPAttnAdapter") -> ScheduleBatch:
         idle_batch = ScheduleBatch.init_new(
             [],
             self.req_to_token_pool,

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -230,7 +230,7 @@ class SchedulerPPMixin:
 
                 self.process_prefill_chunk()
                 batch = self.get_new_batch_prefill()
-                batch = self.maybe_prepare_mlp_sync_batch(batch)
+                batch = self.maybe_prepare_mlp_sync_batch(self.dp_attn_adapter, batch)
                 self.mbs[mb_id] = batch
                 self.running_mbs[mb_id] = self.running_batch