From dd77542360f1ad96557749cdec110a2e7cf5682e Mon Sep 17 00:00:00 2001 From: zyzshishui Date: Mon, 9 Mar 2026 05:59:52 +0000 Subject: [PATCH] 1 --- .../distributed/device_communicators/custom_all_reduce.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index b6ca22ca4b79..5e852a080402 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -4,6 +4,7 @@ import logging import os from contextlib import contextmanager +from functools import partial from typing import Any, List, Optional, Union import torch @@ -495,7 +496,11 @@ def dispatch_custom_allreduce(): ) logger.info("[AR] Using AiterCustomAllreduce (AMD default)") - return AiterCustomAllreduce + tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get() + return partial( + AiterCustomAllreduce, + enable_register_for_capturing=not tms_cudagraph, + ) except ImportError as e: logger.warning( "[AR] Aiter custom all-reduce not available; "