minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Fix] Add _MOE_TP to graph_capture for MoE models with ep>1 (#21907)

Browse Source 
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Kangyan-Zhou
2026-04-03 02:33:16 -07:00
committed by GitHub
co-authored by Claude Opus 4.6 Baizhou Zhang
parent ac593fed90
commit 56ac9c9932
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/parallel_state.py
+6 -5
View File
@@ -1532,11 +1532,12 @@ def graph_capture(stream: Optional[torch.cuda.Stream] = None):
with get_tp_group().graph_capture( with get_tp_group().graph_capture(
stream=stream stream=stream
) as context, get_pp_group().graph_capture(context): ) as context, get_pp_group().graph_capture(context):
moe_ep = _MOE_EP with contextlib.ExitStack() as stack:
if moe_ep is not None and moe_ep is not _TP: seen = {id(_TP)}
with moe_ep.graph_capture(context): for group in (_MOE_EP, _MOE_TP):
yield context if group is not None and id(group) not in seen:
else: seen.add(id(group))
stack.enter_context(group.graph_capture(context))
yield context yield context
python/sglang/srt/server_args.py
+3
View File
@@ -1133,6 +1133,9 @@ class ServerArgs:
# 16. Expert distribution recorder # 16. Expert distribution recorder
if self.enable_eplb or self.expert_distribution_recorder_mode is not None: if self.enable_eplb or self.expert_distribution_recorder_mode is not None:
self.disable_piecewise_cuda_graph = True self.disable_piecewise_cuda_graph = True
# 17. Context parallel
if self.attn_cp_size > 1:
self.disable_piecewise_cuda_graph = True
def _handle_gpu_memory_settings(self, gpu_mem): def _handle_gpu_memory_settings(self, gpu_mem):
""" """