From 8f40b5eb3f59e031bfa9fd161a240aaf50444172 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 6 Jul 2026 13:01:53 -0700 Subject: [PATCH] When attention TP for linear and full attention, use Flashinfer allreduce fusion (#29699) Co-authored-by: Brayden Zhong --- .../attention/hybrid_linear_attn_backend.py | 2 + .../srt/layers/attention/mamba/mamba.py | 5 +- python/sglang/srt/models/nemotron_h.py | 74 +++++++++++++++---- 3 files changed, 64 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index efefc6bb3..101b50287 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -742,6 +742,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase): forward_batch: ForwardBatch, mup_vector: Optional[torch.Tensor] = None, use_triton_causal_conv: bool = False, + should_allreduce_fusion: bool = False, ): assert isinstance(self.forward_metadata, Mamba2Metadata) # Page-major stores state strided; only the stride-aware Triton causal-conv @@ -759,6 +760,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase): forward_batch=forward_batch, mup_vector=mup_vector, use_triton_causal_conv=use_triton_causal_conv, + should_allreduce_fusion=should_allreduce_fusion, ) if forward_batch.mamba_track_mask is not None: diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py index 80ec14481..7e6c072a5 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba.py +++ b/python/sglang/srt/layers/attention/mamba/mamba.py @@ -448,6 +448,7 @@ class MambaMixer2(torch.nn.Module): forward_batch: ForwardBatch, mup_vector: Optional[torch.Tensor] = None, use_triton_causal_conv: bool = False, + should_allreduce_fusion: bool = False, ): # Returns the projected result. When `output` is given it is also # written into that buffer (required by the cuda-graph split ops, which @@ -760,7 +761,9 @@ class MambaMixer2(torch.nn.Module): # norm usage hidden_states = self.norm(preallocated_ssm_out, gate) - mixer_out, _ = self.out_proj(hidden_states) + mixer_out, _ = self.out_proj( + hidden_states, skip_all_reduce=should_allreduce_fusion + ) if output is not None: output[:padded_num_tokens].copy_(mixer_out) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 3eb70777f..1fc5fe2f0 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -495,11 +495,18 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer): ) self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) - self.layer_communicator = make_layer_communicator(self.norm, for_attn=True) + self.layer_communicator = make_layer_communicator( + self.norm, + for_attn=True, + is_last_layer=layer_idx == len(config.hybrid_override_pattern) - 1, + ) self._set_prev_layer_is_attn(config, layer_idx) def _forward_mamba( - self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + should_allreduce_fusion: bool = False, ) -> torch.Tensor: """Core Mamba forward logic, called directly or via split op.""" original_num_tokens = hidden_states.shape[0] @@ -517,6 +524,7 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer): output=None, forward_batch=forward_batch, use_triton_causal_conv=True, + should_allreduce_fusion=should_allreduce_fusion, ) return pad_to_original_num_tokens(output, original_num_tokens) @@ -544,18 +552,30 @@ class NemotronHMambaDecoderLayer(NemotronHAttnLikeDecoderLayer): self.norm, hidden_states, residual ) + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + if is_in_breakable_cuda_graph(): output = torch.empty_like(hidden_states) - breakable_nemotron_mamba2_with_output(hidden_states, output, self.layer_id) - return output, residual - - if is_in_tc_piecewise_cuda_graph(): + breakable_nemotron_mamba2_with_output( + hidden_states, output, self.layer_id, should_allreduce_fusion + ) + elif is_in_tc_piecewise_cuda_graph(): output = torch.empty_like(hidden_states) - nemotron_mamba2_with_output(hidden_states, output, self.layer_id) - return output, residual + nemotron_mamba2_with_output( + hidden_states, output, self.layer_id, should_allreduce_fusion + ) else: - output = self._forward_mamba(hidden_states, forward_batch) - return output, residual + output = self._forward_mamba( + hidden_states, forward_batch, should_allreduce_fusion + ) + + if should_allreduce_fusion: + output._sglang_needs_allreduce_fusion = True + return output, residual class NemotronHAttention(nn.Module): @@ -625,13 +645,18 @@ class NemotronHAttention(nn.Module): ) def forward( - self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + should_allreduce_fusion: bool = False, ) -> torch.Tensor: if not is_dp_attention_enabled(): qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) attn_output = self.attn.forward(q, k, v, forward_batch) - output, _ = self.o_proj(attn_output) + output, _ = self.o_proj( + attn_output, skip_all_reduce=should_allreduce_fusion + ) return output padded_shape = hidden_states.shape[0] @@ -682,7 +707,11 @@ class NemotronHAttentionDecoderLayer(NemotronHAttnLikeDecoderLayer): ) self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) - self.layer_communicator = make_layer_communicator(self.norm, for_attn=True) + self.layer_communicator = make_layer_communicator( + self.norm, + for_attn=True, + is_last_layer=layer_idx == len(config.hybrid_override_pattern) - 1, + ) self._set_prev_layer_is_attn(config, layer_idx) def forward( @@ -705,9 +734,19 @@ class NemotronHAttentionDecoderLayer(NemotronHAttnLikeDecoderLayer): self.norm, hidden_states, residual ) - hidden_states = self.mixer.forward( - hidden_states=hidden_states, forward_batch=forward_batch + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) ) + + hidden_states = self.mixer.forward( + hidden_states=hidden_states, + forward_batch=forward_batch, + should_allreduce_fusion=should_allreduce_fusion, + ) + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True return hidden_states, residual @@ -1168,6 +1207,7 @@ def nemotron_mamba2_with_output( hidden_states: torch.Tensor, output: torch.Tensor, layer_id: int, + should_allreduce_fusion: bool = False, ) -> None: """Split op for Mamba2 forward in piecewise CUDA graph mode.""" context = get_tc_piecewise_forward_context() @@ -1187,7 +1227,9 @@ def nemotron_mamba2_with_output( if hidden_states.shape[0] != num_actual_tokens: hidden_states = hidden_states[:num_actual_tokens] - ret = mamba_layer._forward_mamba(hidden_states, forward_batch) + ret = mamba_layer._forward_mamba( + hidden_states, forward_batch, should_allreduce_fusion + ) # Copy result back; output may be larger (padded) so only fill actual tokens output[:num_actual_tokens].view(ret.shape).copy_(ret)