[JAX] Integrate BF16 Grouped GEMM with on-device group sizes

[jberchtold-nvidia]

Description

Integrate new BF16 grouped GEMM from TE common/cuBLASLt that supports on-device group sizes without a D2H memcpy and stream sync. This grouped GEMM is faster and CUDA-graph safe.

Also fixes #2659

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

• Added a new cuda-graph-safe grouped GEMM to TE/JAX that is automatically used as the backend when the input data is bf16 (no scaling recipe) and not bias is required.
• Exposed a new make_ragged_dot_cls for easy integration into existing models. This will be most useful when quantization is supported and storing recipe state is required

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 integrates a new CUDA-graph-safe BF16 grouped GEMM backend (GroupedGemmV2FFI / nvte_grouped_gemm) into the JAX TE layer, replacing the older multi-stream path that required a D2H memcpy and stream synchronisation for plain BF16 / NO_SCALING inputs. It also exposes a make_grouped_dense_cls factory for easy drop-in use in existing Flax MoE models, and fixes a dimension-swap bug in the C++ grouped GEMM test.

Key changes:

• New C++ backend (GroupedGemmV2FFI) dispatches to nvte_grouped_gemm (cuBLAS 13.2+) and is registered with FFI_CudaGraph_Traits, making it safe for CUDA graph capture.
• Automatic dispatch in Python: _can_use_v2_grouped_gemm selects the V2 path for bfloat16 + NO_SCALING + no-bias; otherwise falls back to the legacy multi-stream path.
• _v2_grouped_gemm_available is cached at module import time via a lightweight probe, avoiding per-trace overhead.
• make_grouped_dense_cls provides a make_dot_general_cls-style factory; however, the internal Flax module is registered under "ragged_dot" instead of "grouped_dense", creating a checkpoint-path inconsistency that should be fixed.
• GroupedGemmV2FFI carries over several dead variables and an unreachable FP8 architecture check from the legacy path that were not cleaned up during the port, which will generate compiler warnings.
• del kwargs in te_grouped_dot_general silently discards any keyword arguments instead of validating or forwarding them.
• Bug fixes: jnp.empty → jnp.ones for the uninitialized scale buffer in grouped_quantize; M/N dimension swap fixed in the C++ test; d_rows/d_cols swap fixed in the setup kernel.

Confidence Score: 3/5

• Safe for BF16 inference/training but has a checkpoint-breaking naming bug in make_grouped_dense_cls that should be fixed before wider adoption.
• The core CUDA-graph-safe GEMM path and dispatch logic appear correct and the PR includes test coverage. However, the "ragged_dot" internal module name in make_grouped_dense_cls is a logic-level inconsistency that will silently corrupt checkpoint parameter paths for any model using that factory. Additionally, GroupedGemmV2FFI has accumulated a significant number of dead variables and an unreachable code block that will produce compiler warnings and hinder maintainability. These issues lower confidence from an otherwise well-structured feature addition.
• transformer_engine/jax/flax/module.py (ragged_dot naming bug) and transformer_engine/jax/csrc/extensions/gemm.cpp (dead variables, unreachable FP8 check).

Important Files Changed

