From 8acb0270fdc2aaf6cdc29e8d96fe82083cfdcfc4 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Sun, 10 May 2026 19:57:46 -0400 Subject: [PATCH] Followup fix for Custom AR V2 in non NVL scenarios (#24742) Co-authored-by: b8zhong --- .../custom_all_reduce_v2.py | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 11599552f..01f162ff7 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -93,24 +93,26 @@ class CustomAllReduceV2: @contextmanager def capture(self): + if self.disabled: + yield + return try: self.obj.set_cuda_graph_capture(True) yield finally: self.obj.set_cuda_graph_capture(False) - if not self.disabled: - # cannot call when graph is capturing - assert ( - torch.cuda.is_current_stream_capturing() == False - ), "Cannot register graph inputs while capturing CUDA graph" - pairs = self.obj.share_graph_inputs() - handles = [handle for _, handle in pairs] - offsets = [offset for offset, _ in pairs] - handles_all = self._share_list(handles) - offsets_all = self._share_list(offsets) - result = [list(zip(o, h)) for o, h in zip(offsets_all, handles_all)] - self.obj.register_inputs(result) - log_info_on_rank0(logger, f"Registering {len(pairs)} cuda graph addresses") + # cannot call when graph is capturing + assert ( + torch.cuda.is_current_stream_capturing() == False + ), "Cannot register graph inputs while capturing CUDA graph" + pairs = self.obj.share_graph_inputs() + handles = [handle for _, handle in pairs] + offsets = [offset for offset, _ in pairs] + handles_all = self._share_list(handles) + offsets_all = self._share_list(offsets) + result = [list(zip(o, h)) for o, h in zip(offsets_all, handles_all)] + self.obj.register_inputs(result) + log_info_on_rank0(logger, f"Registering {len(pairs)} cuda graph addresses") def should_custom_ar(self, inp: torch.Tensor) -> bool: """Check if the input tensor is suitable for custom all-reduce."""