tests/jax/test_custom_call_compute.py · 397c4be6629d9e6345355c623e9a75c5c59a5356 · OpenDAS / TransformerEngine

[JAX] Fix grouped GEMM error on CUDA 12.9.1 & later (#1925) · 3c4dfffb

Hua Huang authored Jul 09, 2025



* Fix JAX grouped gemm error on CUDA 12.9.1 & later by using 16B alignment for scale ptr
Signed-off-by: Hua Huang <huah@nvidia.com>

* Pad MXFP8 scales with 2*-127 instead of NaNs
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Hua Huang <huah@nvidia.com>

3c4dfffb

test_custom_call_compute.py 56.1 KB

Replace test_custom_call_compute.py