Filename	Overview
transformer_engine/jax/csrc/extensions/gemm.cpp	Adds new GroupedGemmV2FFI and supporting JAXX_GroupedTensorWrapper class for CUDA-graph-safe grouped GEMM. The implementation has several dead variables (lhs_ptr, rhs_ptr, out_ptr, num_math_sm, grad, accumulate, use_split_accumulator, bias_shape, out_dtype_bytes) and an unreachable FP8 arch check that were carried over from the legacy path. A redundant lhs_is_trans = true assignment in the wgrad branch is also present.
transformer_engine/jax/flax/module.py	Adds make_grouped_dense_cls factory function and n_groups support in quantizer-set generation. Has two bugs: (1) the internal Flax module is registered as "ragged_dot" instead of "grouped_dense", causing checkpoint path inconsistency; (2) **kwargs are silently discarded rather than validated or forwarded.
transformer_engine/common/gemm/cublaslt_grouped_gemm.cu	Adds compute_grouped_tensor_offset device helper with cumulative-sum for per-tensor dims, exposes nvte_get_grouped_gemm_setup_workspace_size, and adds nvte_convert_int32_to_int64 utility kernel. Also fixes d_rows/d_cols swap. Code is in the correct preprocessor guards; the cuBLAS stubs are properly placed in the #else block.
transformer_engine/jax/cpp_extensions/gemm.py	Integrates the V2 CUDA-graph-safe backend into the JAX primitive layer; adds _can_use_v2_grouped_gemm dispatch helper and _v2_grouped_gemm_available cached flag at import time. The dual-backend GroupedGemmPrimitive cleanly dispatches to either te_grouped_gemm_ffi or te_grouped_gemm_v2_ffi depending on dtype/scaling-mode.
transformer_engine/jax/cpp_extensions/quantization.py	Fixes uninitialized scale buffer (jnp.empty → jnp.ones) and improves assertion message formatting. Both changes are correct and safe.
tests/cpp/operator/test_grouped_gemm.cu	Fixes M/N dimension swap in test shape construction for transposed GEMM cases; shapes now match the expected A[N,K]/B[K,M] layout that corresponds to the fix in #2659.

Sequence Diagram

sequenceDiagram
    participant PY as grouped_gemm() [Python]
    participant PRIM as GroupedGemmPrimitive
    participant V2FFI as GroupedGemmV2FFI [C++]
    participant V1FFI as GroupedGemmFFI [C++]
    participant NVTE as nvte_grouped_gemm [cuBLAS 13.2+]
    participant MULTI as nvte_multi_tensor_gemm [legacy]

    PY->>PY: _can_use_v2_grouped_gemm(scaling_mode, dtype, has_bias)
    alt BF16 + NO_SCALING + no bias + cuBLAS≥13.2
        PY->>PRIM: bind(use_v2_ffi=True, alpha, beta)
        PRIM->>V2FFI: te_grouped_gemm_v2_ffi (FFI_CudaGraph_Traits)
        V2FFI->>V2FFI: nvte_convert_int32_to_int64(group_sizes)
        V2FFI->>V2FFI: build JAXX_GroupedTensorWrapper objects
        V2FFI->>NVTE: nvte_grouped_gemm(rhs, lhs, out, α, β, ws_setup, ws_cublas)
        NVTE-->>V2FFI: result (CUDA-graph safe)
        V2FFI-->>PY: output
    else FP8 / MXFP8 / has_bias / old cuBLAS
        PY->>PRIM: bind(use_v2_ffi=False, group_offset)
        PRIM->>V1FFI: te_grouped_gemm_ffi (multi-stream)
        V1FFI->>V1FFI: D2H memcpy + stream sync for group_sizes
        V1FFI->>MULTI: nvte_multi_tensor_gemm(per-stream)
        MULTI-->>V1FFI: result
        V1FFI-->>PY: output
    end

Loading

_{Last reviewed commit: 505feac}

[greptile-apps]

_{8 files reviewed, 8 comments}

_{Edit Code Review Agent Settings | Greptile}

[jberchtold-nvidia]

/te-ci L0 jax

[jberchtold-nvidia]

/te-ci L0 jax

[jberchtold-nvidia]

/te-ci

[greptile-apps]

