grouped gemm: accept canonical (sum_m,k)/(l,n,k) layouts and flat SF buffers

[hwanseoc]

Description

Implements the canonical-layout proposal for the contiguous grouped GEMM SwiGLU forward and dSwiGLU backward APIs (Slack thread): callers hand cuDNN the buffers they already own, in natural row-major form, instead of pre-permuting every operand into kernel-facing cute views. This removes the per-call layout gymnastics TE does today (e.g. sf.view(dtype=e8m0).view(1, m//128, k//128, 32, 4, 4).permute(3, 4, 1, 5, 2, 0) in grouped_mlp.py), and it makes every input of these two kernels expressible as a dense row-major array — the prerequisite for a jax_api.py for grouped/swiglu + grouped/dswiglu (follow-up PR; today they are the only grouped variants without one, blocking TE JAX MoE).

New accepted signatures (additive — every pre-permuted form keeps working unchanged)

grouped_gemm_swiglu_wrapper_sm100 / GroupedGemmSwigluSm100 and grouped_gemm_dswiglu_wrapper_sm100 / GroupedGemmDswigluSm100 now also accept, per operand and independently:

operand	kernel-facing form (unchanged)	canonical form (new)
A	(valid_m, k, 1) k-major	(valid_m, k) row-major
B	(n, k, l) k-major strided view	(l, n, k) C-contiguous
C (dswiglu input)	(valid_m, 2n, 1) n-major	(valid_m, 2n) row-major
SFA/SFB	MMA-tiled 6-D (32, 4, mn/128, 4, rest_k, l) strided view	any dense C-contiguous buffer with the same element count: flat 1-D, or physical (l, mn/128, rest_k, 32, 4, 4)
prob	(valid_m, 1, 1) float32	(valid_m,) float32 or bfloat16
alpha	required	optional; None → cached ones
stream	optional (already defaulted)	unchanged

When A is canonical (2-D), outputs come back natural-shaped: c (m, n), d/d_col (m, n/2) (dswiglu: d_row/d_col (m, 2n), dprob (m,)) row-major, and sfd_row/sfd_col as C-contiguous physical (1, mn/128, rest, 32, 4, 4) buffers.

What FE derives internally

• Flat SF buffers become raw pointers. Both kernels already discard the incoming SF layout on device — they rebuild it from the A/B/D shapes via tile_atom_to_shape_SF and consume only the base pointer (grouped_gemm_swiglu_quant.py:588). A canonical SF buffer therefore compiles as a flat 1-D dynamic-length fake tensor; no MMA-permuted view is ever materialized, host- or device-side. Kernel bodies unchanged.
• A/B/C/prob normalize to the kernel-facing views with zero-copy unsqueeze/permute at the API boundary (grouped/canonical.py); form detection is by rank/stride, so the two forms coexist and cache separately (compile signatures differ; the wrapper cache keys carry the form).
• bf16 prob: the only kernel change in this PR is an explicit .to(cutlass.Float32) at the prob load site in both kernels (no-op for fp32 prob).

Compat

• Pre-permuted-view callers (TE PyT path today) are bit-for-bit unaffected: same descriptors, same compiled signatures, same cache keys, same outputs.
• Degenerate case where a legacy 6-D SF view is also C-contiguous (all tiled dims unit-sized): both interpretations address identical memory, so routing it down the flat path is behavior-preserving.
• JAX arrays are still rejected at these entry points until the follow-up jax_api.py lands (the eager wrapper remains torch-lazy).

Host latency at Jeremy's DSv3 shape

benchmark/gemm/bench_grouped_gemm_canonical_host_latency.py, (sum_m, n, k) = (24576, 7168, 2048), 8 experts, MXFP8, first dim overallocated 1.5–4×, B200, GPU idle before each timed call. "Legacy" includes the TE-style per-call view/permute gymnastics the old contract forces on the caller.

overalloc	tensor_m	legacy p50 / p90 (µs)	canonical p50 / p90 (µs)
1.5×	36864	55.4 / 60.0	52.5 / 58.3
2.0×	49152	56.1 / 60.4	52.9 / 61.2
4.0×	98304	55.7 / 60.4	53.4 / 58.2

The wrapper host cost is dominated by descriptor validation + tvm-ffi marshalling (see #627/#589), so retiring the caller-side view/permute is worth ~3 µs/call (~5%) here; what it doesn't capture is TE's surrounding per-call weight copy()/swizzle staging and the code it deletes on the TE side. The primary win of this PR is the contract: every input is now a plain dense row-major buffer.

Testing

• New test/python/fe_api/grouped_gemm/test_grouped_gemm_canonical_layouts.py (9 tests, L0): canonical (physical and flat SF) vs legacy runs on identical device buffers must match bitwise, for swiglu (fp8 + fp4) and dswiglu (fp8); bf16-prob parity; alpha-default parity; natural output shape checks.
• Full legacy suites pass unchanged: test_grouped_gemm_swiglu.py, test_grouped_gemm_dswiglu.py, and the JAX-rejection tests. Drive-by: the swiglu wrapper cache key now includes the SF dtype (dswiglu already did); previously e8m0-vs-e4m3 runs of the same config collided on one compiled signature, which surfaced as Mismatched Tensor ... expected dtype=float8_e8m0fnu errors (18 pre-existing test skips on develop now pass).

Out of scope (follow-ups)

• jax_api.py for grouped/swiglu + grouped/dswiglu on cudnn.jax.call — this PR makes the representation expressible; the bridge is a separate small PR.
• The same canonical acceptance for the unified grouped/glu/dglu entry points and the bias (n, l) repack, if TE wants it there too.

[hwanseoc]

@coderabbitai ignore

[coderabbitai]

Note

Reviews paused

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

✅ Action performed

Reviews paused.