Add GPU multi-head attention benchmark (NNlib vs LuxLib vs NNkernels)

[CarloLucibello]

Adds a standalone GPU benchmark (benchmark/cuda_attention/) comparing three implementations of the same scaled dot-product (multi-head) attention, across sizes from tiny up to a single Llama3-8B attention layer, timing both the forward and the full Zygote gradient (fwd+bwd).

Library	Function	Approach	Memory
NNlib	dot_product_attention	materialized: full score tensor → softmax → batched_mul	O(L²)
LuxLib	scaled_dot_product_attention	materialized softmax path (own layout)	O(L²)
NNkernels	flash_attention	fused FlashAttention; scores never materialized	O(L)

All three compute identical math (scale 1/√E); the script builds all three from one source tensor since each wants a different q,k,v layout.

Usage

It's CLI-driven, so a single dtype/impl combination can be measured without rerunning the whole suite:

julia --project=. bench_attention.jl                                   # full suite
julia --project=. bench_attention.jl --dtypes f32 --impls flash        # one slice
julia --project=. bench_attention.jl --dtypes f16,bf16 --impls nnlib,lux --sizes tiny,small
julia --project=. bench_attention.jl --help        # --dtypes/--impls/--causal/--sizes/--seconds, --list

Timing is minimum of CUDA.@sync(op) over BenchmarkTools.run(...; seconds, evals=1) — evals=1 is per-sample (many samples within the budget), one clean synced GPU execution each, matching cuda_softmax/bench_softmax.jl.

---

Test machine

• GPU: NVIDIA GeForce RTX 5090, 32 GiB (Blackwell, sm_120), driver 595.71.05 / CUDA driver 13.2
• CPU: AMD Ryzen Threadripper PRO 9955WX (16 cores / 32 threads)
• RAM: 122 GiB
• OS: Ubuntu 26.04 LTS (x86_64)
• Software: Julia, CUDA.jl 6.2 / cuDNN 6.2, NNlib 0.9.36, LuxLib 1.15.9, NNkernels 0.2.0 (CUDA 5.11 numbers in results_rtx5090_cuda511.txt)

Results — RTX 5090 (Blackwell, sm_120), CUDA.jl 6.2 / cuDNN 6.2

n/s = path can't run that config; OOM = out of GPU memory; <impl>/fl = time relative to flash (>1 → flash faster). Times in ms, lower is better. Configs: last three are a Llama3-8B layer (E=128, H=32) at growing context.

Float16

Forward

config	L	nnlib	lux	flash	nnlib/fl	lux/fl
tiny	128	0.791	0.856	0.053	15.0×	16.2×
small	512	1.526	1.531	0.402	3.8×	3.8×
gpt2-ish	1024	1.157	1.155	0.922	1.25×	1.25×
llama3 L=2k	2048	1.511	1.518	8.110	0.19×	0.19×
llama3 L=4k	4096	4.138	4.160	29.82	0.14×	0.14×
llama3 L=8k	8192	n/s	n/s	114.0	—	—

Full gradient (fwd+bwd)

config	L	nnlib	lux	flash	nnlib/fl	lux/fl
tiny	128	2.490	3.155	0.290	8.6×	10.9×
small	512	4.726	5.457	3.227	1.5×	1.7×
gpt2-ish	1024	3.765	4.880	11.67	0.32×	0.42×
llama3 L=2k	2048	4.282	4.460	167.4	0.03×	0.03×
llama3 L=4k	4096	11.27	11.42	674.7	0.02×	0.02×
llama3 L=8k	8192	n/s	n/s	2723	—	—

BFloat16 — runs on the materialized paths only

NNlib/LuxLib work under CUDA 6.2 (they were n/s on 5.11 — cuDNN had no bf16 softmax). Flash is n/s: NNkernels' _flash_attention dispatches only on T<:Union{Float16,Float32} at the source level, independent of CUDA version.

config	L	nnlib (fwd)	lux (fwd)	nnlib (grad)	lux (grad)	flash
tiny	128	0.786	0.791	2.796	3.172	n/s
small	512	1.717	1.561	5.335	5.508	n/s
gpt2-ish	1024	1.264	1.270	4.068	4.722	n/s
llama3 L=2k	2048	1.686	1.688	4.688	4.845	n/s
llama3 L=4k	4096	4.494	4.510	12.13	12.23	n/s
llama3 L=8k	8192	n/s	n/s	n/s	n/s	n/s

Float32

Forward

config	L	nnlib	lux	flash	nnlib/fl	lux/fl
tiny	128	0.050	0.051	0.068	0.73×	0.75×
small	512	0.243	0.242	0.806	0.30×	0.30×
gpt2-ish	1024	0.743	0.746	2.238	0.33×	0.33×
llama3 L=2k	2048	2.338	2.353	18.33	0.13×	0.13×
llama3 L=4k	4096	8.759	8.787	69.50	0.13×	0.13×
llama3 L=8k	8192	n/s	n/s	270.5	—	—

Full gradient (fwd+bwd)

config	L	nnlib	lux	flash	nnlib/fl	lux/fl
tiny	128	0.281	0.895	0.621	0.45×	1.44×
small	512	0.766	1.027	8.702	0.09×	0.12×
gpt2-ish	1024	2.014	2.098	35.54	0.06×	0.06×
llama3 L=2k	2048	6.522	6.670	467.6	0.01×	0.01×
llama3 L=4k	4096	24.61	24.52	1891	0.01×	0.01×
llama3 L=8k	8192	OOM	OOM	7552	—	—

(Causal variants are in the README / raw output; causal roughly halves flash's work and slightly slows the materialized path.)

Summary

• NNlib ≈ LuxLib everywhere — same materialized algorithm, within noise. LuxLib carries a bit more AD overhead at tiny sizes.
• Flash wins only at small Float16 sizes (3–17× faster at tiny/small), then loses from gpt2-ish up: at Llama sizes it's ~5–6× slower on the forward and ~50–100× slower on the gradient. For Float32 it's slower at every size. Two causes, one confirmed in source: NNkernels only takes the WMMA tensor-core path when sizeof(T)==2 (so Float32 always uses the scalar mma!), and even Float16 (which does use WMMA) trails cuDNN here — its sm_70/80-era wmma intrinsics aren't tuned for Blackwell sm_120. On the Ampere/Ada cards NNkernels' CI targets this likely looks different.
• Flash's real advantage is memory: it's the only path that completes L=8192 — the materialized (8192,8192,32) score tensor has 2³¹ elements, overflowing cuDNN's int32 softmax descriptor (n/s), and OOMs in Float32-causal.

CUDA.jl 6.2 vs 5.11

• BFloat16 became usable on the materialized paths (headline of the upgrade).
• FP32 materialized path sped up ~20–25% at large sizes (e.g. llama3 L=4k FP32 gradient 31.9 → 24.6 ms).
• ⚠️ NNkernels 0.2.0 is unusable under CUDA.jl 6.x out of the box: it pins CUDA = "5", and even with that relaxed its NNkernelsCUDAExt hits UndefVarError at runtime because CUDA.jl 6.0 renamed CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK (dropped the CU_ prefix). The flash columns above required a locally-dev-ed NNkernels with a compat bump + that one-line API fix (documented in the README). Worth an upstream NNkernels issue/PR.

Both raw logs (results_rtx5090_cuda511.txt, results_rtx5090_cuda62.txt) and full per-causal tables are included in the folder.

🤖 Generated with Claude Code