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

[Fix] Enable graph capture and MSCCL++ for attention TP groups (#31629)

Browse Source 
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
This commit is contained in:
Caio Rocha
2026-07-27 15:56:41 -07:00
committed by GitHub
co-authored by Cheng Wan
parent 3005af0941
commit 5a46e16f01
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/parallel_state.py
+1 -2
View File
@@ -1971,7 +1971,7 @@ def graph_capture(stream=None):
): ):
with contextlib.ExitStack() as stack: with contextlib.ExitStack() as stack:
seen = {id(_TP), id(_PP)} 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: if group is not None and id(group) not in seen:
seen.add(id(group)) seen.add(id(group))
stack.enter_context(group.graph_capture(context)) stack.enter_context(group.graph_capture(context))
@@ -2404,7 +2404,6 @@ def initialize_model_parallel(
get_world_group().local_rank, get_world_group().local_rank,
backend, backend,
use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP or enable_symm_mem, use_pynccl=SYNC_TOKEN_IDS_ACROSS_TP or enable_symm_mem,
use_mscclpp_allreduce=False,
use_custom_allreduce=False, use_custom_allreduce=False,
use_torch_symm_mem_allreduce=False, use_torch_symm_mem_allreduce=False,
use_message_queue_broadcaster=envs.SGLANG_USE_MESSAGE_QUEUE_BROADCASTER.get(), use_message_queue_broadcaster=envs.SGLANG_USE_MESSAGE_QUEUE_BROADCASTER.get(),