attnax 0.2.0

Composable attention and transformer components for JAX.

Attnax

Attention kernels and transformer components for JAX.

Installation | Quick start | Documentation | Examples | Citation

Overview

Attnax is built on JAX and Flax and provides:

• Attention kernels as pure JAX functions sharing a single AttentionFn protocol: standard_attention, memory_efficient_attention, flash_attention, linear_attention, ring_attention, pallas_flash_attention, paged_attention, lite_attention.
• ScoreMod / MaskMod constructors for ALiBi, sliding window, prefix-LM, document masks, and arbitrary additive biases — composed with compose_score_mods.
• MultiHeadAttention with MHA, GQA, MQA, RoPE, sliding window, and optional contiguous or paged KV caching (KVLayerCache, PagedKVCache).
• EncoderBlock, DecoderBlock, TransformerEncoder, TransformerDecoder, VisionTransformer, FeedForward, MixtureOfExperts, RMSNorm, RoPE and the usual positional embeddings.

Documentation on Attnax can be found at attnax.readthedocs.io.

Installation

pip install attnax

From source:

git clone https://github.com/glibtkachenko/attnax.git
cd attnax
pip install -e .

Requires Python 3.10+, JAX 0.10.0+, and Flax 0.12.7+.

Quick start

Attention kernels are pure JAX functions:

import jax, jax.numpy as jnp
from attnax import standard_attention

q = jax.random.normal(jax.random.key(0), (1, 4, 64, 32))
k = jax.random.normal(jax.random.key(1), (1, 4, 64, 32))
v = jax.random.normal(jax.random.key(2), (1, 4, 64, 32))
out = standard_attention(q, k, v)

Biases compose as ScoreMods:

from attnax import alibi_mod, compose_score_mods, sliding_window_mod

mod = compose_score_mods(
    alibi_mod(num_heads=4),
    sliding_window_mod(window_size=128, causal=True),
)
out = standard_attention(q, k, v, score_mod=mod)

Any kernel matching AttentionFn plugs into MultiHeadAttention:

import flax.nnx as nnx
from attnax import MultiHeadAttention, pallas_flash_attention

attn = MultiHeadAttention(
    nnx.Rngs(0),
    num_heads=8,
    in_features=512,
    attention_fn=pallas_flash_attention,
)

A full transformer stack:

from attnax import TransformerConfig, TransformerEncoder

config = TransformerConfig(
    vocab_size=32000, d_model=512, num_heads=8, num_layers=6,
)
model = TransformerEncoder(nnx.Rngs(0), config)
y = model(jnp.ones((2, 16), dtype=jnp.int32), deterministic=True)

See the getting-started notebook for a walkthrough covering score-mods, custom kernels, KV caching, paged caching, Mixture-of-Experts, the Vision Transformer, and training.

Citing Attnax

@software{attnax2025github,
  author = {Glib Tkachenko},
  title = {{Attnax}: Attention Kernels and Transformer Components for {JAX}},
  url = {https://github.com/glibtkachenko/attnax},
  version = {0.2.0},
  year = {2025},
}