[JAX][Core] Fix Grouped GEMM cuBLAS version and SM arch checks

[jberchtold-nvidia]

Description

Fixes SM arch check for TE/JAX Grouped GEMM. Also fixes cuBLAS version check in TE common to cuBLAS 13.3+ as cuBLAS 13.2 contains a wgrad numerical error if any groups have size 0.

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

• Fix SM arch check for TE/JAX grouped GEMM
• Fix cuBLAS version check in TE common to be cuBLAS 13.3+ instead of cuBLAS 13.2+

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

[jberchtold-nvidia]

/te-ci

[greptile-apps]

Greptile Summary

This PR fixes two related bugs in Transformer Engine's grouped GEMM support: (1) the JAX Python-side capability check in _can_use_v2_grouped_gemm was missing an SM arch guard, so on Hopper (SM90) and older GPUs the v2 code path could be selected and subsequently fail in the C++ layer; (2) the cuBLAS minimum version requirement was raised from 13.2 to 13.3 everywhere (compile-time guard, runtime check, and all fallback stub messages) because cuBLAS 13.2 contains a wgrad numerical correctness bug when any group has k=0.

Key changes:

• Adds CUBLAS_GROUPED_GEMM_VERSION 130300 macro separate from CUBLAS_MXFP8_GROUPED_GEMM_VERSION, and uses it consistently for the #if compile-time guard and the runtime check_grouped_gemm_requirements call.
• All stub error messages in the #else branch updated from "13.2+" / "CUDA 13.1" to "13.3+" / "CUDA 13.3".
• _can_use_v2_grouped_gemm in the JAX extension now short-circuits to False on get_device_compute_capability(0) < 100, matching the C++ runtime guard.
• One stale comment in gemm.py (line 75) still reads "cuBLAS < 13.2" and was not updated to reflect the new 13.3 requirement.

Confidence Score: 4/5

• This PR is safe to merge; changes are correctness fixes with no functional regressions for supported configurations.
• Both fixes are well-scoped and targeted: the SM arch check in Python directly mirrors the existing C++ runtime guard, and the version bump is clearly justified by the described cuBLAS 13.2 bug. All previously stale error strings have been updated consistently. The only remaining issue is a one-line stale comment and a pre-existing hardcoded device-0 pattern that is unlikely to cause problems in real deployments.
• transformer_engine/jax/cpp_extensions/gemm.py — stale cuBLAS version comment at line 75 and hardcoded device-0 index at line 1941.

Important Files Changed

Filename	Overview
transformer_engine/common/gemm/cublaslt_grouped_gemm.cu	Correctly introduces CUBLAS_GROUPED_GEMM_VERSION (130300) macro separate from CUBLAS_MXFP8_GROUPED_GEMM_VERSION, updates the compile-time guard, runtime version check, and all stub error messages to consistently require cuBLAS 13.3+. All previously flagged stale "13.2" and "CUDA 13.1" strings are fixed.
transformer_engine/jax/cpp_extensions/gemm.py	Adds the missing SM100+ compute capability guard to _can_use_v2_grouped_gemm so Hopper and older architectures fall back to the v1 path. One stale comment at line 75 still references cuBLAS 13.2 instead of 13.3; the device-0 hardcoding is a pre-existing limitation but is worth noting.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[grouped_gemm called] --> B{_can_use_v2_grouped_gemm?}
    B --> C{_v2_grouped_gemm_available\ncompile-time cuBLAS >= 13.3?}
    C -- No --> V1[Use v1 legacy path\nnvte_multi_tensor_gemm]
    C -- Yes --> D{get_device_compute_capability 0\n>= SM100?}
    D -- No\nHopper / older --> V1
    D -- Yes --> E{scaling_mode == NO_SCALING\ndtype == bfloat16\nno bias?}
    E -- No --> V1
    E -- Yes --> F[Use v2 path\nnvte_grouped_gemm]
    F --> G{C++ runtime checks}
    G --> H{sm_arch current_device\n>= 100?}
    H -- No --> ERR1[NVTE_CHECK error:\nrequires SM100+]
    H -- Yes --> I{cublas_version\n>= 130300?}
    I -- No --> ERR2[NVTE_CHECK error:\nrequires cuBLAS 13.3+]
    I -- Yes --> OK[Execute grouped GEMM]

Loading

_{Last reviewed commit: f17e674}

[ksivaman]

LGTM

[ksivaman]

/te-ci

[greptile-apps]

Hardcoded device index 0 may misidentify architecture in multi-GPU systems

get_device_compute_capability(0) always queries device 0, regardless of which GPU JAX is actually dispatching work to. In a heterogeneous multi-GPU system — for example, one where device 0 is SM100+ (Blackwell) but the active computation device is SM90 (Hopper) — the check would incorrectly allow the v2 path, causing the C++ runtime guard in check_grouped_gemm_requirements to raise an error. The reverse (device 0 is SM90, active device is SM100+) would just silently fall back to v1, wasting performance.

While this matches the existing pattern used at line 88 (get_cublas_workspace_size_bytes), that context is less likely to cause a hard failure. For correctness it would be safer to iterate over all JAX devices and take the minimum compute capability, or find the current active device. That said, this is a pre-existing limitation of the helper function itself and extremely unlikely in practice since mixed SM90/SM100 node configurations are rare.

[jberchtold-nvidia]

/te-ci

[jberchtold-nvidia]

/te-ci