Attention performance¶

probjax's attention kernels (mha, flex_attention and the default dot_product_attention, which every probjax transformer uses as attention_fn) compared with jax.nn.dot_product_attention (XLA and cuDNN backends) and PyTorch scaled_dot_product_attention pinned to its flash, memory-efficient and math backends. GPU: NVIDIA GeForce RTX 3090 (24 GB), JAX 0.11.2, torch 2.14.1+cu130. Setup and methodology are on the overview.

Inputs are [batch, seq, heads, head_dim] with batch 4 and 8 heads unless stated otherwise; torch receives the transposed layout it expects. fwd is one forward call; fwd+bwd is the forward plus the gradient with respect to queries, keys and values. Times are medians in milliseconds. OOM: out of device memory. fail†: XLA found no kernel configuration that fits during autotuning (NOT_FOUND), so the case did not run. unsup.: the implementation rejects the shape or backend. error: another failure (the JSON row has the message). —: not measured. cuDNN and torch flash have no fp32 path.

When to use which¶

• On an H100 (sm_90) the defaults are device-specific: large bf16/fp16 calls with no mask, a KVLenMask or equal-length causal masking run cuDNN (1.0–3.6× faster fwd+bwd there) or, where cuDNN cannot take them, the Mosaic GPU kernel; and backward_pass_impl="auto" resolves to the split backward. Rules and measurements: H100 page. The rest of this page is the RTX 3090.

• Default (dot_product_attention, implementation=None). On NVIDIA GPUs it runs the fused flex_attention kernel whenever that is exact (rules below); flex_attention and mha are the same kernel. Its memory grows linearly in T: fwd+bwd ran up to T = 65,536 (the largest length tried) at batch 4 and 8 heads, dense attention only up to T = 4,096.

• jax.nn.dot_product_attention(implementation='cudnn') is faster in the bf16 configurations listed under where an alternative is faster; use it there if you need neither probjax masks nor fp32.
• Dense attention (jax.nn XLA, torch math, probjax implementation='xla') materializes the T × T scores. In this grid jax.nn XLA fails from T = 8192 and torch math from T = 4096 (fwd+bwd); it is only competitive for short sequences.
• Backward pass. backward_pass_impl='auto' picks the kernel. On GPUs with about 100 KB of shared memory per block (RTX 30/40 class) it uses the single-pass backward ('triton_onepass': dK/dV per KV block, dQ accumulated with atomic adds, as in FlashAttention-2) for head_dim 32 and 64 and for bf16 head_dim 128 at the lengths where it is faster, otherwise the two-pass kernels ('triton_fused', 'triton_split'). The single-pass dQ is summed in a nondeterministic order, so gradients are not bitwise reproducible across runs; select a two-pass kernel if you need that. It also keeps a float32 dQ buffer, which adds about a quarter to the backward's temporaries for 16-bit inputs.
• Decoding with a KV cache (few queries, long cache): the default selects the split-KV kernel automatically; see decoding.
• CPU: probjax runs an XLA formulation whose memory is linear in T; torch's native CPU kernel is faster (CPU).

Time¶

bf16, batch 4, 8 heads, RTX 3090. Top row forward, bottom row forward + backward; T = 600, 1500 and 3000 test lengths that are not a multiple of the tile size.

fp32, batch 4, 8 heads, RTX 3090. fp32 matmuls run in TF32 (JAX default precision).

fwd+bwd time of each implementation divided by the probjax default at the same shape; below the line is faster than probjax.

Where an alternative is faster¶

Configurations (T ≥ 512) in which another implementation is more than 5% faster than the probjax default, computed from gpu_attention.json:

• bf16, head_dim 64, non-causal, forward + backward: jax.nn cuDNN at T = 4096, 8192; torch SDPA flash at T = 2048, 4096, 8192. Largest gap: torch SDPA flash is 1.18× faster at T=8192.
• bf16, head_dim 64, causal, forward + backward: torch SDPA flash at T = 4096, 8192. Largest gap: torch SDPA flash is 1.16× faster at T=8192.
• bf16, head_dim 128, non-causal, forward + backward: jax.nn cuDNN at T = 1024, 2048, 4096, 8192; torch SDPA flash at T = 1024, 2048, 4096, 8192. Largest gap: torch SDPA flash is 1.24× faster at T=8192.
• bf16, head_dim 128, causal, forward + backward: jax.nn cuDNN at T = 1024, 2048, 4096, 8192; torch SDPA flash at T = 2048, 4096, 8192. Largest gap: torch SDPA flash is 1.28× faster at T=8192.
• fp32, head_dim 128, non-causal, forward + backward: jax.nn XLA at T = 512, 1024, 2048, 4096. Largest gap: jax.nn XLA is 1.27× faster at T=1024.
• bf16, head_dim 64, non-causal, forward: torch SDPA flash at T = 512, 8192. Largest gap: torch SDPA flash is 1.17× faster at T=512.
• bf16, head_dim 64, causal, forward: torch SDPA flash at T = 512. Largest gap: torch SDPA flash is 1.07× faster at T=512.
• bf16, head_dim 128, non-causal, forward: jax.nn cuDNN at T = 4096, 8192; torch SDPA flash at T = 512, 1024, 2048, 4096, 8192. Largest gap: torch SDPA flash is 1.11× faster at T=512.
• bf16, head_dim 128, causal, forward: jax.nn cuDNN at T = 2048, 4096, 8192; torch SDPA flash at T = 512, 1024, 2048, 4096, 8192. Largest gap: torch SDPA flash is 1.26× faster at T=512.

