From 4f3dc1ef5b621a3d0f3827be7a42a368c4c48d35 Mon Sep 17 00:00:00 2001 From: Yuzhen Zhou <82826991+zyzshishui@users.noreply.github.com> Date: Sun, 22 Feb 2026 23:27:31 -0800 Subject: [PATCH] [ROCm] Use unreg path for custom all-reduce during CUDA graph capture (#19162) --- .../srt/distributed/device_communicators/custom_all_reduce.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index b761bb69b..b6ca22ca4 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -411,6 +411,8 @@ class CustomAllreduce: if self._IS_CAPTURING: if torch.cuda.is_current_stream_capturing(): if _is_hip: + if self.tms_cudagraph: + return self.all_reduce_unreg(input) return self.all_reduce_reg(input) else: return self.all_reduce(input, registered=not self.tms_cudagraph)