kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-11 09:47:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

[DP-Attn] Clarify MLP sync / idle batch preparation logic (#12843)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-08 23:23:14 +08:00
committed by GitHub
parent 6fee2c535c
commit 243ea585fc
7 changed files with 229 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/bench_one_batch.py
View File

@@ -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(),

34
python/sglang/srt/disaggregation/decode.py
View File

@@ -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

6
python/sglang/srt/disaggregation/prefill.py
View File

@@ -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

158
python/sglang/srt/managers/scheduler.py
View File

@@ -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."""

185
python/sglang/srt/managers/scheduler_dp_attn_mixin.py Normal file
View File

@@ -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

13
python/sglang/srt/managers/scheduler_pp_mixin.py
View File

@@ -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

4
python/sglang/srt/two_batch_overlap.py
View File

@@ -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