tests/jax/test_custom_call_compute.py · b10f436aa28ec8d885eb0d9bf134c320ffb6353d · OpenDAS / TransformerEngine

[JAX] Remove cudaGraph compatible trait from GroupedGemmFFI and GroupedQuantizeFFI (#2048) · 9f9b4816

Phuong Nguyen authored Aug 08, 2025



* rm cudaGraph compatible trait from GroupedGEMM and groupedQuantize
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* add grouped_gemm jitting in the unit test
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

9f9b4816

test_custom_call_compute.py 56.5 KB

Replace test_custom_call_compute.py