jax
bb454d20 - [Mosaic GPU] Fix a race in test_remote_async_copy

Commit

177 days ago

[Mosaic GPU] Fix a race in test_remote_async_copy Previously the kernel in the test would exit without observing the completion of the current GPU's and the peer GPU's writes to each other's memory, which occasionally led to failing correctness checks. We fix the race by adding a system memory barrier to observe the completion of this GPU's write to the peer GPU's memory, and a semaphore that allows to observe the peer's GPU write to this GPU's memory.

References

#30159 - [Mosaic GPU] Fix a race in test_remote_async_copy

#31381 - Remove old ROCm build code

#31720 - Fix ann_test.py numerical bug in target reshape

#31768 - [ROCm] Support lowering through PJRT_Triton_Extension

#32115 - Relax version requirements for ROCm Jax Plugin wheels

#33157 - Resolve undefined behavior in bitshift unit test

#33186 - Make nvidia version data optional for ROCm builds

#579 - Create rocm-test-requirements.txt

#580 - GESVDJ support for ROCm GPUs in JAX

#581 - Fix/pallas tests shared memory

#584 - Use plain bazel to test jax, use hermetic rocm dependency