[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]

Confidence Score: 5/5

[Medium impact] The reviewed changes appear safe to merge after the stated prerequisite PR.

Summary

Adds experimental cuBLASLt strided batched GEMM and the fusible PyTorch operation BatchedLinear.

• Batched GEMM multiplies separate matrices without rearranging their inputs.
• BatchedLinear applies one learned transformation to each input group.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[BatchedLinear] --> B[Prepare input and weight]
  B --> C[PyTorch GEMM binding]
  C --> D{MXFP8 inputs?}
  D -->|Yes| E[Pack scales by batch]
  E --> F[cuBLASLt batched GEMM]
  D -->|No| F
  F --> G[Add bias and restore output shape]

Loading

_{Reviews (12) · Last reviewed commit: "Pass 3D operands from BatchedLinear to s..." · Reviewed by 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).

[greptile-apps]

Comments Outside Diff

These findings could not be posted inline.

•  cuBLASLt handle leak on any API failure after partial descriptor creation transformer_engine/common/gemm/cublaslt_gemm.cu:1148 ▶

  Adesc, Bdesc, Cdesc, Ddesc, operationDesc, and preference are created incrementally. If any NVTE_CHECK_CUBLAS call after the first cublasLtMatrixLayoutCreate throws - for example in set_strided_batch, in the MXFP8 cublasLtMatmulDescSetAttribute block, in cublasLtMatmulAlgoGetHeuristic, or in cublasLtMatmul - the already-allocated handles are never passed to their corresponding Destroy functions. The cleanup block at the bottom of the function is only reached on the happy path.

  The MXFP8 version checks at the top of the function were correctly moved before descriptor creation, so those specific paths are safe. The leak window is every NVTE_CHECK_CUBLAS call from the first descriptor creation onward. Adding RAII wrappers (or a try/catch with a cleanup lambda) to the new function would prevent the problem here.

•  C-matrix always aliased to D - undocumented design restriction transformer_engine/pytorch/cpp_extensions/gemm.py:1491 ▶

  The public C-level API nvte_cublas_gemm_strided_batched accepts independent C, ldc, and stridec parameters (useful for a fused D = A*B + beta*C with distinct accumulation buffers). This Python wrapper always passes out for both C and D, permanently losing that flexibility. A comment in the docstring noting that accumulation is always done in-place (beta is applied to out itself) would prevent surprises for future callers who discover the separate C parameter in the C++ header.

  Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[timmoon10]

There are a lot of redundant params here. m, n, k, lda, stridea, etc are contained in the dimensions of A and B. Also, it seems this is a general pattern that we could apply in other situation. We should avoid hard-coding this to only handle MLA.

[yaox12]

I kept ld/stride explicit on purpose. The main use case is an input activation stored contiguously as [S, B, G, D] with the GEMM batch dim G interleaved, i.e. ld = G*D and stride = D. NVTETensor has no strides, so this layout can't be inferred from shapes. Copying to [G, S*B, D] would defeat the purpose of the strided kernel.

m/n/k could be derived from shapes for the batch-major case, but not for the interleaved case without extra conventions. I kept them explicit so the API mirrors cuBLASLt's strided-batch API. Nothing in the API is MLA-specific; it is a generic strided batched GEMM, and the header docs now describe it that way. I'm happy to revisit this if NVTETensor gets stride support.

[timmoon10]

Logically, the inputs to batched GEMM are 3D tensors, where the 3 dims are interpreted as the batch dim column dim, and row dim. We can figure everything out based on the batch dim and trans*:

• Batch dim tells us how to flatten to 3D. batch_dim=0 means we flatten [G, ..., D] to [G, flat, D]. batch_dim=-2 means we flatten [..., G, D] to [flat, G, D].
• trans* tells us whether the col dim or row dim comes first. transa=False and transb=True means the last dim is k.

So we can make this significantly more logical and less redundant by replacing all of the stride and dim stuff with batch_dim_a and batch_dim_b.

For comparison, plain GEMM operates on 2D data and we implicitly flatten [..., D] to [flat, D]. Batched GEMM adds an extra dim, so one extra parameter seems fair.

[yaox12]

Done. nvte_cublas_gemm_strided_batched and nvte_cublas_gemm_strided_batched_mxfp8 now take batch_dim_a, batch_dim_b, and batch_dim_d instead of dims, leading dims, and strides. batch_dim_d applies to both C and D, since the fprop and dgrad outputs are also interleaved. Each tensor is flattened to [G, R, C] (batch_dim=0) or [R, G, C] (batch_dim=-2). m/n/k, ld, and stride are derived from the shapes and trans*, with consistency checks across A, B, and D.

One detail in the PyTorch binding: with batch_dim=-2, column-wise MXFP8 data has to be quantized as a 2D [R, G*C] view, otherwise 32-row scaling blocks would mix batch entries. The binding therefore treats 2D MXFP8 data with batch_dim=-2 as that view and takes G from the output.

[timmoon10]

Logically, the inputs to batched GEMM are 3D tensors, where the 3 dims are interpreted as the batch dim column dim, and row dim. We can figure everything out based on the batch dim and trans*:

• Batch dim tells us how to flatten to 3D. batch_dim=0 means we flatten [G, ..., D] to [G, flat, D]. batch_dim=-2 means we flatten [..., G, D] to [flat, G, D].
• trans* tells us whether the col dim or row dim comes first. transa=False and transb=True means the last dim is k.

So we can make this significantly more logical and less redundant by replacing all of the stride and dim stuff with batch_dim_a and batch_dim_b.

For comparison, plain GEMM operates on 2D data and we implicitly flatten [..., D] to [flat, D]. Batched GEMM adds an extra dim, so one extra parameter seems fair.

[timmoon10]

We should add a fallback implementation for unsupported quantization recipes. MXFP8 support is an internal detail, and if users need to be aware of this before calling this API then it significantly reduces its usefulness (e.g. for experimental or newly-developed recipes). It's fine if the fallback is slow, but it should work.

There are two natural fallbacks:

• Apply quantizers without checking the recipe, and if it's not MXFP8 then dequantize to BF16 before calling batched GEMM (preferably in cpp_extensions.strided_batched_gemm). This is somewhat easier since the op can use the same logic for every quantization format, but it adds unnecessary dequantize kernels.
• Reorder input tensor if batch_dim=-2, then slice. Independently call quantize and GEMM for each slice. This requires some MXFP8-awareness in the op impl, but it means the fallback impl matches what we were doing before this PR.

[yaox12]

I went with the first option. BatchedLinear now accepts any recipe and applies the recipe's quantizers as usual. For MXFP8, the data is quantized so its scales can be packed by batch. For other formats, the op only quantizes row-wise, and cpp_extensions.strided_batched_gemm dequantizes and runs the high-precision batched GEMM, using your suggested fallback. There is one special case: if the weight is quantized in a non-MXFP8 format but compute is MXFP8, the weight is dequantized and requantized to MXFP8. The fusible-op tests now cover the whole _quantization_list, and test_deferred_param_init_quantized includes batched_linear for all recipes.

[timmoon10]

This has 384 test cases, which is just about manageable. If we add more parametization in the future, we should split this up into multiple subtests so that the growth rate changes from 2^n to 2*n.

[yaox12]

Agreed. With _quantization_list it now collects about 770 cases, though many are skipped (e.g. quantized FP32), and the batched-linear tests run in about 10 s on B200. I'll split it into subtests if we add more parametrization.