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

[anijain2305]

Summary

GroupedExperts._grouped_mm is the seam low-precision variants override to swap in a scaled grouped GEMM. It currently receives B_t — the already-transposed weight. Pass the weight in its stored orientation instead, and do the transpose inside the seam.

# before
def _grouped_mm(self, *, A, B_t, offs):
    return torch._grouped_mm(A, B_t, offs=offs)

# after
def _grouped_mm(self, *, A, weight_ENK, offs):
    return torch._grouped_mm(A, weight_ENK.bfloat16().transpose(-2, -1), offs=offs)

Why

B_t is enough for a variant that quantizes on every call, which is what the float8 and MXFP8 converters do today. It is not enough for one whose quantized representation is owned by the weight's FSDP unshard lifetime: that representation is keyed off the stored (experts, out_features, in_features) orientation, and the transpose has already erased which parameter it came from.

Split out of #4344, which needs this, so the mechanical refactor is reviewable on its own.

Scope

No behavior change — the base implementation transposes exactly as before.

Callers updated: the common, gpt_oss, and kimi_k3 experts, the fused SwiGLU override, and the float8 converter's override. The fused SwiGLU case reshapes its 4D w13 parameter into the (E, 2F, D) grouped weight first, which the seam then transposes.

[tianyu-l]

this rename is confusing, because it conflicts with existing suffix vocab https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe.py#L25

[tianyu-l]

begging for a non-conflicting name and then we can land

[tianyu-l]

@claude The PR summary states the motivation, which I couldn't visualize / digest. Could you give an example where previous representation is problematic.

[claude]

Claude finished @tianyu-l's task in 1m 28s —— View job

---

Concrete example: why B_t breaks the MXFP8 weight cache

Here's the scenario from #4344 that the current B_t seam can't support.

The float8 rowwise path (works fine with B_t) re-quantizes the weight on every call:

def _grouped_mm(self, *, A, B_t, offs):
    # quantize B_t right here, every forward. Nothing is cached, so
    # receiving the already-transposed view is all we need.
    return _quantize_then_scaled_grouped_mm(A, B_t, ...)

The MXFP8 path amortizes the weight quantization over the FSDP unshard window. MXFP8 weight quantization is expensive, so instead of redoing it per grouped-mm call, it quantizes each expert weight once when FSDP all-gathers (unshards) the parameter, and stashes that MXFP8 representation for the lifetime of the unshard. Conceptually:

# FSDP all-gather post-hook, once per unshard:
#   w1_EFD  ->  mxfp8(w1_EFD)   # (E, out=F, in=D), block-scaled along the K=D axis
# cached, keyed on the parameter, reused by every grouped_mm in the window.

Now look at what the seam receives in each design:

# core forward, BEFORE this PR: caller transposes, seam gets B_t
B_t = w1_EFD.transpose(-2, -1)      # (E, D, F), a fresh view
self._grouped_mm(A=..., B_t=B_t, offs=...)

Inside the MXFP8 seam I now want to fetch the cached representation:

def _grouped_mm(self, *, A, B_t, offs):
    mx = self._mxfp8_cache[???]     # I need w1_EFD, but I only have B_t
    ...

Two things have gone wrong, both because the transpose happened before the seam:

1. Identity is gone. B_t is a new view produced by transpose; it is not the parameter object the cache was keyed on. self._mxfp8_cache[B_t] misses, and there's no robust way to recover w1_EFD from B_t. The transpose "erased which parameter it came from."

2. Orientation is gone. MXFP8 is block-scaled along a specific axis. The cached representation was quantized for the stored (E, out=F, in=D) layout (scales along the K=D contraction axis). B_t is (E, D, F) — you can't cheaply reinterpret an already-quantized MXFP8 tensor into the transposed orientation, because the scale blocks would be along the wrong axis. So even ignoring (1), quantizing B_t as-is wouldn't match what was cached.

The only escape with the old seam is to re-quantize inside _grouped_mm every call — which throws away the entire point of amortizing over the unshard window.

With this PR, the seam receives weight_ENK in its stored orientation, so it can look the parameter up in the cache and apply the transpose after fetching the quantized weight (in whatever form the scaled grouped GEMM expects):

def _grouped_mm(self, *, A, weight_ENK, offs):
    mx = self._mxfp8_cache[weight_ENK]   # keyed on the stored (E, F, D) parameter
    return _scaled_grouped_mm(A, mx.transpose(-2, -1), offs=offs)

So the short version: B_t is sufficient for a "quantize every call" variant, but a "quantize once per unshard and cache" variant needs the parameter in its stored orientation to (a) find its cache entry and (b) match the axis its block-scales were computed along. The base implementation still transposes exactly as before, so there's no behavior change for the existing paths.

(Separately, I saw @tianyu-l's inline note about the weight_ENK suffix colliding with the N = routed tokens legend at moe.py:25 — happy to rename, e.g. weight_ENK -> weight_EOI (out/in) or fold out_features/in_features into the name, if you want that in this PR too.)