Releases: nshepperd/flash_attn_jax
Releases · nshepperd/flash_attn_jax
v0.2.2
v0.2.1
Bump minor version to 0.2.1.
v0.2.0
Expanded vmap support for flash_mha. Vmapping q but not k,v reduces t…
v0.1.0a3
Try cibuildwheel.
v0.1.0a2
Try cibuildwheel.
v0.1.0a1
Set CUDA_HOME for the sdist
v0.1.0
Try release with github actions and rebased repo.
v2.5.5a2
Implement ring attention backward pass. More tests.
v2.5.5a1
Merge up to v2.5.5
v2.5.0a4
Implement custom sharding for flash_mha, to allow efficient multi-gpu…