From 5a46e16f01e213d141a69dcc732d334cd03913fc Mon Sep 17 00:00:00 2001 From: Caio Rocha <164253795+caiocbr@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:56:41 -0700 Subject: [PATCH] [Fix] Enable graph capture and MSCCL++ for attention TP groups (#31629) Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com> --- python/sglang/srt/distributed/parallel_state.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 88cdb7485..2d3f8932a 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1971,7 +1971,7 @@ def graph_capture(stream=None): ): with contextlib.ExitStack() as stack: seen = {id(_TP), id(_PP)} - for group in (_DCP, _MOE_EP, _MOE_TP): + for group in (_DCP, _ATTN_TP, _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)) @@ -2404,7 +2404,6 @@ def initialize_model_parallel( get_world_group().local_rank, backend, use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP or enable_symm_mem, - use_mscclpp_allreduce=False, use_custom_allreduce=False, use_torch_symm_mem_allreduce=False, use_message_queue_broadcaster=envs.SGLANG_USE_MESSAGE_QUEUE_BROADCASTER.get(),