From 834400705f2de2378a327121340f57e324ca5a36 Mon Sep 17 00:00:00 2001
From: YAMY <74099316+YAMY1234@users.noreply.github.com>
Date: Fri, 21 Aug 2026 15:39:44 -0700
Subject: [PATCH] perf: overlap Qwen shared expert with DeepEP routed experts
(#34938)
---
.../docs/references/environment_variables.mdx | 5 +++++
python/sglang/srt/environ.py | 1 +
python/sglang/srt/models/qwen2_moe.py | 22 +++++++++++++++++++
3 files changed, 28 insertions(+)
diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx
index ea09a2fad..2d79ff276 100644
--- a/docs/docs/references/environment_variables.mdx
+++ b/docs/docs/references/environment_variables.mdx
@@ -331,6 +331,11 @@ SGLang supports various environment variables that can be used to configure its
| 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"` |
+
+ | `SGLANG_ENABLE_QWEN_DEEPEP_SHARED_OVERLAP` |
+ Run the Qwen shared expert on an alternate CUDA stream while DeepEP executes the routed experts. |
+ `"true"` |
+
SGLANG_DISABLED_MODEL_ARCHS |
Comma-separated list of model architectures to disable from auto-registration. |
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index af5f70c7b..af10cf771 100644
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -1028,6 +1028,7 @@ class Envs:
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)
+ SGLANG_ENABLE_QWEN_DEEPEP_SHARED_OVERLAP = EnvBool(True)
# Force dynamic Waterfill with runtime EP all-reduce instead of the default
# static local-batch path.
SGLANG_DISABLE_STATIC_WATERFILL = EnvBool(False)
diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py
index 120b49124..c699ec897 100644
--- a/python/sglang/srt/models/qwen2_moe.py
+++ b/python/sglang/srt/models/qwen2_moe.py
@@ -91,6 +91,9 @@ from sglang.srt.model_executor.cuda_graph_config import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+ is_in_breakable_cuda_graph,
+)
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.runtime_context import get_exec, get_forward, get_parallel
from sglang.srt.utils import (
@@ -533,6 +536,16 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
and envs.SGLANG_NPU_USE_MULTI_STREAM.get()
and forward_batch.forward_mode.is_cuda_graph()
)
+ enable_cuda_shared_overlap = (
+ _is_cuda
+ and envs.SGLANG_ENABLE_QWEN_DEEPEP_SHARED_OVERLAP.get()
+ # Breakable CUDA graph joins side streams before the eager DeepEP
+ # break, so this path cannot overlap the two expert computations.
+ and not is_in_breakable_cuda_graph()
+ and self.alt_stream is not None
+ and self.shared_expert is not None
+ and hidden_states.shape[0] > 0
+ )
shared_output = None
if hidden_states.shape[0] > 0:
# router_logits: (num_tokens, n_experts)
@@ -541,6 +554,13 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
shared_output = shared_expert_on_independent_stream(
hidden_states.clone(), self._forward_shared_experts
)
+ elif enable_cuda_shared_overlap:
+ current_stream = torch.cuda.current_stream()
+ self.alt_stream.wait_stream(current_stream)
+ with torch.cuda.stream(self.alt_stream):
+ shared_output = self._forward_shared_experts(hidden_states)
+ shared_output.record_stream(self.alt_stream)
+ shared_event = self.alt_stream.record_event()
else:
shared_output = self._forward_shared_experts(hidden_states)
topk_output = self.topk(
@@ -563,6 +583,8 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
)
if enable_dual_stream:
wait_share_stream()
+ elif enable_cuda_shared_overlap:
+ torch.cuda.current_stream().wait_event(shared_event)
if shared_output is not None:
final_hidden_states.add_(shared_output)