From e2d33531f396018ee2dc1a361c67b85223870590 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Thu, 22 Jan 2026 13:31:57 +0800 Subject: [PATCH] [Kernel] Little refactor of flashinfer allreduce norm fusion (#17474) --- python/sglang/srt/layers/communicator.py | 44 +++++++++++------------- 1 file changed, 21 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 1636ed706..0115a189c 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/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 - and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE + apply_flashinfer_allreduce_fusion(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 - # We set the max token num to 128 for allreduce fusion with min-latency case(use_oneshot=True). - 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 + if apply_flashinfer_allreduce_fusion(hidden_states.shape[0]) and hasattr( + layernorm, "forward_with_allreduce_fusion" ): hidden_states, residual = layernorm.forward_with_allreduce_fusion( hidden_states, residual