From 243ea585fc2f2ac8a519a36a9f537e2acde01ee3 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 8 Nov 2025 23:23:14 +0800 Subject: [PATCH] [DP-Attn] Clarify MLP sync / idle batch preparation logic (#12843) --- python/sglang/bench_one_batch.py | 6 +- python/sglang/srt/disaggregation/decode.py | 34 ++-- python/sglang/srt/disaggregation/prefill.py | 6 +- python/sglang/srt/managers/scheduler.py | 158 ++------------- .../srt/managers/scheduler_dp_attn_mixin.py | 185 ++++++++++++++++++ .../sglang/srt/managers/scheduler_pp_mixin.py | 13 +- python/sglang/srt/two_batch_overlap.py | 4 +- 7 files changed, 229 insertions(+), 177 deletions(-) create mode 100644 python/sglang/srt/managers/scheduler_dp_attn_mixin.py diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index ed4ff9cbd..25b16d310 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -68,7 +68,7 @@ from sglang.srt.distributed.parallel_state import destroy_distributed_environmen from sglang.srt.entrypoints.engine import _set_envs_and_config from sglang.srt.layers.moe import initialize_moe_config from sglang.srt.managers.schedule_batch import Req, ScheduleBatch -from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.sampling.sampling_params import SamplingParams @@ -391,15 +391,13 @@ def decode(input_token_ids, batch, model_runner): def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner): if require_mlp_sync(model_runner.server_args): - Scheduler.prepare_mlp_sync_batch_raw( + prepare_mlp_sync_batch_raw( batch, dp_size=model_runner.server_args.dp_size, attn_tp_size=1, tp_group=model_runner.tp_group, get_idle_batch=None, disable_cuda_graph=model_runner.server_args.disable_cuda_graph, - spec_algorithm=SpeculativeAlgorithm.NONE, - speculative_num_draft_tokens=None, require_mlp_tp_gather=require_mlp_tp_gather(model_runner.server_args), disable_overlap_schedule=model_runner.server_args.disable_overlap_schedule, offload_tags=set(), diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 60bf58090..328a22adf 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -25,7 +25,7 @@ import time from collections import deque from dataclasses import dataclass from http import HTTPStatus -from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union +from typing import TYPE_CHECKING, List, Optional, Type, Union import torch from torch.distributed import ProcessGroup @@ -63,7 +63,7 @@ from sglang.srt.tracing.trace import ( trace_slice_batch, trace_slice_end, ) -from sglang.srt.utils import get_int_env_var, require_mlp_sync +from sglang.srt.utils import get_int_env_var from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter logger = logging.getLogger(__name__) @@ -833,8 +833,6 @@ class SchedulerDisaggregationDecodeMixin: batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch - prepare_mlp_sync_flag = require_mlp_sync(self.server_args) - if batch: # Generate fake extend output. if batch.forward_mode.is_extend(): @@ -843,15 +841,15 @@ class SchedulerDisaggregationDecodeMixin: batch.reqs, any(req.return_logprob for req in batch.reqs) ) trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs) - if prepare_mlp_sync_flag: - self._prepare_idle_batch_and_run(None) + if self.require_mlp_sync: + self._prepare_idle_batch_and_run() else: - if prepare_mlp_sync_flag: + if self.require_mlp_sync: self.prepare_mlp_sync_batch(batch) result = self.run_batch(batch) self.process_batch_result(batch, result) - elif prepare_mlp_sync_flag: - batch, _ = self._prepare_idle_batch_and_run(None) + elif self.require_mlp_sync: + batch, _ = self._prepare_idle_batch_and_run() queue_size = ( len(self.waiting_queue) @@ -882,8 +880,6 @@ class SchedulerDisaggregationDecodeMixin: self.cur_batch = batch last_batch_in_queue = False - prepare_mlp_sync_flag = require_mlp_sync(self.server_args) - batch_result = None if batch: # Generate fake extend output. @@ -893,23 +889,23 @@ class SchedulerDisaggregationDecodeMixin: batch.reqs, any(req.return_logprob for req in batch.reqs) ) trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs) - if prepare_mlp_sync_flag: + if self.require_mlp_sync: batch_, batch_result = self._prepare_idle_batch_and_run( - None, delay_process=True + delay_process=True ) if batch_: self.result_queue.append((batch_.copy(), batch_result)) last_batch_in_queue = True else: - if prepare_mlp_sync_flag: + if self.require_mlp_sync: self.prepare_mlp_sync_batch(batch) batch_result = self.run_batch(batch) self.result_queue.append((batch.copy(), batch_result)) last_batch_in_queue = True - elif prepare_mlp_sync_flag: + elif self.require_mlp_sync: batch, batch_result = self._prepare_idle_batch_and_run( - None, delay_process=True + delay_process=True ) if batch: self.result_queue.append((batch.copy(), batch_result)) @@ -936,8 +932,8 @@ class SchedulerDisaggregationDecodeMixin: self.last_batch = batch self.last_batch_in_queue = last_batch_in_queue - def _prepare_idle_batch_and_run(self: Scheduler, batch, delay_process=False): - batch = self.prepare_mlp_sync_batch(batch) + def _prepare_idle_batch_and_run(self: Scheduler, delay_process=False): + batch = self.prepare_mlp_sync_batch(None) result = None if batch: result = self.run_batch(batch) @@ -947,7 +943,7 @@ class SchedulerDisaggregationDecodeMixin: def get_next_disagg_decode_batch_to_run( self: Scheduler, - ) -> Optional[Tuple[ScheduleBatch, bool]]: + ) -> Optional[ScheduleBatch]: """Create fake completed prefill if possible and merge with running batch""" # Merge the prefill batch into the running batch last_batch = self.last_batch diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 5b2e5cb41..044234e85 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -54,7 +54,7 @@ from sglang.srt.mem_cache.memory_pool import ( SWAKVPool, ) from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end -from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj, require_mlp_sync +from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj if TYPE_CHECKING: from torch.distributed import ProcessGroup @@ -326,7 +326,7 @@ class SchedulerDisaggregationPrefillMixin: attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()} trace_event_batch("schedule", batch.reqs, attrs=attrs) - if require_mlp_sync(self.server_args): + if self.require_mlp_sync: batch = self.prepare_mlp_sync_batch(batch) self.cur_batch = batch @@ -361,7 +361,7 @@ class SchedulerDisaggregationPrefillMixin: attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()} trace_event_batch("schedule", batch.reqs, attrs=attrs) - if require_mlp_sync(self.server_args): + if self.require_mlp_sync: batch = self.prepare_mlp_sync_batch(batch) self.cur_batch = batch diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 92f4e1025..7125d3310 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -128,6 +128,7 @@ from sglang.srt.managers.schedule_policy import ( PrefillAdder, SchedulePolicy, ) +from sglang.srt.managers.scheduler_dp_attn_mixin import SchedulerDPAttnMixin from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker from sglang.srt.managers.scheduler_metrics_mixin import ( RECORD_STEP_TIME, @@ -166,7 +167,6 @@ from sglang.srt.tracing.trace import ( trace_slice_end, trace_slice_start, ) -from sglang.srt.two_batch_overlap import TboDPAttentionPreparer from sglang.srt.utils import ( DynamicGradMode, broadcast_pyobj, @@ -181,7 +181,6 @@ from sglang.srt.utils import ( numa_bind_to_node, point_to_point_pyobj, require_mlp_sync, - require_mlp_tp_gather, set_gpu_proc_affinity, set_random_seed, suppress_other_loggers, @@ -218,6 +217,7 @@ class Scheduler( SchedulerMultiplexMixin, SchedulerRuntimeCheckerMixin, SchedulerPPMixin, + SchedulerDPAttnMixin, ): """A scheduler that manages a tensor parallel GPU worker.""" @@ -523,6 +523,9 @@ class Scheduler( # Init overlap self.init_overlap() + # Init mlp sync flag + self.require_mlp_sync = require_mlp_sync(server_args) + # Init request dispatcher self._request_dispatcher = TypeBasedDispatcher( [ @@ -1667,13 +1670,14 @@ class Scheduler( new_batch = self.get_new_batch_prefill() - need_dp_attn_preparation = require_mlp_sync(self.server_args) - - if need_dp_attn_preparation and not self.spec_algorithm.is_none(): - # In speculative decoding, prefill batches and decode batches cannot be processed in the same DP attention group. - # We prepare idle batches in advance to skip preparing decode batches when there are prefill batches in the group. + need_mlp_sync = self.require_mlp_sync + if need_mlp_sync and not self.spec_algorithm.is_none(): + # NOTE: This branch makes sure prefill and decode batches will not be mixed when spec and dp-attn is enabled. + # 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.prepare_mlp_sync_batch(new_batch) - need_dp_attn_preparation = new_batch is None + need_mlp_sync = new_batch is None if new_batch is not None: # Run prefill first if possible @@ -1687,7 +1691,7 @@ class Scheduler( ret = None # Handle DP attention - if need_dp_attn_preparation: + if need_mlp_sync: ret = self.prepare_mlp_sync_batch(ret) if ret: @@ -2108,142 +2112,6 @@ class Scheduler( self.return_health_check_ct -= 1 self.send_to_tokenizer.send_output(HealthCheckOutput()) - def prepare_mlp_sync_batch(self, local_batch: ScheduleBatch): - return self.prepare_mlp_sync_batch_raw( - local_batch, - dp_size=self.server_args.dp_size, - attn_tp_size=self.attn_tp_size, - tp_group=self.tp_group, - get_idle_batch=self.get_idle_batch, - disable_cuda_graph=self.server_args.disable_cuda_graph, - spec_algorithm=self.spec_algorithm, - speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, - 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 prepare_mlp_sync_batch_raw( - local_batch: ScheduleBatch, - dp_size, - attn_tp_size: int, - tp_group, - get_idle_batch, - disable_cuda_graph: bool, - spec_algorithm, - speculative_num_draft_tokens, - 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: - 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 - if local_batch.return_logprob: - 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, - ) - ) - else: - # When return_logprob = False, only need last token per request - num_tokens_for_logprob = local_batch.batch_size() - - if local_batch is None or local_batch.forward_mode.is_decode_or_idle(): - can_cuda_graph = 1 - else: - can_cuda_graph = 0 - - is_extend_in_batch = ( - local_batch.forward_mode.is_extend() if local_batch else False - ) - - tbo_preparer = TboDPAttentionPreparer() - if len(offload_tags) == 0 and disable_overlap_schedule: - group = tp_group.device_group - device = tp_group.device - else: - group = tp_group.cpu_group - device = "cpu" - - local_info = torch.tensor( - [ - num_tokens, - can_cuda_graph, - num_tokens_for_logprob, - is_extend_in_batch, - *tbo_preparer.prepare_all_gather( - local_batch, - ), - ], - dtype=torch.int64, - device=device, - ) - global_info = torch.empty( - (dp_size, attn_tp_size, 6), - dtype=torch.int64, - device=device, - ) - torch.distributed.all_gather_into_tensor( - global_info.flatten(), - local_info, - group=group, - ) - global_num_tokens = global_info[:, 0, 0].tolist() - can_cuda_graph = min(global_info[:, 0, 1].tolist()) - global_num_tokens_for_logprob = global_info[:, 0, 2].tolist() - is_extend_in_batch = global_info[:, 0, 3].tolist() - - tbo_split_seq_index, global_forward_mode = tbo_preparer.compute_output( - global_info[:, :, 4:6] - ) - - if local_batch is None and max(global_num_tokens) > 0: - local_batch = get_idle_batch() - - if local_batch is not None: - # TODO: handle the case when moe_dense_tp_size != 1 - if not require_mlp_tp_gather: - local_batch.global_num_tokens = [num_tokens] - local_batch.global_num_tokens_for_logprob = [num_tokens_for_logprob] - else: - local_batch.global_num_tokens = global_num_tokens - local_batch.global_num_tokens_for_logprob = ( - global_num_tokens_for_logprob - ) - local_batch.is_extend_in_batch = any(is_extend_in_batch) - local_batch.tbo_split_seq_index = tbo_split_seq_index - local_batch.global_forward_mode = global_forward_mode - - # Check forward mode for cuda graph - if not disable_cuda_graph: - local_batch.can_run_dp_cuda_graph = can_cuda_graph - - return local_batch - - def get_idle_batch(self): - 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 - def move_ready_grammar_requests(self): """Move requests whose grammar objects are ready from grammar_queue to waiting_queue.""" diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py new file mode 100644 index 000000000..99001d7ad --- /dev/null +++ b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable + +import torch + +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.two_batch_overlap import TboDPAttentionPreparer +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 + + +@dataclass +class MLPSyncBatchInfo: + dp_size: int + tp_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 + + 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 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, 6), + dtype=torch.int64, + device=device, + ) + + torch.distributed.all_gather_into_tensor( + global_info_tensor.flatten(), + local_info_tensor, + group=group, + ) + + tp0_info = global_info_tensor[:, 0, :] + self.tp0_info = tp0_info + self.global_num_tokens = tp0_info[:, 0].tolist() + self.global_num_tokens_for_logprob = tp0_info[:, 1].tolist() + self.can_cuda_graph = bool(tp0_info[:, 2].min().item()) + self.is_extend_in_batch = bool(tp0_info[:, 3].max().item()) + + +def prepare_mlp_sync_batch_raw( + local_batch: ScheduleBatch, + dp_size: int, + attn_tp_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: + 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 + if local_batch.return_logprob: + 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, + ) + ) + else: + # When return_logprob = False, only need last token per request + num_tokens_for_logprob = local_batch.batch_size() + + if local_batch is None or local_batch.forward_mode.is_decode_or_idle(): + can_cuda_graph = 1 + else: + can_cuda_graph = 0 + + is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False + + tbo_preparer = TboDPAttentionPreparer() + if len(offload_tags) == 0 and disable_overlap_schedule: + 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, + 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, + ) + mlp_sync_info.all_gather(device=device, group=group) + + tbo_split_seq_index, global_forward_mode = tbo_preparer.compute_output( + mlp_sync_info.tp0_info[:, 4:6], + ) + + if local_batch is None and max(mlp_sync_info.global_num_tokens) > 0: + local_batch = get_idle_batch() + + if local_batch is not None: + # TODO: handle the case when moe_dense_tp_size != 1 + if not require_mlp_tp_gather: + local_batch.global_num_tokens = [num_tokens] + local_batch.global_num_tokens_for_logprob = [num_tokens_for_logprob] + else: + local_batch.global_num_tokens = mlp_sync_info.global_num_tokens + local_batch.global_num_tokens_for_logprob = ( + mlp_sync_info.global_num_tokens_for_logprob + ) + local_batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch + local_batch.tbo_split_seq_index = tbo_split_seq_index + local_batch.global_forward_mode = global_forward_mode + + # Check forward mode for cuda graph + if not disable_cuda_graph: + local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph + + return local_batch + + +class SchedulerDPAttnMixin: + def prepare_mlp_sync_batch(self: Scheduler, local_batch: ScheduleBatch): + return prepare_mlp_sync_batch_raw( + local_batch, + dp_size=self.server_args.dp_size, + attn_tp_size=self.attn_tp_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 get_idle_batch(self: Scheduler) -> 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 c49ae937f..d14a31b85 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -1,10 +1,15 @@ -from typing import List, Optional +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.model_executor.forward_batch_info import PPProxyTensors -from sglang.srt.utils import DynamicGradMode, point_to_point_pyobj, require_mlp_sync +from sglang.srt.utils import DynamicGradMode, point_to_point_pyobj + +if TYPE_CHECKING: + from sglang.srt.managers.scheduler import Scheduler class SchedulerPPMixin: @@ -132,7 +137,7 @@ class SchedulerPPMixin: self.self_check_during_idle() @DynamicGradMode() - def event_loop_pp_disagg_prefill(self): + def event_loop_pp_disagg_prefill(self: Scheduler): """ An event loop for the prefill server in pipeline parallelism. @@ -238,7 +243,7 @@ class SchedulerPPMixin: self.process_prefill_chunk() batch = self.get_new_batch_prefill() - if require_mlp_sync(self.server_args): + if self.require_mlp_sync: batch = self.prepare_mlp_sync_batch(batch) mbs[mb_id] = batch diff --git a/python/sglang/srt/two_batch_overlap.py b/python/sglang/srt/two_batch_overlap.py index a07c112fe..e72420407 100644 --- a/python/sglang/srt/two_batch_overlap.py +++ b/python/sglang/srt/two_batch_overlap.py @@ -407,8 +407,8 @@ class TboDPAttentionPreparer: return local_can_run_tbo, local_forward_mode def compute_output(self, partial_global_info): - local_can_run_tbo_aggregated = min(partial_global_info[:, 0, 0].tolist()) - forward_modes = partial_global_info[:, 0, 1].tolist() + local_can_run_tbo_aggregated = min(partial_global_info[:, 0].tolist()) + forward_modes = partial_global_info[:, 1].tolist() global_forward_mode, forward_mode_agree = self._compute_global_forward_mode( forward_modes