Pass the expert weight, not its transpose, to the grouped-mm seam

torchtitan/components/quantization/float8.py

@@ -198,13 +198,16 @@ def __init__(self, config: Config):
 
             self._float8_op_config = Float8TrainingOpConfig()
 
-        def _grouped_mm(self, *, A, B_t, offs):
+        def _grouped_mm(self, *, A, weight_EOI, offs):
             from torchao.prototype.moe_training.utils import (
                 _quantize_then_scaled_grouped_mm,
             )
 
             return _quantize_then_scaled_grouped_mm(
-                A, B_t, config=self._float8_op_config, offs=offs
+                A,
+                weight_EOI.bfloat16().transpose(-2, -1),
+                config=self._float8_op_config,
+                offs=offs,
             )
 
     Float8GroupedExperts.__name__ = f"Float8{parent_cls.__name__}"

torchtitan/models/common/moe.py

@@ -96,32 +96,30 @@ def forward(
                 spmd.mutate_type(offsets_E, axis, src=spmd.P, dst=spmd.V)
 
         h_RF = F.silu(
-            self._grouped_mm(
-                A=x_RD.bfloat16(),
-                B_t=w1_EFD.bfloat16().transpose(-2, -1),
-                offs=offsets_E,
-            )
+            self._grouped_mm(A=x_RD.bfloat16(), weight_EOI=w1_EFD, offs=offsets_E)
         )
         h_RF = h_RF * self._grouped_mm(
-            A=x_RD.bfloat16(),
-            B_t=w3_EFD.bfloat16().transpose(-2, -1),
-            offs=offsets_E,
+            A=x_RD.bfloat16(), weight_EOI=w3_EFD, offs=offsets_E
         )
-        return self._grouped_mm(
-            A=h_RF, B_t=w2_EDF.bfloat16().transpose(-2, -1), offs=offsets_E
-        ).type_as(x_RD)
+        return self._grouped_mm(A=h_RF, weight_EOI=w2_EDF, offs=offsets_E).type_as(x_RD)
 
     def _grouped_mm(
-        self, *, A: torch.Tensor, B_t: torch.Tensor, offs: torch.Tensor
+        self, *, A: torch.Tensor, weight_EOI: torch.Tensor, offs: torch.Tensor
     ) -> torch.Tensor:
-        """Grouped matmul of ``A @ B_t`` with per-expert token offsets.
-
-        Overridable seam for low-precision variants (e.g. the MXFP8 converter
-        swaps this for a dynamically-quantized scaled grouped GEMM). Keeping the
-        op here -- rather than behind a tensor-subclass ``__torch_function__`` --
-        means it is captured by FX tracers such as graph_trainer's make_fx path.
+        """Grouped matmul of ``A @ weight_EOI.transpose(-2, -1)``.
+
+        ``weight_EOI`` is the grouped expert weight in its stored
+        ``(experts, out_features, in_features)`` orientation; the transpose to
+        the grouped-GEMM right operand happens here. Overridable seam for
+        low-precision variants (e.g. the MXFP8 converter swaps this for a
+        scaled grouped GEMM). Variants receive the weight rather than its
+        transpose because a quantized representation may be owned by the
+        weight's FSDP unshard lifetime and is keyed off the stored orientation.
+        Keeping the op here -- rather than behind a tensor-subclass
+        ``__torch_function__`` -- means it is captured by FX tracers such as
+        graph_trainer's make_fx path.
         """
-        return torch._grouped_mm(A, B_t, offs=offs)
+        return torch._grouped_mm(A, weight_EOI.bfloat16().transpose(-2, -1), offs=offs)
 
 
 class RoutedExperts(Module):

torchtitan/models/gpt_oss/moe.py

@@ -162,7 +162,7 @@ def forward(
         # G = gate+up dimension (2*F)
         h_RG = self._grouped_mm(
             A=x_RD.bfloat16(),
-            B_t=mlp1_weight_EGD.transpose(-2, -1).bfloat16(),
+            weight_EOI=mlp1_weight_EGD,
             offs=offsets_E,
         )
 
@@ -175,9 +175,7 @@ def forward(
         h_RG = h_RG + b1_RG.to(h_RG.dtype)
 
         h_RF = swiglu(h_RG, limit=self.swiglu_limit)
-        h_RD = self._grouped_mm(
-            A=h_RF, B_t=mlp2_weight_EDF.transpose(-2, -1).bfloat16(), offs=offsets_E
-        )
+        h_RD = self._grouped_mm(A=h_RF, weight_EOI=mlp2_weight_EDF, offs=offsets_E)
 
         # Apply custom autograd function to scale bias in forward but not in backward
         b2 = torch.cat(

torchtitan/models/kimi_k3/moe.py

@@ -89,20 +89,20 @@ def forward(
 
         gate_RF = self._grouped_mm(
             A=x_RD.bfloat16(),
-            B_t=w1_EFD.bfloat16().transpose(-2, -1),
+            weight_EOI=w1_EFD,
             offs=offsets_E,
         )
         up_RF = self._grouped_mm(
             A=x_RD.bfloat16(),
-            B_t=w3_EFD.bfloat16().transpose(-2, -1),
+            weight_EOI=w3_EFD,
             offs=offsets_E,
         )
 
         h_RF = _situ_glu(gate_RF, up_RF, self.beta, self.linear_beta)
 
         return self._grouped_mm(
             A=h_RF,
-            B_t=w2_EDF.bfloat16().transpose(-2, -1),
+            weight_EOI=w2_EDF,
             offs=offsets_E,
         ).type_as(x_RD)
 

torchtitan/overrides/fused_swiglu.py

@@ -647,15 +647,15 @@ def forward(
         E, F, _, D = w13.shape
         offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32)
 
-        w13_E_D_2F = w13.bfloat16().reshape(E, F * 2, D).transpose(-2, -1)
+        # The fused parameter stores gate and up interleaved as (E, F, 2, D);
+        # the grouped GEMM consumes them as one (E, 2F, D) expert weight.
+        w13_E_2F_D = w13.bfloat16().reshape(E, F * 2, D)
         gate_up_R2F = self._grouped_mm(
-            A=x_RD.bfloat16(), B_t=w13_E_D_2F, offs=offsets_E
+            A=x_RD.bfloat16(), weight_EOI=w13_E_2F_D, offs=offsets_E
         )
         gate_RF, up_RF = gate_up_R2F.reshape(-1, F, 2).unbind(-1)
         h_RF = silu_and_mul_op(gate_RF, up_RF, offsets_E)
-        return self._grouped_mm(
-            A=h_RF, B_t=w2_EDF.bfloat16().transpose(-2, -1), offs=offsets_E
-        ).type_as(x_RD)
+        return self._grouped_mm(A=h_RF, weight_EOI=w2_EDF, offs=offsets_E).type_as(x_RD)
 
     @staticmethod
     def _split_w13_on_save(module, state_dict, prefix, local_metadata) -> None: