Permute-Free Grouped GEMM for MoE (bf16, gfx950)

[sudhu2k]

Description

This PR adds a permute-free execution path for MoE GroupedLinear on ROCm. Instead of
physically permuting activations into expert-contiguous order before each grouped GEMM (and
un-permuting after), the gather is folded into the GEMM: the kernel reads each expert's rows
directly from the token-ordered activation buffer using a precomputed route list. This eliminates
the permute/unpermute activation copies, lets SwiGLU be treated as "just an op" (recomputed in the
backward rather than checkpointed), matches the MegaMOE reference kernel numerically, beats the
other bf16 grouped-GEMM backends currently in TE, and is fully CUDA-graph capturable (the
routing/align path is sync-free).

Opt-in via NVTE_PERMUTE_FREE_GROUPED_GEMM=1 (ROCm builds only). ~6.9k LOC across 16 files; new
code is self-contained under transformer_engine/pytorch/moe/ and
transformer_engine/pytorch/flydsl_kernels/, wired into GroupedLinear.

---

1. Motivation: what's expensive about the standard MoE path

A conventional MoE layer does, per micro-batch:

router → permute (scatter tokens into expert-contiguous order)
       → FC1 grouped GEMM  → SwiGLU activation
       → FC2 grouped GEMM
       → unpermute (scatter results back to token order) + combine

Two sources of avoidable activation memory / bandwidth:

1. Permute/unpermute copies. The permuted activation ([num_routes, hidden]) is a full extra
   materialization of the activations, written before FC1 and read/scattered after FC2 — plus the
   symmetric copies in the backward.
2. SwiGLU checkpointing. The standard path saves the F-wide activated tensor (and/or the 2F
   pre-activation) for the backward, so the activation output lives across the whole fwd→bwd
   boundary.

For large expert counts and top-k this permute + activation traffic is a meaningful fraction of MoE
step time and peak memory.

---

2. Core idea: gather-in-GEMM over a route-list layout

Rather than moving data, we build a compact route list once and let the GEMM kernels do the
gather while they load their A tiles:

• The activation buffer stays in received-token order ([num_recv_tokens, hidden]), never
  permuted.
• A sync-free "align" pass builds an expert-sorted, block-padded index (sorted_slot_ids) that
  maps each padded GEMM slot → the received-token row to gather.
• The forward GEMM (A = activations, B = [E, N, K] weights) reads A[sorted_slot_ids[slot], :]
  inside the tile load — the gather is free (it rides the existing global loads), and no permuted
  copy is ever written.
• The token combine (and the FC1 input-grad reduction) is done by a contention-free
  gather-combine (each output token pulls its own routes) instead of an atomic scatter.

Because the whole routing/align path is over-allocated to static, shape-derived upper bounds and
carries the real extents as device scalars (no .item(), no host sync), the entire flow is
CUDA-graph capturable.

---

3. End-to-end flow

The same built routing metadata drives both FCs; a route_space flag selects direction (FC1 vs
FC2), so FC1 and FC2 share a single align build (dataclasses.replace(meta, route_space=True)).

                         received-token order                route/padded order
                         [num_recv, hidden]                  [em_max, ...]

FC1 (route_space=False)  x ──gather-in-GEMM (A gathered by sorted_slot_ids)──▶ [em_max, 2F]  (raw [gate|up])
                                                                                    │
SwiGLU (standalone op)   act(gate)*up*prob  ── recomputed, F-wide transient ◀───────┘  (only 2F preact persists)
                                                                                    │
FC2 (route_space=True)   [em_max, F] ──gather-in-GEMM (A read by route pos)──▶ [em_max, out]
                                                                                    │
combine                  gather-combine (each token pulls its routes) ──▶ [num_recv, out]

• FC1 (route_space=False) — input in token order. Forward gathers per expert
  (index_a_by_route_pos=False) into the padded [em_max, 2F] raw [gate|up] buffer. Dgrad
  combines the input gradient back to token rows via the contention-free gather-combine.
