fix: Graceful fallback to CustomAllReduce when full_nvlink is not True (#25650)

python/sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -338,7 +338,10 @@ class CustomAllreduce:
         self.close()
 
 
-def dispatch_custom_allreduce():
+def dispatch_custom_allreduce(
+    group: ProcessGroup,
+    device: torch.device,
+):
     """Return the CustomAllreduce class to use (aiter on ROCm if enabled).
 
     On AMD with 1-stage AR enabled, use sglang's CustomAllreduce.
@@ -350,8 +353,12 @@ def dispatch_custom_allreduce():
     ``nnodes > 1`` since custom AR is intra-node only.
     """
     if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
-        from .custom_all_reduce_v2 import CustomAllReduceV2
+        from .custom_all_reduce_v2 import (
+            CustomAllReduceV2,
+            can_use_custom_all_reduce_v2,
+        )
 
+        if can_use_custom_all_reduce_v2(group=group, device=device):
             logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)")
             return CustomAllReduceV2
 

python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py

@@ -38,15 +38,9 @@ class CustomAllReduceV2:
         max_pull_blocks: Optional[int] = None,
         max_push_blocks: Optional[int] = None,
     ) -> None:
-        _init_config()
+        _maybe_init_config()
         self.disabled = True
-        full_nvlink = can_use_custom_all_reduce_with_nvlink(
+        if not can_use_custom_all_reduce_v2(group=group, device=device):
-            group=group,
-            device=device,
-            supported_world_size=list(THRESHOLD_2_SHOT_MAP.keys()),
-            cls_name="CustomAllReduceV2",
-        )
-        if full_nvlink != True:
             return
 
         self.group = group
@@ -172,8 +166,10 @@ class CustomAllReduceV2:
         self.close()
 
 
-def _init_config():
+def _maybe_init_config():
     global THRESHOLD_2_SHOT_MAP
+    if THRESHOLD_2_SHOT_MAP:
+        return
     KB, MB = 1024, 1024 * 1024
 
     if is_sm100_supported():
@@ -201,4 +197,19 @@ def _init_config():
     # TODO: tune on more GPUs, e.g A100
 
 
+def can_use_custom_all_reduce_v2(
+    group: ProcessGroup,
+    device: torch.device,
+) -> bool:
+    # call _maybe_init_config() to ensure THRESHOLD_2_SHOT_MAP is initialized, since can_use_custom_all_reduce_v2 can be called before CustomAllReduceV2 is initialized
+    _maybe_init_config()
+    full_nvlink = can_use_custom_all_reduce_with_nvlink(
+        group=group,
+        device=device,
+        supported_world_size=list(THRESHOLD_2_SHOT_MAP.keys()),
+        cls_name="CustomAllReduceV2",
+    )
+    return full_nvlink is True
+
+
 THRESHOLD_2_SHOT_MAP: Dict[int, ModeConfig] = {}

python/sglang/srt/distributed/parallel_state.py

@@ -384,7 +384,10 @@ class GroupCoordinator:
         if use_custom_allreduce and self.world_size > 1:
             # Initialize a custom fast all-reduce implementation.
             try:
-                CAClass = dispatch_custom_allreduce()
+                CAClass = dispatch_custom_allreduce(
+                    group=self.cpu_group,
+                    device=self.device,
+                )
                 self.ca_comm = CAClass(
                     group=self.cpu_group,
                     device=self.device,