From 1f61e2407b7641c7865c2c87094f5f2308da3ada Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Sat, 22 Aug 2026 03:14:59 -0700
Subject: [PATCH 1/9] perf(moe): add frozen torch grouped GEMM backend

Signed-off-by: Taufeeque <taufeeque@far.ai>
(cherry picked from commit 193d5fe2bf3abf9f987015e3633faf331f94cf02)
---
 .../core/extensions/transformer_engine.py     | 106 ++++++++++++++
 .../core/transformer/transformer_config.py    |  32 ++++
 .../moe/test_frozen_torch_grouped_mlp.py      | 137 ++++++++++++++++++
 3 files changed, 275 insertions(+)
 create mode 100644 tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 7f7424e6c02..bbdaf39e6fe 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -11,6 +11,7 @@
 import re
 import warnings
 from contextlib import contextmanager, nullcontext
+from itertools import accumulate
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, cast
 
 import torch
@@ -2358,6 +2359,7 @@ def __init__(
                 name (str | None): module instance name passed top-down from its paranet module
             """
             self.config = config
+            self.expert_gemm_backend = config.moe_expert_gemm_backend
 
             # TE returns a zero length Tensor when bias=False and
             # return_bias=True, but we prefer None.  So in that case we
@@ -2632,6 +2634,107 @@ def _split_grouped_checkpoint_tensor(
                 f"into {self.num_gemms} GEMM shards."
             )
 
+        def _torch_grouped_weight_layout_is_current(self) -> bool:
+            """Return whether each expert parameter still views the contiguous buffer."""
+            grouped_weight = self._buffers.get('_torch_grouped_weight')
+            if grouped_weight is None or grouped_weight.shape[0] != self.num_gemms:
+                return False
+
+            first_weight = self.weight0
+            last_weight = getattr(self, f'weight{self.num_gemms - 1}')
+            if (
+                grouped_weight.device != first_weight.device
+                or grouped_weight.dtype != first_weight.dtype
+                or grouped_weight.shape[1:] != first_weight.shape
+                or last_weight.device != first_weight.device
+            ):
+                return False
+
+            storage_pointer = grouped_weight.untyped_storage().data_ptr()
+            expert_elements = first_weight.numel()
+            return (
+                first_weight.untyped_storage().data_ptr() == storage_pointer
+                and first_weight.storage_offset() == 0
+                and first_weight.stride() == grouped_weight[0].stride()
+                and last_weight.untyped_storage().data_ptr() == storage_pointer
+                and last_weight.storage_offset() == (self.num_gemms - 1) * expert_elements
+                and last_weight.stride() == grouped_weight[-1].stride()
+            )
+
+        @torch.no_grad()
+        def prepare_torch_grouped_mm(self) -> int:
+            """Relocate frozen BF16 expert weights into one contiguous allocation.
+
+            The existing per-expert ``Parameter`` objects become views of the new allocation, so
+            parameter names and checkpoint structure remain unchanged. The backing allocation is a
+            non-persistent buffer and therefore does not add a checkpoint entry.
+
+            Returns:
+                Number of bytes in the contiguous backing allocation.
+            """
+            if self.expert_gemm_backend != 'torch':
+                raise RuntimeError(
+                    "prepare_torch_grouped_mm requires moe_expert_gemm_backend='torch'"
+                )
+            if not hasattr(torch, '_grouped_mm'):
+                raise RuntimeError("this PyTorch build does not provide torch._grouped_mm")
+            if self.use_bias:
+                raise RuntimeError("torch grouped expert GEMM does not support bias")
+            if getattr(self, 'single_grouped_weight', False):
+                raise RuntimeError(
+                    "torch grouped expert GEMM expects per-GEMM weight parameters, but "
+                    "moe_single_grouped_weight=True makes TE hold one grouped weight tensor"
+                )
+
+            weights = [getattr(self, f'weight{index}') for index in range(self.num_gemms)]
+            if any(weight.requires_grad for weight in weights):
+                raise RuntimeError("torch grouped expert GEMM requires frozen base weights")
+            if any(weight.device.type != 'cuda' for weight in weights):
+                raise RuntimeError("torch grouped expert GEMM requires CUDA weights")
+            if any(weight.dtype != torch.bfloat16 for weight in weights):
+                raise RuntimeError("torch grouped expert GEMM requires BF16 weights")
+            if any(weight.shape != weights[0].shape for weight in weights[1:]):
+                raise RuntimeError("torch grouped expert GEMM requires uniform expert shapes")
+
+            if not self._torch_grouped_weight_layout_is_current():
+                grouped_weight = torch.stack([weight.detach() for weight in weights]).contiguous()
+                for index, weight in enumerate(weights):
+                    weight.data = grouped_weight[index]
+                if '_torch_grouped_weight' in self._buffers:
+                    self._buffers['_torch_grouped_weight'] = grouped_weight
+                else:
+                    self.register_buffer('_torch_grouped_weight', grouped_weight, persistent=False)
+
+            grouped_weight = self._buffers['_torch_grouped_weight']
+            return grouped_weight.numel() * grouped_weight.element_size()
+
+        @property
+        def torch_grouped_mm_prepared(self) -> bool:
+            """Return whether the torch grouped GEMM backing allocation is ready."""
+            return self._torch_grouped_weight_layout_is_current()
+
+        def _torch_grouped_mm_forward(self, x, m_splits):
+            """Run the frozen expert base branch with ``torch._grouped_mm``."""
+            if not self._torch_grouped_weight_layout_is_current():
+                self.prepare_torch_grouped_mm()
+
+            x_2d = x.reshape(-1, x.shape[-1])
+            input_rows = sum(m_splits)
+            if input_rows != x_2d.shape[0]:
+                raise RuntimeError(
+                    f"expert splits sum to {input_rows}, but input has {x_2d.shape[0]} rows"
+                )
+
+            grouped_weight = self._buffers['_torch_grouped_weight']
+            if input_rows == 0:
+                output = x_2d.matmul(grouped_weight[0].transpose(0, 1))
+            else:
+                offsets = torch.tensor(
+                    list(accumulate(m_splits)), device=x.device, dtype=torch.int32
+                )
+                output = torch._grouped_mm(x_2d, grouped_weight.transpose(1, 2), offs=offsets)
+            return output.reshape(*x.shape[:-1], output.shape[-1]), None
+
         def finish_init(self, quantization_config: QuantizationConfig):
             """Post-init of quantization override"""
             if quantization_config is None:
@@ -2647,6 +2750,9 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool:
 
         def forward(self, x, m_splits):
             """Forward."""
+            if self.expert_gemm_backend == 'torch':
+                return self._torch_grouped_mm_forward(x, m_splits)
+
             _is_first_microbatch = (
                 None if self.disable_parameter_transpose_cache else self.is_first_microbatch
             )
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 66cfaded213..6812ef1880a 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -843,6 +843,14 @@ class TransformerConfig(ModelParallelConfig):
     parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``
     and ``add_bias_linear=True``."""
 
+    moe_expert_gemm_backend: Literal['transformer_engine', 'torch'] = 'transformer_engine'
+    """Backend for grouped expert linear layers during training.
+
+    ``torch`` uses ``torch._grouped_mm`` for frozen BF16, bias-free expert weights. It is intended
+    for parameter-efficient fine-tuning where only the adapter branch is trainable. The default
+    ``transformer_engine`` backend supports trainable expert weights and other precisions.
+    """
+
     moe_aux_loss_coeff: Union[float, List[float]] = 0.0
     """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended.
     If a list of load balancing types is provided for `moe_router_load_balancing_type`,
@@ -1551,6 +1559,30 @@ def __post_init__(self):
         if self.num_moe_experts is not None and self.num_moe_experts <= 0:
             raise ValueError("num_moe_experts must be non-negative.")
 
+        if self.moe_expert_gemm_backend not in ('transformer_engine', 'torch'):
+            raise ValueError(
+                "moe_expert_gemm_backend must be 'transformer_engine' or 'torch', "
+                f"got {self.moe_expert_gemm_backend!r}"
+            )
+        if self.moe_expert_gemm_backend == 'torch':
+            if self.num_moe_experts is None:
+                raise ValueError("moe_expert_gemm_backend='torch' requires num_moe_experts")
+            if not self.moe_grouped_gemm:
+                raise ValueError("moe_expert_gemm_backend='torch' requires moe_grouped_gemm=True")
+            if not self.bf16 or self.params_dtype != torch.bfloat16:
+                raise ValueError("moe_expert_gemm_backend='torch' requires BF16 parameters")
+            if self.add_bias_linear:
+                raise ValueError("moe_expert_gemm_backend='torch' does not support expert bias")
+            if self.fp8 or self.fp4:
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' does not support FP8 or FP4 experts"
+                )
+            if self.moe_single_grouped_weight:
+                raise ValueError(
+                    "moe_expert_gemm_backend='torch' requires per-GEMM expert weight parameters "
+                    "and is incompatible with moe_single_grouped_weight=True"
+                )
+
         if self.num_moe_experts is not None and self.moe_ffn_hidden_size is None:
             self.moe_ffn_hidden_size = self.ffn_hidden_size
             warnings.warn("moe_ffn_hidden_size is not set, using ffn_hidden_size instead.")
diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
new file mode 100644
index 00000000000..714a657a571
--- /dev/null
+++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
@@ -0,0 +1,137 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_layer_with_transformer_engine_submodules,
+)
+from megatron.core.transformer.module import Float16Module
+from megatron.core.transformer.moe.moe_layer import MoELayer
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import is_te_min_version
+from tests.unit_tests.test_utilities import Utils
+
+
+def _config(*, backend: str = "transformer_engine") -> TransformerConfig:
+    return TransformerConfig(
+        num_layers=1,
+        hidden_size=16,
+        num_attention_heads=4,
+        num_moe_experts=4,
+        moe_grouped_gemm=True,
+        moe_expert_gemm_backend=backend,
+        add_bias_linear=False,
+        gated_linear_unit=False,
+        activation_func=F.gelu,
+        bias_activation_fusion=False,
+        bf16=True,
+        params_dtype=torch.bfloat16,
+        moe_router_load_balancing_type="sinkhorn",
+        moe_router_topk=1,
+    )
+
+
+@pytest.mark.parametrize(
+    "override,error",
+    [
+        ({"num_moe_experts": None}, "requires num_moe_experts"),
+        ({"moe_grouped_gemm": False}, "requires moe_grouped_gemm=True"),
+        ({"bf16": False, "params_dtype": torch.float32}, "requires BF16 parameters"),
+        ({"add_bias_linear": True}, "does not support expert bias"),
+    ],
+)
+def test_torch_grouped_expert_gemm_config_validation(override, error):
+    kwargs = {
+        "num_layers": 1,
+        "hidden_size": 16,
+        "num_attention_heads": 4,
+        "num_moe_experts": 2,
+        "moe_grouped_gemm": True,
+        "moe_expert_gemm_backend": "torch",
+        "add_bias_linear": False,
+        "bf16": True,
+        "params_dtype": torch.bfloat16,
+    }
+    kwargs.update(override)
+
+    with pytest.raises(ValueError, match=error):
+        TransformerConfig(**kwargs)
+
+
+@pytest.mark.skipif(
+    not is_te_min_version("1.9.0.dev0")
+    or not torch.cuda.is_available()
+    or not hasattr(torch, "_grouped_mm")
+    or torch.cuda.get_device_capability()[0] < 10,
+    reason="torch grouped expert GEMM requires TE and a compatible CUDA device",
+)
+class TestFrozenTorchGroupedMLP:
+    def setup_method(self):
+        Utils.initialize_model_parallel(1, 1)
+
+    def teardown_method(self):
+        Utils.destroy_model_parallel()
+
+    @staticmethod
+    def _model(config: TransformerConfig) -> MoELayer:
+        model = MoELayer(
+            config,
+            get_gpt_layer_with_transformer_engine_submodules(
+                config.num_moe_experts, moe_grouped_gemm=True
+            ).mlp.submodules,
+        )
+        return Float16Module(config, model).module.cuda()
+
+    def test_forward_input_gradient_and_checkpoint_parity(self):
+        baseline = self._model(_config())
+        torch_model = self._model(_config(backend="torch"))
+        torch_model.load_state_dict(baseline.state_dict())
+        torch_linears = (torch_model.experts.linear_fc1, torch_model.experts.linear_fc2)
+        for linear in torch_linears:
+            for parameter in linear.parameters():
+                parameter.requires_grad = False
+        state_keys_before = set(torch_model.state_dict())
+
+        baseline_input = torch.rand(
+            (32, 2, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True
+        )
+        torch_input = baseline_input.detach().clone().requires_grad_(True)
+        baseline_output, _ = baseline(baseline_input)
+        torch_output, _ = torch_model(torch_input)
+        grad_output = torch.randn_like(baseline_output)
+        baseline_output.backward(grad_output)
+        torch_output.backward(grad_output)
+
+        torch.testing.assert_close(torch_output, baseline_output, rtol=0.02, atol=0.02)
+        torch.testing.assert_close(torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02)
+        assert set(torch_model.state_dict()) == state_keys_before
+        assert not any("_torch_grouped_weight" in key for key in state_keys_before)
+        for linear in torch_linears:
+            assert linear.torch_grouped_mm_prepared
+            assert (
+                linear._buffers["_torch_grouped_weight"].untyped_storage().data_ptr()
+                == linear.weight0.untyped_storage().data_ptr()
+            )
+
+    def test_no_tokens(self):
+        config = deepcopy(_config(backend="torch"))
+        model = self._model(config)
+        for parameter in model.experts.parameters():
+            parameter.requires_grad = False
+        hidden_states = torch.empty(
+            (0, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True
+        )
+
+        output, _ = model.experts(
+            hidden_states,
+            tokens_per_expert=torch.zeros(4, dtype=torch.int32, device="cuda"),
+            permuted_probs=torch.empty(0, dtype=torch.float32, device="cuda"),
+        )
+        output.sum().backward()
+
+        assert output.shape == (0, 16)
+        assert hidden_states.grad is not None

From dcf358be9ba3956a6f9a46a2859b8451f46e933a Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Sat, 22 Aug 2026 03:20:58 -0700
Subject: [PATCH 2/9] test(moe): cover zero-token grouped expert splits

Signed-off-by: Taufeeque <taufeeque@far.ai>
(cherry picked from commit ecd9ca99d2346c01e5b967c41675aef09e85d90c)
---
 .../moe/test_frozen_torch_grouped_mlp.py      | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
index 714a657a571..4d3ee6b93b8 100644
--- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
+++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
@@ -135,3 +135,31 @@ def test_no_tokens(self):
 
         assert output.shape == (0, 16)
         assert hidden_states.grad is not None
+
+    def test_zero_token_groups(self):
+        baseline = self._model(_config())
+        torch_model = self._model(_config(backend="torch"))
+        torch_model.load_state_dict(baseline.state_dict())
+        for parameter in torch_model.experts.parameters():
+            parameter.requires_grad = False
+
+        baseline_input = torch.rand(
+            (5, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True
+        )
+        torch_input = baseline_input.detach().clone().requires_grad_(True)
+        tokens_per_expert = torch.tensor([2, 0, 3, 0], device="cuda")
+        permuted_probs = torch.rand(5, dtype=torch.float32, device="cuda")
+        baseline_output, _ = baseline.experts(
+            baseline_input, tokens_per_expert, permuted_probs
+        )
+        torch_output, _ = torch_model.experts(
+            torch_input, tokens_per_expert, permuted_probs
+        )
+        grad_output = torch.randn_like(baseline_output)
+        baseline_output.backward(grad_output)
+        torch_output.backward(grad_output)
+
+        torch.testing.assert_close(torch_output, baseline_output, rtol=0.02, atol=0.02)
+        torch.testing.assert_close(
+            torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02
+        )

From 356b80d44a91619b69f07b53e6ae278450ede3b3 Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Sat, 22 Aug 2026 03:23:34 -0700
Subject: [PATCH 3/9] fix(moe): validate every grouped expert weight view

Signed-off-by: Taufeeque <taufeeque@far.ai>
(cherry picked from commit f2a9b74de1dfd52407d8481bcd0fdb0a05b13503)
---
 .../core/extensions/transformer_engine.py     | 23 +++++++++++--------
 .../moe/test_frozen_torch_grouped_mlp.py      | 15 ++++++++++++
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index bbdaf39e6fe..80f6b8eb2c7 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -2641,25 +2641,28 @@ def _torch_grouped_weight_layout_is_current(self) -> bool:
                 return False
 
             first_weight = self.weight0
-            last_weight = getattr(self, f'weight{self.num_gemms - 1}')
             if (
                 grouped_weight.device != first_weight.device
                 or grouped_weight.dtype != first_weight.dtype
                 or grouped_weight.shape[1:] != first_weight.shape
-                or last_weight.device != first_weight.device
             ):
                 return False
 
             storage_pointer = grouped_weight.untyped_storage().data_ptr()
             expert_elements = first_weight.numel()
-            return (
-                first_weight.untyped_storage().data_ptr() == storage_pointer
-                and first_weight.storage_offset() == 0
-                and first_weight.stride() == grouped_weight[0].stride()
-                and last_weight.untyped_storage().data_ptr() == storage_pointer
-                and last_weight.storage_offset() == (self.num_gemms - 1) * expert_elements
-                and last_weight.stride() == grouped_weight[-1].stride()
-            )
+            for index in range(self.num_gemms):
+                weight = getattr(self, f'weight{index}')
+                grouped_weight_view = grouped_weight[index]
+                if (
+                    weight.device != grouped_weight.device
+                    or weight.dtype != grouped_weight.dtype
+                    or weight.shape != grouped_weight_view.shape
+                    or weight.untyped_storage().data_ptr() != storage_pointer
+                    or weight.storage_offset() != index * expert_elements
+                    or weight.stride() != grouped_weight_view.stride()
+                ):
+                    return False
+            return True
 
         @torch.no_grad()
         def prepare_torch_grouped_mm(self) -> int:
diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
index 4d3ee6b93b8..c9970dda8ee 100644
--- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
+++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
@@ -136,6 +136,21 @@ def test_no_tokens(self):
         assert output.shape == (0, 16)
         assert hidden_states.grad is not None
 
+    def test_reprepares_after_middle_expert_weight_replacement(self):
+        model = self._model(_config(backend="torch"))
+        linear = model.experts.linear_fc1
+        for parameter in linear.parameters():
+            parameter.requires_grad = False
+        linear.prepare_torch_grouped_mm()
+        replacement = torch.randn_like(linear.weight1)
+
+        linear.weight1.data = replacement
+
+        assert not linear.torch_grouped_mm_prepared
+        linear.prepare_torch_grouped_mm()
+        assert linear.torch_grouped_mm_prepared
+        torch.testing.assert_close(linear.weight1, replacement)
+
     def test_zero_token_groups(self):
         baseline = self._model(_config())
         torch_model = self._model(_config(backend="torch"))

From 814469a43145e09464c2738ed534b8ad3a3f329e Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Sat, 22 Aug 2026 07:24:29 -0700
Subject: [PATCH 4/9] perf(moe): checkpoint expert FC1 activation

Signed-off-by: Taufeeque <taufeeque@far.ai>
(cherry picked from commit 931658048abb094f7c66a7bf943ee298780f449f)
---
 megatron/core/transformer/moe/experts.py      | 99 ++++++++++++-------
 .../core/transformer/transformer_config.py    | 38 ++++++-
 .../moe/test_frozen_torch_grouped_mlp.py      | 85 +++++++++++++++-
 3 files changed, 184 insertions(+), 38 deletions(-)

diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
index 59f8deffeca..c47c372b546 100644
--- a/megatron/core/transformer/moe/experts.py
+++ b/megatron/core/transformer/moe/experts.py
@@ -269,7 +269,19 @@ def __init__(
             self.config.recompute_granularity == 'selective'
             and "moe_act" in self.config.recompute_modules
         )
-        if self.activation_recompute and (self.config.fp8 or self.config.fp4):
+        self.expert_fc1_activation_recompute = (
+            self.config.recompute_granularity == 'selective'
+            and "expert_fc1_act" in self.config.recompute_modules
+        )
+        if self.expert_fc1_activation_recompute and (
+            self.offload_expert_fc1 or self.offload_moe_act
+        ):
+            raise ValueError(
+                "expert_fc1_act recomputation cannot be combined with expert activation offload."
+            )
+        if (self.activation_recompute or self.expert_fc1_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)
@@ -789,21 +801,6 @@ def forward(
             # Probs already applied, so reset to 1.
             permuted_probs = torch.ones_like(permuted_probs)
 
-        expert_fc1_manager = off_interface(
-            self.offload_expert_fc1, permuted_local_hidden_states, "expert_fc1"
-        )
-        with expert_fc1_manager as permuted_local_hidden_states:
-            fc1_output, bias_parallel = apply_module(self.linear_fc1)(
-                permuted_local_hidden_states, tokens_per_expert
-            )
-        fc1_output = expert_fc1_manager.group_offload(
-            fc1_output,
-            forced_released_tensors=[permuted_local_hidden_states],
-            delay_offload=self.config.delay_offload_until_cuda_graph,
-        )
-
-        moe_act_manager = off_interface(self.offload_moe_act, fc1_output, "moe_act")
-
         def bias_act_func(intermediate_parallel, bias_parallel, permuted_probs):
 
             # Whether activation function is interleaved GLU
@@ -879,26 +876,60 @@ def glu(x):
                 intermediate_parallel = intermediate_parallel.to(original_dtype)
             return intermediate_parallel
 
-        if self.activation_recompute:
-            self.activation_checkpoint = tensor_parallel.CheckpointWithoutOutput()
-            with moe_act_manager as fc1_output:
-                bias_act_output = self.activation_checkpoint.checkpoint(
-                    bias_act_func, fc1_output, bias_parallel, permuted_probs
+        if self.expert_fc1_activation_recompute:
+            # Discard the FC1 output as well as the activation output, and recompute both in the
+            # backward pass. The FC1 output is the largest expert activation, so checkpointing the
+            # pair saves more memory than checkpointing the activation alone.
+            expert_fc1_activation_checkpoint = tensor_parallel.CheckpointWithoutOutput(
+                fp8=self.config.fp8 or self.config.fp4
+            )
+
+            def expert_fc1_activation(hidden_states, expert_probs):
+                fc1_output, bias_parallel = apply_module(self.linear_fc1)(
+                    hidden_states, tokens_per_expert
                 )
+                return bias_act_func(fc1_output, bias_parallel, expert_probs)
+
+            bias_act_output = expert_fc1_activation_checkpoint.checkpoint(
+                expert_fc1_activation, permuted_local_hidden_states, permuted_probs
+            )
+            output, output_bias = apply_module(self.linear_fc2)(bias_act_output, tokens_per_expert)
+            expert_fc1_activation_checkpoint.discard_output_and_register_recompute(output)
         else:
-            with moe_act_manager as fc1_output:
-                bias_act_output = bias_act_func(fc1_output, bias_parallel, permuted_probs)
-        output, output_bias = apply_module(self.linear_fc2)(bias_act_output, tokens_per_expert)
-        if self.activation_recompute:
-            self.activation_checkpoint.discard_output_and_register_recompute(output)
-
-        # Delay the offload of the moe act until after the linear_fc2 has been computed
-        # to make sure the fc1_output is reloaded to GPU before recomputing moe_act.
-        output = moe_act_manager.group_offload(
-            output,
-            forced_released_tensors=[fc1_output],
-            delay_offload=self.config.delay_offload_until_cuda_graph,
-        )
+            expert_fc1_manager = off_interface(
+                self.offload_expert_fc1, permuted_local_hidden_states, "expert_fc1"
+            )
+            with expert_fc1_manager as permuted_local_hidden_states:
+                fc1_output, bias_parallel = apply_module(self.linear_fc1)(
+                    permuted_local_hidden_states, tokens_per_expert
+                )
+            fc1_output = expert_fc1_manager.group_offload(
+                fc1_output,
+                forced_released_tensors=[permuted_local_hidden_states],
+                delay_offload=self.config.delay_offload_until_cuda_graph,
+            )
+
+            moe_act_manager = off_interface(self.offload_moe_act, fc1_output, "moe_act")
+            if self.activation_recompute:
+                self.activation_checkpoint = tensor_parallel.CheckpointWithoutOutput()
+                with moe_act_manager as fc1_output:
+                    bias_act_output = self.activation_checkpoint.checkpoint(
+                        bias_act_func, fc1_output, bias_parallel, permuted_probs
+                    )
+            else:
+                with moe_act_manager as fc1_output:
+                    bias_act_output = bias_act_func(fc1_output, bias_parallel, permuted_probs)
+            output, output_bias = apply_module(self.linear_fc2)(bias_act_output, tokens_per_expert)
+            if self.activation_recompute:
+                self.activation_checkpoint.discard_output_and_register_recompute(output)
+
+            # Delay the offload of the moe act until after the linear_fc2 has been computed
+            # to make sure the fc1_output is reloaded to GPU before recomputing moe_act.
+            output = moe_act_manager.group_offload(
+                output,
+                forced_released_tensors=[fc1_output],
+                delay_offload=self.config.delay_offload_until_cuda_graph,
+            )
         output = self._apply_bias(output, output_bias, tokens_per_expert, permuted_probs)
 
         # upad and concat the output
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 6812ef1880a..6bc07f4859c 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -545,10 +545,11 @@ class TransformerConfig(ModelParallelConfig):
 
     recompute_modules: Optional[List[str]] = None
     """The submodules to recompute.
-    choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe",
+    choices: "core_attn", "expert_fc1_act", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe",
     "shared_experts", "gdn_norm_out".
     default: ["core_attn"].
     "core_attn": recompute the core attention part of the transformer layer.
+    "expert_fc1_act": recompute grouped expert FC1 and its activation together.
     "moe_act": recompute the MoE MLP activation function.
     "layernorm": recompute the input_layernorm and pre_mlp_layernorm.
     "mla_up_proj": recompute the MLA up projection and RoPE applying parts.
@@ -556,7 +557,8 @@ class TransformerConfig(ModelParallelConfig):
     "moe": recompute the MoE layer.
     "shared_experts": recompute the shared experts in the MoE layer.
     "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all.
-    "moe_act", "layernorm", "mla_up_proj", and "gdn_norm_out" use output-discarding checkpointing,
+    "expert_fc1_act", "moe_act", "layernorm", "mla_up_proj", and "gdn_norm_out" use
+    output-discarding checkpointing,
     "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing.
     """
 
@@ -1825,6 +1827,7 @@ def __post_init__(self):
             if len(self.recompute_modules) > 0:
                 allowed_modules = {
                     "core_attn",
+                    "expert_fc1_act",
                     "moe_act",
                     "layernorm",
                     "mla_up_proj",
@@ -1844,6 +1847,31 @@ def __post_init__(self):
                     "moe_act in recompute_modules is only supported with moe_grouped_gemm."
                 )
 
+            if "expert_fc1_act" in self.recompute_modules:
+                if not self.moe_grouped_gemm:
+                    raise ValueError(
+                        "expert_fc1_act in recompute_modules requires moe_grouped_gemm."
+                    )
+                if self.transformer_impl != "transformer_engine":
+                    raise ValueError(
+                        "expert_fc1_act in recompute_modules requires transformer_engine."
+                    )
+                if self.use_transformer_engine_op_fuser:
+                    raise ValueError(
+                        "expert_fc1_act in recompute_modules is not supported by the "
+                        "Transformer Engine op fuser path."
+                    )
+                conflicting_modules = {"moe", "moe_act"}.intersection(self.recompute_modules)
+                if conflicting_modules:
+                    raise ValueError(
+                        "expert_fc1_act cannot be combined with MoE recompute modules: "
+                        f"{sorted(conflicting_modules)}"
+                    )
+                if self.fp8 or self.fp4:
+                    raise ValueError(
+                        "expert_fc1_act in recompute_modules currently supports BF16/FP16 only."
+                    )
+
             if "mla_up_proj" in self.recompute_modules and not self.multi_latent_attention:
                 raise ValueError(
                     "mla_up_proj in recompute_modules is only supported with "
@@ -1899,6 +1927,8 @@ def __post_init__(self):
                 raise ValueError(
                     "Do not set --moe-layer-recompute with full recompute granularity. "
                 )
+            if "expert_fc1_act" in self.recompute_modules:
+                raise ValueError("expert_fc1_act cannot be combined with moe_layer_recompute.")
             self.recompute_granularity = "selective"
             if "moe" not in self.recompute_modules:
                 self.recompute_modules.append("moe")
@@ -1907,6 +1937,10 @@ def __post_init__(self):
             assert (
                 not self.cpu_offloading
             ), "fine_grained_activation_offloading cannot be enabled with cpu_offloading."
+            if "expert_fc1_act" in self.recompute_modules:
+                raise ValueError(
+                    "expert_fc1_act cannot be combined with fine-grained activation offloading."
+                )
             assert self.offload_modules is not None and len(self.offload_modules) > 0
             allowed_modules = {
                 "core_attn",
diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
index c9970dda8ee..c6068965c3f 100644
--- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
+++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
@@ -16,7 +16,9 @@
 from tests.unit_tests.test_utilities import Utils
 
 
-def _config(*, backend: str = "transformer_engine") -> TransformerConfig:
+def _config(
+    *, backend: str = "transformer_engine", expert_fc1_act_recompute: bool = False
+) -> TransformerConfig:
     return TransformerConfig(
         num_layers=1,
         hidden_size=16,
@@ -32,6 +34,8 @@ def _config(*, backend: str = "transformer_engine") -> TransformerConfig:
         params_dtype=torch.bfloat16,
         moe_router_load_balancing_type="sinkhorn",
         moe_router_topk=1,
+        recompute_granularity="selective" if expert_fc1_act_recompute else None,
+        recompute_modules=["expert_fc1_act"] if expert_fc1_act_recompute else None,
     )
 
 
@@ -62,6 +66,45 @@ def test_torch_grouped_expert_gemm_config_validation(override, error):
         TransformerConfig(**kwargs)
 
 
+@pytest.mark.parametrize("conflict", ["moe", "moe_act"])
+def test_expert_fc1_activation_recompute_rejects_nested_moe_recompute(conflict):
+    with pytest.raises(ValueError, match="cannot be combined"):
+        TransformerConfig(
+            num_layers=1,
+            hidden_size=16,
+            num_attention_heads=4,
+            num_moe_experts=2,
+            moe_grouped_gemm=True,
+            recompute_granularity="selective",
+            recompute_modules=["expert_fc1_act", conflict],
+        )
+
+
+def test_expert_fc1_activation_recompute_requires_grouped_transformer_engine():
+    with pytest.raises(ValueError, match="requires moe_grouped_gemm"):
+        TransformerConfig(
+            num_layers=1,
+            hidden_size=16,
+            num_attention_heads=4,
+            num_moe_experts=2,
+            moe_grouped_gemm=False,
+            recompute_granularity="selective",
+            recompute_modules=["expert_fc1_act"],
+        )
+
+    with pytest.raises(ValueError, match="requires transformer_engine"):
+        TransformerConfig(
+            num_layers=1,
+            hidden_size=16,
+            num_attention_heads=4,
+            num_moe_experts=2,
+            moe_grouped_gemm=True,
+            transformer_impl="local",
+            recompute_granularity="selective",
+            recompute_modules=["expert_fc1_act"],
+        )
+
+
 @pytest.mark.skipif(
     not is_te_min_version("1.9.0.dev0")
     or not torch.cuda.is_available()
@@ -86,7 +129,7 @@ def _model(config: TransformerConfig) -> MoELayer:
         )
         return Float16Module(config, model).module.cuda()
 
-    def test_forward_input_gradient_and_checkpoint_parity(self):
+    def test_forward_and_input_gradient_parity(self):
         baseline = self._model(_config())
         torch_model = self._model(_config(backend="torch"))
         torch_model.load_state_dict(baseline.state_dict())
@@ -117,6 +160,44 @@ def test_forward_input_gradient_and_checkpoint_parity(self):
                 == linear.weight0.untyped_storage().data_ptr()
             )
 
+    def test_expert_fc1_activation_recompute_parity_and_call_counts(self):
+        reference = self._model(_config(backend="torch"))
+        recompute = self._model(_config(backend="torch", expert_fc1_act_recompute=True))
+        recompute.load_state_dict(reference.state_dict())
+        for model in (reference, recompute):
+            for parameter in model.experts.parameters():
+                parameter.requires_grad = False
+
+        calls = {"fc1": 0, "fc2": 0}
+
+        def count_fc1(*_args):
+            calls["fc1"] += 1
+
+        def count_fc2(*_args):
+            calls["fc2"] += 1
+
+        recompute.experts.linear_fc1.register_forward_pre_hook(count_fc1)
+        recompute.experts.linear_fc2.register_forward_pre_hook(count_fc2)
+
+        reference_input = torch.rand(
+            (5, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True
+        )
+        recompute_input = reference_input.detach().clone().requires_grad_(True)
+        tokens_per_expert = torch.tensor([2, 0, 3, 0], device="cuda")
+        reference_probs = torch.rand(5, device="cuda", requires_grad=True)
+        recompute_probs = reference_probs.detach().clone().requires_grad_(True)
+
+        reference_output, _ = reference.experts(reference_input, tokens_per_expert, reference_probs)
+        recompute_output, _ = recompute.experts(recompute_input, tokens_per_expert, recompute_probs)
+        grad_output = torch.randn_like(reference_output)
+        reference_output.backward(grad_output)
+        recompute_output.backward(grad_output)
+
+        torch.testing.assert_close(recompute_output, reference_output)
+        torch.testing.assert_close(recompute_input.grad, reference_input.grad)
+        torch.testing.assert_close(recompute_probs.grad, reference_probs.grad)
+        assert calls == {"fc1": 2, "fc2": 1}
+
     def test_no_tokens(self):
         config = deepcopy(_config(backend="torch"))
         model = self._model(config)

From f8bb00bd1787379a2da7b74f8a6ddabed59aa299 Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Sat, 22 Aug 2026 05:57:44 -0700
Subject: [PATCH 5/9] perf(moe): fuse shared expert LoRA projection

Signed-off-by: Taufeeque <taufeeque@far.ai>
(cherry picked from commit 2a529bc8647c4e03af01ab5bab7702e213756ba3)
---
 .../core/extensions/transformer_engine.py     | 229 ++++++++++++++++++
 .../moe/test_frozen_torch_grouped_mlp.py      | 163 +++++++++++++
 2 files changed, 392 insertions(+)

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 80f6b8eb2c7..42859d058da 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -124,6 +124,63 @@ def _set_expert_parameter_attributes(
             param.partition_stride = 1
 
 
+def _run_torch_grouped_mm(x: Tensor, weight: Tensor, offsets: Tensor) -> Tensor:
+    """Run grouped matrix multiplication with a CPU reference for unit tests."""
+    if x.device.type == 'cuda':
+        return torch._grouped_mm(x, weight, offs=offsets)
+
+    outputs = []
+    start = 0
+    for group, end_tensor in enumerate(offsets):
+        end = int(end_tensor)
+        outputs.append(x[start:end].matmul(weight[group]))
+        start = end
+    if not outputs:
+        return x.new_empty((0, weight.shape[-1]))
+    return torch.cat(outputs, dim=0)
+
+
+class _TorchGroupedFusedLoRA(torch.autograd.Function):
+    """Differentiate a frozen grouped GEMM with a shared low-rank branch."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        augmented_weight: Tensor,
+        lora_a: Tensor,
+        lora_b: Tensor,
+        offsets: Tensor,
+        output_features: int,
+        scale: float,
+        counters: Dict[str, int],
+    ) -> Tensor:
+        with torch.no_grad():
+            augmented_weight[:, output_features:, :].copy_(lora_a)
+        augmented_output = _run_torch_grouped_mm(x, augmented_weight.transpose(1, 2), offsets)
+        base_output = augmented_output[:, :output_features]
+        low_rank = augmented_output[:, output_features:]
+        output = base_output + scale * low_rank.matmul(lora_b.transpose(0, 1))
+
+        ctx.save_for_backward(x, lora_b, low_rank, offsets)
+        ctx.augmented_weight = augmented_weight
+        ctx.scale = scale
+        ctx.counters = counters
+        counters["forward_calls"] += 1
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: Tensor):
+        x, lora_b, low_rank, offsets = ctx.saved_tensors
+        grad_low_rank = ctx.scale * grad_output.matmul(lora_b)
+        grad_augmented_output = torch.cat((grad_output, grad_low_rank), dim=-1)
+        grad_x = _run_torch_grouped_mm(grad_augmented_output, ctx.augmented_weight, offsets)
+        grad_lora_a = grad_low_rank.transpose(0, 1).matmul(x)
+        grad_lora_b = ctx.scale * grad_output.transpose(0, 1).matmul(low_rank)
+        ctx.counters["backward_calls"] += 1
+        return grad_x, None, grad_lora_a, grad_lora_b, None, None, None, None
+
+
 class TransformerEngineConfigType(enum.Enum):
     """Configuration object types in config dictionary"""
 
