chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Single Batch Overlap for MoE Models (#9660)

Browse Source 
Co-authored-by: Cheng Wan <wan4ch@gmail.com>
Co-authored-by: Zqy11 <841971412@qq.com>
Co-authored-by: AniZpZ <aniz1905@gmail.com>
Co-authored-by: TianyuZhang1214 <tianyuzhang1214@gmail.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
This commit is contained in:
Sulfur6-L8972
2025-12-04 02:07:42 +08:00
committed by GitHub
parent 974c562a25
commit 20aad5b5ab
10 changed files with 226 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
docker/Dockerfile
View File

@@ -7,6 +7,7 @@ ARG BRANCH_TYPE=remote
ARG GRACE_BLACKWELL=0
ARG GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2
ARG HOPPER_SBO_DEEPEP_COMMIT=9f2fc4b3182a51044ae7ecb6610f7c9c3258c4d6
ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee
ARG BUILD_AND_DOWNLOAD_PARALLEL=8
ARG SGL_KERNEL_VERSION=0.3.18.post2
@@ -149,6 +150,12 @@ RUN set -eux; \
git checkout ${GRACE_BLACKWELL_DEEPEP_BRANCH} && \
sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \
cd .. ; \
elif [ "$HOPPER_SBO" = "1" ]; then \
git clone https://github.com/deepseek-ai/DeepEP.git -b antgroup-opt && \
cd DeepEP && \
git checkout ${HOPPER_SBO_DEEPEP_COMMIT} && \
sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \
cd .. ; \
else \
wget -q https://${GITHUB_ARTIFACTORY}/deepseek-ai/DeepEP/archive/${DEEPEP_COMMIT}.zip && \
unzip ${DEEPEP_COMMIT}.zip && rm ${DEEPEP_COMMIT}.zip && mv DeepEP-${DEEPEP_COMMIT} DeepEP && cd DeepEP && \

45
python/sglang/srt/batch_overlap/single_batch_overlap.py
View File

@@ -21,28 +21,37 @@ import torch
from sglang.srt.layers.moe import get_moe_runner_backend
from sglang.srt.layers.moe.utils import is_sbo_enabled
from sglang.srt.utils import get_int_env_var
from sglang.srt.utils import get_int_env_var, is_blackwell
class SboFlags:
# TODO may have: "enable_dispatch_shared_one_stream_overlap", "enable_dispatch_gateup_gemm_two_stream_overlap", ...
# TODO may have: "enable_dispatch_gateup_gemm_two_stream_overlap", ...
@classmethod
def enable_combine_down_gemm_two_stream_overlap(cls):
return (
is_sbo_enabled()
# currently only cutedsl backend supports it
and get_moe_runner_backend().is_flashinfer_cutedsl()
and (
get_moe_runner_backend().is_flashinfer_cutedsl()
or (get_moe_runner_backend().is_deep_gemm() and not is_blackwell())
)
)
@classmethod
def enable_combine_shared_two_stream_overlap(cls):
return is_sbo_enabled()
return is_sbo_enabled() and not cls.enable_dispatch_shared_one_stream_overlap()
@classmethod
def enable_dispatch_shared_one_stream_overlap(cls):
return is_sbo_enabled() and not is_blackwell()
@classmethod
def fuse_shared_experts_inside_sbo(cls):
# TODO after antgroup's PR, should be `... or cls.enable_dispatch_shared_one_stream_overlap()`
return cls.enable_combine_shared_two_stream_overlap()
return (
cls.enable_combine_shared_two_stream_overlap()
or cls.enable_dispatch_shared_one_stream_overlap()
)
@dataclass
@@ -51,9 +60,10 @@ class CombineOverlapArgs:
overlap: bool
stream: torch.cuda.Stream
wait_event: torch.cuda.Event
num_sms: int
num_sms: Optional[int] = None
signal: Optional[torch.Tensor] = None
threshold: int = 0
block_m: Optional[int] = 64
threshold: Optional[int] = 0
@dataclass
@@ -77,7 +87,9 @@ def compute_overlap_args(dispatch_output, alt_stream):
total_num_sms = torch.cuda.get_device_properties(
device="cuda"
).multi_processor_count
communicate_num_sms = get_int_env_var("SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS", 32)
communicate_num_sms = get_int_env_var(
"SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS", 32 if is_blackwell() else 3
)
compute_num_sms = total_num_sms - communicate_num_sms
assert alt_stream is not None
@@ -96,9 +108,18 @@ def compute_overlap_args(dispatch_output, alt_stream):
if SboFlags.enable_combine_down_gemm_two_stream_overlap():
# TODO use zero_allocator to remove this `torch.zeros` call
# NOTE ours v2 use uint32 not int32 currently
combine_signal = torch.zeros(
num_local_experts, dtype=torch.uint32, device=hidden_states.device
)
if is_blackwell():
combine_signal = torch.zeros(
num_local_experts, dtype=torch.uint32, device=hidden_states.device
)
else:
MIN_BLOCK_M = 64
combine_signal_size = num_local_experts * (
(num_tokens_static + MIN_BLOCK_M - 1) // MIN_BLOCK_M
)
combine_signal = torch.zeros(
combine_signal_size, dtype=torch.int32, device=hidden_states.device
)
down_gemm_overlap_args = DownGemmOverlapArgs(
signal=combine_signal,

6
python/sglang/srt/batch_overlap/two_batch_overlap.py
View File

@@ -1009,6 +1009,12 @@ class MaybeTboDeepEPDispatcher(BaseDispatcher):
def combine_b(self, **kwargs):
return self._execute("combine_b", **kwargs)
def register_deepep_dispatch_hook(self, hook):
handle_list = []
for inner in self._inners:
handle_list.append(inner.register_deepep_dispatch_hook(hook))
return handle_list
def set_quant_config(self, quant_config: dict):
super().set_quant_config(quant_config)
for inner in self._inners:

31
python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py
View File

@@ -1,6 +1,6 @@
import logging
from contextlib import contextmanager
from typing import Tuple
from typing import Any, Optional, Tuple
import torch
@@ -29,6 +29,8 @@ def grouped_gemm_nt_f8f8bf16_masked(
out: torch.Tensor,
masked_m: torch.Tensor,
expected_m: int,
overlap_args: Optional[Any] = None,
max_block_n: int = 256,
):
num_groups, _, k = lhs[0].shape
_, n, _ = rhs[0].shape
@@ -40,13 +42,26 @@ def grouped_gemm_nt_f8f8bf16_masked(
with compile_utils.deep_gemm_execution_hook(
expected_m, n, k, num_groups, kernel_type
):
deep_gemm.fp8_m_grouped_gemm_nt_masked(
lhs,
rhs,
out,
masked_m,
expected_m,
)
with configure_deep_gemm_num_sms(
overlap_args.num_sms if overlap_args is not None else None
):
return deep_gemm.fp8_m_grouped_gemm_nt_masked(
lhs,
rhs,
out,
masked_m,
expected_m,
**(
dict(
enable_overlap=True,
max_block_n=max_block_n,
signal=overlap_args.signal,
)
if overlap_args is not None
else {}
),
)
def grouped_gemm_nt_f8f8bf16_contig(

19
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -268,6 +268,9 @@ class FusedMoE(torch.nn.Module):
self.down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = None
self.meta_overlap_args: Optional[dict] = None
if self.quant_method is not None and hasattr(self.quant_method, "runner"):
self.runner = self.quant_method.runner
def _load_per_tensor_weight_scale(
self,
shard_id: str,
@@ -1010,12 +1013,20 @@ class FusedMoE(torch.nn.Module):
def set_overlap_args(
self, down_gemm_overlap_args: DownGemmOverlapArgs, meta_overlap_args: dict
):
self.down_gemm_overlap_args = down_gemm_overlap_args
self.meta_overlap_args = meta_overlap_args
if hasattr(self, "runner"):
self.runner.set_overlap_args(down_gemm_overlap_args, meta_overlap_args)
else:
# TODO: remove this branch after MoE refactor
self.down_gemm_overlap_args = down_gemm_overlap_args
self.meta_overlap_args = meta_overlap_args
def clear_overlap_args(self) -> None:
self.down_gemm_overlap_args = None
self.meta_overlap_args = None
if hasattr(self, "runner"):
self.runner.clear_overlap_args()
else:
# TODO: remove this branch after MoE refactor
self.down_gemm_overlap_args = None
self.meta_overlap_args = None
class FlashInferFusedMoE(FusedMoE):

23
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
View File

@@ -40,6 +40,7 @@ if not (_is_npu or _is_hip):
_MASKED_GEMM_FAST_ACT = get_bool_env_var("SGLANG_MASKED_GEMM_FAST_ACT")
_DEEPGEMM_ON_H20 = get_bool_env_var("SGLANG_DEEPGEMM_ON_H20")
# TODO(kaixih@nvidia): ideally we should merge this logic into
@@ -315,13 +316,33 @@ class DeepGemmRunnerCore(MoeRunnerCore):
down_output = torch.empty(
(num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16
)
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
down_gemm_overlap_args = running_state.get("down_gemm_overlap_args", None)
if down_gemm_overlap_args is None:
gemm_overlap_args_dict = {}
else:
down_gemm_overlap_args.start_event.record()
max_block_n = (
160 if (_DEEPGEMM_ON_H20 and runner_input.expected_m <= 64) else 256
)
gemm_overlap_args_dict = {
"overlap_args": down_gemm_overlap_args,
"max_block_n": max_block_n,
}
deep_gemm_return_value = deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
(down_input, down_input_scale),
(w2_weight, w2_scale),
down_output,
masked_m,
expected_m,
**gemm_overlap_args_dict,
)
meta_overlap_args = running_state.get("meta_overlap_args", None)
if meta_overlap_args is not None:
block_m, threshold = deep_gemm_return_value
meta_overlap_args["block_m"] = block_m
meta_overlap_args["threshold"] = threshold
return down_output

24
python/sglang/srt/layers/moe/moe_runner/runner.py
View File

@@ -2,7 +2,7 @@ from __future__ import annotations
import logging
import os
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional
from sglang.srt.layers.moe.moe_runner.base import (
FusedOpPool,
@@ -15,6 +15,7 @@ from sglang.srt.layers.moe.moe_runner.triton_kernels import TritonKernelsRunnerC
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
if TYPE_CHECKING:
from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
from sglang.srt.layers.moe.moe_runner.base import MoeQuantInfo
from sglang.srt.layers.moe.token_dispatcher.base import CombineInput, DispatchOutput
from sglang.srt.layers.moe.utils import MoeRunnerBackend
@@ -42,10 +43,14 @@ class MoeRunner:
a2a_backend_name = get_moe_a2a_backend().value
runner_backend_name = runner_backend.value
# TODO(cwan): add a server argument to disable fused func
self.fused_func = FusedOpPool.get_fused_func(
a2a_backend_name, runner_backend_name
)
self.down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = None
self.meta_overlap_args: Optional[dict] = None
SGLANG_CI_DISABLE_MOE_FUSED_FUNC = os.environ.get(
"SGLANG_CI_DISABLE_MOE_FUSED_FUNC", "0"
)
@@ -69,6 +74,11 @@ class MoeRunner:
)
running_state = {}
if self.down_gemm_overlap_args is not None:
running_state["down_gemm_overlap_args"] = self.down_gemm_overlap_args
if self.meta_overlap_args is not None:
running_state["meta_overlap_args"] = self.meta_overlap_args
runner_input = self.pre_permute_func(
dispatch_output, quant_info, self.config, running_state
)
@@ -84,3 +94,15 @@ class MoeRunner:
)
return combine_input
def set_overlap_args(
self, down_gemm_overlap_args: DownGemmOverlapArgs, meta_overlap_args: dict
):
assert self.fused_func is None, "Fused func is not supported for overlap args"
self.down_gemm_overlap_args = down_gemm_overlap_args
self.meta_overlap_args = meta_overlap_args
def clear_overlap_args(self) -> None:
assert self.fused_func is None, "Fused func is not supported for overlap args"
self.down_gemm_overlap_args = None
self.meta_overlap_args = None

10
python/sglang/srt/layers/moe/token_dispatcher/base.py
View File

@@ -49,7 +49,7 @@ class _RemovableDispatcherHandle:
del hooks_dict[self.id]
class _DispatcherBaseHooks:
class DispatcherBaseHooks:
def __init__(self):
self.hook_dict = OrderedDict[int, Callable]()
@@ -63,7 +63,7 @@ class _DispatcherBaseHooks:
raise NotImplementedError("This method should be overridden by subclasses")
class _PreDispatchHooks(_DispatcherBaseHooks):
class _PreDispatchHooks(DispatcherBaseHooks):
def __call__(
self,
@@ -78,7 +78,7 @@ class _PreDispatchHooks(_DispatcherBaseHooks):
return hidden_states, topk_output
class _PostDispatchHooks(_DispatcherBaseHooks):
class _PostDispatchHooks(DispatcherBaseHooks):
def __call__(
self, dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
@@ -90,7 +90,7 @@ class _PostDispatchHooks(_DispatcherBaseHooks):
return dispatch_output
class _PreCombineHooks(_DispatcherBaseHooks):
class _PreCombineHooks(DispatcherBaseHooks):
def __call__(
self, dispatcher: BaseDispatcher, combine_input: CombineInput
@@ -102,7 +102,7 @@ class _PreCombineHooks(_DispatcherBaseHooks):
return combine_input
class _PostCombineHooks(_DispatcherBaseHooks):
class _PostCombineHooks(DispatcherBaseHooks):
def __call__(
self, dispatcher: BaseDispatcher, hidden_states: torch.Tensor

44
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
View File

@@ -13,6 +13,7 @@ from sglang.srt.layers.moe.token_dispatcher.base import (
BaseDispatcherConfig,
CombineInput,
CombineInputFormat,
DispatcherBaseHooks,
DispatchOutput,
DispatchOutputFormat,
)
@@ -26,6 +27,7 @@ from sglang.srt.layers.moe.utils import (
from sglang.srt.utils import (
get_bool_env_var,
get_int_env_var,
is_blackwell,
is_hip,
is_npu,
load_json_config,
@@ -58,6 +60,13 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
logger = logging.getLogger(__name__)
class DeepEPPDispatchHooks(DispatcherBaseHooks):
def __call__(self, dispatcher: BaseDispatcher):
for hook_fun in self.hook_dict.values():
hook_fun(dispatcher)
class DeepEPNormalDispatchOutput(NamedTuple):
"""DeepEP normal dispatch output."""
@@ -660,12 +669,31 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
):
buffer = self._get_buffer()
overlap_args = self.overlap_args
meta_overlap_args = self.meta_overlap_args
ctx = nullcontext()
if overlap_args is not None:
overlap_args.stream.wait_event(overlap_args.wait_event)
ctx = torch.cuda.stream(overlap_args.stream)
if is_blackwell():
overlap_args_dict = dict(
overlap=overlap_args.overlap,
src_signals=overlap_args.signal,
src_signal_expect_value=overlap_args.threshold,
)
else:
overlap_args_dict = dict(
overlap=overlap_args.overlap,
packed_recv_count=self.packed_recv_count,
comp_signal=overlap_args.signal,
block_m=meta_overlap_args["block_m"],
threshold=meta_overlap_args["threshold"],
num_sms=overlap_args.num_sms,
)
else:
overlap_args_dict = {}
with ctx:
combined_hidden_states, event, hook = buffer.low_latency_combine(
x=hidden_states,
@@ -674,15 +702,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
handle=self.handle,
async_finish=not self.return_recv_hook,
return_recv_hook=self.return_recv_hook,
**(
dict(
overlap=overlap_args.overlap,
src_signals=overlap_args.signal,
src_signal_expect_value=overlap_args.threshold,
)
if overlap_args is not None
else {}
),
**overlap_args_dict,
)
self.packed_recv_count = self.handle = None
@@ -749,6 +769,7 @@ class DeepEPDispatcher(BaseDispatcher):
)
self._stage = _Stage.INITIAL
self._deepep_dispatch_hooks = DeepEPPDispatchHooks()
def dispatch(
self,
@@ -756,6 +777,8 @@ class DeepEPDispatcher(BaseDispatcher):
topk_output: TopKOutput,
) -> DispatchOutput:
self.dispatch_a(hidden_states, topk_output)
if self._deepep_dispatch_hooks is not None:
self._deepep_dispatch_hooks(self)
ret = self.dispatch_b()
return ret
@@ -844,3 +867,6 @@ class DeepEPDispatcher(BaseDispatcher):
self._low_latency_dispatcher.clear_overlap_args()
if self.deepep_mode.enable_normal():
self._normal_dispatcher.clear_overlap_args()
def register_deepep_dispatch_hook(self, hook):
return self._deepep_dispatch_hooks.register_hook(hook)

60
python/sglang/srt/models/deepseek_v2.py
View File

@@ -30,7 +30,10 @@ from torch import nn
from transformers import PretrainedConfig
from sglang.srt.batch_overlap.single_batch_overlap import SboFlags, compute_overlap_args
from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.batch_overlap.two_batch_overlap import (
MaybeTboDeepEPDispatcher,
model_forward_maybe_tbo,
)
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.configs.model_config import (
get_nsa_index_head_dim,
@@ -962,6 +965,13 @@ class DeepseekV2MoE(nn.Module):
) -> torch.Tensor:
shared_output = None
sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn
sbo_overlap_dispatch_flag = (
sbo_enabled_flag and SboFlags.enable_dispatch_shared_one_stream_overlap()
)
sbo_overlap_combine_flag = (
sbo_enabled_flag and SboFlags.enable_combine_shared_two_stream_overlap()
)
if hidden_states.shape[0] > 0:
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states, forward_batch=forward_batch)
@@ -978,8 +988,52 @@ class DeepseekV2MoE(nn.Module):
else:
topk_output = self.topk.empty_topk_output(hidden_states.device)
# SBO is not yet implemented for NextN
if sbo_enabled_flag:
if sbo_overlap_dispatch_flag:
shared_output = None
def _deepep_dispatch_hook(dispatcher: BaseDispatcher):
nonlocal shared_output
shared_output = self._forward_shared_experts(hidden_states)
for handle in deepep_dispatch_hook_handle:
handle.remove()
def _post_dispatch_hook(
dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
):
combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
compute_overlap_args(dispatch_output, self.alt_stream)
)
dispatcher.set_overlap_args(
combine_overlap_args=combine_overlap_args,
meta_overlap_args=meta_overlap_args,
)
self.experts.set_overlap_args(
down_gemm_overlap_args=down_gemm_overlap_args,
meta_overlap_args=meta_overlap_args,
)
post_dispatch_hook_handle.remove()
def _post_combine_hook(
dispatcher: BaseDispatcher, hidden_states: torch.Tensor
):
dispatcher.clear_overlap_args()
self.experts.clear_overlap_args()
post_combine_hook_handle.remove()
assert isinstance(self.experts.dispatcher, MaybeTboDeepEPDispatcher)
deepep_dispatch_hook_handle = (
self.experts.dispatcher.register_deepep_dispatch_hook(
_deepep_dispatch_hook
)
)
post_dispatch_hook_handle = (
self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
)
post_combine_hook_handle = (
self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
)
elif sbo_overlap_combine_flag:
shared_output = None
def _post_dispatch_hook(