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

Followup fix for Custom AR V2 in non NVL scenarios (#24742)

Browse Source 
Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
Brayden Zhong
2026-05-10 16:57:46 -07:00
committed by GitHub
co-authored by b8zhong
parent b202778aa5
commit 8acb0270fd
1 changed files with 15 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py
+3 -1
View File
@@ -93,12 +93,14 @@ class CustomAllReduceV2:
@contextmanager @contextmanager
def capture(self): def capture(self):
if self.disabled:
yield
return
try: try:
self.obj.set_cuda_graph_capture(True) self.obj.set_cuda_graph_capture(True)
yield yield
finally: finally:
self.obj.set_cuda_graph_capture(False) self.obj.set_cuda_graph_capture(False)
if not self.disabled:
# cannot call when graph is capturing # cannot call when graph is capturing
assert ( assert (
torch.cuda.is_current_stream_capturing() == False torch.cuda.is_current_stream_capturing() == False