From 913b6bd7c336b82cae35f708cbd3c7fb0455abf3 Mon Sep 17 00:00:00 2001 From: Alex Magro Date: Wed, 17 Jun 2026 17:41:03 +0000 Subject: [PATCH 1/2] Use fused pad+MCT when available --- .../core/extensions/transformer_engine.py | 12 +++++- megatron/core/transformer/moe/experts.py | 40 ++++++++++++++++--- 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 3ee4fc30488..d2a975dcae7 100755 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1743,7 +1743,8 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool: self.te_quant_params, self.training, is_context_quantized ) - def forward(self, x, m_splits, m_splits_gpu=None): + def forward(self, x, m_splits, m_splits_gpu=None, actual_m_splits=None, + unpad_output=False): """Forward.""" _is_first_microbatch = ( None if self.disable_parameter_transpose_cache else self.is_first_microbatch @@ -1754,8 +1755,15 @@ def forward(self, x, m_splits, m_splits_gpu=None): # Added in TE commit: https://github.com/ROCm/TransformerEngine/commit/2776c33 if is_te_min_version("2.7.0", check_equality=False): if class_has_method_param(te.pytorch.GroupedLinear, "forward", "m_splits_tensor"): + fused_padding_kwargs = {} + if class_has_method_param( + te.pytorch.GroupedLinear, "forward", "actual_m_splits" + ): + fused_padding_kwargs["actual_m_splits"] = actual_m_splits + fused_padding_kwargs["unpad_output"] = unpad_output with quant_context: - out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch, m_splits_tensor=m_splits_gpu) + out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch, + m_splits_tensor=m_splits_gpu, **fused_padding_kwargs) else: warnings.warn( "Transformer Engine is missing `m_splits_tensor` parameter in GroupedLinear.forward; " diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 21efa6e7390..7c3d5431ded 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -59,6 +59,16 @@ HAVE_TE = False +is_rocm = hasattr(torch.version, 'hip') and torch.version.hip is not None +HAVE_FUSED_PADDING_MCT = False +if HAVE_TE and is_rocm: + HAVE_FUSED_PADDING_MCT = hasattr(Fp8Padding, 'compute_padded_splits') + if not HAVE_FUSED_PADDING_MCT: + logging.getLogger(__name__).info( + "Fp8Padding.compute_padded_splits not found; " + "using unfused BF16 padding path for MoE." + ) + logger = logging.getLogger(__name__) @@ -643,14 +653,20 @@ def forward( output (torch.Tensor): The output of the local experts. """ tokens_per_expert_gpu = None + actual_tokens_per_expert = None if isinstance(tokens_per_expert, tuple): tokens_per_expert, tokens_per_expert_gpu = tokens_per_expert tokens_per_expert = tokens_per_expert.tolist() if self.config.fp8 or self.config.fp4: actual_tokens_per_expert = tokens_per_expert - permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( - permuted_local_hidden_states, tokens_per_expert - ) + if HAVE_FUSED_PADDING_MCT: + # ROCm: fuse padding into cast_transpose — compute padded + # splits without allocating a BF16 intermediate. + tokens_per_expert = self.quantization_padding.compute_padded_splits(tokens_per_expert) + else: + permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( + permuted_local_hidden_states, tokens_per_expert + ) permuted_probs, _ = self.quantization_padding( permuted_probs.unsqueeze(-1), actual_tokens_per_expert ) @@ -667,11 +683,12 @@ def forward( # Probs already applied, so reset to 1. permuted_probs = torch.ones_like(permuted_probs) + _fused_kwargs = {"actual_m_splits": actual_tokens_per_expert} if HAVE_FUSED_PADDING_MCT else {} with off_interface( self.offload_expert_fc1, permuted_local_hidden_states, "expert_fc1" ) as permuted_local_hidden_states: fc1_output, bias_parallel = self.linear_fc1( - permuted_local_hidden_states, tokens_per_expert, tokens_per_expert_gpu + permuted_local_hidden_states, tokens_per_expert, tokens_per_expert_gpu, **_fused_kwargs, ) if self.offload_expert_fc1: fc1_output = off_interface.group_commit( @@ -748,7 +765,18 @@ def glu(x): with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: bias_act_output = bias_act_func(fc1_output, bias_parallel, permuted_probs) - output, output_bias = self.linear_fc2(bias_act_output, tokens_per_expert, tokens_per_expert_gpu) + _can_unpad_fc2 = HAVE_FUSED_PADDING_MCT and not self.config.add_bias_linear + if _can_unpad_fc2: + _fc2_fused_kwargs = { + "actual_m_splits": actual_tokens_per_expert, + "unpad_output": True, + } + else: + _fc2_fused_kwargs = {} + output, output_bias = self.linear_fc2( + bias_act_output, tokens_per_expert, tokens_per_expert_gpu, + **_fc2_fused_kwargs, + ) if self.activation_recompute: self.activation_checkpoint.discard_output_and_register_recompute(output) @@ -761,7 +789,7 @@ def glu(x): output = self._apply_bias(output, output_bias, tokens_per_expert, permuted_probs) # upad and concat the output - if self.config.fp8 or self.config.fp4: + if (self.config.fp8 or self.config.fp4) and not _can_unpad_fc2: output = self.quantization_unpadding(output, actual_tokens_per_expert) output_bias = None From 8c84746d233ebc90a425f9a156083aa4ec27d0a4 Mon Sep 17 00:00:00 2001 From: Alex Magro Date: Fri, 26 Jun 2026 18:01:20 +0000 Subject: [PATCH 2/2] Fix PR comments --- megatron/core/extensions/transformer_engine.py | 7 +++++-- megatron/core/transformer/moe/experts.py | 10 +++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index d2a975dcae7..5c282c9c82d 100755 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1756,11 +1756,14 @@ def forward(self, x, m_splits, m_splits_gpu=None, actual_m_splits=None, if is_te_min_version("2.7.0", check_equality=False): if class_has_method_param(te.pytorch.GroupedLinear, "forward", "m_splits_tensor"): fused_padding_kwargs = {} - if class_has_method_param( + if actual_m_splits is not None and class_has_method_param( te.pytorch.GroupedLinear, "forward", "actual_m_splits" ): fused_padding_kwargs["actual_m_splits"] = actual_m_splits - fused_padding_kwargs["unpad_output"] = unpad_output + if unpad_output and class_has_method_param( + te.pytorch.GroupedLinear, "forward", "unpad_output" + ): + fused_padding_kwargs["unpad_output"] = unpad_output with quant_context: out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch, m_splits_tensor=m_splits_gpu, **fused_padding_kwargs) diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 7c3d5431ded..25992e98839 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -59,18 +59,18 @@ HAVE_TE = False +logger = logging.getLogger(__name__) + is_rocm = hasattr(torch.version, 'hip') and torch.version.hip is not None HAVE_FUSED_PADDING_MCT = False if HAVE_TE and is_rocm: HAVE_FUSED_PADDING_MCT = hasattr(Fp8Padding, 'compute_padded_splits') if not HAVE_FUSED_PADDING_MCT: - logging.getLogger(__name__).info( + logger.warning( "Fp8Padding.compute_padded_splits not found; " "using unfused BF16 padding path for MoE." ) -logger = logging.getLogger(__name__) - class GroupedMLP(MegatronModule): """An efficient implementation of the Experts layer using GroupedGEMM. @@ -683,7 +683,7 @@ def forward( # Probs already applied, so reset to 1. permuted_probs = torch.ones_like(permuted_probs) - _fused_kwargs = {"actual_m_splits": actual_tokens_per_expert} if HAVE_FUSED_PADDING_MCT else {} + _fused_kwargs = {"actual_m_splits": actual_tokens_per_expert} if HAVE_FUSED_PADDING_MCT and actual_tokens_per_expert is not None else {} with off_interface( self.offload_expert_fc1, permuted_local_hidden_states, "expert_fc1" ) as permuted_local_hidden_states: @@ -765,7 +765,7 @@ def glu(x): with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: bias_act_output = bias_act_func(fc1_output, bias_parallel, permuted_probs) - _can_unpad_fc2 = HAVE_FUSED_PADDING_MCT and not self.config.add_bias_linear + _can_unpad_fc2 = HAVE_FUSED_PADDING_MCT and actual_tokens_per_expert is not None and not self.config.add_bias_linear if _can_unpad_fc2: _fc2_fused_kwargs = { "actual_m_splits": actual_tokens_per_expert,