[TRTLLM-6445] feat: Enable AllReduce associated fusion patterns in Llama3/4.

Signed-off-by: Yukun He <[email protected]>

Author: hyukn

tensorrt_llm/_torch/models/modeling_llama.py

@@ -1,4 +1,5 @@
 import copy
+import os
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
@@ -337,7 +338,7 @@ def forward(
         assert shared_output.size() == routed_output.size(
         ), f'unmatched tensor shape'
         final_hidden_states = shared_output + routed_output
-        if not self.enable_attention_dp and self.mapping.tp_size > 1:
+        if not self.enable_attention_dp and self.mapping.has_tp():
             final_hidden_states = self.all_reduce(
                 final_hidden_states, all_reduce_params=final_all_reduce_params)
 
@@ -367,9 +368,6 @@ def __init__(
         self.fusion_config = EagerFusionConfig()
         # self.fusion_config.PRE_MOE_FUSION = model_config.mapping.has_tp(
         # )
-        # TODO: re-enable these fusions
-        self.fusion_config.PRE_MOE_FUSION = False
-        self.fusion_config.POST_MLP_FUSION = False
 
         nope_layer = config.no_rope_layers[layer_idx] == 0
         attention_chunk_size = getattr(config, "attention_chunk_size",
@@ -387,6 +385,20 @@ def __init__(
         self.is_mlp_layer = (layer_idx +
                              1) % config.interleave_moe_layer_step != 0
 
+        self.enable_fusion = os.environ.get(
+            "TRTLLM_LLAMA_EAGER_FUSION_DISABLED", "0") == "0"
+
+        if self.is_nvfp4:
+            self.pre_feed_forward_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4
+            self.post_feed_forward_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4
+        # TODO: enable fp8 quant fusion later
+        # elif self.is_fp8_quant:
+        #     self.pre_feed_forward_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8
+        #     self.post_feed_forward_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8
+        else:
+            self.pre_feed_forward_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM
+            self.post_feed_forward_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM
+
         if self.is_mlp_layer:
             self.feed_forward = GatedMLP(
                 hidden_size=config.hidden_size,
@@ -399,8 +411,11 @@ def __init__(
                 layer_idx=layer_idx,
             )
 
-            # self.fusion_config.POST_MLP_FUSION = model_config.mapping.has_tp(
-            # )
+            self.fusion_config.PRE_MLP_FUSION = model_config.mapping.has_tp(
+            ) and not self.enable_attention_dp and self.enable_fusion
+            self.fusion_config.POST_MLP_FUSION = model_config.mapping.has_tp(
+            ) and not self.enable_attention_dp and self.enable_fusion
+
         else:
             self.feed_forward = Llama4MoE(
                 num_experts=config.num_local_experts,
@@ -413,8 +428,10 @@ def __init__(
                 dtype=config.torch_dtype,
                 layer_idx=layer_idx)
 
-            # self.fusion_config.POST_MOE_FUSION = model_config.mapping.has_tp(
-            # )
+            self.fusion_config.PRE_MOE_FUSION = model_config.mapping.has_tp(
+            ) and not self.enable_attention_dp and self.enable_fusion
+            self.fusion_config.POST_MOE_FUSION = model_config.mapping.has_tp(
+            ) and not self.enable_attention_dp and self.enable_fusion
 
         self.input_layernorm = RMSNorm(hidden_size=config.hidden_size,
                                        eps=config.rms_norm_eps,
@@ -432,6 +449,17 @@ def __init__(
 
         self.moe_allreduce = MoEAllReduce(self.mapping)
 
+        self.disable_attn_allreduce = (self.fusion_config.PRE_MOE_FUSION
+                                       or self.fusion_config.PRE_MLP_FUSION
+                                       or self.mapping.tp_size == 1
+                                       or self.enable_attention_dp)
+        self.disable_feed_forward_allreduce = (
+            self.fusion_config.POST_MOE_FUSION
+            or self.fusion_config.POST_MLP_FUSION or self.mapping.tp_size == 1
+            or self.enable_attention_dp)
+
+        print(f"init Llama4DecoderLayer")
+
     def forward(
         self,
         position_ids: torch.IntTensor,
@@ -461,34 +489,43 @@ def forward(
             position_ids=position_ids,
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
-            all_reduce_params=AllReduceParams(enable_allreduce=not (
-                self.fusion_config.PRE_MOE_FUSION or self.mapping.tp_size == 1
-                or self.enable_attention_dp)),
+            all_reduce_params=AllReduceParams(
+                enable_allreduce=not self.disable_attn_allreduce),
             **kwargs,
         )
 
-        if self.fusion_config.PRE_MOE_FUSION:
-            hidden_states, residual = self.all_reduce(
+        if self.is_nvfp4 or self.is_fp8_quant:
+            scale = self.self_attn.qkv_proj.input_scale
+        else:
+            scale = None
+
+        if self.fusion_config.PRE_MLP_FUSION or self.fusion_config.PRE_MOE_FUSION:
+            allreduce_output = self.all_reduce(
                 hidden_states,
                 all_reduce_params=AllReduceParams(
-                    fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+                    fusion_op=self.pre_feed_forward_fusion_op,
                     residual=residual,
                     norm_weight=self.post_attention_layernorm.weight,
+                    scale=scale,
                     eps=self.post_attention_layernorm.variance_epsilon,
                 ))
         else:
             # Fully Connected
-            hidden_states, residual = self.post_attention_layernorm(
+            allreduce_output = self.post_attention_layernorm(
                 hidden_states, residual)
 
+        if self.is_nvfp4:
+            act_fp4, act_sf, residual = allreduce_output
+            hidden_states = Fp4QuantizedTensor(act_fp4, act_sf)
+        else:
+            hidden_states, residual = allreduce_output
+
         hidden_states = self.feed_forward(
             hidden_states,
             all_rank_num_tokens=attn_metadata.all_rank_num_tokens,
             all_rank_max_num_tokens=attn_metadata.all_rank_max_num_tokens,
-            final_all_reduce_params=AllReduceParams(enable_allreduce=not (
-                self.fusion_config.POST_MOE_FUSION
-                or self.fusion_config.POST_MLP_FUSION
-                or self.mapping.tp_size == 1 or self.enable_attention_dp)),
+            final_all_reduce_params=AllReduceParams(
+                enable_allreduce=not self.disable_feed_forward_allreduce),
             cutlass_min_latency_mode=cutlass_min_latency_mode,
         )
 
@@ -500,16 +537,23 @@ def forward(
             spec_metadata.maybe_capture_hidden_states(self.layer_idx,
                                                       hidden_states, residual)
 
-        if (self.fusion_config.POST_MOE_FUSION
+        if (
+                self.fusion_config.POST_MOE_FUSION
                 or self.fusion_config.POST_MLP_FUSION
-            ) and self.next_layer_layernorm is not None:
+        ) and self.next_layer_layernorm is not None and self.next_attn is not None:
+            # Get the scale for the next allreduce fusion op
+            if self.is_nvfp4 or self.is_fp8_quant:
+                scale = self.next_attn.qkv_proj.input_scale
+            else:
+                scale = None
+
             if cutlass_min_latency_mode:
                 shared_output = hidden_states[0]
                 hidden_states_activated_experts = hidden_states[1]
                 num_activated_experts_per_node = hidden_states[2]
                 experts_to_token_score = hidden_states[3]
 
-                hidden_states, residual = self.moe_allreduce(
+                allreduce_output = self.moe_allreduce(
                     residual,
                     self.next_layer_layernorm.weight,
                     device_num_experts=num_activated_experts_per_node,
@@ -519,18 +563,30 @@ def forward(
                     eps=self.next_layer_layernorm.variance_epsilon,
                 )
             else:
-                hidden_states, residual = self.all_reduce(
-                    hidden_states,
-                    all_reduce_params=AllReduceParams(
-                        fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
-                        residual=residual,
-                        norm_weight=self.next_layer_layernorm.weight,
-                        eps=self.next_layer_layernorm.variance_epsilon,
-                    ))
+                if self.fusion_config.POST_MLP_FUSION or self.fusion_config.POST_MOE_FUSION:
+                    allreduce_output = self.all_reduce(
+                        hidden_states,
+                        all_reduce_params=AllReduceParams(
+                            fusion_op=self.post_feed_forward_fusion_op,
+                            residual=residual,
+                            norm_weight=self.next_layer_layernorm.weight,
+                            scale=scale,
+                            eps=self.next_layer_layernorm.variance_epsilon,
+                        ))
+                else:
+                    raise ValueError("Unknown fusion config")
         elif self.next_layer_layernorm:
-            hidden_states, residual = self.next_layer_layernorm(
+            print(f"{self.layer_idx}, {self.next_layer_layernorm}")
+            allreduce_output = self.next_layer_layernorm(
                 hidden_states, residual)
 
+        print(f"in forward")
+        if self.is_nvfp4:
+            act_fp4, act_sf, residual = allreduce_output
+            hidden_states = Fp4QuantizedTensor(act_fp4, act_sf)
+        else:
+            hidden_states, residual = allreduce_output
+
         return hidden_states, residual
 
 
@@ -544,6 +600,14 @@ def __init__(
         super().__init__()
         config = model_config.pretrained_config
         self.layer_idx = layer_idx
+        self.mapping = model_config.mapping
+        self.enable_attention_dp = model_config.mapping.enable_attention_dp
+        self.is_quanted = model_config.quant_config and model_config.quant_config.quant_mode.has_any_quant(
+        )
+        self.is_fp8_quant = self.is_quanted and model_config.quant_config.quant_mode.has_fp8_qdq(
+        )
+        self.is_nvfp4 = self.is_quanted and model_config.quant_config.quant_mode.has_nvfp4(
+        )
 
         self.self_attn = LlamaAttention(
             model_config,
@@ -566,11 +630,43 @@ def __init__(
                                                 eps=config.rms_norm_eps,
                                                 dtype=config.torch_dtype)
 
+        self.all_reduce = AllReduce(mapping=model_config.mapping)
+
+        self.next_layer_layernorm: RMSNorm = None
+        self.next_attn: LlamaAttention = None
+
         self.attention_mask = PredefinedAttentionMask.CAUSAL
         # If the model is being used as an encoder model (prefill only) we use a full attention mask
         if not model_config.is_generation:
             self.attention_mask = PredefinedAttentionMask.FULL
 
+        self.enable_fusion = os.environ.get(
+            "TRTLLM_LLAMA_EAGER_FUSION_DISABLED", "0") == "0"
+        self.PRE_MLP_FUSION = self.mapping.has_tp(
+        ) and not self.enable_attention_dp and self.enable_fusion
+        self.POST_MLP_FUSION = self.mapping.has_tp() and self.enable_fusion
+
+        if self.is_nvfp4:
+            self.pre_mlp_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4
+            self.post_mlp_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4
+        # TODO: enable fp8 quant fusion later
+        # elif self.is_fp8_quant:
+        #     self.pre_mlp_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8
+        #     self.post_mlp_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8
+        else:
+            self.pre_mlp_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM
+            self.post_mlp_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM
+
+        # TODO: Disable this to avoid large accuracy drop
+        self.POST_MLP_FUSION = False
+
+        self.disable_attn_allreduce = (self.PRE_MLP_FUSION
+                                       or self.mapping.tp_size == 1
+                                       or self.enable_attention_dp)
+        self.disable_mlp_allreduce = (self.POST_MLP_FUSION
+                                      or self.mapping.tp_size == 1
+                                      or self.enable_attention_dp)
+
     def forward(
         self,
         position_ids: torch.IntTensor,
@@ -583,30 +679,78 @@ def forward(
         if residual is None:
             residual = hidden_states
             hidden_states = self.input_layernorm(hidden_states)
-        else:
-            hidden_states, residual = self.input_layernorm(
-                hidden_states, residual)
 
         # Self Attention
         hidden_states = self.self_attn(
             position_ids=position_ids,
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
             attention_mask=self.attention_mask,
+            all_reduce_params=AllReduceParams(
+                enable_allreduce=not self.disable_attn_allreduce),
             **kwargs,
         )
-
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(
-            hidden_states, residual)
-        hidden_states = self.mlp(hidden_states, **kwargs)
+        if self.PRE_MLP_FUSION:
+            if self.is_nvfp4:
+                scale = self.mlp.gate_up_proj.input_scale
+            else:
+                scale = None
+            all_reduce_output = self.all_reduce(
+                hidden_states,
+                all_reduce_params=AllReduceParams(
+                    fusion_op=self.pre_mlp_fusion_op,
+                    residual=residual,
+                    norm_weight=self.post_attention_layernorm.weight,
+                    scale=scale,
+                    eps=self.post_attention_layernorm.variance_epsilon,
+                ))
+            if self.is_nvfp4:
+                act_fp4, act_sf, residual = all_reduce_output
+                hidden_states = Fp4QuantizedTensor(act_fp4, act_sf)
+            else:
+                hidden_states, residual = all_reduce_output
+        else:
+            hidden_states, residual = self.post_attention_layernorm(
+                hidden_states, residual)
+
+        hidden_states = self.mlp(
+            hidden_states,
+            final_all_reduce_params=AllReduceParams(
+                enable_allreduce=not self.disable_mlp_allreduce),
+            **kwargs,
+        )
+
         if spec_metadata is not None:
             # We save the hidden states in the spec metadata here. In _prepare_draft_tokens,
             # PyExecutor will extract these from the model engine's spec metadata.
             # They will be passed to the draft model engine on the first draft iteration.
             # TODO: can we support multiple model outputs instead?
             spec_metadata.maybe_capture_hidden_states(self.layer_idx,
                                                       hidden_states, residual)
+        if self.POST_MLP_FUSION and self.next_attn is not None:
+            if self.is_nvfp4:
+                scale = self.next_attn.qkv_proj.input_scale
+            else:
+                scale = None
+            all_reduce_output = self.all_reduce(
+                hidden_states,
+                all_reduce_params=AllReduceParams(
+                    fusion_op=self.post_mlp_fusion_op,
+                    residual=residual,
+                    norm_weight=self.next_layer_layernorm.weight,
+                    scale=scale,
+                    eps=self.next_layer_layernorm.variance_epsilon,
+                ))
+            if self.is_nvfp4:
+                act_fp4, act_sf, residual = all_reduce_output
+                hidden_states = Fp4QuantizedTensor(act_fp4, act_sf)
+            else:
+                hidden_states, residual = all_reduce_output
+        elif self.next_layer_layernorm:
+            hidden_states, residual = self.next_layer_layernorm(
+                hidden_states, residual)
+
         return hidden_states, residual
 
 
@@ -727,7 +871,7 @@ def __init__(self, model_config: ModelConfig[LlamaConfig]):
 
         if self.has_custom_embed_tokens:
             with torch.no_grad():
-                if model_config.mapping.tp_size > 1:
+                if model_config.mapping.has_tp():
                     weight = split_matrix_tp(
                         weight,
                         model_config.mapping.tp_size,
@@ -775,7 +919,6 @@ def forward(
                 lora_params=lora_params,
             )
 
-        hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
 
@@ -788,6 +931,18 @@ def __init__(
     ):
         super().__init__(LlamaModel(model_config), model_config)
 
+    def load_weights(self, weights: Dict):
+        super().load_weights(weights)
+
+        for idx, layer in enumerate(
+                self.model.layers[:self.config.num_hidden_layers]):
+            if idx == self.config.num_hidden_layers - 1:
+                layer.next_layer_layernorm = self.model.norm
+            else:
+                layer.next_layer_layernorm = self.model.layers[
+                    idx + 1].input_layernorm
+                layer.next_attn = self.model.layers[idx + 1].self_attn
+
 
 class Llama4InputProcessor(InputProcessor):