@@ -2664,6 +2721,37 @@ def _torch_grouped_weight_layout_is_current(self) -> bool:
                     return False
             return True
 
+        def _torch_grouped_fused_lora_layout_is_current(self, rank: int) -> bool:
+            """Return whether expert parameters view the augmented grouped buffer."""
+            augmented_weight = self._buffers.get("_torch_grouped_fused_lora_weight")
+            if augmented_weight is None or augmented_weight.shape[0] != self.num_gemms:
+                return False
+
+            first_weight = self.weight0
+            expected_shape = (self.num_gemms, first_weight.shape[0] + rank, first_weight.shape[1])
+            if (
+                augmented_weight.device != first_weight.device
+                or augmented_weight.dtype != first_weight.dtype
+                or augmented_weight.shape != expected_shape
+            ):
+                return False
+
+            storage_pointer = augmented_weight.untyped_storage().data_ptr()
+            group_elements = augmented_weight.shape[1] * augmented_weight.shape[2]
+            for index in range(self.num_gemms):
+                weight = getattr(self, f"weight{index}")
+                weight_view = augmented_weight[index, : first_weight.shape[0], :]
+                if (
+                    weight.device != augmented_weight.device
+                    or weight.dtype != augmented_weight.dtype
+                    or weight.shape != weight_view.shape
+                    or weight.untyped_storage().data_ptr() != storage_pointer
+                    or weight.storage_offset() != index * group_elements
+                    or weight.stride() != weight_view.stride()
+                ):
+                    return False
+            return True
+
         @torch.no_grad()
         def prepare_torch_grouped_mm(self) -> int:
             """Relocate frozen BF16 expert weights into one contiguous allocation.
@@ -2703,6 +2791,7 @@ def prepare_torch_grouped_mm(self) -> int:
                 grouped_weight = torch.stack([weight.detach() for weight in weights]).contiguous()
                 for index, weight in enumerate(weights):
                     weight.data = grouped_weight[index]
+                self._buffers.pop("_torch_grouped_fused_lora_weight", None)
                 if '_torch_grouped_weight' in self._buffers:
                     self._buffers['_torch_grouped_weight'] = grouped_weight
                 else:
@@ -2711,6 +2800,146 @@ def prepare_torch_grouped_mm(self) -> int:
             grouped_weight = self._buffers['_torch_grouped_weight']
             return grouped_weight.numel() * grouped_weight.element_size()
 
+        @torch.no_grad()
+        def prepare_torch_grouped_mm_fused_lora(self, lora_a: Tensor) -> Dict[str, Any]:
+            """Prepare augmented frozen expert weights for shared LoRA-A fusion."""
+            if self.expert_gemm_backend != "torch":
+                raise RuntimeError("fused expert LoRA requires moe_expert_gemm_backend='torch'")
+            if not hasattr(torch, "_grouped_mm"):
+                raise RuntimeError("this PyTorch build does not provide torch._grouped_mm")
+            if self.use_bias:
+                raise RuntimeError("fused expert LoRA does not support expert bias")
+
+            weights = [getattr(self, f"weight{index}") for index in range(self.num_gemms)]
+            if any(weight.requires_grad for weight in weights):
+                raise RuntimeError("fused expert LoRA requires frozen base weights")
+            if any(weight.device.type != "cuda" for weight in weights):
+                raise RuntimeError("fused expert LoRA requires CUDA weights")
+            if any(weight.dtype != torch.bfloat16 for weight in weights):
+                raise RuntimeError("fused expert LoRA requires BF16 base weights")
+            if any(weight.shape != weights[0].shape for weight in weights[1:]):
+                raise RuntimeError("fused expert LoRA requires uniform expert shapes")
+            if lora_a.device != weights[0].device or lora_a.dtype != torch.bfloat16:
+                raise RuntimeError("fused expert LoRA requires a BF16 LoRA-A on the base device")
+            if lora_a.ndim != 2 or lora_a.shape[1] != weights[0].shape[1]:
+                raise RuntimeError(
+                    f"LoRA-A shape {tuple(lora_a.shape)} does not match base input "
+                    f"features {weights[0].shape[1]}"
+                )
+
+            memory_allocated_before = torch.cuda.memory_allocated(lora_a.device)
+            memory_allocated_at_coexistence = memory_allocated_before
+            peak_allocated_before = torch.cuda.max_memory_allocated(lora_a.device)
+            rank = lora_a.shape[0]
+            if not self._torch_grouped_fused_lora_layout_is_current(rank):
+                output_features, input_features = weights[0].shape
+                augmented_weight = weights[0].new_empty(
+                    (self.num_gemms, output_features + rank, input_features)
+                )
+                memory_allocated_at_coexistence = torch.cuda.memory_allocated(lora_a.device)
+                for index, weight in enumerate(weights):
+                    augmented_weight[index, :output_features, :].copy_(weight)
+                    weight.data = augmented_weight[index, :output_features, :]
+                augmented_weight[:, output_features:, :].copy_(lora_a)
+                self._buffers.pop("_torch_grouped_weight", None)
+                if "_torch_grouped_fused_lora_weight" in self._buffers:
+                    self._buffers["_torch_grouped_fused_lora_weight"] = augmented_weight
+                else:
+                    self.register_buffer(
+                        "_torch_grouped_fused_lora_weight", augmented_weight, persistent=False
+                    )
+
+            augmented_weight = self._buffers["_torch_grouped_fused_lora_weight"]
+            if "_torch_grouped_weight" in self._buffers:
+                raise RuntimeError("fused expert LoRA retained the separate grouped base buffer")
+            if not self._torch_grouped_fused_lora_layout_is_current(rank):
+                raise RuntimeError(
+                    "fused expert LoRA parameters do not share the augmented backing storage"
+                )
+            base_storage_bytes = sum(weight.numel() * weight.element_size() for weight in weights)
+            resident_storage_bytes = augmented_weight.numel() * augmented_weight.element_size()
+            memory_allocated_after = torch.cuda.memory_allocated(lora_a.device)
+            return {
+                "base_storage_bytes": base_storage_bytes,
+                "augmentation_storage_bytes": resident_storage_bytes - base_storage_bytes,
+                "resident_storage_bytes": resident_storage_bytes,
+                "duplicate_base_buffer_present": False,
+                "base_parameters_share_augmented_storage": True,
+                "cuda_memory_allocated_before": memory_allocated_before,
+                "cuda_memory_allocated_at_coexistence": memory_allocated_at_coexistence,
+                "cuda_relocation_peak_delta": (
+                    memory_allocated_at_coexistence - memory_allocated_before
+                ),
+                "cuda_memory_allocated_after": memory_allocated_after,
+                "cuda_memory_allocated_delta": memory_allocated_after - memory_allocated_before,
+                "cuda_peak_allocated_before": peak_allocated_before,
+                "cuda_peak_allocated_after": torch.cuda.max_memory_allocated(lora_a.device),
+            }
+
+        def torch_grouped_mm_fused_lora_forward(
+            self,
+            x: Tensor,
+            m_splits: List[int],
+            *,
+            lora_a: Tensor,
+            lora_b: Tensor,
+            scale: float,
+            adapter_enabled: bool,
+            counters: Dict[str, int],
+        ) -> Tuple[Tensor, None]:
+            """Run the frozen grouped base and shared LoRA branch in one grouped GEMM."""
+            if not self._torch_grouped_fused_lora_layout_is_current(lora_a.shape[0]):
+                self.prepare_torch_grouped_mm_fused_lora(lora_a)
+            x_2d = x.reshape(-1, x.shape[-1])
+            input_rows = sum(m_splits)
+            if input_rows != x_2d.shape[0]:
+                raise RuntimeError(
+                    f"expert splits sum to {input_rows}, but input has {x_2d.shape[0]} rows"
+                )
+
+            augmented_weight = self._buffers["_torch_grouped_fused_lora_weight"]
+            output_features = self.weight0.shape[0]
+            if lora_b.device != x.device or lora_b.dtype != torch.bfloat16:
+                raise RuntimeError("fused expert LoRA requires a BF16 LoRA-B on the input device")
+            if lora_b.shape != (output_features, lora_a.shape[0]):
+                raise RuntimeError(
+                    f"LoRA-B shape {tuple(lora_b.shape)} does not match "
+                    f"({output_features}, {lora_a.shape[0]})"
+                )
+
+            if input_rows == 0:
+                base_output = x_2d.matmul(self.weight0.transpose(0, 1))
+                if adapter_enabled:
+                    output = base_output + scale * x_2d.matmul(lora_a.transpose(0, 1)).matmul(
+                        lora_b.transpose(0, 1)
+                    )
+                    counters["forward_calls"] += 1
+                else:
+                    output = base_output
+                    counters["base_only_forward_calls"] += 1
+            else:
+                offsets = torch.tensor(
+                    list(accumulate(m_splits)), device=x.device, dtype=torch.int32
+                )
+                if adapter_enabled:
+                    output = _TorchGroupedFusedLoRA.apply(
+                        x_2d,
+                        augmented_weight,
+                        lora_a,
+                        lora_b,
+                        offsets,
+                        output_features,
+                        scale,
+                        counters,
+                    )
+                else:
+                    augmented_output = torch._grouped_mm(
+                        x_2d, augmented_weight.transpose(1, 2), offs=offsets
+                    )
+                    output = augmented_output[:, :output_features]
+                    counters["base_only_forward_calls"] += 1
+            return output.reshape(*x.shape[:-1], output.shape[-1]), None
+
         @property
         def torch_grouped_mm_prepared(self) -> bool:
             """Return whether the torch grouped GEMM backing allocation is ready."""
diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
index c6068965c3f..ddb486a3cb3 100644
--- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
+++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
@@ -6,6 +6,10 @@
 import torch
 import torch.nn.functional as F
 
+from megatron.core.extensions.transformer_engine import (
+    _TorchGroupedFusedLoRA,
+    _run_torch_grouped_mm,
+)
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_layer_with_transformer_engine_submodules,
 )
@@ -259,3 +263,162 @@ def test_zero_token_groups(self):
         torch.testing.assert_close(
             torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02
         )
+
+
+def test_fused_expert_lora_matches_shared_adapter_and_gradients_on_cpu():
+    torch.manual_seed(7)
+    group_sizes = [2, 0, 3]
+    offsets = torch.tensor(group_sizes, dtype=torch.int32).cumsum(0)
+    base_weight = torch.randn(3, 7, 5, dtype=torch.double)
+    reference_a = torch.randn(3, 5, dtype=torch.double, requires_grad=True)
+    reference_b = torch.randn(7, 3, dtype=torch.double, requires_grad=True)
+    reference_input = torch.randn(5, 5, dtype=torch.double, requires_grad=True)
+    fused_input = reference_input.detach().clone().requires_grad_(True)
+    fused_a = reference_a.detach().clone().requires_grad_(True)
+    fused_b = reference_b.detach().clone().requires_grad_(True)
+    scale = 2.5
+    base_output = _run_torch_grouped_mm(
+        reference_input, base_weight.transpose(1, 2), offsets
+    )
+    reference = base_output + scale * reference_input.matmul(
+        reference_a.transpose(0, 1)
+    ).matmul(reference_b.transpose(0, 1))
+    augmented_weight = base_weight.new_empty((3, 10, 5))
+    augmented_weight[:, :7, :].copy_(base_weight)
+    augmented_weight[:, 7:, :].copy_(fused_a)
+    counters = {
+        "forward_calls": 0,
+        "backward_calls": 0,
+        "base_only_forward_calls": 0,
+    }
+    fused = _TorchGroupedFusedLoRA.apply(
+        fused_input,
+        augmented_weight,
+        fused_a,
+        fused_b,
+        offsets,
+        7,
+        scale,
+        counters,
+    )
+    grad_output = torch.randn_like(reference)
+    reference_grads = torch.autograd.grad(
+        reference, (reference_input, reference_a, reference_b), grad_output
+    )
+    fused_grads = torch.autograd.grad(
+        fused, (fused_input, fused_a, fused_b), grad_output
+    )
+
+    torch.testing.assert_close(fused, reference)
+    for actual, expected in zip(fused_grads, reference_grads):
+        torch.testing.assert_close(actual, expected)
+    assert counters == {
+        "forward_calls": 1,
+        "backward_calls": 1,
+        "base_only_forward_calls": 0,
+    }
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or not hasattr(torch, "_grouped_mm")
+    or torch.cuda.get_device_capability()[0] < 10,
+    reason="production fused expert LoRA parity requires a Blackwell CUDA device",
+)
+@pytest.mark.parametrize(
+    ("input_features", "output_features"),
+    [(1024, 2688), (2688, 1024)],
+    ids=["super-fc1", "super-fc2"],
+)
+def test_fused_expert_lora_super_shapes_over_repeated_updates(input_features, output_features):
+    generator = torch.Generator(device="cuda").manual_seed(20260822)
+    group_sizes = [0 if index % 17 == 0 else 1 + (index * 13) % 4 for index in range(128)]
+    offsets = torch.tensor(group_sizes, device="cuda", dtype=torch.int32).cumsum(0)
+    rows = sum(group_sizes)
+    rank = 32
+    scale = 64.0 / rank
+    base_weight = torch.randn(
+        (len(group_sizes), output_features, input_features),
+        device="cuda",
+        dtype=torch.bfloat16,
+        generator=generator,
+    )
+    reference_a = torch.randn(
+        (rank, input_features),
+        device="cuda",
+        dtype=torch.bfloat16,
+        generator=generator,
+        requires_grad=True,
+    )
+    reference_b = torch.randn(
+        (output_features, rank),
+        device="cuda",
+        dtype=torch.bfloat16,
+        generator=generator,
+        requires_grad=True,
+    )
+    fused_a = reference_a.detach().clone().requires_grad_(True)
+    fused_b = reference_b.detach().clone().requires_grad_(True)
+    augmented_weight = base_weight.new_empty(
+        (len(group_sizes), output_features + rank, input_features)
+    )
+    augmented_weight[:, :output_features, :].copy_(base_weight)
+    augmented_weight[:, output_features:, :].copy_(fused_a)
+    reference_optimizer = torch.optim.SGD([reference_a, reference_b], lr=0.01)
+    fused_optimizer = torch.optim.SGD([fused_a, fused_b], lr=0.01)
+    counters = {"forward_calls": 0, "backward_calls": 0, "base_only_forward_calls": 0}
+
+    for _ in range(3):
+        reference_optimizer.zero_grad()
+        fused_optimizer.zero_grad()
+        reference_inputs = []
+        fused_inputs = []
+        reference_outputs = []
+        fused_outputs = []
+        for _ in range(2):
+            reference_input = torch.randn(
+                (rows, input_features),
+                device="cuda",
+                dtype=torch.bfloat16,
+                generator=generator,
+                requires_grad=True,
+            )
+            fused_input = reference_input.detach().clone().requires_grad_(True)
+            base_output = torch._grouped_mm(
+                reference_input, base_weight.transpose(1, 2), offs=offsets
+            )
+            reference_output = base_output + scale * reference_input.matmul(
+                reference_a.transpose(0, 1)
+            ).matmul(reference_b.transpose(0, 1))
+            fused_output = _TorchGroupedFusedLoRA.apply(
+                fused_input,
+                augmented_weight,
+                fused_a,
+                fused_b,
+                offsets,
+                output_features,
+                scale,
+                counters,
+            )
+            torch.testing.assert_close(fused_output, reference_output, rtol=0.02, atol=0.02)
+            reference_inputs.append(reference_input)
+            fused_inputs.append(fused_input)
+            reference_outputs.append(reference_output)
+            fused_outputs.append(fused_output)
+
+        grad_output = torch.randn(
+            reference_outputs[0].shape, device="cuda", dtype=torch.bfloat16, generator=generator
+        )
+        torch.autograd.backward(reference_outputs, [grad_output, grad_output])
+        torch.autograd.backward(fused_outputs, [grad_output, grad_output])
+        for fused_input, reference_input in zip(fused_inputs, reference_inputs):
+            torch.testing.assert_close(fused_input.grad, reference_input.grad, rtol=0.02, atol=0.02)
+        torch.testing.assert_close(fused_a.grad, reference_a.grad, rtol=0.02, atol=0.02)
+        torch.testing.assert_close(fused_b.grad, reference_b.grad, rtol=0.02, atol=0.02)
+        reference_optimizer.step()
+        fused_optimizer.step()
+        torch.testing.assert_close(fused_a, reference_a, rtol=0.02, atol=0.02)
+        torch.testing.assert_close(fused_b, reference_b, rtol=0.02, atol=0.02)
+
+    assert any(size == 0 for size in group_sizes)
+    assert counters == {"forward_calls": 6, "backward_calls": 6, "base_only_forward_calls": 0}

From 47bfd8946d8d0364842b17603fa924fe6e50850a Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Sat, 22 Aug 2026 06:55:23 -0700
Subject: [PATCH 6/9] test(moe): preserve grouped offsets dtype

Signed-off-by: Taufeeque <taufeeque@far.ai>
(cherry picked from commit a9c866bfc4102d44da7a6942430d888f1780a543)
---
 .../transformer/moe/test_frozen_torch_grouped_mlp.py          | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
index ddb486a3cb3..bdffec2b696 100644
--- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
+++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
@@ -333,7 +333,9 @@ def test_fused_expert_lora_matches_shared_adapter_and_gradients_on_cpu():
 def test_fused_expert_lora_super_shapes_over_repeated_updates(input_features, output_features):
     generator = torch.Generator(device="cuda").manual_seed(20260822)
     group_sizes = [0 if index % 17 == 0 else 1 + (index * 13) % 4 for index in range(128)]
-    offsets = torch.tensor(group_sizes, device="cuda", dtype=torch.int32).cumsum(0)
+    offsets = torch.tensor(group_sizes, device="cuda").cumsum(
+        0, dtype=torch.int32
+    )
     rows = sum(group_sizes)
     rank = 32
     scale = 64.0 / rank

From d5ff43c2ce940a07478f47fead33c78e33d06c6d Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Sat, 22 Aug 2026 07:00:50 -0700
Subject: [PATCH 7/9] test(moe): allow BF16 fused input rounding

Signed-off-by: Taufeeque <taufeeque@far.ai>
(cherry picked from commit b64bdd8f54ee42e31d5cb25b7a743f56cb89ef6f)
---
 .../transformer/moe/test_frozen_torch_grouped_mlp.py        | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
index bdffec2b696..12d155a0329 100644
--- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
+++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
@@ -333,9 +333,7 @@ def test_fused_expert_lora_matches_shared_adapter_and_gradients_on_cpu():
 def test_fused_expert_lora_super_shapes_over_repeated_updates(input_features, output_features):
     generator = torch.Generator(device="cuda").manual_seed(20260822)
     group_sizes = [0 if index % 17 == 0 else 1 + (index * 13) % 4 for index in range(128)]
-    offsets = torch.tensor(group_sizes, device="cuda").cumsum(
-        0, dtype=torch.int32
-    )
+    offsets = torch.tensor(group_sizes, device="cuda").cumsum(0, dtype=torch.int32)
     rows = sum(group_sizes)
     rank = 32
     scale = 64.0 / rank
@@ -414,7 +412,7 @@ def test_fused_expert_lora_super_shapes_over_repeated_updates(input_features, ou
         torch.autograd.backward(reference_outputs, [grad_output, grad_output])
         torch.autograd.backward(fused_outputs, [grad_output, grad_output])
         for fused_input, reference_input in zip(fused_inputs, reference_inputs):
-            torch.testing.assert_close(fused_input.grad, reference_input.grad, rtol=0.02, atol=0.02)
+            torch.testing.assert_close(fused_input.grad, reference_input.grad, rtol=0.02, atol=1.0)
         torch.testing.assert_close(fused_a.grad, reference_a.grad, rtol=0.02, atol=0.02)
         torch.testing.assert_close(fused_b.grad, reference_b.grad, rtol=0.02, atol=0.02)
         reference_optimizer.step()

From 8d9ad06ba1bcfa4b5cf5358d92d5693b69c32bf4 Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Mon, 24 Aug 2026 15:43:35 -0700
Subject: [PATCH 8/9] fix(recompute): keep chunk outputs differentiable with a
 frozen embedding

Re-entrant activation checkpointing only attaches a grad_fn when some tensor
input requires grad. Under adapter-only training the embedding is frozen, so
every checkpointed chunk output carried no grad_fn and the adapters inside the
chunks received no gradient. Detach and re-enable grad on the block input
before the chunk loop.

The rest of the source commits added a MambaStack checkpointed forward, which
upstream has since absorbed: MambaStack moved to
megatron/core/models/hybrid/hybrid_block.py and routes full-granularity
recompute through megatron/core/recompute.py::checkpointed_forward.

Signed-off-by: Taufeeque <taufeeque@far.ai>
(cherry picked from commit 91c8b4180)
(cherry picked from commit abd235c8d)
---
 megatron/core/recompute.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/megatron/core/recompute.py b/megatron/core/recompute.py
index bd0d1bcb3b2..fb68354781c 100644
--- a/megatron/core/recompute.py
+++ b/megatron/core/recompute.py
@@ -154,6 +154,13 @@ def chunk_runner(start: int, end: int, use_checkpoint: bool):
             if (start + layer_offset) in extract_layer_indices:
                 intermediate_hidden_states.append(hidden_states)
 
+    if not hidden_states.requires_grad:
+        # Re-entrant checkpointing only attaches a grad_fn when some tensor input
+        # requires grad. With a frozen embedding (adapter-only training) every chunk
+        # output would otherwise carry no grad_fn and the adapters inside the chunks
+        # would receive no gradient.
+        hidden_states = hidden_states.detach().requires_grad_(True)
+
     if self.config.recompute_method == 'uniform':
         # Uniformly divide the total number of layers and checkpoint
         # the input activation of each divided chunk.

From e8c17b92b5c1900b973ad2e8629aa46bcb287885 Mon Sep 17 00:00:00 2001
From: Taufeeque <taufeeque@far.ai>
Date: Mon, 24 Aug 2026 15:44:36 -0700
Subject: [PATCH 9/9] style(moe): match repo black and pylint conventions in
 the ported MoE stack

The fork commits were written against an older checkout and predate this
repo's black profile (line length 100, --skip-magic-trailing-comma) and its
pylint docstring gate on megatron/core.

Signed-off-by: Taufeeque <taufeeque@far.ai>
---
 .../core/extensions/transformer_engine.py     |  2 +
 .../moe/test_frozen_torch_grouped_mlp.py      | 45 +++++--------------
 2 files changed, 12 insertions(+), 35 deletions(-)

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 42859d058da..b7959047062 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -155,6 +155,7 @@ def forward(
         scale: float,
         counters: Dict[str, int],
     ) -> Tensor:
+        """Forward pass."""
         with torch.no_grad():
             augmented_weight[:, output_features:, :].copy_(lora_a)
         augmented_output = _run_torch_grouped_mm(x, augmented_weight.transpose(1, 2), offsets)
@@ -171,6 +172,7 @@ def forward(
 
     @staticmethod
     def backward(ctx, grad_output: Tensor):
+        """Backward pass."""
         x, lora_b, low_rank, offsets = ctx.saved_tensors
         grad_low_rank = ctx.scale * grad_output.matmul(lora_b)
         grad_augmented_output = torch.cat((grad_output, grad_low_rank), dim=-1)
diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
index 12d155a0329..5e7b056113e 100644
--- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
+++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py
@@ -249,20 +249,14 @@ def test_zero_token_groups(self):
         torch_input = baseline_input.detach().clone().requires_grad_(True)
         tokens_per_expert = torch.tensor([2, 0, 3, 0], device="cuda")
         permuted_probs = torch.rand(5, dtype=torch.float32, device="cuda")
-        baseline_output, _ = baseline.experts(
-            baseline_input, tokens_per_expert, permuted_probs
-        )
-        torch_output, _ = torch_model.experts(
-            torch_input, tokens_per_expert, permuted_probs
-        )
+        baseline_output, _ = baseline.experts(baseline_input, tokens_per_expert, permuted_probs)
+        torch_output, _ = torch_model.experts(torch_input, tokens_per_expert, permuted_probs)
         grad_output = torch.randn_like(baseline_output)
         baseline_output.backward(grad_output)
         torch_output.backward(grad_output)
 
         torch.testing.assert_close(torch_output, baseline_output, rtol=0.02, atol=0.02)
-        torch.testing.assert_close(
-            torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02
-        )
+        torch.testing.assert_close(torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02)
 
 
 def test_fused_expert_lora_matches_shared_adapter_and_gradients_on_cpu():
@@ -277,46 +271,27 @@ def test_fused_expert_lora_matches_shared_adapter_and_gradients_on_cpu():
     fused_a = reference_a.detach().clone().requires_grad_(True)
     fused_b = reference_b.detach().clone().requires_grad_(True)
     scale = 2.5
-    base_output = _run_torch_grouped_mm(
-        reference_input, base_weight.transpose(1, 2), offsets
+    base_output = _run_torch_grouped_mm(reference_input, base_weight.transpose(1, 2), offsets)
+    reference = base_output + scale * reference_input.matmul(reference_a.transpose(0, 1)).matmul(
+        reference_b.transpose(0, 1)
     )
-    reference = base_output + scale * reference_input.matmul(
-        reference_a.transpose(0, 1)
-    ).matmul(reference_b.transpose(0, 1))
     augmented_weight = base_weight.new_empty((3, 10, 5))
     augmented_weight[:, :7, :].copy_(base_weight)
     augmented_weight[:, 7:, :].copy_(fused_a)
-    counters = {
-        "forward_calls": 0,
-        "backward_calls": 0,
-        "base_only_forward_calls": 0,
-    }
+    counters = {"forward_calls": 0, "backward_calls": 0, "base_only_forward_calls": 0}
     fused = _TorchGroupedFusedLoRA.apply(
-        fused_input,
-        augmented_weight,
-        fused_a,
-        fused_b,
-        offsets,
-        7,
-        scale,
-        counters,
+        fused_input, augmented_weight, fused_a, fused_b, offsets, 7, scale, counters
     )
     grad_output = torch.randn_like(reference)
     reference_grads = torch.autograd.grad(
         reference, (reference_input, reference_a, reference_b), grad_output
     )
-    fused_grads = torch.autograd.grad(
-        fused, (fused_input, fused_a, fused_b), grad_output
-    )
+    fused_grads = torch.autograd.grad(fused, (fused_input, fused_a, fused_b), grad_output)
 
     torch.testing.assert_close(fused, reference)
     for actual, expected in zip(fused_grads, reference_grads):
         torch.testing.assert_close(actual, expected)
-    assert counters == {
-        "forward_calls": 1,
-        "backward_calls": 1,
-        "base_only_forward_calls": 0,
-    }
+    assert counters == {"forward_calls": 1, "backward_calls": 1, "base_only_forward_calls": 0}
 
 
 @pytest.mark.skipif(