Skip to content

Use fused pad+MCT when available #142

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alextmagro wants to merge 2 commits into
base: rocm_dev
Choose a base branch
from
Open

Use fused pad+MCT when available #142

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
913b6bd
Use fused pad+MCT when available
alextmagro Jun 17, 2026
8c84746
Fix PR comments
alextmagro Jun 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions megatron/core/extensions/transformer_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -1754,8 +1755,18 @@ 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 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
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)
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; "
Expand Down
40 changes: 34 additions & 6 deletions megatron/core/transformer/moe/experts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,16 @@

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:
logger.warning(
"Fp8Padding.compute_padded_splits not found; "
"using unfused BF16 padding path for MoE."
)


class GroupedMLP(MegatronModule):
"""An efficient implementation of the Experts layer using GroupedGEMM.
Expand Down Expand Up @@ -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)
Comment on lines +662 to +665

@sudhu2k sudhu2k Jun 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't have a fp8 triton groupedgemm kernel for now and we would fall back to hipblasLT Multistream GEMM. But I guess it would be future proof to do this and the better way would be to pad on the gpu tensors directly instead of creating a new tensor from a list.

@alextmagro alextmagro Jun 26, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think maybe this should be postponed until we have intentions to do a grouped gemm in triton.

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
)
Expand All @@ -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 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:
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(
Expand Down Expand Up @@ -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 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,
"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)

Expand All @@ -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
Expand Down
Loading