From ec962fd12da14787e54d6995e8bc9084b4902477 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 28 Aug 2026 22:20:35 -0700 Subject: [PATCH] [Fix] Fall back to the process-group broadcast for DSA topk when PyNCCL is absent (#36963) --- .../sglang/srt/layers/attention/dsa/dsa_indexer.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index f07732065..4f630c124 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -154,11 +154,14 @@ def _broadcast_indexer_topk_from_rank0_impl(topk_indices: torch.Tensor) -> None: if group.world_size == 1: return - if topk_indices.device.type == "cuda" and torch.cuda.is_current_stream_capturing(): - if group.pynccl_comm is None: - raise RuntimeError( - "SGLANG_DSA_TOPK_BROADCAST requires PyNCCL during CUDA graph capture." - ) + # PyNCCL is the faster path under capture, but it is not a precondition: + # a split attn-TP group is built without one (parallel_state.py), and the + # process-group broadcast captures and replays correctly. + 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): group.pynccl_comm.broadcast(topk_indices, src=0) else: