Flash Attention - Jax
Implementation of Flash Attention in Jax. It will likely not be as performant as with the official CUDA version, given lack of ability for fine memory management. But just for educational purposes as well as to see how clever XLA compiler is (or is not).
Install
$ pip install flash-attention-jax
Usage
from jax import random
from flash_attention_jax import flash_attention
rng_key = random.PRNGKey(42)
q = random.normal(rng_key, (131072, 512))
k = random.normal(rng_key, (131072, 512))
v = random.normal(rng_key, (131072, 512))
mask = random.randint(rng_key, (131072,), 0, 2)
out, _ = flash_attention(q, k, v, mask)
out.shape # (131072, 512)
Quick sanity check
from flash_attention_jax import plain_attention, flash_attention, value_and_grad_difference
diff, (dq_diff, dk_diff, dv_diff) = value_and_grad_difference(
plain_attention,
flash_attention,
seed = 42
)
print('shows differences between normal and flash attention for output, dq, dk, dv')
print(f'o: {diff}') # < 1e-4
print(f'dq: {dq_diff}') # < 1e-6
print(f'dk: {dk_diff}') # < 1e-6
print(f'dv: {dv_diff}') # < 1e-6
Autoregressive Flash Attention - GPT-like decoder attention
from jax import random
from flash_attention_jax import causal_flash_attention
rng_key = random.PRNGKey(42)
q = random.normal(rng_key, (131072, 512))
k = random.normal(rng_key, (131072, 512))
v = random.normal(rng_key, (131072, 512))
out, _ = causal_flash_attention(q, k, v)
out.shape # (131072, 512)
Using in Pytorch
You'll have to have both Jax and Pytorch installed
First install jax2torch
$ pip install jax2torch
Then
import torch
from jax2torch import jax2torch
from flash_attention_jax import flash_attention
q = torch.randn(131072, 512).cuda()
k = torch.randn(131072, 512).cuda()
v = torch.randn(131072, 512).cuda()
mask = torch.ones((131072,)).cuda() # dlpack does not support boolean types, use 1s and 0s
torch_flash_attention = jax2torch(flash_attention)
out = torch_flash_attention(q, k, v, mask)
out.shape # (131072, 512)
Todo
- figure out issue with jit and static argnums
- comment with references to paper algorithms and explanations
- make sure it can work one-headed key / values, as in PaLM
Citations
@article{Dao2022FlashAttentionFA,
title = {FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness},
author = {Tri Dao and Daniel Y. Fu and Stefano Ermon and Atri Rudra and Christopher R'e},
journal = {ArXiv},
year = {2022},
volume = {abs/2205.14135}
}
@article{Rabe2021SelfattentionDN,
title = {Self-attention Does Not Need O(n2) Memory},
author = {Markus N. Rabe and Charles Staats},
journal = {ArXiv},
year = {2021},
volume = {abs/2112.05682}
}