From 20aad5b5abcad8179563161aedefcdb152c4b70a Mon Sep 17 00:00:00 2001 From: Sulfur6-L8972 Date: Thu, 4 Dec 2025 02:07:42 +0800 Subject: [PATCH] Single Batch Overlap for MoE Models (#9660) Co-authored-by: Cheng Wan Co-authored-by: Zqy11 <841971412@qq.com> Co-authored-by: AniZpZ Co-authored-by: TianyuZhang1214 Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com> --- docker/Dockerfile | 7 +++ .../srt/batch_overlap/single_batch_overlap.py | 45 ++++++++++---- .../srt/batch_overlap/two_batch_overlap.py | 6 ++ .../layers/deep_gemm_wrapper/entrypoint.py | 31 +++++++--- .../srt/layers/moe/fused_moe_triton/layer.py | 19 ++++-- .../srt/layers/moe/moe_runner/deep_gemm.py | 23 ++++++- .../srt/layers/moe/moe_runner/runner.py | 24 +++++++- .../srt/layers/moe/token_dispatcher/base.py | 10 ++-- .../srt/layers/moe/token_dispatcher/deepep.py | 44 +++++++++++--- python/sglang/srt/models/deepseek_v2.py | 60 ++++++++++++++++++- 10 files changed, 226 insertions(+), 43 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 7121a4d58..c6348364c 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -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 && \ diff --git a/python/sglang/srt/batch_overlap/single_batch_overlap.py b/python/sglang/srt/batch_overlap/single_batch_overlap.py index 7b324c37d..fc88ff27d 100644 --- a/python/sglang/srt/batch_overlap/single_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/single_batch_overlap.py @@ -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, diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index becb54e44..4d5746d26 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -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: diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index bf2ab4800..88d0a959b 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -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( diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index cb8b41ba5..763414cc2 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -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): diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 436f0364a..f60a428ef 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 294b90e81..fa0fd2559 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py index 48ebf22cf..06e2e2e5d 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/base.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index c73f2ce8b..0b0766271 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -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) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 31348b5be..29a896649 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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(