_{9 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

O(n²) complexity in parallel kernel. Each of the n threads calls this function with a different idx, and for case 2 (per-tensor dims without explicit offsets), thread idx performs a sequential loop from 0 to idx-1. This creates O(1 + 2 + ... + n) = O(n²) total work across all threads.

For large numbers of groups, consider either:

1. Computing offsets once on CPU and passing them explicitly
2. Using a parallel prefix sum (scan) to compute cumulative offsets
3. Documenting this limitation if group counts are expected to be small

[jberchtold-nvidia]

/te-ci

[greptile-apps]

make_ragged_dot_cls returns an instance, not a class, contradicting its name and docstring.

The function is documented as "Creates a ragged dot (grouped GEMM) class" and follows the make_*_cls naming convention, but the trailing () instantiates the result. This is inconsistent with make_dot_general_cls (line 1443), which returns the class itself:

# make_dot_general_cls returns a class:
DotGeneral = make_dot_general_cls(recipe)
result = DotGeneral()(x, kernel)          # instantiate → call

# make_ragged_dot_cls (as written) returns an instance:
ragged_dot = make_ragged_dot_cls(None)    # already an instance!
result = ragged_dot(x, kernel, group_sizes)  # just call

This inconsistency will confuse users who expect API parity with make_dot_general_cls. Either:

1. Remove the trailing () so the function returns the class (consistent with the name/docs), or
2. Rename to make_ragged_dot to accurately reflect that it returns a ready-to-use module instance.

Suggested change

	return wrap_function_in_te_state_module(
	te_grouped_dot_general, quantization_recipe, "ragged_dot"
	)()
	return wrap_function_in_te_state_module(
	te_grouped_dot_general, quantization_recipe, "ragged_dot"
	)

[phu0ngng]

Should we call it grouped_dense here to be consistent?

[greptile-apps]

Good catch! There's an inconsistency in terminology. The function calls grouped_dense internally (and the underlying implementation in dense.py is named grouped_dense), but the wrapper is called make_ragged_dot_cls.

Two options for consistency:

1. Rename to make_grouped_dense_cls - matches the underlying implementation and would align with the grouped_dense function already exported from the JAX API
2. Keep make_ragged_dot_cls - if "ragged dot" is intended as the user-facing terminology for this feature (though it's currently not used elsewhere in the codebase)

Which terminology do you want to standardize on - "ragged_dot" or "grouped_dense"?

[jberchtold-nvidia]

Updated to rename to grouped_dense

[greptile-apps]

Fragile capability detection via exception message matching.

Using try/except with a substring match ("cublas" in str(e).lower()) for runtime capability detection is fragile:

1. Substring fragility: If the error message ever changes (e.g., capitalization, wording), the if branch fails silently and the detection breaks completely, changing the fallback behavior.
2. Called every JAX trace: This check runs on every invocation of _can_use_cuda_graphable_grouped_gemm during abstract evaluation, adding unnecessary overhead for what is a build-time constant.

A more robust approach would be to cache the result at module import time:

# Computed once at import; `None` signals "not yet tested".
_CUDA_GRAPHABLE_GROUPED_GEMM_AVAILABLE: Optional[bool] = None

def _check_cuda_graphable_grouped_gemm_available() -> bool:
    global _CUDA_GRAPHABLE_GROUPED_GEMM_AVAILABLE
    if _CUDA_GRAPHABLE_GROUPED_GEMM_AVAILABLE is None:
        try:
            get_grouped_gemm_setup_workspace_size(1)
            _CUDA_GRAPHABLE_GROUPED_GEMM_AVAILABLE = True
        except RuntimeError as e:
            if "cublas" in str(e).lower():
                _CUDA_GRAPHABLE_GROUPED_GEMM_AVAILABLE = False
            else:
                raise
    return _CUDA_GRAPHABLE_GROUPED_GEMM_AVAILABLE

[phu0ngng]

I'm not sure about this helper.

Why don't we simply check for NO_SCALING, then go with the v2?

Also, we should be able to check it inside primitive methods instead of passing it as an argument so that when we eventually remove the multi-stream approach, we don't need to change the primitive inputs.

[greptile-apps]

Good points on both fronts.

1. Simplifying the check:
You're right - if we're only supporting BF16 + NO_SCALING in v2 initially, the try/except adds complexity for runtime capability detection that may not be needed. The cublas version check could be done once at module import or handled by the FFI call itself (which already has #if CUBLAS_VERSION >= 130200 guards).

2. Checking inside primitives:
This is a better design. Instead of passing use_cuda_graphable_ffi through all the primitive signatures, you could:

• Check scaling_mode, dtype, has_bias directly inside lowering() to decide which FFI to call
• Or make it purely a dispatch decision in the C++ layer based on the inputs

This way when you eventually remove the multi-stream path, you just delete the old FFI handler and the dispatch logic - no need to change primitive inputs/outputs or all the call sites.

The runtime check (if still needed) could be a module-level _HAS_CUDA_GRAPHABLE_GEMM cached at import time rather than being checked on every trace.

[jberchtold-nvidia]

Why don't we simply check for NO_SCALING, then go with the v2?

In addition to NO_SCALING, we need to check the cuBLAS version. In fact, more than that, we need to check the cuBLAS version TE was built with, which may not be the same depending on how the user installed TE. Otherwise, we will dispatch to v2 but it will fail during lowering because the C++ functions raise an error in TE core based on the cuBLAS version it was compiled with. In this case, instead of erroring, I'd rather keep the old behavior and fall back to the slower implementation.

Also, we should be able to check it inside primitive methods instead of passing it as an argument so that when we eventually remove the multi-stream approach, we don't need to change the primitive inputs.

I was planning to do this originally but ran into two issues:

1. I needed to check which backend we're using (original or V2) in multiple places in the primitive. The abstract requires it and the lowering requires it. In the lowering, we don't have easy access to dtype as it's not a jnp.dtype it's some mlir dtype that isn't as easy to convert.
2. The MLIR ctx bakes in information about the inputs. You can't simply slice args in lowering and pick and choose which primitive to dispatch to. If you do, it'll give an MLIR error that the ctx and args don't match. As a result, I'd have to update the original FFI to accept alpha and beta which would break FFI compatibility. Which is why I've gone with the solution of passing additional_arg_0 and additional_arg_1 which needs to happen outside the primitive. This isn't ideal, but was required to get around this fixed arg limitation. Additionally, the static arg indices are fixed at primitive registration and currently the original FFI and the V2 FFI take a different number of args, so for the original FFI I have to use an unused addtional_arg_1 in the outer impl.

As a result of both of these, it's simpler to just check once outside the primitive whether we should use the original or V2, then pass that as a bool flag. Given these restrictions of 1. and 2., I'm not sure this consolidated primitive is a cleaner solution than the split primitives (unless I can break FFI compatibility to unify the FFIs to make the original FFI take alpha and beta).

[phu0ngng]

Register nvte_ API with NVTE_API_CALL.

[jberchtold-nvidia]

Done

[phu0ngng]

Remove these three attributes as they are not needed.

[jberchtold-nvidia]

Removed has_bias and use_async_d2h_group_sizes. The is_grouped_dense_wgrad is still required currently but I agree in future we should refactor this logic

[phu0ngng]

Should we simply call it v2 for now and eventually rename it back to GroupedGEMM?

Besides, if it is possible, could we mark it as an experimental API and subject to changes?

[jberchtold-nvidia]

Sounds good, renamed to V2

[phu0ngng]

I'm not sure about this helper.

Why don't we simply check for NO_SCALING, then go with the v2?

Also, we should be able to check it inside primitive methods instead of passing it as an argument so that when we eventually remove the multi-stream approach, we don't need to change the primitive inputs.

[phu0ngng]

Should we call it grouped_dense here to be consistent?

[phu0ngng]

Nit: could we add get i.e., nvte_get_grouped_gemm_setup_workspace_size? :D

[jberchtold-nvidia]

Updated

[greptile-apps]

Additional Comments (1)

transformer_engine/common/gemm/cublaslt_grouped_gemm.cu
Inconsistent CUDA version in upgrade hint

The error message says "Please upgrade to CUDA 13.1 or newer", but the corresponding stub for nvte_grouped_gemm_setup_workspace_size (line 673 in the same #else block) says "Please upgrade to CUDA 13.2 or newer". Both stubs guard the same CUBLAS_VERSION >= 130200 requirement, so the version suggested to users is contradictory depending on which function they call first. Consider aligning both messages to the same version string.

  NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.2+, but compile-time cuBLAS version is ",
             CUBLAS_VERSION, ". Please upgrade to CUDA 13.2 or newer.");

[greptile-apps]

Additional Comments (2)

transformer_engine/common/gemm/cublaslt_grouped_gemm.cu
O(n²) work in device kernel for case 2

The compute_grouped_tensor_offset helper is called from setup_grouped_gemm_kernel which launches one thread per group. For case 2 (per-tensor dims, no explicit offset array), thread idx serially loops from 0 to idx-1 to accumulate the cumulative sum. The total work across all n threads is 0 + 1 + 2 + … + (n-1) = O(n²).

For a typical MoE with 64 or 128 experts this is tolerable, but for larger group counts the kernel becomes the bottleneck. Note that case 2 is reachable from GroupedGemmV2FFI whenever first_dims is set without offsets, which is exactly the is_grouped_dense_wgrad path (set_group_sizes_only populates first_dims but not offsets).

Consider either:

1. Computing offsets once with a parallel prefix-sum (scan) kernel before launching setup_grouped_gemm_kernel, or
2. Documenting an explicit upper bound on supported group counts for this path.

---

transformer_engine/jax/csrc/extensions/gemm.cpp
Redundant assignment to lhs_is_trans

lhs_is_trans is a bool function parameter. The NVTE_CHECK three lines above already asserts that lhs_is_trans == true, so the explicit re-assignment on this line is dead code that may confuse future readers into thinking the caller can pass lhs_is_trans = false here.

    // lhs_is_trans is already guaranteed true by the NVTE_CHECK above.
    auto lhs_tensor = make_grouped_tensor(lhs_data, lhs_sinv, scaling_mode, num_gemms, lhsShape);

[jberchtold-nvidia]

/te-ci

[greptile-apps]

Inconsistent internal module name "ragged_dot" in make_grouped_dense_cls

The wrap_function_in_te_state_module call passes "ragged_dot" as the Flax module name, but the public-facing function was renamed to make_grouped_dense_cls. This string becomes the Flax variable-collection name (used in init / apply param dicts and serialised checkpoints). Any model that saves weights with this wrapper will have "ragged_dot" in its checkpoint paths, which is both confusing and will silently break loading if the name is ever corrected in a follow-up PR.

Suggested change

	return wrap_function_in_te_state_module(
	te_grouped_dot_general, quantization_recipe, "ragged_dot"
	)()
	return wrap_function_in_te_state_module(
	te_grouped_dot_general, quantization_recipe, "grouped_dense"
	)()

[greptile-apps]

del kwargs silently discards user-provided arguments

Any keyword argument passed to the returned module's __call__ (e.g. precision, preferred_element_type, or future contracting_dims) is silently swallowed rather than forwarded to grouped_dense or rejected with an informative error. Users who accidentally pass such arguments will get incorrect behaviour with no warning. Consider at minimum asserting that kwargs is empty:

def te_grouped_dot_general(generate_quantizer_set, x, kernel, group_sizes, **kwargs):
    assert not kwargs, f"Unexpected keyword arguments: {list(kwargs.keys())}"

[greptile-apps]

Multiple dead variables and unreachable block in GroupedGemmV2FFI

Several variables computed in this function are never consumed by the final nvte_grouped_gemm call (the V2 API takes grouped-tensor wrappers, not raw pointers or per-stream bookkeeping), and will generate -Wunused-variable compiler warnings:

• lhs_ptr, rhs_ptr (lines 585–586) — raw pointers, never passed anywhere
• lhs_sinv_ptr, rhs_sinv_ptr (lines 587–588) — same
• out_ptr (line 618) — same
• num_math_sm, grad, accumulate, use_split_accumulator, bias_shape (lines 680–684) — multi-stream bookkeeping that is unused in the V2 path
• out_dtype_bytes (line 652) — never referenced

Additionally, the FP8 architecture check (lines 687–691) is unreachable: the NVTE_CHECK(scaling_mode == JAXX_Scaling_Mode::NO_SCALING, ...) on line 606 guarantees is_fp8_gemm is always false in this function. The whole block can be removed or marked with a // TODO: wire up when FP8 is added comment alongside the other FP8 placeholders.

These should be cleaned up to avoid misleading future contributors and to keep compiler warnings clean.

[greptile-apps]

Redundant lhs_is_trans = true assignment

The NVTE_CHECK(lhs_is_trans && !rhs_is_trans, ...) guard on line 707 already guarantees that lhs_is_trans is true when execution reaches line 717. The reassignment is dead code and may confuse readers into thinking the shape construction depends on overriding a possibly-false value.

Suggested change

	//// LHS
	NVTEShape lhsShape{.data = {k, m}, .ndim = 2};
	lhs_is_trans = true;
	//// LHS
	NVTEShape lhsShape{.data = {k, m}, .ndim = 2};
	// lhs_is_trans is already verified true by the NVTE_CHECK above.