From 6ea491e4392d8cb4bcf38c21430eef594ed62eb6 Mon Sep 17 00:00:00 2001
From: Baizhou Zhang <sobereddiezhang@gmail.com>
Date: Wed, 21 Jan 2026 02:56:55 +0800
Subject: [PATCH] Overlap shared experts with deepep dispatch for single batch
 overlap on Blackwell (#17289)

---
 docs/references/environment_variables.md      |  1 +
 .../srt/batch_overlap/single_batch_overlap.py |  6 ++-
 python/sglang/srt/environ.py                  |  1 +
 python/sglang/srt/models/deepseek_v2.py       | 46 +++++++++++++++++++
 4 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md
index ac8a41dbe..e9dc3b6fa 100644
--- a/docs/references/environment_variables.md
+++ b/docs/references/environment_variables.md
@@ -62,6 +62,7 @@ SGLang supports various environment variables that can be used to configure its
 | `SGLANG_DEEPEP_BF16_DISPATCH` | Use Bfloat16 for dispatch | `"false"` |
 | `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each GPU | `"128"` |
 | `SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS` | Number of SMs used for DeepEP combine when single batch overlap is enabled | `"32"` |
+| `SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO` | Run shared experts on an alternate stream when single batch overlap is enabled on GB200. When not setting this flag, shared experts and down gemm will be overlapped with DeepEP combine together. | `"false"` |
 
 ## NSA Backend Configuration (For DeepSeek V3.2)
 
diff --git a/python/sglang/srt/batch_overlap/single_batch_overlap.py b/python/sglang/srt/batch_overlap/single_batch_overlap.py
index 3209f7d5b..815ba8715 100644
--- a/python/sglang/srt/batch_overlap/single_batch_overlap.py
+++ b/python/sglang/srt/batch_overlap/single_batch_overlap.py
@@ -41,7 +41,11 @@ class SboFlags:
 
     @classmethod
     def enable_combine_shared_two_stream_overlap(cls):
-        return is_sbo_enabled() and not cls.enable_dispatch_shared_one_stream_overlap()
+        return (
+            is_sbo_enabled()
+            and not cls.enable_dispatch_shared_one_stream_overlap()
+            and not envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get()
+        )
 
     @classmethod
     def enable_dispatch_shared_one_stream_overlap(cls):
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index 5b66db887..f962b4a89 100644
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -341,6 +341,7 @@ class Envs:
     SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False)
     SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
     SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32)
+    SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False)
 
     # NSA Backend
     SGLANG_NSA_FUSE_TOPK = EnvBool(True)
diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index b5d274a3c..b6986977a 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -882,6 +882,52 @@ class DeepseekV2MoE(nn.Module):
             post_combine_hook_handle = (
                 self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
             )
+        elif envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get():
+            # On GB200: Shared experts overlapped on alt_stream, down gemm overlapped with DeepEP Combine
+
+            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 _pre_combine_hook(
+                dispatcher: BaseDispatcher, combine_input: CombineInput
+            ):
+                if (
+                    e := dispatcher.meta_overlap_args.get("record_event_after_down")
+                ) is not None:
+                    e.record()
+                pre_combine_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()
+
+            post_dispatch_hook_handle = (
+                self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
+            )
+            pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
+                _pre_combine_hook
+            )
+            post_combine_hook_handle = (
+                self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
+            )
 
         final_hidden_states = self.experts(
             hidden_states=hidden_states,