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): """