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

benchmark/gemm/bench_grouped_gemm_canonical_host_latency.py

@@ -0,0 +1,110 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Host-latency comparison for the grouped GEMM SwiGLU wrapper: legacy pre-permuted
+call (including the TransformerEngine-style per-call view/permute gymnastics the
+legacy contract forces on the caller) vs the canonical natural-layout call.
+
+DSv3 fc1 shape from MLPerf MoE: (sum_m, n, k) = (24576, 7168, 2048), MXFP8 inputs,
+first dim overallocated 1.5x-4x for EP routing slack.
+
+Usage: python bench_grouped_gemm_canonical_host_latency.py
+"""
+
+import time
+
+import torch
+
+from cudnn import grouped_gemm_swiglu_wrapper_sm100
+from cudnn.api_base import ceil_div
+
+VALID_M, N, K = 24576, 7168, 2048
+EXPERTS = 8
+SF_VEC = 32
+WARMUP, ITERS = 20, 200
+
+
+def make_buffers(tensor_m):
+    dev = "cuda"
+    rest_k = ceil_div(ceil_div(K, SF_VEC), 4)
+    # Natural (canonical) buffers, as a framework owns them.
+    a = torch.randint(0, 200, (tensor_m, K), dtype=torch.uint8, device=dev).view(torch.float8_e4m3fn)
+    b = torch.randint(0, 200, (EXPERTS, N, K), dtype=torch.uint8, device=dev).view(torch.float8_e4m3fn)
+    sfa = torch.randint(118, 132, (1, ceil_div(tensor_m, 128), rest_k, 32, 4, 4), dtype=torch.uint8, device=dev).view(torch.float8_e8m0fnu)
+    sfb = torch.randint(118, 132, (EXPERTS, ceil_div(N, 128), rest_k, 32, 4, 4), dtype=torch.uint8, device=dev).view(torch.float8_e8m0fnu)
+    group = VALID_M // EXPERTS
+    offsets = torch.arange(group, VALID_M + 1, group, dtype=torch.int32, device=dev)
+    alpha = torch.ones(EXPERTS, dtype=torch.float32, device=dev)
+    prob = torch.rand(tensor_m, dtype=torch.float32, device=dev)
+    norm_const = torch.tensor([0.01], dtype=torch.float32, device=dev)
+    return dict(a=a, b=b, sfa=sfa, sfb=sfb, offsets=offsets, alpha=alpha, prob=prob, norm_const=norm_const)
+
+
+def call_legacy(buf):
+    # TE-style per-call layout gymnastics required by the legacy contract
+    # (see transformer_engine grouped_mlp.py: 6-D SF view+permute, B permute).
+    m = buf["a"].shape[0]
+    a3d = buf["a"].view(m, K, 1)
+    b_nkl = buf["b"].permute(1, 2, 0)
+    sfa6d = buf["sfa"].view(torch.float8_e8m0fnu).view(1, ceil_div(m, 128), ceil_div(ceil_div(K, SF_VEC), 4), 32, 4, 4).permute(3, 4, 1, 5, 2, 0)
+    sfb6d = buf["sfb"].view(torch.float8_e8m0fnu).view(EXPERTS, ceil_div(N, 128), ceil_div(ceil_div(K, SF_VEC), 4), 32, 4, 4).permute(3, 4, 1, 5, 2, 0)
+    prob3d = buf["prob"].view(m, 1, 1)
+    return grouped_gemm_swiglu_wrapper_sm100(
+        a_tensor=a3d,
+        b_tensor=b_nkl,
+        sfa_tensor=sfa6d,
+        sfb_tensor=sfb6d,
+        padded_offsets=buf["offsets"],
+        alpha_tensor=buf["alpha"],
+        norm_const_tensor=buf["norm_const"],
+        prob_tensor=prob3d,
+        d_dtype=torch.float8_e4m3fn,
+        sf_vec_size=SF_VEC,
+    )
+
+
+def call_canonical(buf):
+    return grouped_gemm_swiglu_wrapper_sm100(
+        a_tensor=buf["a"],
+        b_tensor=buf["b"],
+        sfa_tensor=buf["sfa"],
+        sfb_tensor=buf["sfb"],
+        padded_offsets=buf["offsets"],
+        alpha_tensor=None,
+        norm_const_tensor=buf["norm_const"],
+        prob_tensor=buf["prob"],
+        d_dtype=torch.float8_e4m3fn,
+        sf_vec_size=SF_VEC,
+    )
+
+
+def bench(fn, buf):
+    for _ in range(WARMUP):
+        fn(buf)
+    torch.cuda.synchronize()
+    times = []
+    for _ in range(ITERS):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        fn(buf)
+        times.append(time.perf_counter() - t0)
+    torch.cuda.synchronize()
+    times.sort()
+    n = len(times)
+    return times[n // 2] * 1e6, times[int(n * 0.9)] * 1e6
+
+
+def main():
+    torch.manual_seed(0)
+    print(f"grouped_gemm_swiglu host latency, (sum_m, n, k)=({VALID_M}, {N}, {K}), {EXPERTS} experts, MXFP8")
+    print(f"{'overalloc':>9} | {'tensor_m':>8} | {'legacy p50/p90 (us)':>22} | {'canonical p50/p90 (us)':>22}")
+    for factor in (1.5, 2.0, 4.0):
+        tensor_m = ceil_div(int(VALID_M * factor), 256) * 256
+        buf = make_buffers(tensor_m)
+        legacy_p50, legacy_p90 = bench(call_legacy, buf)
+        canon_p50, canon_p90 = bench(call_canonical, buf)
+        print(f"{factor:>8}x | {tensor_m:>8} | {legacy_p50:>10.1f} / {legacy_p90:>7.1f} | {canon_p50:>11.1f} / {canon_p90:>7.1f}")
+
+
+if __name__ == "__main__":
+    main()

docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md

@@ -363,6 +363,25 @@ Returns a `TupleDict` - a dictionary-like object that also supports tuple unpack
 - `C`, `D_row`, and `D_col` must be **N-major** (contiguous along N dimension)
 - All tensors must be **16-byte aligned** along the contiguous dimension
 
+### Canonical layouts (additive)
+
+Each input is also accepted in its natural row-major form, normalized internally;
+the pre-permuted kernel-facing forms above keep working unchanged:
+
+- `A`: `(valid_m, K)` row-major
+- `B`: `(L, N, K)` C-contiguous
+- `C`: `(valid_m, 2N)` row-major
+- `SFA`/`SFB`: any dense C-contiguous buffer with the MMA-tiled element count,
+  e.g. flat 1-D or the physical `(L, ceil(mn/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4)`
+  allocation. The kernel rebuilds the MMA-tiled SF layouts from the GEMM shapes and
+  reads only the base pointer.
+- `prob`: `(valid_m,)`, `float32` or `bfloat16`
+- `alpha_tensor` may be omitted (defaults to cached ones)
+
+When `A` is canonical (2-D), the wrapper returns natural-shaped outputs:
+`d_row`/`d_col (valid_m, 2N)` row-major, `dprob (valid_m,)`, and
+`sfd_row`/`sfd_col` as C-contiguous physical `(1, ceil(mn/128), rest, 32, 4, 4)` buffers.
+
 ### Data Types
 
 #### Input/Weight Types (ab_dtype)

docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md

@@ -340,6 +340,24 @@ Returns a `TupleDict` - a dictionary-like object that also supports tuple unpack
 - `C`, `D`, and `D_col` must be **N-major** (contiguous along N dimension)
 - All tensors must be **16-byte aligned** along the contiguous dimension
 
+### Canonical layouts (additive)
+
+Each input is also accepted in its natural row-major form, normalized internally;
+the pre-permuted kernel-facing forms above keep working unchanged:
+
+- `A`: `(valid_m, K)` row-major
+- `B`: `(L, N, K)` C-contiguous
+- `SFA`/`SFB`: any dense C-contiguous buffer with the MMA-tiled element count,
+  e.g. flat 1-D or the physical `(L, ceil(mn/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4)`
+  allocation — no `.view().permute()` gymnastics required. The kernel rebuilds the
+  MMA-tiled SF layouts from the GEMM shapes and reads only the base pointer.
+- `prob`: `(valid_m,)`, `float32` or `bfloat16`
+- `alpha_tensor` may be omitted (defaults to cached ones)
+
+When `A` is canonical (2-D), the wrapper returns natural-shaped outputs:
+`c (valid_m, N)`, `d`/`d_col (valid_m, N/2)` row-major, and `sfd_row`/`sfd_col` as
+C-contiguous physical `(1, ceil(mn/128), rest, 32, 4, 4)` buffers.
+
 ### Data Types
 
 #### Input/Weight Types (ab_dtype)

python/cudnn/gemm/cutedsl/grouped/canonical.py

@@ -0,0 +1,103 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Host-side helpers for canonical (natural row-major) grouped GEMM tensor layouts.
+
+The contiguous grouped GEMM kernels historically required callers to pre-permute
+every operand into the kernel-facing form: A/C/D as (m, x, 1) with a trailing unit
+L mode, B as (n, k, l) k-major strided views, prob as (m, 1, 1), and the block
+scale factors as 6-D MMA-tiled (32, 4, mn//128, 4, rest_k, l) strided views of a
+dense buffer. These helpers additively accept the natural buffers instead --
+A (sum_m, k) row-major, B (l, n, k) row-major, prob (sum_m,), and flat/dense
+C-contiguous scale-factor buffers -- and normalize them to the kernel-facing form
+with zero-copy views. Kernel-facing inputs pass through unchanged, so existing
+callers are unaffected.
+
+The scale-factor kernels rebuild the MMA-tiled SF layouts from the A/B/D shapes on
+device and consume only the SF base pointers, so a canonical (C-contiguous) SF
+buffer is compiled as a flat 1-D tensor: no MMA-permuted view is ever materialized.
+"""
+
+from __future__ import annotations
+
+import cutlass.cute as cute
+
+_cache_of_alpha_ones = {}
+
+
+def default_alpha_ones(l: int, device):
+    """Cached all-ones per-group scale for callers that don't scale per group."""
+    import torch
+
+    key = (l, str(device))
+    alpha = _cache_of_alpha_ones.get(key)
+    if alpha is None:
+        alpha = torch.ones(l, dtype=torch.float32, device=device)
+        _cache_of_alpha_ones[key] = alpha
+    return alpha
+
+
+def unsqueeze_l_dim(tensor):
+    """Canonical (m, x) row-major -> kernel-facing (m, x, 1); 3-D passes through."""
+    if tensor is not None and tensor.ndim == 2:
+        return tensor.unsqueeze(-1)
+    return tensor
+
+
+def is_canonical_b(tensor) -> bool:
+    """True for a canonical (l, n, k) row-major weight tensor.
+
+    The kernel-facing forms keep a stride-1 k mode at dim 1 (k-major (n, k, l)) or a
+    stride-1 n mode at dim 0 (n-major), so a stride-1 innermost dim 2 identifies the
+    canonical form.
+    """
+    if tensor is None or tensor.ndim != 3:
+        return False
+    stride = tensor.stride()
+    return stride[2] == 1 and stride[1] != 1 and stride[0] != 1
+
+
+def to_kernel_b(tensor):
+    """Canonical (l, n, k) row-major -> kernel-facing (n, k, l); other forms pass through."""
+    if is_canonical_b(tensor):
+        return tensor.permute(1, 2, 0)
+    return tensor
+
+
+def to_kernel_prob(tensor):
+    """Canonical (m,) -> kernel-facing (m, 1, 1); other ranks pass through."""
+    if tensor is not None and tensor.ndim == 1:
+        return tensor.view(-1, 1, 1)
+    return tensor
+
+
+def is_flat_sf(tensor) -> bool:
+    """True when a scale-factor tensor is a dense C-contiguous buffer (canonical form).
+
+    The legacy MMA-tiled 6-D views are non-contiguous except for degenerate unit-dim
+    cases where both interpretations address identical memory.
+    """
+    return tensor is not None and tensor.is_contiguous()
+
+
+def to_kernel_sf(tensor, flat: bool):
+    """Flatten a canonical scale-factor buffer to 1-D; legacy MMA views pass through."""
+    if tensor is None or not flat:
+        return tensor
+    return tensor if tensor.ndim == 1 else tensor.view(-1)
+
+
+def make_flat_sf_fake(api, desc):
+    """Fake cute tensor for a flat (1-D, dynamic-length) scale-factor buffer.
+
+    The kernels rebuild the SF layout from the GEMM operand shapes and read only the
+    base pointer, so the compiled signature needs nothing beyond dtype and a dynamic
+    length (always a multiple of the 32x4x4 = 512-element SF atom).
+    """
+    if desc is None:
+        return None
+    return api._make_fake_cute_tensor(
+        dtype=desc.dtype,
+        shape=(cute.sym_int(divisibility=512),),
+        stride=(1,),
+    )