[JAX] Hopper BF16 grouped GEMM v2 support - #3083
Conversation
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
|
/te-ci L1 jax |
Greptile SummaryThis PR enables the JAX V2 grouped-GEMM path for non-quantized BF16 workloads on Hopper and adapts alpha/beta metadata to the architecture-specific native interface.
Confidence Score: 4/5The PR is not yet safe to merge because mixed Hopper/Blackwell processes can fail grouped-GEMM native validation on Blackwell devices. Alpha/beta cardinality is selected using the minimum capability across the process, while the native implementation validates it using the GPU executing each call; an SM90/SM100 process therefore supplies singleton buffers to an SM100 kernel that requires per-group values. The earlier bias-plus-ragged-last-dim finding also remains outstanding because the current routing still does not reject that combination before V2 execution. Files Needing Attention: transformer_engine/jax/cpp_extensions/gemm.py Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[JAX grouped_gemm] --> B{Minimum visible SM}
B -->|SM90 present| C[Singleton alpha and beta]
B -->|All SM100+| D[Per-group alpha and beta]
C --> E[FFI dispatch on selected GPU]
D --> E
E --> F{Current device SM}
F -->|SM90| G[Native expects singleton]
F -->|SM100| H[Native expects per-group]
C -. heterogeneous SM90/SM100 mismatch .-> H
Reviews (3): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if out_last_dims is not None: | ||
| raise NotImplementedError("V2 grouped GEMM bias is not supported for ragged last dims") |
There was a problem hiding this comment.
When has_bias=True and out_last_dims is not None, _is_v2_grouped_gemm_supported still returns True for BF16 on SM90+ (no gate for this combination), so the full V2 kernel is dispatched before _add_grouped_gemm_bias raises NotImplementedError. The check should be moved upstream — either into _is_v2_grouped_gemm_supported (returning False to fall back to V1) or as an early guard in grouped_gemm before the FFI bind.
There was a problem hiding this comment.
This will be especially bad if NVTE_JAX_ENFORCE_V2_GROUPED_GEMM=1 because v1 fallback is disabled
There was a problem hiding this comment.
Exactly right — that makes the bug significantly worse. With NVTE_JAX_ENFORCE_V2_GROUPED_GEMM=1, _should_enforce_v2_grouped_gemm() returns True, use_v2_ffi is forced True, and the (bias + ragged last dims) combination crashes inside _add_grouped_gemm_bias with NotImplementedError — with no V1 fallback available and no useful error message referencing the env var.
The clean fix is to gate this combination in _is_v2_grouped_gemm_supported. The function already accepts has_bias; it just needs out_last_dims (or a boolean flag) added so it can return early with a clear message:
| if out_last_dims is not None: | |
| raise NotImplementedError("V2 grouped GEMM bias is not supported for ragged last dims") | |
| def _is_v2_grouped_gemm_supported( | |
| scaling_mode: ScalingMode, | |
| dtype: jnp.dtype, | |
| has_bias: bool, | |
| lhs_shape: Tuple[int, ...], | |
| rhs_shape: Tuple[int, ...], | |
| lhs_axis_boundary: int, | |
| rhs_axis_boundary: int, | |
| has_ragged_last_dims: bool = False, | |
| ) -> Tuple[bool, str]: |
Then add a guard near the top of the function (after the SM90 check, before the BF16 fast-path):
if has_bias and has_ragged_last_dims:
return (
False,
"V2 grouped GEMM does not support bias with ragged last dims.",
)And pass has_ragged_last_dims=(out_last_dims is not None) at the call site in grouped_gemm.
This way, when NVTE_JAX_ENFORCE_V2_GROUPED_GEMM=1 is set, the enforce check fires with the correct reason rather than letting the kernel launch and then crashing in Python post-processing.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
|
/te-ci |
| if out_last_dims is not None: | ||
| raise NotImplementedError("V2 grouped GEMM bias is not supported for ragged last dims") |
There was a problem hiding this comment.
This will be especially bad if NVTE_JAX_ENFORCE_V2_GROUPED_GEMM=1 because v1 fallback is disabled
| return (out_2d + bias_rows).reshape(out_shape) | ||
|
|
||
| bias_shape = (num_gemms,) + (1,) * (out.ndim - 2) + (n_dim,) | ||
| return out + bias_2d.reshape(bias_shape) |
There was a problem hiding this comment.
Correct me if I am wrong but I think the test file is currently having bias=None, which will not test this path?
a0b17d8 to
77f7098
Compare
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
e42bc59 to
d264a62
Compare
for more information, see https://pre-commit.ci
|
/te-ci L0 jax |
|
|
||
| @cache | ||
| def _v2_grouped_gemm_supports_per_group_alpha_beta() -> bool: | ||
| """Whether nvte_grouped_gemm accepts per-group alpha/beta on all visible devices.""" |
There was a problem hiding this comment.
If a process exposes both Hopper and Blackwell GPUs, this process-wide minimum-capability check selects singleton alpha/beta buffers for every device. The native grouped-GEMM implementation validates their size using the device that executes the kernel, and Blackwell requires one value per grouped GEMM. A call dispatched to the Blackwell GPU therefore fails native validation. Select the alpha/beta representation per executing device or make the native contract accept the process-wide choice.
Knowledge Base Used: JAX runtime and public API
There was a problem hiding this comment.
Do we even support this use case in TE in general?
There was a problem hiding this comment.
Not explicitly, at least not in the JAX TE path. The code assumes one effective device capability for a compiled computation—feature gates use get_min_device_compute_capability(), and grouped GEMM has no per-executing-device dispatch. Supported deployments therefore appear to assume a homogeneous local mesh; mixed Hopper/Blackwell devices are outside the documented/implemented contract, even though there is no explicit homogeneity check.
I agree this should not be treated as a required fix for this PR. The remaining choice is whether to add an explicit rejection for mixed-capability meshes so they fail clearly, rather than implying that the process-wide minimum-capability decision supports them.
2 participants
Description
Adds support for Hopper BF16 grouped GEMM
Type of change
Changes
Checklist: