[Fix] Add fallback for flashinfer allreduce fusion (#20384)

python/sglang/srt/layers/communicator.py

@@ -50,6 +50,7 @@ from sglang.srt.layers.dp_attention import (
     is_allocation_symmetric,
     is_dp_attention_enabled,
 )
+from sglang.srt.layers.flashinfer_comm_fusion import is_flashinfer_allreduce_unavailable
 from sglang.srt.layers.moe import (
     get_moe_a2a_backend,
     should_use_flashinfer_cutlass_moe_fp4_allgather,
@@ -100,6 +101,7 @@ def apply_flashinfer_allreduce_fusion(batch_size: int):
         and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
         and not is_dp_attention_enabled()
         and get_global_server_args().enable_flashinfer_allreduce_fusion
+        and not is_flashinfer_allreduce_unavailable()
     )
 
 

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -14,6 +14,7 @@ logger = logging.getLogger(__name__)
 
 _flashinfer_comm = None
 _workspace_manager = None
+_flashinfer_allreduce_unavailable = False
 
 if is_flashinfer_available():
     try:
@@ -24,17 +25,23 @@ if is_flashinfer_available():
         ):
             _flashinfer_comm = comm
         else:
+            _flashinfer_allreduce_unavailable = True
             logger.warning(
                 "flashinfer.comm unified allreduce_fusion API is not available, "
                 "falling back to standard implementation"
             )
     except ImportError:
+        _flashinfer_allreduce_unavailable = True
         logger.warning(
             "flashinfer.comm is not available, falling back to standard "
             "implementation"
         )
 
 
+def is_flashinfer_allreduce_unavailable() -> bool:
+    return _flashinfer_allreduce_unavailable
+
+
 class FlashInferWorkspaceManager:
     def __init__(self):
         self.workspace = None
@@ -57,7 +64,7 @@ class FlashInferWorkspaceManager:
         """Initialize workspace"""
         if _flashinfer_comm is None:
             logger.warning(
-                "FlashInfer comm not available, skipping workspace " "initialization"
+                "FlashInfer comm not available, skipping workspace initialization"
             )
             return
 
@@ -73,7 +80,12 @@ class FlashInferWorkspaceManager:
                 force_oneshot_support=bool(use_oneshot),
             )
         except Exception as e:
-            logger.warning(f"Failed to initialize FlashInfer workspace: {e}")
+            global _flashinfer_allreduce_unavailable
+            _flashinfer_allreduce_unavailable = True
+            logger.warning(
+                f"Failed to initialize FlashInfer workspace: {e}. "
+                "Disabling flashinfer allreduce fusion permanently."
+            )
             self.workspace = None
             self.initialized = False
             return
@@ -140,6 +152,9 @@ def ensure_workspace_initialized(
     use_oneshot: Optional[bool] = None,
 ):
     """Ensure workspace is initialized"""
+    if _flashinfer_allreduce_unavailable:
+        return False
+
     if not is_flashinfer_available() or _flashinfer_comm is None:
         return False
 
@@ -220,7 +235,7 @@ def flashinfer_allreduce_residual_rmsnorm(
     """
     if not is_flashinfer_available() or _flashinfer_comm is None:
         logger.debug(
-            "FlashInfer not available, falling back to standard " "implementation"
+            "FlashInfer not available, falling back to standard implementation"
         )
         return None, None
 

python/sglang/srt/layers/layernorm.py

@@ -327,6 +327,12 @@ class RMSNorm(MultiPlatformOp):
                     )
                     if fused_result is not None:
                         return fused_result
+                    else:
+                        logger.warning(
+                            "AITER fused AR+RMSNorm failed, falling back to standard implementation"
+                        )
+                        x = tensor_model_parallel_all_reduce(x)
+                        return self.forward(x, residual, None)
                 else:
                     fused_result = flashinfer_allreduce_residual_rmsnorm(
                         input_tensor=x,
@@ -336,14 +342,12 @@ class RMSNorm(MultiPlatformOp):
                     )
                     if fused_result[0] is not None:
                         return fused_result
-
+                    else:
-                # For AITER route, preserve correctness when fused path is unavailable.
+                        logger.warning(
-                if (
+                            "FlashInfer allreduce fusion failed, falling back to standard implementation"
-                    _use_aiter
+                        )
-                    and get_global_server_args().enable_aiter_allreduce_fusion
+                        x = tensor_model_parallel_all_reduce(x)
-                ):
+                        return self.forward(x, residual, None)
-                    x = tensor_model_parallel_all_reduce(x)
-                    return self.forward(x, residual, None)
 
         return self.forward(x, residual, post_residual_addition)