chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Baizhou Zhang
2026-01-22 13:31:57 +08:00
committed by GitHub
parent 71482dd171
commit e2d33531f3
1 changed files with 21 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
python/sglang/srt/layers/communicator.py
View File

@@ -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