From 050f108c29f0e581b602af41601d1f448d70a091 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Sun, 21 Dec 2025 09:34:04 +0800 Subject: [PATCH] Optimize Bailing-MoE with FlashInfer Fused All-Reduce (#15526) Co-authored-by: luoyuan.luo --- python/sglang/srt/models/bailing_moe.py | 78 ++++++++++++++++++------- 1 file changed, 58 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 02048850d..3fed3528d 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -19,7 +19,7 @@ # limitations under the License. """SGLang BailingMoE model.""" import logging -from typing import Iterable, Optional, Tuple, Union +from typing import Iterable, List, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -54,7 +54,11 @@ from sglang.srt.layers.linear import ( RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor -from sglang.srt.layers.moe import get_deepep_mode, get_moe_a2a_backend +from sglang.srt.layers.moe import ( + get_deepep_mode, + get_moe_a2a_backend, + should_use_flashinfer_cutlass_moe_fp4_allgather, +) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.token_dispatcher import DeepEPDispatcher @@ -124,6 +128,7 @@ class BailingMoEMLP(nn.Module): self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, + should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ) -> torch.Tensor: if (self.tp_size == 1) and hidden_states.shape[0] == 0: @@ -132,7 +137,7 @@ class BailingMoEMLP(nn.Module): gate_up, _ = self.gate_up_proj(hidden_states) hidden_states = self.act_fn(gate_up) hidden_states, _ = self.down_proj( - hidden_states, skip_all_reduce=use_reduce_scatter + hidden_states, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter ) return hidden_states @@ -301,10 +306,15 @@ class BailingMoESparseMoeBlock(nn.Module): self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, + should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ) -> torch.Tensor: if not get_moe_a2a_backend().is_deepep(): - return self.forward_normal(hidden_states, use_reduce_scatter) + return self.forward_normal( + hidden_states, + should_allreduce_fusion, + use_reduce_scatter, + ) else: return self.forward_deepep(hidden_states, forward_batch) @@ -344,6 +354,7 @@ class BailingMoESparseMoeBlock(nn.Module): def forward_normal( self, hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ) -> torch.Tensor: num_tokens, hidden_size = hidden_states.shape @@ -364,7 +375,12 @@ class BailingMoESparseMoeBlock(nn.Module): if self.num_shared_experts > 0: final_hidden_states = final_hidden_states + shared_output - if self.tp_size > 1 and not use_reduce_scatter: + if ( + self.tp_size > 1 + and not should_allreduce_fusion + and not use_reduce_scatter + and not should_use_flashinfer_cutlass_moe_fp4_allgather() + ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_size) @@ -560,6 +576,7 @@ class BailingMoEBlock(nn.Module): alt_stream: Optional[torch.cuda.Stream] = None, ): super().__init__() + self.config = config hidden_size = config.hidden_size self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps) @@ -625,6 +642,7 @@ class BailingMoEBlock(nn.Module): input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, allow_reduce_scatter=True, + is_last_layer=(self.layer_id == self.config.num_hidden_layers - 1), ) def _is_layer_sparse( @@ -640,18 +658,23 @@ class BailingMoEBlock(nn.Module): hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], + captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, ) -> torch.Tensor: - hidden_states, residual = self.layer_communicator.prepare_attn( - hidden_states=hidden_states, - residual=residual, - forward_batch=forward_batch, + hidden_states, residual = ( + self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( + hidden_states, + residual, + forward_batch, + captured_last_layer_outputs=captured_last_layer_outputs, + ) ) - hidden_states = self.attention( - positions=positions, - hidden_states=hidden_states, - forward_batch=forward_batch, - ) + if hidden_states.shape[0] != 0: + hidden_states = self.attention( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states=hidden_states, @@ -659,19 +682,28 @@ class BailingMoEBlock(nn.Module): forward_batch=forward_batch, ) + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + # For DP with padding, reduce scatter can be used instead of all-reduce. use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( forward_batch ) - hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) - - hidden_states, residual = self.layer_communicator.postprocess_layer( - hidden_states=hidden_states, - residual=residual, - forward_batch=forward_batch, + hidden_states = self.mlp( + hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter ) + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True + else: + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + return hidden_states, residual @@ -739,6 +771,7 @@ class BailingMoEModel(nn.Module): hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] + aux_hidden_states = [] for i in range(self.start_layer, self.end_layer): with get_global_expert_distribution_recorder().with_current_layer(i): layer = self.layers[i] @@ -747,6 +780,11 @@ class BailingMoEModel(nn.Module): hidden_states, forward_batch, residual, + captured_last_layer_outputs=( + aux_hidden_states + if getattr(layer, "_is_layer_to_capture", False) + else None + ), ) if not self.pp_group.is_last_rank: return PPProxyTensors(