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

[Fix] Fall back to the process-group broadcast for DSA topk when PyNCCL is absent (#36963)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-28 22:20:35 -07:00
committed by GitHub
parent 3904b309db
commit ec962fd12d
1 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+8 -5
View File
@@ -154,11 +154,14 @@ def _broadcast_indexer_topk_from_rank0_impl(topk_indices: torch.Tensor) -> None:
if group.world_size == 1: if group.world_size == 1:
return return
if topk_indices.device.type == "cuda" and torch.cuda.is_current_stream_capturing(): # PyNCCL is the faster path under capture, but it is not a precondition:
if group.pynccl_comm is None: # a split attn-TP group is built without one (parallel_state.py), and the
raise RuntimeError( # process-group broadcast captures and replays correctly.
"SGLANG_DSA_TOPK_BROADCAST requires PyNCCL during CUDA graph capture." if (
) topk_indices.device.type == "cuda"
and torch.cuda.is_current_stream_capturing()
and group.pynccl_comm is not None
):
with group.pynccl_comm.change_state(enable=True): with group.pynccl_comm.change_state(enable=True):
group.pynccl_comm.broadcast(topk_indices, src=0) group.pynccl_comm.broadcast(topk_indices, src=0)
else: else: