EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Do layernorm before allgather for DP attention (#8631)

Browse Source
This commit is contained in:
Trevor Morris
2025-08-03 00:53:08 -07:00
committed by GitHub
parent f7b2853ff8
commit 32f2815451
1 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -404,14 +404,24 @@ class CommunicateWithAllReduceAndLayerNormFn:
if context.attn_dp_size != 1: if context.attn_dp_size != 1:
if context.attn_tp_rank == 0: if context.attn_tp_rank == 0:
hidden_states += residual hidden_states += residual
# Perform layernorm on smaller data before comm. Only valid when attn_tp_size is 1 (tp_size == dp_size)
use_layer_norm_before_gather = context.attn_tp_size == 1
if use_layer_norm_before_gather:
residual.copy_(hidden_states)
if hidden_states.shape[0] != 0:
hidden_states = layernorm(hidden_states)
hidden_states, local_hidden_states = ( hidden_states, local_hidden_states = (
forward_batch.gathered_buffer, forward_batch.gathered_buffer,
hidden_states, hidden_states,
) )
dp_gather_partial(hidden_states, local_hidden_states, forward_batch) dp_gather_partial(hidden_states, local_hidden_states, forward_batch)
dp_scatter(residual, hidden_states, forward_batch)
if hidden_states.shape[0] != 0: if not use_layer_norm_before_gather:
hidden_states = layernorm(hidden_states) dp_scatter(residual, hidden_states, forward_batch)
if hidden_states.shape[0] != 0:
hidden_states = layernorm(hidden_states)
else: else:
# According to the discussion in https://github.com/flashinfer-ai/flashinfer/issues/1223#issuecomment-3047256465 # 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). # We set the max token num to 128 for allreduce fusion with min-latency case(use_oneshot=True).