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

[JAX] Async issuing D2H memcpy for grouped_gemm group_sizes array (#2213) · af2a0c16

Hua Huang authored Oct 08, 2025



* Try async copy of grouped GEMM group_sizes data
Signed-off-by: Hua Huang <huah@nvidia.com>

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

af2a0c16

test_custom_call_compute.py 59.4 KB

Replace test_custom_call_compute.py