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

fuse allreduce and residual_rmsnorm (#8731)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-08-11 13:50:53 -07:00
committed by GitHub
parent 8c07fabda7
commit 44e86480e8
8 changed files with 135 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+49 -34
View File
@@ -212,7 +212,7 @@ class DeepseekV2MLP(nn.Module):
self,
x,
forward_batch=None,
can_fuse_mlp_allreduce: bool = False,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
):
if (self.tp_size == 1) and x.shape[0] == 0:
@@ -221,7 +221,7 @@ class DeepseekV2MLP(nn.Module):
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x, skip_all_reduce=can_fuse_mlp_allreduce or use_reduce_scatter
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return x
@@ -448,7 +448,7 @@ class DeepseekV2MoE(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
can_fuse_mlp_allreduce: bool = False,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if not self._enable_deepep_moe:
@@ -459,11 +459,11 @@ class DeepseekV2MoE(nn.Module):
and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD
):
return self.forward_normal_dual_stream(
hidden_states, can_fuse_mlp_allreduce, use_reduce_scatter
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else:
return self.forward_normal(
hidden_states, can_fuse_mlp_allreduce, use_reduce_scatter
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else:
return self.forward_deepep(hidden_states, forward_batch)
@@ -471,7 +471,7 @@ class DeepseekV2MoE(nn.Module):
def forward_normal_dual_stream(
self,
hidden_states: torch.Tensor,
can_fuse_mlp_allreduce: bool = False,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
@@ -500,20 +500,20 @@ class DeepseekV2MoE(nn.Module):
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
final_hidden_states = final_hidden_states_out
sm.tag(final_hidden_states)
if self.tp_size > 1 and not can_fuse_mlp_allreduce and not use_reduce_scatter:
if self.tp_size > 1 and not should_allreduce_fusion and not use_reduce_scatter:
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states
def forward_normal(
self,
hidden_states: torch.Tensor,
can_fuse_mlp_allreduce: bool = False,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if hasattr(self, "shared_experts") and use_intel_amx_backend(
self.shared_experts.gate_up_proj
):
return self.forward_cpu(hidden_states, can_fuse_mlp_allreduce)
return self.forward_cpu(hidden_states, should_allreduce_fusion)
shared_output = self._forward_shared_experts(hidden_states)
# router_logits: (num_tokens, n_experts)
@@ -537,12 +537,14 @@ class DeepseekV2MoE(nn.Module):
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
final_hidden_states = final_hidden_states_out
sm.tag(final_hidden_states)
if self.tp_size > 1 and not can_fuse_mlp_allreduce and not use_reduce_scatter:
if self.tp_size > 1 and not should_allreduce_fusion and not use_reduce_scatter:
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states
def forward_cpu(
self, hidden_states: torch.Tensor, can_fuse_mlp_allreduce: bool = False
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
) -> torch.Tensor:
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states)
@@ -593,7 +595,7 @@ class DeepseekV2MoE(nn.Module):
None, # a2_scale
True, # is_vnni
)
if self.tp_size > 1 and not can_fuse_mlp_allreduce:
if self.tp_size > 1 and not should_allreduce_fusion:
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states
@@ -1842,6 +1844,8 @@ class DeepseekV2DecoderLayer(nn.Module):
allow_reduce_scatter=True,
)
self._fuse_allreduce_lookup_table = self._build_fuse_allreduce_lookup_table()
def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool:
return is_nextn or (
self.config.n_routed_experts is not None
@@ -1850,27 +1854,18 @@ class DeepseekV2DecoderLayer(nn.Module):
)
def _should_fuse_mlp_allreduce_with_next_layer(self, forward_batch) -> bool:
"""Check if MLP allreduce can be fused with next layer's add_rmsnorm"""
"""Check if MLP allreduce can be fused with next layer's residual_rmsnorm"""
if (
self.layer_id == self.config.num_hidden_layers - 1
or get_tensor_model_parallel_world_size() <= 1
):
batch_size = (
forward_batch.input_ids.shape[0]
if hasattr(forward_batch, "input_ids")
else 0
)
if batch_size > 128:
return False
if not global_server_args_dict.get("enable_flashinfer_allreduce_fusion", False):
return False
if not _is_sm100_supported or not _is_flashinfer_available:
return False
if hasattr(forward_batch, "input_ids") and (
forward_batch.input_ids.shape[0] == 0
or forward_batch.input_ids.shape[0] > 128
):
return False
return True
return self._fuse_allreduce_lookup_table.get(batch_size, False)
def forward(
self,
@@ -1896,7 +1891,7 @@ class DeepseekV2DecoderLayer(nn.Module):
hidden_states, residual, forward_batch
)
can_fuse_mlp_allreduce = (
should_allreduce_fusion = (
self._should_fuse_mlp_allreduce_with_next_layer(forward_batch)
and not (self.enable_dp_attention and self.speculative_algorithm.is_eagle())
and not self.is_nextn
@@ -1907,13 +1902,13 @@ class DeepseekV2DecoderLayer(nn.Module):
forward_batch
)
hidden_states = self.mlp(
hidden_states, forward_batch, can_fuse_mlp_allreduce, use_reduce_scatter
hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
)
if can_fuse_mlp_allreduce:
if should_allreduce_fusion:
hidden_states._sglang_needs_allreduce_fusion = True
if not can_fuse_mlp_allreduce:
if not should_allreduce_fusion:
hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch
)
@@ -1990,6 +1985,26 @@ class DeepseekV2DecoderLayer(nn.Module):
)
return output
def _build_fuse_allreduce_lookup_table(self):
static_conditions_met = (
self.layer_id != self.config.num_hidden_layers - 1
and get_tensor_model_parallel_world_size() > 1
and global_server_args_dict.get("enable_flashinfer_allreduce_fusion", False)
and _is_sm100_supported
and _is_flashinfer_available
)
if not static_conditions_met:
return {}
lookup_table = {}
for batch_size in range(129): # 0 to 128
is_last_layer = self.layer_id == self.config.num_hidden_layers - 1
should_fuse = batch_size > 0 and batch_size <= 128 and not is_last_layer
lookup_table[batch_size] = should_fuse
return lookup_table
class DeepseekV2Model(nn.Module):
fall_back_to_pt_during_load = False