diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index ec379d53d..4f348583c 100644 --- a/python/sglang/bench_one_batch.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 3e4ed5fe5..258d411ac 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(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 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 96a2b14a4..67ac83bc6 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(self.dp_attn_adapter, batch) + batch = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(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 c309fbc50..7c429020c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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( - self.dp_attn_adapter, new_batch - ) + new_batch = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(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( - self.dp_attn_adapter, ret, need_sync=need_mlp_sync + ret = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch( + ret, need_sync=need_mlp_sync ) # Handle ngram embedding diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 8ff860364..bdda2e058 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py deleted file mode 100644 index b26cd32be..000000000 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ /dev/null @@ -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 diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 9bb5ac984..c57b9c762 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(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 diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 3e6a73c88..4ec25c836 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/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