[Pytorch] Add B200&B300 BF16 Cutlass GroupGEMM for fine-grained MoE and Varlen-K WGrad Grouped GEMM

[alan-hpc]

Description

This PR extends the CUTLASS Group GEMM support added in #2045 to cover the variable-K
(K-grouped / ragged-K) BF16 weight-gradient (wgrad) path of fine-grained MoE models on H100 (SM90).

In expert-parallel MoE training the per-expert token counts — the contraction dimension of the
wgrad GEMM D_i = B_iᵀ @ A_i — are ragged and generally not 128-aligned, which the existing
uniform-K CUTLASS grouped-GEMM fast path from #2045 cannot serve. This PR adds a dedicated path
that handles ragged per-expert token counts directly (SM90 TMA/WGMMA), zero-initializes empty
(K=0) groups, and writes each per-expert D_i in place. Inputs are BF16; output is FP32 (default)
or BF16. The standard uniform-K and Multi-Stream cuBLAS paths are unchanged.

Performance on H100 80GB, BF16, wgrad (D_i = B_iᵀ @ A_i), CUTLASS vs. the Multi-Stream cuBLAS
baseline. Shape is (g, m, n, k[mink, avgk, maxk]): g groups, m = expert dim, n = hidden dim,
k = the per-group routed-token count — the ragged contraction this kernel is built for.

run benchmark with

NVTE_USE_CUTLASS_GROUPED_GEMM=1 python benchmarks/gemm/benchmark_grouped_gemm_fwd_bwd.py --use-cutlass --dtype bf16 --num-experts <E> --ep-size 8 --hidden-dim 2048 --expert-dim 512 [--jagged-splits ...]

Shape(g, m, n, k[mink, avgk, maxk])	TE (cuBLAS, TFLOPs)	Cutlass (TFLOPs)	Speed-Up
(20, 512, 2048, k[3328, 3328, 3328])	445.50	567.65	1.27×
(20, 512, 2048, k[512, 3104, 6016])	444.52	520.92	1.17×
(32, 512, 2048, k[1024, 2048, 3072])	361.13	564.13	1.56×
(32, 512, 2048, k[512, 1024, 1536])	173.50	512.61	2.95×

The gain grows as the per-group K shrinks: small, ragged groups are where the Multi-Stream cuBLAS
per-group launch overhead dominates.

Correctness reuses the existing test harness from #2045 (unchanged in this PR): the parametrized
tests/pytorch/test_grouped_linear.py::test_grouped_gemm with layout=NT (the wgrad case),
use_cutlass=True, dtype=bfloat16 over ragged group splits exercises exactly this path and passes
on SM90.

This path reuses the NVTE_USE_CUTLASS_GROUPED_GEMM toggle introduced in #2045 (default 0):
export NVTE_USE_CUTLASS_GROUPED_GEMM=1 routes the BF16 NT wgrad through CUTLASS, 0 keeps the
Multi-Stream cuBLAS implementation. NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK still warns on fallback.

Type of change

• New feature (non-breaking change which adds functionality)

Changes

• cutlass_grouped_gemm.cuh: add CutlassGroupedGemmWgrad<trans_a, trans_b, ElementD> — an SM90
  grouped-GEMM template specialized for the NT wgrad layout — with explicit instantiations for
  FP32 and BF16 output.
• cutlass_grouped_gemm.cu: add cutlass_grouped_gemm_varlen_k(...). It validates the BF16 NT wgrad
  contract, splits groups into the non-empty set (excluding K=0 groups whose null A/B pointers
  would crash TMA descriptor construction, zero-initializing their outputs when not accumulating),
  and dispatches on output dtype — mirroring the existing cutlass_grouped_gemm call path.
• cublaslt_gemm.cu: wire the path into the nvte_multi_tensor_gemm dispatch
  (uniform-K fast path → K-grouped wgrad → cuBLAS fallback).

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

[greptile-apps]

Greptile Summary

