Use PyTorch CUTLASS for BF16 grouped GEMM on SM100

[wujingyue]

Description

Motivated by NVIDIA/Megatron-LM#7400: on GB200, grouped BF16 GEMM can retain old output values despite beta=0 while NCCL all-gather and reduce-scatter run concurrently. The issue includes a standalone cuBLASLt + NCCL reproducer without PyTorch or TE. The underlying cause remains unresolved.

This PR extends NVTE_USE_CUTLASS_GROUPED_GEMM=1 to the graph-safe BF16 GroupedLinear path on SM100 using torch._grouped_mm. It provides an alternative to the affected cuBLASLt path while preserving independent communication streams and CUDA graph support.

Type of change

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

Changes

• Dispatch supported BF16 forward, dgrad, and overwrite wgrad operations through PyTorch CUTLASS in the shared grouped GEMM wrapper. Expert counts remain on the GPU, and caller-owned output buffers and unused capacity are preserved.
• Retain TE's existing backend for FP32 outputs, fused accumulation, unsupported shapes, and explicit alpha/beta. Adding a rounded BF16 product separately would not preserve fused accumulation precision.
• Reuse the existing CUTLASS flag and extend its existing accuracy test to SM100. The module API also requires NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1; Hopper keeps its existing dispatch. If PyTorch exposes torch.backends.cuda.matmul.prefer_cublaslt_grouped_gemm, it must be False.

The implementation uses temporary weight/output buffers. No performance improvement is claimed.

Validation

GPU validation used GB200, PyTorch 2.14.0a0+4fdf77b940.nv26.08, and TE 2.18.0+27486e03:

• 241 existing tests passed, 19 skipped, covering forward/dgrad/wgrad, bias, empty experts, caller buffers, deferred wgrad, accumulation and MXFP8 fallbacks, and CUDA graph capture/replay. Instrumentation confirmed all three layouts reached _grouped_mm.
• 8-GPU MFSDP v2 + Flex/HybridEP training: 100 eager steps and 100 steps with both training and optimizer CUDA graphs completed with finite gradient norms. Training validation used an earlier revision of the adapter; the existing tests above were rerun after the API cleanup and implementation simplification.
• Isort, Black, Ruff F checks, and git diff --check passed.

This draft ports the change onto main. The new GEMM helper and shared dispatch function are identical to the GPU-tested versions; the port only required resolving imports. A fresh native build/test on this main revision and a physical Hopper run remain outstanding.

Checklist

• Read the contributing guidelines
• Documented the flag and backend limitations
• Extended existing test coverage
• Ran targeted numerical and CUDA graph tests on GB200
• Validate a fresh build of the PR revision
• Run the existing CUTLASS tests on Hopper

[wujingyue]

This looks silly -- why is GroupedTensorStorage used for non-quantized where the only field people care is rowwise_data?

[greptile-apps]

Confidence Score: 4/5

The PR should not merge until scaled bias is made compatible with packed buffers whose capacity exceeds their logical token count.

Findings

1.  Bias scaling breaks capacity ▶

Summary

This PR adds an SM100 BF16 grouped-GEMM route through PyTorch CUTLASS while retaining the existing backend for unsupported configurations.

• Adds TN, NN, and NT dispatch through torch._grouped_mm.
• Preserves caller-owned output buffers and applies grouped bias in Python.
• Extends the existing CUTLASS accuracy test and documents the Blackwell behavior.
• The scaled-bias path does not currently handle packed buffers with unused capacity, and the added test does not prove that the new backend was selected.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  C[Grouped tensor GEMM call] --> G{SM100 BF16 and CUTLASS enabled?}
  G -- No --> TE[Existing TE backend]
  G -- Yes --> MM[torch._grouped_mm]
  MM --> B{Bias provided?}
  B -- No --> W[Write logical rows and preserve capacity]
  B -- Yes --> E[Expand per-group bias across result capacity]
  E --> S{Per-token bias scale?}
  S -- No --> W
  S -- Yes --> M[Scale logical-token vector against capacity-sized bias]
  M --> X[Size mismatch when capacity exceeds logical rows]

Loading

_{Reviews (1) · Last reviewed commit: "Fix grouped GEMM row-index initializatio..."}

[greptile-apps]

Bias scaling breaks capacity

When packed input storage has unused capacity beyond offsets[-1], result and the expanded bias_data include those extra rows, while bias_scale contains only the logical token rows. This multiplication therefore raises a size-mismatch error before the later torch.where can preserve the unused output capacity. Please limit scaling to the logical prefix or otherwise account for the capacity rows.

Knowledge Base Used: Native GEMM and quantization kernels

[greptile-apps]

Comments Outside Diff

These findings sit on lines the diff does not cover, so they could not be posted inline. Each one leaves this list once its file changes.

•  Test permits silent fallback tests/pytorch/test_grouped_linear.py:518 ▶

  The SM100 test only compares numerical output with bias=False and never verifies that torch._grouped_mm was called. Because the wrapper catches NotImplementedError and falls back to the existing backend, this test can pass without exercising the backend it was added to validate. Instrument the dispatch and cover at least the new bias and capacity behavior so regressions do not pass through the fallback.