[Kernel] Little refactor of flashinfer allreduce norm fusion (#17474)

python/sglang/srt/layers/communicator.py

@@ -80,9 +80,26 @@ if _use_aiter and _is_gfx95_supported:
 elif _is_npu:
     from sglang.srt.hardware_backend.npu.cmo import prepare_weight_cache
 
+
+# TODO: According to the discussion in https://github.com/flashinfer-ai/flashinfer/issues/1223#issuecomment-3047256465
+# We set the max token num to 128 for allreduce fusion with min-latency case(use_oneshot=True).
 FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048
 
 
+def apply_flashinfer_allreduce_fusion(batch_size: int):
+    return (
+        # TODO: flashinfer 0.6.1 caused performance regression on sm100 for allreduce fusion
+        # Temporarily disable it on sm100. Add it back after its performance is restored.
+        # Ref: https://github.com/sgl-project/sglang/issues/17237
+        _is_sm90_supported
+        and _is_flashinfer_available
+        and batch_size > 0
+        and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
+        and not is_dp_attention_enabled()
+        and get_global_server_args().enable_flashinfer_allreduce_fusion
+    )
+
+
 class ScatterMode(Enum):
     """
     Suppose we have TP=4, DP=2, enable-dp-attention, and the system handles seq a,b,c,d
@@ -581,24 +598,11 @@ class LayerCommunicator:
             if hasattr(forward_batch, "input_ids")
             else 0
         )
-        if batch_size > FUSE_ALLREDUCE_MAX_BATCH_SIZE:
-            return False
-
-        static_conditions_met = (
-            (not self.is_last_layer)
-            and (self._context.tp_size > 1)
-            and not is_dp_attention_enabled()
-            and get_global_server_args().enable_flashinfer_allreduce_fusion
-            and _is_flashinfer_available
-        )
-
-        if not static_conditions_met:
-            return False
 
         return (
-            batch_size > 0
+            apply_flashinfer_allreduce_fusion(batch_size)
-            and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
             and (not self.is_last_layer)
+            and (self._context.tp_size > 1)
         )
 
 
@@ -796,14 +800,8 @@ class CommunicateWithAllReduceAndLayerNormFn:
                 if hidden_states.shape[0] != 0:
                     hidden_states = layernorm(hidden_states)
         else:
-            # According to the discussion in https://github.com/flashinfer-ai/flashinfer/issues/1223#issuecomment-3047256465
+            if apply_flashinfer_allreduce_fusion(hidden_states.shape[0]) and hasattr(
-            # We set the max token num to 128 for allreduce fusion with min-latency case(use_oneshot=True).
+                layernorm, "forward_with_allreduce_fusion"
-            if (
-                (_is_sm100_supported or _is_sm90_supported)
-                and _is_flashinfer_available
-                and hasattr(layernorm, "forward_with_allreduce_fusion")
-                and get_global_server_args().enable_flashinfer_allreduce_fusion
-                and hidden_states.shape[0] <= 2048
             ):
                 hidden_states, residual = layernorm.forward_with_allreduce_fusion(
                     hidden_states, residual