From f8668d9e78a1a50b8f88faabaea4710acff683c1 Mon Sep 17 00:00:00 2001
From: Baizhou Zhang <sobereddiezhang@gmail.com>
Date: Fri, 13 Mar 2026 01:24:55 -0700
Subject: [PATCH] [Fix] Add fallback for flashinfer allreduce fusion (#20384)

---
 python/sglang/srt/layers/communicator.py      |  2 ++
 .../srt/layers/flashinfer_comm_fusion.py      | 21 ++++++++++++++++---
 python/sglang/srt/layers/layernorm.py         | 20 +++++++++++-------
 3 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py
index 5f7054344..b354d0a04 100644
--- a/python/sglang/srt/layers/communicator.py
+++ b/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()
     )
 
 
diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py
index 9d1f91900..c45ff977f 100644
--- a/python/sglang/srt/layers/flashinfer_comm_fusion.py
+++ b/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
 
diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py
index 383f58399..614be8f0e 100644
--- a/python/sglang/srt/layers/layernorm.py
+++ b/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
-
-                # For AITER route, preserve correctness when fused path is unavailable.
-                if (
-                    _use_aiter
-                    and get_global_server_args().enable_aiter_allreduce_fusion
-                ):
-                    x = tensor_model_parallel_all_reduce(x)
-                    return self.forward(x, residual, None)
+                    else:
+                        logger.warning(
+                            "FlashInfer allreduce fusion failed, falling back to standard implementation"
+                        )
+                        x = tensor_model_parallel_all_reduce(x)
+                        return self.forward(x, residual, None)
 
         return self.forward(x, residual, post_residual_addition)