From 56ac9c9932be44558803d60325a6a175dc6fe04d Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Fri, 3 Apr 2026 02:33:16 -0700 Subject: [PATCH] [Fix] Add _MOE_TP to graph_capture for MoE models with ep>1 (#21907) Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Baizhou Zhang --- python/sglang/srt/distributed/parallel_state.py | 11 ++++++----- python/sglang/srt/server_args.py | 3 +++ 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index a80c6da5d..f0cba2189 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1532,11 +1532,12 @@ def graph_capture(stream: Optional[torch.cuda.Stream] = None): with get_tp_group().graph_capture( stream=stream ) as context, get_pp_group().graph_capture(context): - moe_ep = _MOE_EP - if moe_ep is not None and moe_ep is not _TP: - with moe_ep.graph_capture(context): - yield context - else: + with contextlib.ExitStack() as stack: + seen = {id(_TP)} + for group in (_MOE_EP, _MOE_TP): + if group is not None and id(group) not in seen: + seen.add(id(group)) + stack.enter_context(group.graph_capture(context)) yield context diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 57efd74dd..ec8d39567 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1133,6 +1133,9 @@ class ServerArgs: # 16. Expert distribution recorder if self.enable_eplb or self.expert_distribution_recorder_mode is not None: 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): """