From 99b8f36cb1f1616e4a152af334a86c0cb365ca55 Mon Sep 17 00:00:00 2001 From: Yuzhen Zhou <82826991+zyzshishui@users.noreply.github.com> Date: Tue, 30 Jun 2026 14:53:29 -0700 Subject: [PATCH] Skip custom all-reduce v2 CUDA graph capture with torch memory saver. (#27948) --- .../device_communicators/custom_all_reduce_v2.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 210089633..2cbc5b0fe 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -16,6 +16,7 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_vmm_utils imp VmmGraphInputManager, is_vmm_pointer, ) +from sglang.srt.environ import envs from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) @@ -62,6 +63,7 @@ class CustomAllReduceV2: self.max_size = max(max_pull_size, max_push_size) self.override_shot(None) # set default config based on world size self.override_algo: Optional[AllReduceAlgo] = None + self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get() self.obj = get_custom_all_reduce_cls()( rank=self.rank, world_size=self.world_size, @@ -102,7 +104,7 @@ class CustomAllReduceV2: yield return try: - self.obj.set_cuda_graph_capture(True) + self.obj.set_cuda_graph_capture(not self.tms_cudagraph) yield finally: self.obj.set_cuda_graph_capture(False) @@ -147,7 +149,7 @@ class CustomAllReduceV2: self.obj.set_cuda_graph_capture(False) return self._all_reduce(input) finally: - self.obj.set_cuda_graph_capture(True) + self.obj.set_cuda_graph_capture(not self.tms_cudagraph) return self._all_reduce(input) def close(self):