[Common][PyTorch] Add strided batched GEMM in BF16/MXFP8

[yaox12]

Description

This PR adds experimental cuBLASLt-backed strided batched GEMM support and a new fusible PyTorch operation exported as transformer_engine.pytorch.ops.BatchedLinear.

The main use case is applying one weight matrix per group without rearranging interleaved activations:

[S, B, G, D] @ [G, R, D] -> [S, B, G, R]

Both contiguous [G, ..., D] and [..., G, D] input layouts are supported through explicit leading dimensions and batch strides. MXFP8 scales are packed into the cuBLASLt per-batch layout internally before each GEMM.

Dependency

This PR depends on #3327 for meta-device deferred initialization in the operation fuser and should be merged after it.

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

Changes

• Add an experimental common API for cuBLASLt strided batched GEMM with independent matrix strides.
• Support high-precision and MXFP8 input pairs with high-precision output.
• Add transformer_engine.pytorch.ops.BatchedLinear with high-precision and MXFP8 forward and backward computation.
• Support primary MXFP8 weights, bias and return_bias, save_original_input, custom initialization and RNG tracking, and main-gradient accumulation.
• Add API documentation and PyTorch L0 tests.

Current limitations

• A and B must both be high precision or both be MXFP8.
• MXFP8 operands must use contiguous [G, ..., D] or [..., G, D] layouts.
• Output quantization, fused GEMM epilogues, non-MXFP8 recipes, and tensor parallelism are not supported.

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 adds experimental cuBLASLt-backed strided batched GEMM support (BF16 and MXFP8) and exports it as BatchedLinear, a new fusible PyTorch operation. The primary use case is applying one weight matrix per group over interleaved activations without rearranging the [S, B, G, D] layout.

• C++ core: cublas_gemm_strided_batched in cublaslt_gemm.cu gains a CublasLtDescriptorCleanup RAII guard that correctly destroys all cuBLASLt descriptors on both the happy path and exception unwinding; buffer-span and MXFP8 scale-size validators are added with overflow-checked arithmetic.
• MXFP8 scale packing: Two new CUDA kernels in swizzle.cu (pack_mxfp8_rowwise_scales_for_batched_gemm_kernel / pack_mxfp8_columnwise_scales_for_batched_gemm_kernel) convert per-tensor compact scales into the per-batch GEMM-swizzled layout expected by cuBLASLt, correctly handling both batch-major and interleaved input storage.
• PyTorch layer: BatchedLinear implements forward (TN), dgrad (NN), and wgrad (NT) strided GEMMs with bias, return_bias, save_original_input, main-grad accumulation, deferred meta-device initialisation, and RNG tracker support; comprehensive L0 tests cover numerics, selective gradients, and error paths.

Confidence Score: 5/5

• The change is self-contained, well-tested, and guarded by RAII descriptor cleanup in C++. The two findings are both non-blocking quality observations.
• The core GEMM path is thoroughly exercised by L0 tests covering BF16 and MXFP8 in both batch layouts, forward and backward numerics, error rejection, and edge-case strides. The RAII guard correctly handles descriptor cleanup on failure. Buffer-span validation is present at both the Python and C++ layers. No correctness-affecting defect was found in the changed code paths.
• No files require special attention. The only notes are a missing contiguity assertion in strided_batched_gemm (Python wrapper) and an undocumented asymmetry in zero-dimension handling in cublas_gemm_strided_batched.

Important Files Changed