This PR extends the existing CUTLASS Grouped GEMM path (from #2045) to handle the variable-K (ragged-K) BF16 weight-gradient case that arises in fine-grained MoE training on Hopper (SM90) and now Blackwell (SM100/SM103). It also refactors the host staging buffer from a monolithic 4 MB allocation to a 1024-slot pinned ring buffer (64 KB/slot, 64 MB total) to eliminate the per-call synchronization overhead imposed by pageable host memory.

• Adds CutlassGroupedGemmWgrad<trans_a, trans_b, ElementD, kSm100, kBigN> with SM90 Cooperative (FP32 output) / Pingpong (BF16 output) and SM100 2-SM tile variants, dispatched via a new cutlass_grouped_gemm_varlen_k function that filters out K=0 empty groups before launch.
• Wires the new path into nvte_multi_tensor_gemm behind shape-eligibility guards (is_bf16_wgrad_dtype, is_bf16_wgrad_shape) so mismatched or unsupported shapes still fall back to cuBLAS.
• Also adds forward-path SM100 instantiations for CutlassGroupedGemm, extends the host ring buffer for the forward path, and increases kMaxGroups from 64 to 256 to handle larger expert counts in expert-parallel MoE.

Confidence Score: 5/5

The new varlen-K wgrad dispatch path is well-guarded: shape eligibility is validated before entering the CUTLASS path, K=0 empty groups are correctly excluded with output zero-initialization when not accumulating, and the ring-buffer design properly prevents host-buffer reuse races across concurrent stream launches.

The core correctness logic — group filtering, NT-layout dispatch, SM90/SM100 tile selection, and cuBLAS fallback preservation — is sound. All findings are latent guards that are wrong in principle but cannot trigger given the current kMaxGroups=256 bound (~15 KB per slot, well within the 64 KB ring slot).

cutlass_grouped_gemm.cuh: the ring-buffer slot-size guard, the strict less-than workspace checks, and the two uninstantiated device-path functions whose null problem_sizes_host pointer could crash CUTLASS's scheduler if they are ever wired up without a matching host estimate.

Important Files Changed

Filename	Overview
transformer_engine/common/gemm/cutlass_grouped_gemm.cuh	Core template file adding CutlassGroupedGemmWgrad, SM100 schedule selectors, ring-buffer host workspace, and two new unreachable device-path functions; ring-buffer size guard checks total buffer size instead of per-slot size, and workspace size checks use strict less-than
transformer_engine/common/gemm/cutlass_grouped_gemm.cu	Adds explicit template instantiations for SM100 forward + wgrad variants; adds collect_bf16_wgrad_nt_groups and cutlass_grouped_gemm_varlen_k; correct SM100/SM90 dispatch logic
transformer_engine/common/gemm/cublaslt_gemm.cu	Adds Blackwell detection, is_bf16_wgrad_dtype/shape eligibility guards, and the new else-if branch dispatching to cutlass_grouped_gemm_varlen_k; logic is correct and unguarded shapes fall back to cuBLAS
transformer_engine/common/gemm/cublaslt_grouped_gemm.cu	Adds out_m, out_n, contraction_k fields to GroupedGemmConfig and increases kMaxGroups to 256; fields computed correctly but currently unused (pre-wired for future device-path integration)
tests/pytorch/test_grouped_linear.py	Extends skipif condition to include Blackwell (SM100/SM103) alongside Hopper (SM90); logically correct
transformer_engine/common/CMakeLists.txt	Adds CUDA::nvrtc and CUDA::cuda_driver as public link dependencies; consistent with SM100 CUTLASS requirements

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["nvte_multi_tensor_gemm()"] --> B{is Hopper or Blackwell\nAND use_cutlass?}
    B -- No --> C["cublas_path() fallback"]
    B -- Yes --> D{all_groups_uniform_k128\nAND no epilogue\nAND BF16/FP16 dtype}
    D -- Yes --> E["cutlass_grouped_gemm()\nuniform-K forward SM90+SM100"]
    D -- No --> F{is_bf16_wgrad_dtype\nAND NT layout AND grad\nAND is_bf16_wgrad_shape}
    F -- Yes --> G["cutlass_grouped_gemm_varlen_k()\nragged-K BF16 wgrad"]
    F -- No --> H["warn_fallback then cublas_path()"]
    G --> I["collect_bf16_wgrad_nt_groups()\nfilter K=0 groups\nzero-init empty outputs"]
    I --> J{A_nz empty?}
    J -- Yes --> K["return early all K=0"]
    J -- No --> L{SM100?}
    L -- No --> M["CutlassGroupedGemmWgrad\nSM90 Cooperative FP32\nor Pingpong BF16"]
    L -- Yes --> N{avg_K >= 1536?}
    N -- Yes --> O["CutlassGroupedGemmWgrad\nSM100 256x256 kBigN=true"]
    N -- No --> P["CutlassGroupedGemmWgrad\nSM100 256x128 kBigN=false"]
    M --> Q["getHostWorkspace ring slot\nfill ptrs+shapes cudaMemcpyAsync\nCUTLASS kernel launch"]
    O --> Q
    P --> Q

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["nvte_multi_tensor_gemm()"] --> B{is Hopper or Blackwell\nAND use_cutlass?}
    B -- No --> C["cublas_path() fallback"]
    B -- Yes --> D{all_groups_uniform_k128\nAND no epilogue\nAND BF16/FP16 dtype}
    D -- Yes --> E["cutlass_grouped_gemm()\nuniform-K forward SM90+SM100"]
    D -- No --> F{is_bf16_wgrad_dtype\nAND NT layout AND grad\nAND is_bf16_wgrad_shape}
    F -- Yes --> G["cutlass_grouped_gemm_varlen_k()\nragged-K BF16 wgrad"]
    F -- No --> H["warn_fallback then cublas_path()"]
    G --> I["collect_bf16_wgrad_nt_groups()\nfilter K=0 groups\nzero-init empty outputs"]
    I --> J{A_nz empty?}
    J -- Yes --> K["return early all K=0"]
    J -- No --> L{SM100?}
    L -- No --> M["CutlassGroupedGemmWgrad\nSM90 Cooperative FP32\nor Pingpong BF16"]
    L -- Yes --> N{avg_K >= 1536?}
    N -- Yes --> O["CutlassGroupedGemmWgrad\nSM100 256x256 kBigN=true"]
    N -- No --> P["CutlassGroupedGemmWgrad\nSM100 256x128 kBigN=false"]
    M --> Q["getHostWorkspace ring slot\nfill ptrs+shapes cudaMemcpyAsync\nCUTLASS kernel launch"]
    O --> Q
    P --> Q

Loading

_{Reviews (5): Last reviewed commit: "Merge branch 'NVIDIA:main' into feat/var..." | Re-trigger Greptile}

[ptrendx]

How does this kernel compare performance-wise with the cuBLASLt grouped gemm? Ideally if cuBLAS is better we would like to move towards that solution instead.