Move DP-attention adapter methods to SchedulerDPAttnAdapter (#25612)
This commit is contained in:
fzyzcjy
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user