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

[AMD] Fix AITER custom reduce-scatter CUDA-graph capture crash under torch_memory_saver (#34203)

Browse Source
This commit is contained in:
Zhiyao Jiang
2026-08-10 15:36:15 -07:00
committed by GitHub
parent c80a38edcd
commit 2c72323d90
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/parallel_state.py
+4 -1
View File
@@ -1080,7 +1080,10 @@ class GroupCoordinator:
return False return False
if getattr(ca_comm, "_IS_CAPTURING", False): if getattr(ca_comm, "_IS_CAPTURING", False):
if torch.cuda.is_current_stream_capturing(): if torch.cuda.is_current_stream_capturing():
ca_comm.reduce_scatter(input, output, registered=True) if envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get():
ca_comm.reduce_scatter(input, output, registered=False)
else:
ca_comm.reduce_scatter(input, output, registered=True)
elif is_in_tc_piecewise_cuda_graph(): elif is_in_tc_piecewise_cuda_graph():
ca_comm.reduce_scatter(input, output, registered=False) ca_comm.reduce_scatter(input, output, registered=False)
else: else: