rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00

History

Daniel Suo b3d851d722 Add Jax tracing micro benchmarks.

Add a first benchmark for tracing/lowering pallas splash attention.

Sample results below taken on a GCP n2d-standard-128 instance with 512GB Ram and 128 vCPU AMD EPYC Milan.

---------------------------------------------------------------------------------
Benchmark                                       Time             CPU   Iterations
---------------------------------------------------------------------------------
test_pallas_mqa_splash_attention_trace       39.8 ms         39.8 ms           19
test_pallas_mqa_splash_attention_lower       42.1 ms         41.9 ms           18

PiperOrigin-RevId: 742259409

2025-03-31 07:35:22 -07:00

mosaic

Replace disable_backends with enable_backends on jax_multiplatform_test.

2024-09-27 06:15:31 -07:00

api_benchmark.py

[JAX] Clean up make_array_from_callback_* API benchmarks and add a partially replicated sharding variant

2025-03-13 15:50:46 -07:00

linalg_benchmark.py

Use a Jacobi SVD solver for unbatched SVDs up to 1024x1024 on NVIDIA GPUs.

2023-11-13 12:04:13 -08:00

math_benchmark.py

Update flags to use the ABSL typed flag API.