jax
593ad3e7 - [Mosaic GPU] Add support for 1CTA MMA with M=64

Commit

199 days ago

[Mosaic GPU] Add support for 1CTA MMA with M=64 The TMEM layout is a bit weird (each MMA instruction only accesses half the lanes), which is why the patch is not entirely trivial. There's also another layout we could use, where we don't pack the column tiles and split a 128xN TMEM reference into two 64xN refs, each with half the lanes. It might even be slightly better for performance, as it will allow us to avoid the awkward N split that we have to do right now. PiperOrigin-RevId: 775150187

References

#29602 - [Mosaic GPU] Add support for 1CTA MMA with M=64

#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