• FC2 (route_space=True) — input already in route order (FC1's output). The gated activation
  act(gate)*up[*prob] is applied in a standalone pass into an F-wide transient, then a plain
  gather-GEMM + gather-combine → [num_recv, out]. Dgrad gathers the token-space grad back into the
  compact route buffer.

---

4. Data structures

All routing state lives on MoERoutingMetadata / PermuteFreeMetadata
(transformer_engine/pytorch/moe/moe_routing.py). Everything below is built sync-free and
over-allocated to static bounds; real extents ride as device scalars.

Field	Shape / type	Meaning
routing_map	[num_recv_tokens, num_experts] bool	The input: True where a token feeds a local expert. num_routes = routing_map.sum().
topk	host int (optional)	Upper bound on experts a token can feed. Tightens the padded over-allocation from T*E down to T*min(topk,E) → smaller em_max (and less zero-init).
sorted_slot_ids	[T*min(topk,E)] i32	The route list. Block-padded slot → received-token row to gather (sentinel num_recv_tokens for padding/tail → OOB read returns 0).
expert_ids	[blocks_max] i32	Expert owning each BLOCK_SIZE_M block (-1 past the real count; skipped). Drives the tile→expert map in the GEMM grid.
slot_expert_ids	[T*min(topk,E)] i32	Per-slot expert id (expert_ids[slot // block_size_m]), for the standalone activation kernels.
num_tokens_post_padded	[1] i32 device scalar	Real padded extent em; bounds the kernels without a host sync.
block_start	[num_experts] i32	Per-expert first block index (block units); expert e's slots start at block_start[e]*block_size_m. (Consumed by backward wgrad to locate grad rows — see §6.)
token_routes / token_route_count	[T, min(topk,E)] i32 / [T] i32	Token → its padded slot positions (inverse map). Powers the contention-free gather-combine that replaces the atomic scatter.
block_size_m	host int	BLOCK_SIZE_M the fwd/dgrad layout is padded to.
route_counts / route_within	cached tensors	Block-size-independent scan (per-expert counts + within-expert ranks) shared by the fwd/dgrad and wgrad align builds — computed once per routing map.
wgrad_align	WgradAlign (shared holder)	Separate block-CONTRACT_M(=32) align for the wgrad contraction (sorted_slot_ids, block_start, blocks_per_expert). Shared by reference across the FC1 metadata and its FC2 replace copy so the wgrad align is built once.

Direction tag (PermuteFreeMetadata): route_space (FC1 vs FC2) and an optional activation hint
("silu"/"gelu") consumed on the FC2 direction. Per-route gating probabilities are not stored
here — they need a gradient, so they flow as a separate autograd tensor.

---

5. Kernels

Routing / metadata (Triton, pf_helper_kernels.py)

• route_list_scan — block-size-independent per-expert counts + within-expert ranks (shared across
  FC1/FC2/wgrad).
• route_list_align — builds sorted_slot_ids, expert_ids, block_start, blocks_per_expert,
  and the token_routes inverse map. Fully sync-free / graph-safe.
• route_gather_combine — contention-free per-token reduction over a token's routes (replaces
  atomic scatter for the FC2 combine and FC1 dgrad reduction).
• fused_gated_act_prob_fwd / fused_gated_act_prob_bwd — SwiGLU/GELU recompute + backward (see
  §6). exp2-based silu/gelu, no libdevice dependency.

Grouped GEMMs (FlyDSL, flydsl_kernels/permute_free_grouped_gemm/)

• pf_fwd.py — gather-in-GEMM forward. Tile-centric grid: TILE_TO_GROUP (= expert_ids) maps
  each M-tile to its expert; A rows are gathered by sorted_slot_ids (FC1) or read by route
  position (FC2). MegaMOE-ported fixed tile geometry (32×32×16 MFMA, BLOCK_N=256, GROUP_M=4,
  XCD swizzle over the real tile range).
• pf_dgrad.py — data-gradient GEMM (NN contract against the transposed-weight view).
• pf_wgrad.py — weight-gradient GEMM. Expert-centric grid N×K×E; contracts over the token
  axis using a coalesced LDS fill + hardware transpose-read (ds_read_tr16) so the token slot
  becomes the matrix-core contraction axis with no strided global gather. Each dW tile is owned by
  exactly one workgroup → race-free store; supports bf16 or fp32 output for in-kernel
  main_grad accumulation, and a swap_gather variant for FC2's native dW layout.

Thin dispatch/plumbing: pf_fwd_wrapper.py, pf_wgrad_wrapper.py, tensor_shim.py; orchestration
in permute_free_grouped_gemm.py.

---

6. Why this saves memory (the headline benefits)

Removes the permute activation memory. The activations are never copied into expert-contiguous
order. The GEMM gathers rows on the fly via sorted_slot_ids, so the [num_routes, hidden] permute
buffer (and its unpermute twin) simply doesn't exist — in both forward and backward. The token
combine is a gather (each token pulls its routes) rather than a scatter, so no atomic-scatter
scratch either.

Removes the persisted SwiGLU activation. FC1 emits the raw 2F [gate|up] pre-activation; the
gated activation act(gate)*up*prob is materialized into a short-lived F-wide transient that
FC2 consumes immediately and is never saved. Across the fwd→bwd boundary we persist only the 2F
pre-activation — the F-wide activated tensor that the standard path checkpoints is gone.

Fuses the recompute into the SwiGLU backward — "swiglu is just an op." In the backward,
fused_gated_act_prob_fwd re-materializes the F-wide activation from the saved 2F preact, and
fused_gated_act_prob_bwd computes the grad w.r.t. the preact (and optionally the route-prob grad)
in the same pass. The recompute feeds the unchanged FC2 wgrad, so that kernel stays at full
stored-activation speed while only the 2F preact is kept live. The whole SwiGLU becomes a single
fused op with no separately-checkpointed subgraph.

Parity with the MegaMOE GEMM kernel. The PF fwd GEMM kernels are ported from MegaMOE GEMM kernels, while wgrad is a custom written kernel specifically for Permute free.

Faster than the other TE bf16 grouped-GEMM backends.
benchmarks/microbenchmarks/benchmark_perm_free_grouped_gemm.py compares permute_free against
hipblaslt, ck, and triton on the DeepSeek/Grok/Qwen-style shapes; the permute-free path leads
for bf16 grouped GEMM.

CUDA-graph capturable. The align/routing pass has no host sync — it over-allocates to static,
shape-derived bounds and carries real extents as device scalars — and the GEMMs use fixed geometry
(no first-call autotune benchmarking). The full MoE step can therefore be captured into a HIP/CUDA
graph.

---

7. Integration & enablement

• Wired into transformer_engine/pytorch/module/grouped_linear.py (forward + backward dispatch,
  activation/route-prob handling).
• Public surface exported from transformer_engine/pytorch/moe/__init__.py.
• Enable: NVTE_PERMUTE_FREE_GROUPED_GEMM=1 (ROCm/HIP builds only; gated by IS_HIP_EXTENSION).
  Off by default — zero impact on existing paths when unset.

---

8. Testing

• tests/pytorch/test_perm_free_grouped_linear.py (22 tests): FC1/FC2 fwd, dgrad, wgrad (incl. fp32
  accumulate and recompute-from-preact), the standalone gated-act fwd/bwd for both silu and gelu, the
  full FC1→FC2 gated pipeline, fwd/dgrad/wgrad consistency, gather-combine dispatch, sync-free align,
  and the enable flag.
• Verified on gfx950 (ROCm 7.2, torch 2.8.0+rocm7.2.1): 22 passed.

---

Suggested reading order

1. moe_routing.py — the metadata contract and layout (start here).
2. permute_free_grouped_gemm.py module docstring + permute_free_grouped_gemm_forward /
   _backward — the dispatch/flow.
3. pf_helper_kernels.py — align, gated-act recompute/bwd, gather-combine.
4. flydsl_kernels/permute_free_grouped_gemm/pf_fwd.py, then pf_dgrad.py, then pf_wgrad.py.
5. grouped_linear.py diff — the integration seam.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes