Single Batch Overlap for MoE Models (#9660)

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>

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,

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: