[MAIN][NVFP4] Support NVFP4 MOE with Proper Padding

megatron/core/fp4_utils.py

@@ -48,6 +48,28 @@ def is_nvfp4tensor(tensor: torch.Tensor) -> bool:
     return HAVE_TE_FP4_TENSOR_CLASS and isinstance(tensor, FP4_TENSOR_CLASS)
 
 
+def get_fp4_align_size(fp4_recipe: Fp4Recipe) -> int:
+    """
+    Get the alignment size required for FP4 GEMM.
+    FP4 GEMM requires Blackwell and later architectures.
+
+    The value 32 is a hardware requirement: TMA (Tensor Memory Accelerator) requires
+    a 16-byte aligned address for efficient memory access. Since FP4 uses 4 bits per value,
+    16 bytes (128 bits) corresponds to 32 FP4 values. Therefore, the alignment size for FP4
+    is 32. With this alignment, NVFP4 GEMM can be performed efficiently.
+
+    Note that since we are also random hadamard transform for NVFP4 training, we want
+    fused group nvfp4 quantize plus hadamard transform. Hadamard transform will leverage
+    tensor core instructions for better performance, while group quantize kernels also
+    prefer a more aligned size in token dimension M. Therefore, we apply align size 64
+    here for better performance in MOE.
+
+    Paper link: https://arxiv.org/pdf/2509.25149
+    """
+    # pylint: disable=unused-argument
+    return 64
+
+
 def dequantize_fp4_tensor(fp4_tensor: torch.Tensor) -> torch.Tensor:
     """Dequantize a fp4 tensor to a higher precision tensor."""
     if is_te_min_version("2.7.0.dev0"):

megatron/core/transformer/attention.py

@@ -59,7 +59,7 @@
 
 if not HAVE_FA3:
     try:
-        from flashattn_hopper.flash_attn_interface import _flash_attn_forward
+        from flash_attn_3.flash_attn_interface import _flash_attn_forward
         from flashattn_hopper.flash_attn_interface import (
             flash_attn_with_kvcache as flash_attn3_with_kvcache,
         )
@@ -216,12 +216,17 @@ def __init__(
 
         if (
             HAVE_TE
-            and self.config.fp8
-            and self.config.fp8_recipe != 'delayed'
-            and is_te_min_version("2.6.0dev0")
             and isinstance(self.linear_proj, TELinear)
+            and (
+                (
+                    self.config.fp8
+                    and self.config.fp8_recipe != 'delayed'
+                    and is_te_min_version("2.6.0dev0")
+                )
+                or (self.config.fp4 and is_te_min_version("2.7.0.dev0"))
+            )
         ):
-            # For fp8 training, the output of the fused core_attn is saved by itself, and
+            # For fp8/fp4 training, the output of the fused core_attn is saved by itself, and
             # linear_proj also saves the quantized tensor of this output. Here we set the
             # linear_proj to save the original input tensors to avoid the extra memory usage of
             # the quantized tensor.
@@ -1146,7 +1151,7 @@ def _backward_output_proj(self):
         self.linear_proj.backward_dw()
 
     def set_for_recompute_input_layernorm(self):
-        """Set the attention layer for recompute input_layernorm. Only needed for fp8."""
+        """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4."""
         from megatron.core.extensions.transformer_engine import set_save_original_input
 
         set_save_original_input(self.linear_qkv)

megatron/core/transformer/moe/README.md

@@ -240,7 +240,7 @@ Enable A2A overlap across different batches inspired by the DSv3 DualPipe implme
 | --moe-router-fusion | Enable fusion for MoE TopK routing and aux-loss computation. This is only supported in TransformerEngine 2.7.0 and above. |
 | --moe-router-bias-update-rate | The expert bias is updated based on the number of assigned tokens to each expert in a global batch, where the bias is increased for experts with less assigned tokens and decreased for experts with more assigned tokens. Default is 1e-3 same as that used in DeepSeekV3. |
 | --moe-router-force-load-balancing | (Experimental) Force override routing to balance token distribution using random logits for MoE routers, supporting naive top-k and group-limited top-k. This experimental feature is for benchmarking purposes only! |
-| --moe-router-padding-for-fp8 | Pad the routing_map to make sure the number of tokens each expert received is a multiple of 16/32 for FP8 precision. It is suggested to enable this for dropless training with FP8 precision when num_local_experts > 1. This is a more efficient way to pad for FP8 which eliminates the explicit padding in the GroupedMLP layer. |
+| --moe-router-padding-for-quantization | Pad the routing_map to make sure the number of tokens each expert received is a multiple of 16/32 for FP8/FP4 precision. It is suggested to enable this for dropless training with FP8 precision when num_local_experts > 1. This is a more efficient way to pad for FP8 which eliminates the explicit padding in the GroupedMLP layer. |
 | --moe-aux-loss-coeff | Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended. Default is 0.0. |
 | --moe-z-loss-coeff | Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended. Default is None. |
 | --moe-input-jitter-eps | Add noise to the input tensor by applying jitter with a specified epsilon value. Default is None. |
@@ -469,7 +469,7 @@ Therefore, there are two recommended ways during the first 200 steps to avoid th
 
 **FP8 Training Best Practice**
 - Using latest version of [TransformerEngine](https://github.com/NVIDIA/TransformerEngine).
-- Enable router padding with `--moe-router-padding-for-fp8` to reduce padding overhead.
+- Enable router padding with `--moe-router-padding-for-quantization` to reduce padding overhead.
 - Enable native FP8 weights with `--fp8-param-gather` to reduce weights memory cost.
 
 ### Reference Best Parallel Mapping

megatron/core/transformer/moe/experts.py

@@ -21,7 +21,6 @@
     ShardedTensorFactory,
 )
 from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
-from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.fusions.fused_bias_geglu import quick_gelu, weighted_bias_quick_geglu_impl
 from megatron.core.fusions.fused_bias_swiglu import weighted_bias_swiglu_impl
 from megatron.core.fusions.fused_weighted_squared_relu import weighted_squared_relu_impl
@@ -34,7 +33,10 @@
 from megatron.core.transformer.mlp import MLP, MLPSubmodules, apply_swiglu_sharded_factory
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe import grouped_gemm_util as gg
-from megatron.core.transformer.moe.moe_utils import ProcessGroupCollection
+from megatron.core.transformer.moe.moe_utils import (
+    ProcessGroupCollection,
+    get_align_size_for_quantization,
+)
 from megatron.core.transformer.spec_utils import build_module
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.utils import (
@@ -134,8 +136,10 @@ def glu(x):
             self.config.recompute_granularity == 'selective'
             and "moe_act" in self.config.recompute_modules
         )
-        if self.activation_recompute and self.config.fp8:
-            raise ValueError("moe_act recompute for fp8 cannot work with the legacy GroupedMLP.")
+        if self.activation_recompute and (self.config.fp8 or self.config.fp4):
+            raise ValueError(
+                "moe_act recompute for fp8 or fp4 cannot work with the legacy GroupedMLP."
+            )
 
         @jit_fuser
         def activation_func_with_probs(x, probs):
@@ -809,15 +813,15 @@ def __init__(
             self.config.recompute_granularity == 'selective'
             and "moe_act" in self.config.recompute_modules
         )
-        if self.activation_recompute and self.config.fp8:
+        if self.activation_recompute and (self.config.fp8 or self.config.fp4):
             from megatron.core.extensions.transformer_engine import set_save_original_input
 
             set_save_original_input(self.linear_fc2)
 
-        if self.config.fp8:
-            assert HAVE_TE, "FP8 requires TE."
-            self.fp8_padding = Fp8Padding(self.num_local_experts)
-            self.fp8_unpadding = Fp8Unpadding(self.num_local_experts)
+        if self.config.fp8 or self.config.fp4:
+            assert HAVE_TE, "FP8 and FP4 requires TE."
+            self.quantization_padding = Fp8Padding(self.num_local_experts)
+            self.quantization_unpadding = Fp8Unpadding(self.num_local_experts)
 
     @staticmethod
     def _apply_bias(intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs):
@@ -857,12 +861,12 @@ def forward(
             output (torch.Tensor): The output of the local experts.
         """
         tokens_per_expert = tokens_per_expert.tolist()
-        if self.config.fp8:
+        if self.config.fp8 or self.config.fp4:
             actual_tokens_per_expert = tokens_per_expert
-            permuted_local_hidden_states, tokens_per_expert = self.fp8_padding(
+            permuted_local_hidden_states, tokens_per_expert = self.quantization_padding(
                 permuted_local_hidden_states, tokens_per_expert
             )
-            permuted_probs, _ = self.fp8_padding(
+            permuted_probs, _ = self.quantization_padding(
                 permuted_probs.unsqueeze(-1), actual_tokens_per_expert
             )
         else:
@@ -954,8 +958,8 @@ def glu(x):
             output, output_bias = self.linear_fc2(intermediate_parallel, tokens_per_expert)
 
         # upad and concat the output
-        if self.config.fp8:
-            output = self.fp8_unpadding(output, actual_tokens_per_expert)
+        if self.config.fp8 or self.config.fp4:
+            output = self.quantization_unpadding(output, actual_tokens_per_expert)
 
         output = self._apply_bias(output, output_bias, tokens_per_expert, permuted_probs)
         output_bias = None
@@ -1051,10 +1055,10 @@ def __init__(
             )
             self.local_experts.append(expert)
 
-    def _pad_tensor_for_fp8(self, hidden, probs):
+    def _pad_tensor_for_quantization(self, hidden, probs):
         """Padding tensor shape to multiples of 16/32."""
         actual_num_tokens = hidden.shape[0]
-        divisor = get_fp8_align_size(self.config.fp8_recipe)
+        divisor = get_align_size_for_quantization(self.config)
         padded_num_tokens = ceil(actual_num_tokens / divisor) * divisor - actual_num_tokens
         if padded_num_tokens > 0:
             pad_tensor = torch.zeros(
@@ -1086,8 +1090,8 @@ def forward(
             permuted_probs = torch.ones_like(permuted_probs)
 
         if self.num_local_experts == 1:
-            if self.config.fp8:
-                hidden, probs = self._pad_tensor_for_fp8(
+            if self.config.fp8 or self.config.fp4:
+                hidden, probs = self._pad_tensor_for_quantization(
                     permuted_local_hidden_states, permuted_probs
                 )
                 output, output_bias = self.local_experts[0](hidden, probs)
@@ -1106,8 +1110,8 @@ def forward(
             output_local_list = []
 
             for expert, tokens, probs in zip(self.local_experts, tokens_list, probs_list):
-                if self.config.fp8:
-                    hidden, probs = self._pad_tensor_for_fp8(tokens, probs)
+                if self.config.fp8 or self.config.fp4:
+                    hidden, probs = self._pad_tensor_for_quantization(tokens, probs)
                     output, output_bias = expert(hidden, probs)
                     output = output[: tokens.shape[0]]
                 else:

megatron/core/transformer/moe/moe_layer.py

@@ -199,7 +199,7 @@ def shared_experts_compute(self, hidden_states: torch.Tensor):
         if self.use_shared_expert and not self.shared_expert_overlap:
             # Compute the shared expert separately when not overlapped with communication.
             if self.shared_experts_recompute:
-                if self.config.fp8:
+                if self.config.fp8 or self.config.fp4:
                     shared_expert_output = te_checkpoint(
                         self.shared_experts,
                         False,
@@ -278,7 +278,7 @@ def custom_forward(hidden_states):
             return output, mlp_bias
 
         if self.moe_layer_recompute:
-            if self.config.fp8:
+            if self.config.fp8 or self.config.fp4:
                 output, mlp_bias = te_checkpoint(
                     custom_forward,
                     False,
@@ -300,7 +300,7 @@ def backward_dw(self):
             self.shared_experts.backward_dw()
 
     def set_for_recompute_pre_mlp_layernorm(self):
-        """Set the MoE layer for recompute pre_mlp_layernorm. Only needed for fp8."""
+        """Set the MoE layer for recompute pre_mlp_layernorm. Only needed for fp8/fp4."""
         # If shared_experts_recompute is used, nothing needs to be done because the checkpoint
         # function will save the original input tensors.
         if self.shared_experts is not None and not self.shared_experts_recompute:

megatron/core/transformer/moe/moe_utils.py

@@ -6,7 +6,10 @@
 import torch
 
 from megatron.core import parallel_state
+from megatron.core.fp4_utils import get_fp4_align_size
+from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.transformer_config import TransformerConfig
 
 try:
     import transformer_engine as te  # pylint: disable=unused-import
@@ -1008,6 +1011,15 @@ def router_gating_linear(
     return RouterGatingLinearFunction.apply(inp, weight, bias, router_dtype)
 
 
+def get_align_size_for_quantization(config: TransformerConfig):
+    """Get the alignment size for quantization."""
+    if config.fp8:
+        return get_fp8_align_size(config.fp8_recipe)
+    elif config.fp4:
+        return get_fp4_align_size(config.fp4_recipe)
+    return 16
+
+
 # TODO(Hepteract): delete the usage of the global parallel_state.
 # Initialize process groups with the global parallel_state.
 def get_default_pg_collection():

megatron/core/transformer/moe/shared_experts.py

@@ -62,8 +62,10 @@ def __init__(
         else:
             self.gate_weight = None
 
-        if self.config.fp8 and is_te_min_version("2.6.0dev0"):
-            # For fp8 training, the output of pre_mlp_layernorm is saved by router, and
+        if (self.config.fp8 and is_te_min_version("2.6.0dev0")) or (
+            self.config.fp4 and is_te_min_version("2.7.0.dev0")
+        ):
+            # For fp8/fp4 training, the output of pre_mlp_layernorm is saved by router, and
             # the shared expert linear_fc1 also saves the quantized tensor of this output.
             # Here we set the linear_fc1 to save the original input tensors to avoid the extra
             # memory usage of the quantized tensor.

megatron/core/transformer/moe/token_dispatcher.py

@@ -8,7 +8,6 @@
 
 from megatron.core import utils
 from megatron.core.config import is_experimental_enabled
-from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.fusions.fused_indices_converter import fused_indices_to_multihot
 from megatron.core.fusions.fused_pad_routing_map import fused_pad_routing_map
 from megatron.core.tensor_parallel import (
@@ -25,6 +24,7 @@
 )
 from megatron.core.transformer.moe.moe_utils import (
     ProcessGroupCollection,
+    get_align_size_for_quantization,
     get_capacity,
     maybe_move_tensor_to_cpu,
     pad_routing_map,
@@ -476,7 +476,7 @@ def preprocess(self, routing_map: torch.Tensor) -> torch.Tensor:
 
         if (
             self.config.moe_expert_capacity_factor is not None
-            or self.config.moe_router_padding_for_fp8
+            or self.config.moe_router_padding_for_quantization
         ):
             # When using token dropping or router padding, output size is dynamic.
             # Need to sync output size GPU->CPU before allocating output buffer
@@ -578,8 +578,8 @@ def dispatch_preprocess(
         assert routing_map.dtype == torch.bool, "Expected bool tensor for mask"
         hidden_states = hidden_states.view(-1, self.hidden_shape[-1])
 
-        if self.config.moe_router_padding_for_fp8:
-            pad_multiple = get_fp8_align_size(self.config.fp8_recipe)
+        if self.config.moe_router_padding_for_quantization:
+            pad_multiple = get_align_size_for_quantization(self.config)
             if is_experimental_enabled() and self.config.moe_permute_fusion:
                 self.routing_map = fused_pad_routing_map(self.routing_map, pad_multiple)
             else:
@@ -1002,8 +1002,8 @@ def dispatch(
                     "HybridEP only supports float32 probs, please set --moe-router-dtype=fp32"
                 )
             self.token_probs = self.token_probs.float()  # downcast or upcast
-        if self.config.fp8:
-            self.pad_multiple = get_fp8_align_size(self.config.fp8_recipe)
+        if self.config.fp8 or self.config.fp4:
+            self.pad_multiple = get_align_size_for_quantization(self.config)
         dispatched_hidden, self.dispatched_probs, _, tokens_per_expert, self.handle = (
             hybrid_ep_dispatch(
                 x=hidden_states,
@@ -1224,7 +1224,7 @@ def _pad_routing_map(
         """
         Pad the routing map to the nearest multiple of the pad_multiple.
         """
-        pad_multiple = get_fp8_align_size(self.config.fp8_recipe)
+        pad_multiple = get_align_size_for_quantization(self.config)
 
         num_input_tokens = routing_map.shape[0]
         target_tokens_per_expert = (
@@ -1258,7 +1258,7 @@ def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) ->
             self.dispatched_routing_map, self.dispatched_probs = self._indices_to_multihot(
                 self.dispatched_indices, self.dispatched_probs
             )
-        if self.config.moe_router_padding_for_fp8:
+        if self.config.moe_router_padding_for_quantization:
             self.dispatched_routing_map, self.tokens_per_expert = self._pad_routing_map(
                 self.dispatched_routing_map, self.tokens_per_expert
             )

megatron/core/transformer/multi_latent_attention.py

@@ -177,12 +177,17 @@ def __init__(
 
         if (
             HAVE_TE
-            and self.config.fp8
-            and self.config.fp8_recipe != 'delayed'
-            and is_te_min_version("2.6.0dev0")
             and isinstance(self.linear_proj, TELinear)
+            and (
+                (
+                    self.config.fp8
+                    and self.config.fp8_recipe != 'delayed'
+                    and is_te_min_version("2.6.0dev0")
+                )
+                or (self.config.fp4 and is_te_min_version("2.7.0.dev0"))
+            )
         ):
-            # For fp8 training, the output of the fused core_attn is saved by itself, and
+            # For fp8/fp4 training, the output of the fused core_attn is saved by itself, and
             # linear_proj also saves the quantized tensor of this output. Here we set the
             # linear_proj to save the original input tensors to avoid the extra memory usage of
             # the quantized tensor.
@@ -781,7 +786,8 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
             return query, key, value
 
         if self.recompute_up_proj:
-            self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=self.config.fp8)
+            quantization = self.config.fp8 or self.config.fp4
+            self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization)
             query, key, value = self.qkv_up_checkpoint.checkpoint(
                 qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb
             )
@@ -911,7 +917,7 @@ def _backward_output_proj(self):
         self.linear_proj.backward_dw()
 
     def set_for_recompute_input_layernorm(self):
-        """Set the attention layer for recompute input_layernorm. Only needed for fp8."""
+        """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4."""
         from megatron.core.extensions.transformer_engine import set_save_original_input
 
         if self.config.q_lora_rank is not None:

megatron/core/transformer/multi_token_prediction.py

@@ -570,6 +570,8 @@ def _proj_and_transformer_layer(
             fp8_context = nullcontext()
             transformer_layer_fp8_context = nullcontext()
 
+        # TODO: currently ignoring FP4 in MTP layers because we need more numerical validation
+
         with rng_context:
             with fp8_context:
                 hidden_states = self._concat_embeddings(hidden_states, decoder_input)