fwd+bwd at T=4096 in every configuration (ms):

Memory¶

True device peak of one fwd+bwd call, bf16, head_dim 64, batch 4, 8 heads, non-causal, read in a fresh process, RTX 3090.

Largest power-of-two T whose fwd+bwd ran (the search stopped at 65,536, so that value means at least):

Bytes kept for the backward pass (jax.vjp residuals, MB), batch 2, 8 heads, head_dim 64 (kernel_attn_gpu_residuals.json). The kernels keep the output and the log-sum-exp; dense paths keep the T × T probabilities:

Short sequences and the default dispatch¶

probjax.nn.layers.attention.dot_product_attention is the default attention_fn of Transformer, DiffusionTransformer and the attention conditioners. With implementation=None it runs the fused kernel (flex_attention) on NVIDIA GPUs when that is exact:

• no mask, or a position-only mask (causal, length or window masks); no dense boolean mask,
• no additive bias, no dropout, default precision,
• at least 64 keys.

Otherwise, and on CPU, it uses the Flax einsum. implementation='xla' forces the einsum and implementation='pallas' the kernel. GQA inputs (fewer KV heads than query heads) also go to the kernel.

bf16, batch 64, 8 heads, non-causal, RTX 3090 (kernel_attn_gpu_dispatch.json).

In this sweep the mha kernel is more than 5% faster than the einsum in fwd+bwd at every swept length for head_dim 32 and 64. The default switches to the kernel from 64 keys on. cuDNN pays a fixed launch cost at short lengths.

Tile selection¶

With default block sizes, mha picks tiles from a shared-memory model and a tuned table per (dtype, head_dim). Auto choice against the best of every forced (block, stages, warps) combination, batch 4, 8 heads (gpu_attention_tiles.json); the last column counts configurations that ran and those Triton rejected for exceeding shared memory:

Grouped-query attention and decoding¶

The kernels read KV head h // G directly instead of repeating K and V in memory (flex_attention and dot_product_attention route GQA inputs to them).

32 query heads, head_dim 128, batch 1, Lq = Lkv = T, median of three rounds, RTX 3090 (kernel_attn_ext_gpu.json, suite gqa).

Memory of native GQA against running the kernel on repeated K/V (MB):

Split-KV decoding. For few queries against a long KV cache the regular kernel parallelizes only over batch × heads × query tiles. The split-KV forward splits the keys across programs, packs the query heads of a GQA group into one tile and combines the partial results with their log-sum-exp. It is selected automatically for unmasked inputs with Lq ≤ 64, Lkv ≥ 1024 and at least two splits; split_kv=False disables it. Under jax.grad the regular kernel runs.

bf16, 32 query heads, head_dim 128, median of three rounds, RTX 3090. Top row GQA with 8 KV heads, bottom row 32 KV heads (plain multi-head; batch 8 at 64k keys was not run).

Device peak at batch 8, 64k keys, 8 KV heads, bf16 (fresh process):

CPU¶

On CPU, mha, mla and flex_attention run an XLA formulation: dense for T ≤ 256, a single-tile flash custom_vjp up to 512, and 256 × 256 tiles with block-mask tile skipping above that; the backward recomputes from the log-sum-exp. Masks, biases (with gradients), dropout, GQA without repeating K/V and forward-mode AD are supported. The default dot_product_attention uses the Flax einsum on CPU.

fp32, batch 1, 8 heads, head_dim 64, shared 24-thread host (10–30% run-to-run noise) (kernel_attn-cpu_cpu.json).

Left: XLA temporaries of the non-causal fwd+bwd. Right: compile-time memory (temporaries + arguments + outputs) one fwd+bwd needs, against the host's RAM (kernel_attn-cpu_cpu_memory.json).

At T=4096 (non-causal fwd+bwd) the probjax default takes 721 ms, jax.nn XLA 779 ms and torch 169 ms; torch's native CPU kernel is the fastest option on CPU.

CPU decoding, fp32, 32 query heads, 8 KV heads, head_dim 128, batch 1 (kernel_attn_ext_cpu.json), ms:

How to reproduce¶

# Cross-library grid and the tile sweep (CUDA GPU; torch optional).
python -m docs.scripts.benchmark_gpu_attention --output docs/benchmarks/gpu_attention.json
python -m docs.scripts.benchmark_gpu_attention_tiles --output docs/benchmarks/gpu_attention_tiles.json

# Short-sequence sweep, device peak and residuals.
python -m docs.scripts.benchmark_kernel_attn_gpu --device gpu --suite dispatch \
  --output docs/benchmarks/kernel_attn_gpu_dispatch.json
python -m docs.scripts.benchmark_kernel_attn_gpu --device gpu --suite peak \
  --output docs/benchmarks/kernel_attn_gpu_memory.json

# GQA and split-KV decoding.
python -m docs.scripts.benchmark_kernel_attn_ext --device gpu --suite gqa \
  --output docs/benchmarks/kernel_attn_ext_gpu.json
python -m docs.scripts.benchmark_kernel_attn_ext --device gpu --suite decode \
  --output docs/benchmarks/kernel_attn_ext_gpu.json --append

# CPU.
JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_kernel_attn-cpu \
  --device cpu --output docs/benchmarks/kernel_attn-cpu_cpu.json

# Rebuild the figures and this page from the JSON.
python -m docs.scripts.benchmark_figures