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; "