transformer_engine/jax/csrc/extensions/gemm.cpp · 9f9b48168f106a172b28aeb44cb12f2c2c232181 · 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

gemm.cpp 30.2 KB

Replace gemm.cpp