Filename	Overview
transformer_engine/common/gemm/cublaslt_gemm.cu	Adds cublas_gemm_strided_batched with RAII CublasLtDescriptorCleanup guard (addressing prior descriptor-leak concern), buffer-bounds validation helpers, and nvte_cublas_gemm_strided_batched public wrapper. MXFP8 version checks are correctly placed before any descriptor creation.
transformer_engine/common/swizzle/swizzle.cu	Adds two new CUDA kernels for packing and GEMM-swizzling MXFP8 scales by batch, supporting both batch-major [G,…,D] and interleaved […,G,D] layouts, plus a feature-batched path for column-wise interleaved tensors. Swizzle index arithmetic is consistent with the existing mxfp8_gemm_swizzled_scale_idx function.
transformer_engine/pytorch/csrc/extensions/gemm.cpp	Implements the strided_batched_gemm C++ extension: selects the correct rowwise/columnwise scale buffer per operand, packs compact MXFP8 scales into the per-batch swizzled layout, validates layout constraints, and calls nvte_cublas_gemm_strided_batched. MXFP8 operand contiguity and data-size checks are present.
transformer_engine/pytorch/ops/basic/batched_linear.py	New BatchedLinear operation implementing forward (TN), input-gradient (NN), and weight-gradient (NT) strided batched GEMMs with BF16 and MXFP8 paths, bias, return_bias, main_grad accumulation, save_original_input, deferred meta-device init, and RNG tracker support. Logic is well-tested and consistent with the C++ backend.
transformer_engine/pytorch/cpp_extensions/gemm.py	Adds the strided_batched_gemm Python wrapper that validates input types, resolves beta, and invokes the C++ extension. The workspace size is passed as workspace.shape[0] (correct for a 1D uint8 tensor) but the output tensor is not validated for contiguity.

Sequence Diagram

sequenceDiagram
    participant PY as BatchedLinear.op_forward
    participant GEM as strided_batched_gemm (Python)
    participant CPP as strided_batched_gemm (C++)
    participant PACK as pack_mxfp8_operand_scales
    participant SWIZ as nvte_pack_mxfp8_scales_for_batched_gemm
    participant CUBLASLT as cublas_gemm_strided_batched

    PY->>PY: _quantize_for_batched_gemm(x, w) → compact MXFP8
    PY->>GEM: "strided_batched_gemm(A, B, out, layout="TN", ...)"
    GEM->>CPP: tex.strided_batched_gemm(...)
    CPP->>CPP: detect MXFP8 inputs
    CPP->>PACK: "pack_mxfp8_operand_scales(A, rowwise=transa, ...)"
    PACK->>SWIZ: nvte_pack_mxfp8_scales_for_batched_gemm(compact→swizzled)
    SWIZ-->>PACK: packed scale buffer [G×padded_rows×padded_scale_cols]
    PACK-->>CPP: "PackedMXFP8Operand{tensor, scale_inv}"
    CPP->>PACK: "pack_mxfp8_operand_scales(B, rowwise=!transb, ...)"
    PACK->>SWIZ: nvte_pack_mxfp8_scales_for_batched_gemm / columnwise
    SWIZ-->>PACK: packed scale buffer
    PACK-->>CPP: "PackedMXFP8Operand{tensor, scale_inv}"
    CPP->>CUBLASLT: nvte_cublas_gemm_strided_batched(packed_A, packed_B, ...)
    CUBLASLT->>CUBLASLT: CublasLtDescriptorCleanup RAII guard
    CUBLASLT->>CUBLASLT: CheckStridedBatchBuffer(A,B,C,D)
    CUBLASLT->>CUBLASLT: CheckMXFP8BatchScaleBuffer(A,B)
    CUBLASLT->>CUBLASLT: cublasLtMatmul(strided-batch)
    CUBLASLT-->>CPP: output D (BF16)
    CPP-->>GEM: D tensor
    GEM-->>PY: output
    PY->>PY: _apply_bias(output)

Loading

_{Reviews (7): Last reviewed commit: "add save_original_input" | Re-trigger Greptile}

[yaox12]

/te-ci

[timmoon10]

Mathematically, is this equivalent to GroupedLinear with uniform split sizes (m_splits=[S*B] * G)?

If so, are there perf reasons to support a dedicated kernel? I could imagine that batched GEMM is easier to optimize than grouped GEMM, but grouped GEMM should already be getting a good fraction of the peak compute throughput. I would like to make sure we have a compelling benefit before accepting a major maintenance burden (on the order of Linear and GroupedLinear, especially if we need to support all the weird Mcore-specific integrations).