From 8f37a8a3f3eae0469868f5041478a96256764c33 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 18 May 2026 18:31:45 +0800 Subject: [PATCH] Introduce SchedulerDPAttnAdapter to own DP-attention state (#25611) --- python/sglang/srt/disaggregation/decode.py | 2 +- python/sglang/srt/disaggregation/prefill.py | 2 +- python/sglang/srt/managers/scheduler.py | 25 ++++++++++++++-- .../managers/scheduler_components/dp_attn.py | 30 +++++++++++++++++++ .../srt/managers/scheduler_dp_attn_mixin.py | 21 ++++++++----- .../sglang/srt/managers/scheduler_pp_mixin.py | 2 +- 6 files changed, 69 insertions(+), 13 deletions(-) create mode 100644 python/sglang/srt/managers/scheduler_components/dp_attn.py diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 3c26570ed..3e4ed5fe5 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 5cb6f9efc..96a2b14a4 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 58f1e5787..c309fbc50 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py new file mode 100644 index 000000000..8ff860364 --- /dev/null +++ b/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] diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index b135f2d2d..b26cd32be 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/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: - batch = self.prepare_mlp_sync_batch(batch) + if need_sync if need_sync is not None else self.get_require_mlp_sync(): + 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, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 354f5ef52..9bb5ac984 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/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