DOCSEARCHbyAlgolia
Submit Docs

Flash Attention JAX

repository·main·Indexed 42 days ago

https://github.com/lucidrains/flash-attention-jax

A JAX implementation of the Flash Attention algorithm designed for educational purposes and XLA compiler optimization. Includes the flash_attention function for standard attention and causal_flash_attention for autoregressive modeling, along with a value_and_grad_difference utility to verify implementation accuracy against plain attention.

Tokens
876
Snippets
4
Records
4
Agent score
3%

What's inside Flash Attention JAX

  2. Verify attention implementation with value_and_grad_difference

    main

    You can perform a sanity check to ensure the flash_attention implementation matches a plain_attention implementation by comparing their outputs and gradients using value_and_grad_difference.

    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(f'o: {diff}')       # Should be < 1e-4
print(f'dq: {dq_diff}')   # Should be < 1e-6
print(f'dk: {dk_diff}')   # Should be < 1e-6
print(f'dv: {dv_diff}')   # Should be < 1e-6

  3. Use flash_attention for standard attention

    main

    The flash_attention function implements the Flash Attention algorithm in JAX. It accepts query (q), key (k), value (v), and an optional mask.

    Input Shapes:

    • q, k, v: (batch, heads, seq, dim)
    • mask: (batch, seq)

    Returns:

    • out: The attention output with shape (batch, heads, seq, dim)
    • _: A second return value (unused in basic usage).
    from jax import random
from flash_attention_jax import flash_attention

rng_key = random.PRNGKey(42)

q = random.normal(rng_key, (1, 2, 131072, 512))  # (batch, heads, seq, dim)
k = random.normal(rng_key, (1, 2, 131072, 512))
v = random.normal(rng_key, (1, 2, 131072, 512))
mask = random.randint(rng_key, (1, 131072,), 0, 2) # (batch, seq)

out, _ = flash_attention(q, k, v, mask)

print(out.shape)  # (1, 2, 131072, 512)

  4. Use causal_flash_attention for autoregressive attention

    main

    For GPT-like decoder attention where causal masking is required, use causal_flash_attention. This variant is designed for autoregressive modeling.

    Input Shapes:

    • q, k, v: (seq, dim) (Note: This implementation example uses 3D inputs without batch/head dimensions for the causal variant).

    Returns:

    • out: The attention output with shape (seq, dim)
    • _: A second return value (unused).
    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)

print(out.shape)  # (131072, 512)