Move DP-attention adapter methods to SchedulerDPAttnAdapter (#25612)

python/sglang/bench_one_batch.py

@@ -71,7 +71,7 @@ from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
 from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
-from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw
+from sglang.srt.managers.scheduler_components.dp_attn import prepare_mlp_sync_batch_raw
 from sglang.srt.mem_cache.base_prefix_cache import EvictParams
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner

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(self.dp_attn_adapter, ret)
+        ret = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(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(self.dp_attn_adapter, batch)
+        batch = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(batch)
 
         if batch:
             set_schedule_time_batch(batch)

python/sglang/srt/managers/scheduler.py

@@ -170,7 +170,6 @@ from sglang.srt.managers.scheduler_components.dp_attn import (
 from sglang.srt.managers.scheduler_components.request_receiver import (
     SchedulerRequestReceiver,
 )
-from sglang.srt.managers.scheduler_dp_attn_mixin import SchedulerDPAttnMixin
 from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker
 from sglang.srt.managers.scheduler_output_processor_mixin import (
     SchedulerOutputProcessorMixin,
@@ -328,7 +327,6 @@ class Scheduler(
     SchedulerMultiplexMixin,
     SchedulerRuntimeCheckerMixin,
     SchedulerPPMixin,
-    SchedulerDPAttnMixin,
     SchedulerDllmMixin,
     SchedulerMlxOverlapMixin,
 ):
@@ -2243,9 +2241,7 @@ 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 = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(new_batch)
-                self.dp_attn_adapter, new_batch
-            )
             need_mlp_sync = new_batch is None
 
         if new_batch is not None:
@@ -2263,8 +2259,8 @@ class Scheduler(
                 ret = None
 
         # Handle DP attention and log stats
-        ret = self.maybe_prepare_mlp_sync_batch(
+        ret = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(
-            self.dp_attn_adapter, ret, need_sync=need_mlp_sync
+            ret, need_sync=need_mlp_sync
         )
 
         # Handle ngram embedding

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

@@ -1,20 +1,236 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, Optional
 
+import torch
+
+from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer
 from sglang.srt.configs.model_config import ModelConfig
+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.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.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.observability.metrics_collector import DPCooperationInfo
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.utils.common import require_mlp_tp_gather
 
 if TYPE_CHECKING:
     from sglang.srt.distributed.parallel_state import GroupCoordinator
 
 
+_ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get()
+
+
+@dataclass
+class MLPSyncBatchInfo:
+    dp_size: int
+    tp_size: int
+    cp_size: int
+
+    num_tokens: int
+    num_tokens_for_logprob: int
+    can_cuda_graph: bool
+    is_extend_in_batch: bool
+    local_can_run_tbo: bool
+    local_forward_mode: int
+
+    # some gathered elements
+    tp0_info: torch.Tensor = None
+    global_num_tokens: list[int] = None
+    global_num_tokens_for_logprob: list[int] = None
+    tbo_split_seq_index: torch.Tensor = None
+    global_forward_mode: int = None
+    dp_cooperation_info: Optional[DPCooperationInfo] = None
+
+    def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
+        return torch.tensor(
+            [
+                self.num_tokens,
+                self.num_tokens_for_logprob,
+                int(self.can_cuda_graph),
+                int(self.is_extend_in_batch),
+                int(self.local_can_run_tbo),
+                self.local_forward_mode,
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+    def _get_fallback_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
+        return torch.tensor(
+            [
+                0,  # num_tokens
+                0,  # num_tokens_for_logprob
+                1,  # can_cuda_graph
+                0,  # is_extend_in_batch
+                1,  # local_can_run_tbo
+                ForwardMode.IDLE.value,  # local_forward_mode
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+    def all_gather(self, device, group: torch.distributed.ProcessGroup):
+        local_info_tensor = self._get_local_tensor(device=device)
+        global_info_tensor = torch.empty(
+            (self.dp_size, self.tp_size * self.cp_size, 6),
+            dtype=torch.int64,
+            device=device,
+        )
+
+        torch.distributed.all_gather_into_tensor(
+            global_info_tensor.flatten(),
+            local_info_tensor,
+            group=group,
+        )
+        if device == "cpu":
+            tp_active_ranks = get_tp_group().active_ranks_cpu
+        else:
+            tp_active_ranks = get_tp_group().active_ranks
+
+        # Set fallback values for inactive ranks
+        tp_info = global_info_tensor.view(self.dp_size * self.tp_size * self.cp_size, 6)
+        tp_info[tp_active_ranks == 0] = self._get_fallback_tensor(device=device)
+
+        tp0_info = global_info_tensor[:, 0, :]
+        self.tp0_info = tp0_info
+        # Perform only one Device-to-Host (D2H) memory copy
+        cpu_data = tp0_info[:, :2].cpu()
+        self.global_num_tokens = cpu_data[:, 0].tolist()
+        self.global_num_tokens_for_logprob = cpu_data[:, 1].tolist()
+        self.can_cuda_graph = bool(tp0_info[:, 2].min().item())
+        self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
+        if _ENABLE_METRICS_DP_ATTENTION:
+            self.dp_cooperation_info = DPCooperationInfo.create(tp0_info[:, 5].tolist())
+
+
+def _update_gather_batch(
+    batch: ScheduleBatch,
+    mlp_sync_info: MLPSyncBatchInfo,
+    require_mlp_tp_gather: bool,
+    skip_all_gather=False,
+):
+    # TODO: handle the case when moe_dense_tp_size != 1
+    if not require_mlp_tp_gather:
+        batch.global_num_tokens = [mlp_sync_info.num_tokens]
+        batch.global_num_tokens_for_logprob = [mlp_sync_info.num_tokens_for_logprob]
+    else:
+        batch.global_num_tokens = mlp_sync_info.global_num_tokens
+        batch.global_num_tokens_for_logprob = (
+            mlp_sync_info.global_num_tokens_for_logprob
+        )
+    if not skip_all_gather:
+        batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
+        batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index
+        batch.global_forward_mode = mlp_sync_info.global_forward_mode
+
+    # Check forward mode for cuda graph
+    batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
+
+
+def prepare_mlp_sync_batch_raw(
+    local_batch: ScheduleBatch,
+    dp_size: int,
+    attn_tp_size: int,
+    attn_cp_size: int,
+    tp_group: GroupCoordinator,
+    get_idle_batch: Callable[[], ScheduleBatch],
+    disable_cuda_graph: bool,
+    require_mlp_tp_gather: bool,
+    disable_overlap_schedule: bool,
+    offload_tags: set[str],
+):
+    # Check if other DP workers have running batches
+    if local_batch is None or local_batch.forward_mode.is_prebuilt():
+        num_tokens = 0
+        num_tokens_for_logprob = 0
+    elif local_batch.forward_mode.is_decode():
+        num_tokens = local_batch.batch_size()
+        num_tokens_for_logprob = num_tokens
+    else:
+        num_tokens = local_batch.extend_num_tokens
+        num_tokens_for_logprob = sum(
+            # We should have at least 1 token for sample in every case.
+            max(extend_len - logprob_start_len, 1)
+            for logprob_start_len, extend_len in zip(
+                local_batch.extend_logprob_start_lens,
+                local_batch.extend_lens,
+            )
+        )
+        assert (
+            local_batch.return_logprob
+            or num_tokens_for_logprob == local_batch.batch_size()
+        )
+
+    skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get()
+    can_cuda_graph = (
+        local_batch is None
+        or local_batch.forward_mode.is_decode_or_idle()
+        or local_batch.forward_mode.is_prebuilt()
+    ) and not disable_cuda_graph
+
+    is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
+    if local_batch is not None:
+        local_batch.is_extend_in_batch = is_extend_in_batch
+
+    tbo_preparer = TboDPAttentionPreparer()
+    if len(offload_tags) == 0 and (
+        disable_overlap_schedule
+        or envs.SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH.get()
+    ):
+        group = tp_group.device_group
+        device = tp_group.device
+    else:
+        group = tp_group.cpu_group
+        device = "cpu"
+
+    local_can_run_tbo, local_forward_mode = tbo_preparer.prepare_all_gather(local_batch)
+
+    mlp_sync_info = MLPSyncBatchInfo(
+        dp_size=dp_size,
+        tp_size=attn_tp_size,
+        cp_size=attn_cp_size,
+        num_tokens=num_tokens,
+        num_tokens_for_logprob=num_tokens_for_logprob,
+        can_cuda_graph=can_cuda_graph,
+        is_extend_in_batch=is_extend_in_batch,
+        local_can_run_tbo=local_can_run_tbo,
+        local_forward_mode=local_forward_mode,
+    )
+
+    if not skip_all_gather:
+        mlp_sync_info.all_gather(device=device, group=group)
+
+        mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = (
+            tbo_preparer.compute_output(
+                mlp_sync_info.tp0_info[:, 4:6],
+            )
+        )
+
+    need_idle_batch = skip_all_gather or max(mlp_sync_info.global_num_tokens) > 0
+    if need_idle_batch:
+        batch_to_gather = local_batch
+        if local_batch is None:
+            batch_to_gather = local_batch = get_idle_batch()
+        elif local_batch.forward_mode.is_prebuilt():
+            # NOTE: for prebuilt batch, we add an inner idle batch to run MLP sync
+            batch_to_gather = local_batch.inner_idle_batch = get_idle_batch()
+        _update_gather_batch(
+            batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather
+        )
+
+    if _ENABLE_METRICS_DP_ATTENTION and local_batch is not None:
+        local_batch.dp_cooperation_info = mlp_sync_info.dp_cooperation_info
+
+    return local_batch
+
+
 @dataclass(kw_only=True, slots=True, frozen=True)
 class SchedulerDPAttnAdapter:
     tp_group: "GroupCoordinator"
@@ -28,3 +244,47 @@ class SchedulerDPAttnAdapter:
     enable_overlap: bool
     spec_algorithm: SpeculativeAlgorithm
     get_require_mlp_sync: Callable[[], bool]
+
+    def prepare_mlp_sync_batch(self, 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,
+            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,
+        )
+
+    def maybe_prepare_mlp_sync_batch(
+        self,
+        batch: Optional[ScheduleBatch],
+        need_sync: Optional[bool] = None,
+    ) -> Optional[ScheduleBatch]:
+        """
+        Helper to prepare MLP sync batch for DP attention.
+        Should be called after get_new_batch_prefill().
+
+        Args:
+            batch: The batch to process
+            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.get_require_mlp_sync():
+            batch = self.prepare_mlp_sync_batch(batch)
+        return batch
+
+    def get_idle_batch(self) -> ScheduleBatch:
+        idle_batch = ScheduleBatch.init_new(
+            [],
+            self.req_to_token_pool,
+            self.token_to_kv_pool_allocator,
+            self.tree_cache,
+            self.model_config,
+            self.enable_overlap,
+            self.spec_algorithm,
+        )
+        idle_batch.prepare_for_idle()
+        return idle_batch

python/sglang/srt/managers/scheduler_dp_attn_mixin.py

@@ -1,276 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Optional
-
-import torch
-
-from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer
-from sglang.srt.distributed.parallel_state import get_tp_group
-from sglang.srt.environ import envs
-from sglang.srt.managers.schedule_batch import ScheduleBatch
-from sglang.srt.model_executor.forward_batch_info import ForwardMode
-from sglang.srt.observability.metrics_collector import DPCooperationInfo
-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_components.dp_attn import SchedulerDPAttnAdapter
-
-
-_ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get()
-
-
-@dataclass
-class MLPSyncBatchInfo:
-    dp_size: int
-    tp_size: int
-    cp_size: int
-
-    num_tokens: int
-    num_tokens_for_logprob: int
-    can_cuda_graph: bool
-    is_extend_in_batch: bool
-    local_can_run_tbo: bool
-    local_forward_mode: int
-
-    # some gathered elements
-    tp0_info: torch.Tensor = None
-    global_num_tokens: list[int] = None
-    global_num_tokens_for_logprob: list[int] = None
-    tbo_split_seq_index: torch.Tensor = None
-    global_forward_mode: int = None
-    dp_cooperation_info: Optional[DPCooperationInfo] = None
-
-    def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
-        return torch.tensor(
-            [
-                self.num_tokens,
-                self.num_tokens_for_logprob,
-                int(self.can_cuda_graph),
-                int(self.is_extend_in_batch),
-                int(self.local_can_run_tbo),
-                self.local_forward_mode,
-            ],
-            device=device,
-            dtype=dtype,
-        )
-
-    def _get_fallback_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
-        return torch.tensor(
-            [
-                0,  # num_tokens
-                0,  # num_tokens_for_logprob
-                1,  # can_cuda_graph
-                0,  # is_extend_in_batch
-                1,  # local_can_run_tbo
-                ForwardMode.IDLE.value,  # local_forward_mode
-            ],
-            device=device,
-            dtype=dtype,
-        )
-
-    def all_gather(self, device, group: torch.distributed.ProcessGroup):
-        local_info_tensor = self._get_local_tensor(device=device)
-        global_info_tensor = torch.empty(
-            (self.dp_size, self.tp_size * self.cp_size, 6),
-            dtype=torch.int64,
-            device=device,
-        )
-
-        torch.distributed.all_gather_into_tensor(
-            global_info_tensor.flatten(),
-            local_info_tensor,
-            group=group,
-        )
-        if device == "cpu":
-            tp_active_ranks = get_tp_group().active_ranks_cpu
-        else:
-            tp_active_ranks = get_tp_group().active_ranks
-
-        # Set fallback values for inactive ranks
-        tp_info = global_info_tensor.view(self.dp_size * self.tp_size * self.cp_size, 6)
-        tp_info[tp_active_ranks == 0] = self._get_fallback_tensor(device=device)
-
-        tp0_info = global_info_tensor[:, 0, :]
-        self.tp0_info = tp0_info
-        # Perform only one Device-to-Host (D2H) memory copy
-        cpu_data = tp0_info[:, :2].cpu()
-        self.global_num_tokens = cpu_data[:, 0].tolist()
-        self.global_num_tokens_for_logprob = cpu_data[:, 1].tolist()
-        self.can_cuda_graph = bool(tp0_info[:, 2].min().item())
-        self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
-        if _ENABLE_METRICS_DP_ATTENTION:
-            self.dp_cooperation_info = DPCooperationInfo.create(tp0_info[:, 5].tolist())
-
-
-def _update_gather_batch(
-    batch: ScheduleBatch,
-    mlp_sync_info: MLPSyncBatchInfo,
-    require_mlp_tp_gather: bool,
-    skip_all_gather=False,
-):
-    # TODO: handle the case when moe_dense_tp_size != 1
-    if not require_mlp_tp_gather:
-        batch.global_num_tokens = [mlp_sync_info.num_tokens]
-        batch.global_num_tokens_for_logprob = [mlp_sync_info.num_tokens_for_logprob]
-    else:
-        batch.global_num_tokens = mlp_sync_info.global_num_tokens
-        batch.global_num_tokens_for_logprob = (
-            mlp_sync_info.global_num_tokens_for_logprob
-        )
-    if not skip_all_gather:
-        batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
-        batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index
-        batch.global_forward_mode = mlp_sync_info.global_forward_mode
-
-    # Check forward mode for cuda graph
-    batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
-
-
-def prepare_mlp_sync_batch_raw(
-    local_batch: ScheduleBatch,
-    dp_size: int,
-    attn_tp_size: int,
-    attn_cp_size: int,
-    tp_group: GroupCoordinator,
-    get_idle_batch: Callable[[], ScheduleBatch],
-    disable_cuda_graph: bool,
-    require_mlp_tp_gather: bool,
-    disable_overlap_schedule: bool,
-    offload_tags: set[str],
-):
-    # Check if other DP workers have running batches
-    if local_batch is None or local_batch.forward_mode.is_prebuilt():
-        num_tokens = 0
-        num_tokens_for_logprob = 0
-    elif local_batch.forward_mode.is_decode():
-        num_tokens = local_batch.batch_size()
-        num_tokens_for_logprob = num_tokens
-    else:
-        num_tokens = local_batch.extend_num_tokens
-        num_tokens_for_logprob = sum(
-            # We should have at least 1 token for sample in every case.
-            max(extend_len - logprob_start_len, 1)
-            for logprob_start_len, extend_len in zip(
-                local_batch.extend_logprob_start_lens,
-                local_batch.extend_lens,
-            )
-        )
-        assert (
-            local_batch.return_logprob
-            or num_tokens_for_logprob == local_batch.batch_size()
-        )
-
-    skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get()
-    can_cuda_graph = (
-        local_batch is None
-        or local_batch.forward_mode.is_decode_or_idle()
-        or local_batch.forward_mode.is_prebuilt()
-    ) and not disable_cuda_graph
-
-    is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
-    if local_batch is not None:
-        local_batch.is_extend_in_batch = is_extend_in_batch
-
-    tbo_preparer = TboDPAttentionPreparer()
-    if len(offload_tags) == 0 and (
-        disable_overlap_schedule
-        or envs.SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH.get()
-    ):
-        group = tp_group.device_group
-        device = tp_group.device
-    else:
-        group = tp_group.cpu_group
-        device = "cpu"
-
-    local_can_run_tbo, local_forward_mode = tbo_preparer.prepare_all_gather(local_batch)
-
-    mlp_sync_info = MLPSyncBatchInfo(
-        dp_size=dp_size,
-        tp_size=attn_tp_size,
-        cp_size=attn_cp_size,
-        num_tokens=num_tokens,
-        num_tokens_for_logprob=num_tokens_for_logprob,
-        can_cuda_graph=can_cuda_graph,
-        is_extend_in_batch=is_extend_in_batch,
-        local_can_run_tbo=local_can_run_tbo,
-        local_forward_mode=local_forward_mode,
-    )
-
-    if not skip_all_gather:
-        mlp_sync_info.all_gather(device=device, group=group)
-
-        mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = (
-            tbo_preparer.compute_output(
-                mlp_sync_info.tp0_info[:, 4:6],
-            )
-        )
-
-    need_idle_batch = skip_all_gather or max(mlp_sync_info.global_num_tokens) > 0
-    if need_idle_batch:
-        batch_to_gather = local_batch
-        if local_batch is None:
-            batch_to_gather = local_batch = get_idle_batch()
-        elif local_batch.forward_mode.is_prebuilt():
-            # NOTE: for prebuilt batch, we add an inner idle batch to run MLP sync
-            batch_to_gather = local_batch.inner_idle_batch = get_idle_batch()
-        _update_gather_batch(
-            batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather
-        )
-
-    if _ENABLE_METRICS_DP_ATTENTION and local_batch is not None:
-        local_batch.dp_cooperation_info = mlp_sync_info.dp_cooperation_info
-
-    return local_batch
-
-
-class SchedulerDPAttnMixin:
-    @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=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: "SchedulerDPAttnAdapter",
-        batch: Optional[ScheduleBatch],
-        need_sync: Optional[bool] = None,
-    ) -> Optional[ScheduleBatch]:
-        """
-        Helper to prepare MLP sync batch for DP attention.
-        Should be called after get_new_batch_prefill().
-
-        Args:
-            batch: The batch to process
-            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.get_require_mlp_sync():
-            batch = SchedulerDPAttnMixin.prepare_mlp_sync_batch(self, batch)
-        return batch
-
-    @staticmethod
-    def get_idle_batch(self: "SchedulerDPAttnAdapter") -> ScheduleBatch:
-        idle_batch = ScheduleBatch.init_new(
-            [],
-            self.req_to_token_pool,
-            self.token_to_kv_pool_allocator,
-            self.tree_cache,
-            self.model_config,
-            self.enable_overlap,
-            self.spec_algorithm,
-        )
-        idle_batch.prepare_for_idle()
-        return idle_batch

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(self.dp_attn_adapter, batch)
+                batch = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(batch)
                 self.mbs[mb_id] = batch
                 self.running_mbs[mb_id] = self.running_batch
 

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -87,7 +87,10 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
     s.running_batch.batch_is_full = False
     s.running_batch.reqs = []
     s.get_new_batch_prefill = MagicMock(return_value=None)
-    s.maybe_prepare_mlp_sync_batch = MagicMock(side_effect=lambda batch, **_: batch)
+    s.dp_attn_adapter = MagicMock()
+    s.dp_attn_adapter.maybe_prepare_mlp_sync_batch = MagicMock(
+        side_effect=lambda batch, **_: batch
+    )
     s._maybe_prepare_ngram_embedding = MagicMock(side_effect=lambda batch: batch)
     s.update_running_batch = MagicMock(side_effect=lambda batch: batch)
     s.tree_cache = tree_cache