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

Skip custom all-reduce v2 CUDA graph capture with torch memory saver. (#27948)

Browse Source
This commit is contained in:
Yuzhen Zhou
2026-06-30 14:53:29 -07:00
committed by GitHub
parent bb98629157
commit 99b8f36cb1
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py
+4 -2
View File
@@ -16,6 +16,7 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_vmm_utils imp
VmmGraphInputManager, VmmGraphInputManager,
is_vmm_pointer, is_vmm_pointer,
) )
from sglang.srt.environ import envs
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph, is_in_tc_piecewise_cuda_graph,
) )
@@ -62,6 +63,7 @@ class CustomAllReduceV2:
self.max_size = max(max_pull_size, max_push_size) self.max_size = max(max_pull_size, max_push_size)
self.override_shot(None) # set default config based on world size self.override_shot(None) # set default config based on world size
self.override_algo: Optional[AllReduceAlgo] = None self.override_algo: Optional[AllReduceAlgo] = None
self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
self.obj = get_custom_all_reduce_cls()( self.obj = get_custom_all_reduce_cls()(
rank=self.rank, rank=self.rank,
world_size=self.world_size, world_size=self.world_size,
@@ -102,7 +104,7 @@ class CustomAllReduceV2:
yield yield
return return
try: try:
self.obj.set_cuda_graph_capture(True) self.obj.set_cuda_graph_capture(not self.tms_cudagraph)
yield yield
finally: finally:
self.obj.set_cuda_graph_capture(False) self.obj.set_cuda_graph_capture(False)
@@ -147,7 +149,7 @@ class CustomAllReduceV2:
self.obj.set_cuda_graph_capture(False) self.obj.set_cuda_graph_capture(False)
return self._all_reduce(input) return self._all_reduce(input)
finally: finally:
self.obj.set_cuda_graph_capture(True) self.obj.set_cuda_graph_capture(not self.tms_cudagraph)
return self._all_reduce(input) return self._all_reduce(input)
def close(self): def close(self):