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

[ROCm] Use unreg path for custom all-reduce during CUDA graph capture (#19162)

Browse Source
This commit is contained in:
Yuzhen Zhou
2026-02-22 23:27:31 -08:00
committed by GitHub
parent c3c1532820
commit 4f3dc1ef5b
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
+2
View File
@@ -411,6 +411,8 @@ class CustomAllreduce:
if self._IS_CAPTURING: if self._IS_CAPTURING:
if torch.cuda.is_current_stream_capturing(): if torch.cuda.is_current_stream_capturing():
if _is_hip: if _is_hip:
if self.tms_cudagraph:
return self.all_reduce_unreg(input)
return self.all_reduce_reg(input) return self.all_reduce_reg(input)
else: else:
return self.all_reduce(input, registered=not self.tms_cudagraph) return self.all_reduce(input, registered=not self.tms_cudagraph)