diff --git a/benchmark/cuda_attention/Project.toml b/benchmark/cuda_attention/Project.toml new file mode 100644 index 000000000..21935e7c3 --- /dev/null +++ b/benchmark/cuda_attention/Project.toml @@ -0,0 +1,10 @@ +[deps] +ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63" +BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11" +NNkernels = "1409a70c-e4a7-4ad8-be19-f498502441cf" +NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" diff --git a/benchmark/cuda_attention/README.md b/benchmark/cuda_attention/README.md new file mode 100644 index 000000000..dd2230e52 --- /dev/null +++ b/benchmark/cuda_attention/README.md @@ -0,0 +1,233 @@ +# GPU multi-head attention benchmark — NNlib vs LuxLib vs NNkernels + +Compares three implementations of the same scaled dot-product (multi-head) +attention on the GPU, across input sizes from tiny up to a single Llama3-8B +attention layer, timing both the forward pass and the full **Zygote** gradient +(forward + backward). + +| Library | Function | Approach | Memory | +|---|---|---|---| +| [NNlib](https://github.com/FluxML/NNlib.jl) | `dot_product_attention` | materialized: builds the full score tensor, `softmax`, second `batched_mul` | O(L²) | +| [LuxLib](https://github.com/LuxDL/Lux.jl/tree/main/lib/LuxLib) | `scaled_dot_product_attention` | materialized softmax path (own layout, KV-grouping) | O(L²) | +| [NNkernels](https://github.com/FluxML/NNkernels.jl) | `flash_attention` | fused [FlashAttention](https://arxiv.org/abs/2205.14135); scores never materialized | O(L) | + +All three compute the identical math (scale `1/√E`). They only differ in the +memory layout they want for `q,k,v`, so the script builds all three from one +source tensor: + +``` +NNkernels: (head_dim, seq, nheads, batch) = (E, L, H, B) +NNlib: (head_dim*nheads, seq, batch), nheads=H = (E*H, L, B) +LuxLib: (head_dim, nheads, seq, batch) = (E, H, L, B) (head_dim=1, token_dim=3) +``` + +## Run + +Whole suite (all dtypes × both causal modes × all sizes × all impls): + +```bash +julia --project=. bench_attention.jl +``` + +The script takes CLI flags so a **single dtype/impl combination can be measured +without rerunning the whole suite**: + +```bash +# just NNkernels flash, Float32, non-causal: +julia --project=. bench_attention.jl --dtypes f32 --impls flash --causal false +# the two materialized paths in bf16+f16 on the small configs: +julia --project=. bench_attention.jl --dtypes f16,bf16 --impls nnlib,lux --sizes tiny,small +# quick pass (shorter timing budget): +julia --project=. bench_attention.jl --seconds 0.3 +julia --project=. bench_attention.jl --help # all options +julia --project=. bench_attention.jl --list # available configs / dtypes / impls +# pin a GPU: +CUDA_VISIBLE_DEVICES=GPU- julia --project=. bench_attention.jl ... +``` + +| flag | values | default | +|---|---|---| +| `--dtypes` | `f16`,`bf16`,`f32` (aliases: `float16`/`fp16`/`half`, `bfloat16`, `float32`/`fp32`) | all | +| `--impls` | `nnlib`,`lux`,`flash` | all | +| `--causal` | `false`,`true`, or `both` | both | +| `--sizes` | config names (see `--list`) | all | +| `--seconds` | per-measurement time budget | `1.0` | + +### Timing + +Each cell runs `CUDA.@sync(op)` under `BenchmarkTools.run(...; seconds, evals=1)` +and reports `minimum(trial)` in **ms**. `evals=1` is *per sample*, not total: +within the `seconds` budget BenchmarkTools collects as many samples as fit +(thousands for a sub-ms kernel) and we take the fastest — the standard robust +estimator for GPU kernels. `evals=1` (rather than auto-tuned `evals`) keeps one +clean synchronized GPU execution per sample, matching +[cuda_softmax/bench_softmax.jl](../cuda_softmax/bench_softmax.jl). The only +caveat is the multi-second `llama3 L=8k` flash cells, where a single sample +already exceeds the budget → 1 sample (directional, not tight). + +`n/s` = that path can't run the config (reason echoed once to stderr); +`OOM` = ran out of GPU memory; `/fl` = time relative to flash +(>1 → flash faster). + +Configs (head_dim `E`, seq `L`, heads `H`, batch `B`); the last three are a +single Llama3-8B attention layer (`E=128, H=32`) at growing context: + +``` +tiny E=64 L=128 H=4 B=8 +small E=64 L=512 H=8 B=8 +gpt2-ish E=64 L=1024 H=12 B=4 +llama3 L=2k E=128 L=2048 H=32 B=1 +llama3 L=4k E=128 L=4096 H=32 B=1 +llama3 L=8k E=128 L=8192 H=32 B=1 +``` + +## Environment / setup notes + +The benchmark dev-depends on the local NNlib (`../..`). For the **CUDA.jl 6.2** +run, NNkernels 0.2.0 needed two local patches (it is `dev`-ed from +`~/.julia/dev/NNkernels`), because the registered package predates the CUDA 6 API: + +1. **compat** — `CUDA = "5"` → `CUDA = "5, 6.1"` (the upstream + `dependabot/julia/CUDA-5-and-6.1` branch; admits 6.2). +2. **CUDA 6.0 API rename** — its `NNkernelsCUDAExt` used + `CUDA.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK`, which CUDA.jl 6.0 + dropped the `CU_` prefix from. Without this fix `flash_attention` throws + `UndefVarError` at runtime for **every** dtype (the compat bump alone is not + enough). Patched to pick the available name at load time. + +Neither patch changes kernel behavior; flash numbers match the CUDA 5.11 run. + +## Results — RTX 5090 (Blackwell, sm_120), **CUDA.jl 6.2 / cuDNN 6.2**, NNlib 0.9.36, LuxLib 1.15.9, NNkernels 0.2.0 + +Raw output: [`results_rtx5090_cuda62.txt`](results_rtx5090_cuda62.txt) (CUDA 6.2), +[`results_rtx5090_cuda511.txt`](results_rtx5090_cuda511.txt) (earlier CUDA 5.11). +Times in ms, lower is better. + +### Float16 + +**Forward** — flash wins small, materialized wins from `gpt2-ish` up. + +| 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 — **now runs on the materialized paths under CUDA 6.2** + +Under CUDA 5.11 every bf16 cell was `n/s` (cuDNN had no bf16 softmax). CUDA 6.2's +improved BFloat16 support fixes the NNlib/LuxLib path. **Flash is still `n/s`**: +NNkernels' `_flash_attention` dispatches only on `T<:Union{Float16,Float32}` at +the *source* level — independent of the CUDA version. (bf16 numbers track +fp16 within noise.) + +**Forward** + +| config | L | nnlib | lux | flash | +|---|--:|--:|--:|--:| +| tiny | 128 | 0.786 | 0.791 | n/s | +| small | 512 | 1.717 | 1.561 | n/s | +| gpt2-ish | 1024 | 1.264 | 1.270 | n/s | +| llama3 L=2k | 2048 | 1.686 | 1.688 | n/s | +| llama3 L=4k | 4096 | 4.494 | 4.510 | n/s | +| llama3 L=8k | 8192 | n/s | n/s | n/s | + +**Full gradient (fwd+bwd)** + +| config | L | nnlib | lux | flash | +|---|--:|--:|--:|--:| +| tiny | 128 | 2.796 | 3.172 | n/s | +| small | 512 | 5.335 | 5.508 | n/s | +| gpt2-ish | 1024 | 4.068 | 4.722 | n/s | +| llama3 L=2k | 2048 | 4.688 | 4.845 | n/s | +| llama3 L=4k | 4096 | 12.13 | 12.23 | n/s | +| llama3 L=8k | 8192 | n/s | n/s | n/s | + +Both within noise of the fp16 numbers — bf16 buys the same speed here, just less +precision (8 mantissa bits, so the `lux`-vs-`nnlib` sanity diff is ~1e-2, a +precision/accumulation artifact rather than a bug). + +### 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** | — | — | + +## What changed with CUDA.jl 6.2 (vs 5.11) + +- **BFloat16 became usable** on the NNlib/LuxLib materialized paths (was `n/s` + everywhere). This is the headline of the update. +- **FP32 materialized path sped up ~20–25%** at large sizes from the + cuBLAS/cuDNN bump — e.g. `llama3 L=4k` FP32 gradient **31.9 → 24.6 ms**, + forward **10.5 → 8.8 ms**; `L=2k` gradient **8.4 → 6.5 ms**. Float16 was + already fast and is unchanged. +- **Flash is unchanged** (same kernel) once patched for the CUDA 6.0 API rename — + but note it would have been **completely broken** (`UndefVarError` for all + dtypes) on a stock NNkernels 0.2.0 under CUDA 6.2. + +## Takeaways + +- **NNlib ≈ LuxLib** everywhere — same materialized algorithm, within noise. + LuxLib carries more autodiff overhead at tiny sizes (Float32 gradient `tiny`: + 0.90 vs 0.28 ms). + +- **Flash wins only at small `Float16` sizes, then loses badly at scale on this + GPU.** For `tiny`/`small` `Float16` it is 3–17× faster, but from `gpt2-ish` + upward the materialized path overtakes it, and at Llama sizes flash is **5–6× + slower on the forward and ~50–100× slower on the full gradient** (its backward + kernel is the bottleneck). For `Float32` flash is slower at *every* size. + + Two distinct causes, one confirmed in the NNkernels source: + - **`Float32` is slow by construction.** The kernel only takes the `WMMA` + tensor-core path when `sizeof(T) == 2` (`src/attention.jl:162`); `Float32` + always falls back to the scalar `mma!`, while cuDNN/CUBLAS use FP32 SIMT/TF32 + paths the materialized route rides on. + - **`Float16` *does* use WMMA** (`WMMA.is_available(::CUDABackend)` is hardcoded + `true`), yet still trails at Llama sizes — its sm_70/80-era `wmma` intrinsics + don't map to Blackwell's newest tensor cores, and the tiling/occupancy and + the backward kernel aren't tuned for **sm_120**. On the Ampere/Ada cards + NNkernels' CI targets these numbers would likely look different — **re-run + before drawing hardware-agnostic conclusions.** + +- **Flash's real, layout-independent advantage is memory.** It is the only path + that completes `L=8192`: the materialized score tensor `(8192, 8192, 32)` has + `2³¹` elements, which overflows cuDNN's `int32` softmax descriptor + (`CUDNN_STATUS_NOT_SUPPORTED` → `n/s`), and in `Float32`-causal it exhausts the + 32 GB card (`OOM`). FlashAttention's O(L) memory is what lets it scale to long + context even when it is slower per step here. + +- **Causal masking** roughly halves flash's work (it skips masked tiles) while + making the materialized path slightly *slower* (it builds and applies an + explicit mask), narrowing — but not closing — the gap at large sizes. diff --git a/benchmark/cuda_attention/bench_attention.jl b/benchmark/cuda_attention/bench_attention.jl new file mode 100644 index 000000000..e544ea3df --- /dev/null +++ b/benchmark/cuda_attention/bench_attention.jl @@ -0,0 +1,297 @@ +# GPU multi-head attention benchmark: NNlib vs LuxLib vs NNkernels (Flash Attention). +# +# Three implementations of the SAME scaled dot-product (multi-head) attention +# (scale 1/√E): +# +# * NNlib `dot_product_attention` — materialized path: forms the full +# (kv_len, q_len, nheads, batch) score tensor, `softmax`, second batched_mul. +# Memory is O(seq²). +# * LuxLib `scaled_dot_product_attention` — also a materialized softmax path, +# with its own (E, H, L, B) layout and KV-grouping support. +# * NNkernels `flash_attention` — fused FlashAttention kernel; never +# materializes the scores, so memory is O(seq). +# +# Each library wants a different layout for the SAME data, so we build all three +# from one source tensor: +# * NNkernels: q,k,v = (head_dim, seq, nheads, batch) = (E, L, H, B) +# * NNlib: q,k,v = (head_dim*nheads, seq, batch), nheads=H = (E*H, L, B) +# * LuxLib: q,k,v = (head_dim, nheads, seq, batch) = (E, H, L, B) +# +# Reports forward and full Zygote-gradient (forward+backward) wall time in ms, +# GPU-synchronized. "OOM" = ran out of GPU memory; "n/s" = that path can't run +# the config (echoed once to stderr). +# +# Run (whole suite): +# julia --project=. bench_attention.jl +# Benchmark a single dtype/impl combination without rerunning everything: +# julia --project=. bench_attention.jl --dtypes f32 --impls flash --causal false +# julia --project=. bench_attention.jl --dtypes f16,bf16 --impls nnlib,lux --sizes tiny,small +# Pin a GPU: +# CUDA_VISIBLE_DEVICES=GPU- julia --project=. bench_attention.jl ... +# See all options: +# julia --project=. bench_attention.jl --help + +using CUDA, cuDNN, NNlib, NNkernels, LuxLib, Zygote, BenchmarkTools, Printf, ArgParse +using BFloat16s: BFloat16 +using NNlib: make_causal_mask +using LuxLib: scaled_dot_product_attention + +# --------------------------------------------------------------------------- +# registries: configs, dtypes, implementations +# --------------------------------------------------------------------------- + +# (name, head_dim E, seq_len L, nheads H, batch B). Self-attention: q_len=kv_len=L. +# Climbs from tiny up to a single Llama3-8B attention layer (E=128, H=32). +const CONFIGS = [ + ("tiny", 64, 128, 4, 8), + ("small", 64, 512, 8, 8), + ("gpt2-ish", 64, 1024, 12, 4), + ("llama3 L=2k", 128, 2048, 32, 1), + ("llama3 L=4k", 128, 4096, 32, 1), + ("llama3 L=8k", 128, 8192, 32, 1), +] + +const DTYPE_ALIASES = Dict( + "f16" => Float16, "float16" => Float16, "fp16" => Float16, "half" => Float16, + "bf16" => BFloat16, "bfloat16" => BFloat16, + "f32" => Float32, "float32" => Float32, "fp32" => Float32, +) +const DTYPE_ORDER = (Float16, BFloat16, Float32) # canonical print order + +# An implementation: how to build its inputs from the source (E,L,H,B) tensors, +# and how to run forward / full-gradient on those inputs. +struct Impl + name::String + prep::Function # (qf,kf,vf,causal) -> inputs tuple in this impl's layout + fwd::Function # (inputs) -> output + grad::Function # (inputs) -> gradient w.r.t. q,k,v +end + +# (E,L,H,B) -> (E*H, L, B) +_tonn(x) = reshape(permutedims(x, (1, 3, 2, 4)), size(x, 1) * size(x, 3), size(x, 2), size(x, 4)) +# (E,L,H,B) -> (E,H,L,B) +_tolux(x) = permutedims(x, (1, 3, 2, 4)) + +const IMPLS = Dict( + "nnlib" => Impl("nnlib", + (qf, kf, vf, causal) -> begin + qn, kn, vn = _tonn(qf), _tonn(kf), _tonn(vf) + mask = causal ? make_causal_mask(qn) : nothing + (qn, kn, vn, mask, size(qf, 3)) + end, + I -> dot_product_attention(I[1], I[2], I[3]; nheads=I[5], mask=I[4])[1], + I -> Zygote.gradient((q, k, v) -> + sum(dot_product_attention(q, k, v; nheads=I[5], mask=I[4])[1]), I[1], I[2], I[3])), + "lux" => Impl("lux", + (qf, kf, vf, causal) -> (_tolux(qf), _tolux(kf), _tolux(vf), causal), + I -> scaled_dot_product_attention(I[1], I[2], I[3]; is_causal=I[4])[1], + I -> Zygote.gradient((q, k, v) -> + sum(scaled_dot_product_attention(q, k, v; is_causal=I[4])[1]), I[1], I[2], I[3])), + "flash" => Impl("flash", + (qf, kf, vf, causal) -> (qf, kf, vf, causal), + I -> NNkernels.flash_attention(I[1], I[2], I[3]; causal=I[4]), + I -> Zygote.gradient((q, k, v) -> + sum(NNkernels.flash_attention(q, k, v; causal=I[4])), I[1], I[2], I[3])), +) +const IMPL_ORDER = ["nnlib", "lux", "flash"] # canonical column order + +# --------------------------------------------------------------------------- +# timing +# --------------------------------------------------------------------------- + +# time a thunk in ms (minimum, GPU-synced). Returns a positive time on success, +# NaN for an out-of-memory failure ("OOM"), or -1.0 if that path can't run this +# config ("n/s"). The first reason for each distinct message is echoed to stderr +# so nothing is hidden silently. +const SEEN_FAILURES = Set{String}() +function timed(f, seconds) + try + CUDA.reclaim() + CUDA.@sync f() # warmup / compile + b = @benchmarkable CUDA.@sync($f()) + trial = BenchmarkTools.run(b; seconds=seconds, evals=1) + return minimum(trial).time / 1e6 # ns -> ms + catch e + msg = lowercase(sprint(showerror, e)) + if occursin("out of memory", msg) || occursin("outofgpumemory", msg) || + occursin("out of gpu memory", msg) + return NaN + end + reason = first(split(sprint(showerror, e), '\n')) + if reason ∉ SEEN_FAILURES + push!(SEEN_FAILURES, reason) + @warn "unsupported path → n/s: $reason" + end + return -1.0 + finally + CUDA.reclaim() + end +end + +# --------------------------------------------------------------------------- +# sanity check (only over the selected implementations) +# --------------------------------------------------------------------------- + +# bring each impl's forward output to a common (E,H,L,B) layout for comparison +to_common(::Val{:nnlib}, y, E, L, H, B) = reshape(y, E, H, L, B) +to_common(::Val{:lux}, y, E, L, H, B) = y +to_common(::Val{:flash}, y, E, L, H, B) = permutedims(y, (1, 3, 2, 4)) + +function sanity(T, impls) + E, L, H, B = 32, 64, 4, 2 + qf = CUDA.randn(T, E, L, H, B); kf = CUDA.randn(T, E, L, H, B); vf = CUDA.randn(T, E, L, H, B) + outs = Dict{String,Any}() + for name in impls + im = IMPLS[name] + outs[name] = try + Array(to_common(Val(Symbol(name)), im.fwd(im.prep(qf, kf, vf, false)), E, L, H, B)) + catch + nothing + end + end + ran = [name for name in impls if outs[name] !== nothing] + ref = isempty(ran) ? nothing : outs[first(ran)] + parts = map(impls) do name + d = (outs[name] === nothing || ref === nothing) ? "n/s" : + @sprintf("%.2e", maximum(abs, ref .- outs[name])) + "$name=$d" + end + refname = isempty(ran) ? "—" : first(ran) + @printf("sanity %-9s (vs %s) %s\n", string(T), refname, join(parts, " ")) +end + +# --------------------------------------------------------------------------- +# run + tabulate one (dtype, causal) block +# --------------------------------------------------------------------------- + +fmt(x) = isnan(x) ? "OOM" : x < 0 ? "n/s" : @sprintf("%.3f", x) +ok(x) = !isnan(x) && x > 0 +spd(num, den) = (ok(num) && ok(den)) ? @sprintf("%.2fx", num / den) : "-" + +function print_table(title, impls, rows, key) + haveflash = "flash" in impls && length(impls) > 1 + others = filter(!=("flash"), impls) + header = rpad("config", 13) * lpad("L", 6) * lpad("H", 4) * " |" + for name in impls; header *= lpad(name, 10); end + haveflash && (header *= " |"; for name in others; header *= lpad(name * "/fl", 9); end) + println("\n-- $title --") + println(header) + for r in rows + line = rpad(r.name, 13) * lpad(string(r.L), 6) * lpad(string(r.H), 4) * " |" + for name in impls; line *= lpad(fmt(r.t[name][key]), 10); end + if haveflash + line *= " |" + for name in others; line *= lpad(spd(r.t[name][key], r.t["flash"][key]), 9); end + end + println(line) + end +end + +function run(T, causal, impls, configs, seconds) + @printf("\n%s\n### dtype=%s causal=%s\n%s\n", "#"^104, string(T), causal, "#"^104) + rows = map(configs) do (name, E, L, H, B) + qf = CUDA.randn(T, E, L, H, B); kf = CUDA.randn(T, E, L, H, B); vf = CUDA.randn(T, E, L, H, B) + t = Dict{String,Any}() + for nm in impls + im = IMPLS[nm] + I = try im.prep(qf, kf, vf, causal) catch; nothing end + if I === nothing + t[nm] = (; fwd = -1.0, grad = -1.0) + else + t[nm] = (; fwd = timed(() -> im.fwd(I), seconds), + grad = timed(() -> im.grad(I), seconds)) + end + end + qf = kf = vf = nothing; CUDA.reclaim() + (; name, L, H, t) + end + print_table("FORWARD (ms)", impls, rows, :fwd) + print_table("FULL GRADIENT, fwd+bwd (ms)", impls, rows, :grad) +end + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +function parse_cli(argv) + s = ArgParseSettings(description="GPU multi-head attention benchmark: NNlib vs LuxLib vs NNkernels.") + @add_arg_table! s begin + "--dtypes" + help = "comma list of dtypes: f16,bf16,f32 (aliases ok). Default: all." + default = "f16,bf16,f32" + "--impls" + help = "comma list of implementations: nnlib,lux,flash. Default: all." + default = "nnlib,lux,flash" + "--causal" + help = "comma list of causal modes: false,true (or 'both'). Default: both." + default = "false,true" + "--sizes" + help = "comma list of config names (see --list). Default: all." + default = "all" + "--seconds" + help = "per-measurement time budget for @belapsed." + arg_type = Float64 + default = 1.0 + "--list" + help = "list available configs/dtypes/impls and exit." + action = :store_true + end + return parse_args(argv, s) +end + +splitcsv(s) = strip.(split(s, ','; keepempty=false)) + +function resolve(opts) + dtypes = map(splitcsv(opts["dtypes"])) do d + haskey(DTYPE_ALIASES, lowercase(d)) || error("unknown dtype '$d' (try f16,bf16,f32)") + DTYPE_ALIASES[lowercase(d)] + end + dtypes = filter(in(unique(dtypes)), collect(DTYPE_ORDER)) # canonical order, dedup + + impls = map(lowercase, splitcsv(opts["impls"])) + for im in impls; im in IMPL_ORDER || error("unknown impl '$im' (try nnlib,lux,flash)"); end + impls = filter(in(impls), IMPL_ORDER) + + cstr = lowercase(strip(opts["causal"])) + causals = cstr == "both" ? [false, true] : + [c == "true" for c in splitcsv(cstr)] + causals = unique(causals) + + names = lowercase(strip(opts["sizes"])) == "all" ? first.(CONFIGS) : splitcsv(opts["sizes"]) + configs = filter(c -> c[1] in names, CONFIGS) + isempty(configs) && error("no configs matched --sizes=$(opts["sizes"])") + + return dtypes, impls, causals, configs, opts["seconds"] +end + +function main(argv) + opts = parse_cli(argv) + if opts["list"] + println("configs: ", join(first.(CONFIGS), ", ")) + println("dtypes: f16, bf16, f32") + println("impls: ", join(IMPL_ORDER, ", ")) + return + end + @assert CUDA.functional() "CUDA is not functional" + dtypes, impls, causals, configs, seconds = resolve(opts) + + println("GPU: ", CUDA.name(CUDA.device())) + println("NNlib: ", pkgversion(NNlib)) + println("LuxLib: ", pkgversion(LuxLib)) + println("NNkernels: ", pkgversion(NNkernels)) + println("CUDA.jl: ", pkgversion(CUDA)) + println("selection: dtypes=", join(string.(dtypes), ","), + " impls=", join(impls, ","), + " causal=", join(string.(causals), ","), + " sizes=", join(first.(configs), ",")) + + for T in dtypes; sanity(T, impls); end + for T in dtypes, causal in causals + run(T, causal, impls, configs, seconds) + end + if "flash" in impls && length(impls) > 1 + println("\n/fl = time relative to flash (>1 → flash faster).") + end +end + +main(ARGS) diff --git a/benchmark/cuda_attention/results_rtx5090_cuda511.txt b/benchmark/cuda_attention/results_rtx5090_cuda511.txt new file mode 100644 index 000000000..f065a3d69 --- /dev/null +++ b/benchmark/cuda_attention/results_rtx5090_cuda511.txt @@ -0,0 +1,148 @@ +┌ Warning: unsupported path → n/s: CUDNNError: CUDNN_STATUS_NOT_SUPPORTED (code 3000) +└ @ Main ~/.julia/dev/NNlib/benchmark/cuda_attention/bench_attention.jl:70 +┌ Warning: unsupported path → n/s: MethodError: no method matching cudnnDataType(::Type{BFloat16}) +└ @ Main ~/.julia/dev/NNlib/benchmark/cuda_attention/bench_attention.jl:70 +┌ Warning: unsupported path → n/s: MethodError: no method matching _flash_attention(::CuArray{BFloat16, 4, CUDA.DeviceMemory}, ::CuArray{BFloat16, 4, CUDA.DeviceMemory}, ::CuArray{BFloat16, 4, CUDA.DeviceMemory}, ::Nothing; causal::Bool, kpad_mask::Nothing) +└ @ Main ~/.julia/dev/NNlib/benchmark/cuda_attention/bench_attention.jl:70 +GPU: NVIDIA GeForce RTX 5090 +NNlib: 0.9.36 +LuxLib: 1.15.9 +NNkernels: 0.2.0 +CUDA.jl: 5.11.1 +sanity Float16 vs-nnlib=0.00e+00 vs-lux=2.20e-03 vs-flash=2.93e-03 +sanity BFloat16 vs-nnlib=n/s vs-lux=n/s vs-flash=n/s +sanity Float32 vs-nnlib=0.00e+00 vs-lux=2.98e-07 vs-flash=5.96e-07 + +######################################################################################################## +### dtype=Float16 causal=false +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.868 0.949 0.056 | 15.64x 17.08x +small 512 8 | 1.483 1.491 0.419 | 3.54x 3.56x +gpt2-ish 1024 12 | 1.236 1.236 0.954 | 1.30x 1.30x +llama3 L=2k 2048 32 | 1.465 1.466 8.209 | 0.18x 0.18x +llama3 L=4k 4096 32 | 4.025 4.036 30.502 | 0.13x 0.13x +llama3 L=8k 8192 32 | n/s n/s 116.576 | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 3.051 3.455 0.292 | 10.44x 11.82x +small 512 8 | 4.677 5.505 3.319 | 1.41x 1.66x +gpt2-ish 1024 12 | 4.125 4.434 12.175 | 0.34x 0.36x +llama3 L=2k 2048 32 | 4.161 4.256 169.732 | 0.02x 0.03x +llama3 L=4k 4096 32 | 10.840 10.983 689.596 | 0.02x 0.02x +llama3 L=8k 8192 32 | n/s n/s 2770.579 | - - + +######################################################################################################## +### dtype=Float16 causal=true +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.867 0.876 0.056 | 15.54x 15.68x +small 512 8 | 1.643 1.497 0.256 | 6.43x 5.86x +gpt2-ish 1024 12 | 1.668 1.617 0.626 | 2.67x 2.58x +llama3 L=2k 2048 32 | 2.368 2.386 4.318 | 0.55x 0.55x +llama3 L=4k 4096 32 | 7.660 7.733 15.729 | 0.49x 0.49x +llama3 L=8k 8192 32 | n/s n/s 59.858 | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 2.880 3.588 0.257 | 11.20x 13.96x +small 512 8 | 5.871 6.143 1.975 | 2.97x 3.11x +gpt2-ish 1024 12 | 7.049 7.556 6.762 | 1.04x 1.12x +llama3 L=2k 2048 32 | 10.863 11.456 87.665 | 0.12x 0.13x +llama3 L=4k 4096 32 | 36.566 36.772 345.198 | 0.11x 0.11x +llama3 L=8k 8192 32 | n/s n/s 1386.710 | - - + +######################################################################################################## +### dtype=BFloat16 causal=false +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | n/s n/s n/s | - - +small 512 8 | n/s n/s n/s | - - +gpt2-ish 1024 12 | n/s n/s n/s | - - +llama3 L=2k 2048 32 | n/s n/s n/s | - - +llama3 L=4k 4096 32 | n/s n/s n/s | - - +llama3 L=8k 8192 32 | n/s n/s n/s | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | n/s n/s n/s | - - +small 512 8 | n/s n/s n/s | - - +gpt2-ish 1024 12 | n/s n/s n/s | - - +llama3 L=2k 2048 32 | n/s n/s n/s | - - +llama3 L=4k 4096 32 | n/s n/s n/s | - - +llama3 L=8k 8192 32 | n/s n/s n/s | - - + +######################################################################################################## +### dtype=BFloat16 causal=true +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | n/s n/s n/s | - - +small 512 8 | n/s n/s n/s | - - +gpt2-ish 1024 12 | n/s n/s n/s | - - +llama3 L=2k 2048 32 | n/s n/s n/s | - - +llama3 L=4k 4096 32 | n/s n/s n/s | - - +llama3 L=8k 8192 32 | n/s n/s n/s | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | n/s n/s n/s | - - +small 512 8 | n/s n/s n/s | - - +gpt2-ish 1024 12 | n/s n/s n/s | - - +llama3 L=2k 2048 32 | n/s n/s n/s | - - +llama3 L=4k 4096 32 | n/s n/s n/s | - - +llama3 L=8k 8192 32 | n/s n/s n/s | - - + +######################################################################################################## +### dtype=Float32 causal=false +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.050 0.050 0.062 | 0.80x 0.81x +small 512 8 | 0.276 0.275 0.730 | 0.38x 0.38x +gpt2-ish 1024 12 | 0.790 0.791 1.984 | 0.40x 0.40x +llama3 L=2k 2048 32 | 2.759 2.768 16.098 | 0.17x 0.17x +llama3 L=4k 4096 32 | 10.522 10.528 61.818 | 0.17x 0.17x +llama3 L=8k 8192 32 | n/s n/s 243.137 | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.259 0.892 0.535 | 0.48x 1.67x +small 512 8 | 0.853 1.065 7.314 | 0.12x 0.15x +gpt2-ish 1024 12 | 2.126 2.222 30.084 | 0.07x 0.07x +llama3 L=2k 2048 32 | 8.351 8.418 397.375 | 0.02x 0.02x +llama3 L=4k 4096 32 | 31.885 31.902 1611.004 | 0.02x 0.02x +llama3 L=8k 8192 32 | n/s n/s 6435.178 | - - + +######################################################################################################## +### dtype=Float32 causal=true +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.054 0.063 0.062 | 0.88x 1.03x +small 512 8 | 0.394 0.404 0.435 | 0.91x 0.93x +gpt2-ish 1024 12 | 1.257 1.267 1.106 | 1.14x 1.15x +llama3 L=2k 2048 32 | 3.699 3.726 8.449 | 0.44x 0.44x +llama3 L=4k 4096 32 | 14.328 14.394 32.048 | 0.45x 0.45x +llama3 L=8k 8192 32 | n/s n/s 123.786 | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.334 1.006 0.387 | 0.86x 2.60x +small 512 8 | 2.247 2.682 4.049 | 0.55x 0.66x +gpt2-ish 1024 12 | 6.204 6.399 15.672 | 0.40x 0.41x +llama3 L=2k 2048 32 | 18.978 19.100 194.069 | 0.10x 0.10x +llama3 L=4k 4096 32 | 75.572 74.532 799.728 | 0.09x 0.09x +llama3 L=8k 8192 32 | OOM OOM 3218.478 | - - + +nnlib/fl, lux/fl = time relative to flash (>1 → flash faster). diff --git a/benchmark/cuda_attention/results_rtx5090_cuda62.txt b/benchmark/cuda_attention/results_rtx5090_cuda62.txt new file mode 100644 index 000000000..fa9a88d1e --- /dev/null +++ b/benchmark/cuda_attention/results_rtx5090_cuda62.txt @@ -0,0 +1,147 @@ +┌ Warning: unsupported path → n/s: CUDNNError: CUDNN_STATUS_NOT_SUPPORTED (code 3000) +└ @ Main ~/.julia/dev/NNlib/benchmark/cuda_attention/bench_attention.jl:123 +┌ Warning: unsupported path → n/s: MethodError: no method matching _flash_attention(::CuArray{BFloat16, 4, CUDACore.DeviceMemory}, ::CuArray{BFloat16, 4, CUDACore.DeviceMemory}, ::CuArray{BFloat16, 4, CUDACore.DeviceMemory}, ::Nothing; causal::Bool, kpad_mask::Nothing) +└ @ Main ~/.julia/dev/NNlib/benchmark/cuda_attention/bench_attention.jl:123 +GPU: NVIDIA GeForce RTX 5090 +NNlib: 0.9.36 +LuxLib: 1.15.9 +NNkernels: 0.2.0 +CUDA.jl: 6.2.0 +selection: dtypes=Float16,BFloat16,Float32 impls=nnlib,lux,flash causal=false,true sizes=tiny,small,gpt2-ish,llama3 L=2k,llama3 L=4k,llama3 L=8k +sanity Float16 (vs nnlib) nnlib=0.00e+00 lux=1.95e-03 flash=2.44e-03 +sanity BFloat16 (vs nnlib) nnlib=0.00e+00 lux=7.81e-02 flash=n/s +sanity Float32 (vs nnlib) nnlib=0.00e+00 lux=3.20e-07 flash=5.89e-07 + +######################################################################################################## +### dtype=Float16 causal=false +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.791 0.856 0.053 | 14.98x 16.20x +small 512 8 | 1.526 1.531 0.402 | 3.80x 3.81x +gpt2-ish 1024 12 | 1.157 1.155 0.922 | 1.25x 1.25x +llama3 L=2k 2048 32 | 1.511 1.518 8.110 | 0.19x 0.19x +llama3 L=4k 4096 32 | 4.138 4.160 29.822 | 0.14x 0.14x +llama3 L=8k 8192 32 | n/s n/s 114.017 | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 2.490 3.155 0.290 | 8.57x 10.86x +small 512 8 | 4.726 5.457 3.227 | 1.46x 1.69x +gpt2-ish 1024 12 | 3.765 4.880 11.667 | 0.32x 0.42x +llama3 L=2k 2048 32 | 4.282 4.460 167.350 | 0.03x 0.03x +llama3 L=4k 4096 32 | 11.265 11.423 674.672 | 0.02x 0.02x +llama3 L=8k 8192 32 | n/s n/s 2723.243 | - - + +######################################################################################################## +### dtype=Float16 causal=true +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.862 0.796 0.050 | 17.17x 15.86x +small 512 8 | 1.693 1.539 0.250 | 6.78x 6.16x +gpt2-ish 1024 12 | 1.513 1.523 0.596 | 2.54x 2.55x +llama3 L=2k 2048 32 | 2.409 2.430 4.291 | 0.56x 0.57x +llama3 L=4k 4096 32 | 7.765 7.834 15.550 | 0.50x 0.50x +llama3 L=8k 8192 32 | n/s n/s 58.736 | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 2.587 3.291 0.238 | 10.85x 13.80x +small 512 8 | 5.900 6.096 1.789 | 3.30x 3.41x +gpt2-ish 1024 12 | 6.746 7.477 5.679 | 1.19x 1.32x +llama3 L=2k 2048 32 | 10.737 11.042 85.648 | 0.13x 0.13x +llama3 L=4k 4096 32 | 36.107 36.452 336.786 | 0.11x 0.11x +llama3 L=8k 8192 32 | n/s n/s 1364.371 | - - + +######################################################################################################## +### dtype=BFloat16 causal=false +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.786 0.791 n/s | - - +small 512 8 | 1.717 1.561 n/s | - - +gpt2-ish 1024 12 | 1.264 1.270 n/s | - - +llama3 L=2k 2048 32 | 1.686 1.688 n/s | - - +llama3 L=4k 4096 32 | 4.494 4.510 n/s | - - +llama3 L=8k 8192 32 | n/s n/s n/s | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 2.796 3.172 n/s | - - +small 512 8 | 5.335 5.508 n/s | - - +gpt2-ish 1024 12 | 4.068 4.722 n/s | - - +llama3 L=2k 2048 32 | 4.688 4.845 n/s | - - +llama3 L=4k 4096 32 | 12.132 12.232 n/s | - - +llama3 L=8k 8192 32 | n/s n/s n/s | - - + +######################################################################################################## +### dtype=BFloat16 causal=true +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.881 0.811 n/s | - - +small 512 8 | 1.559 1.569 n/s | - - +gpt2-ish 1024 12 | 1.614 1.684 n/s | - - +llama3 L=2k 2048 32 | 2.583 2.594 n/s | - - +llama3 L=4k 4096 32 | 8.125 8.178 n/s | - - +llama3 L=8k 8192 32 | n/s n/s n/s | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 2.595 3.303 n/s | - - +small 512 8 | 5.445 6.177 n/s | - - +gpt2-ish 1024 12 | 7.270 8.143 n/s | - - +llama3 L=2k 2048 32 | 10.868 11.307 n/s | - - +llama3 L=4k 4096 32 | 37.047 37.396 n/s | - - +llama3 L=8k 8192 32 | n/s n/s n/s | - - + +######################################################################################################## +### dtype=Float32 causal=false +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.050 0.051 0.068 | 0.73x 0.75x +small 512 8 | 0.243 0.242 0.806 | 0.30x 0.30x +gpt2-ish 1024 12 | 0.743 0.746 2.238 | 0.33x 0.33x +llama3 L=2k 2048 32 | 2.338 2.353 18.331 | 0.13x 0.13x +llama3 L=4k 4096 32 | 8.759 8.787 69.503 | 0.13x 0.13x +llama3 L=8k 8192 32 | n/s n/s 270.467 | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.281 0.895 0.621 | 0.45x 1.44x +small 512 8 | 0.766 1.027 8.702 | 0.09x 0.12x +gpt2-ish 1024 12 | 2.014 2.098 35.535 | 0.06x 0.06x +llama3 L=2k 2048 32 | 6.522 6.670 467.639 | 0.01x 0.01x +llama3 L=4k 4096 32 | 24.613 24.521 1890.846 | 0.01x 0.01x +llama3 L=8k 8192 32 | n/s n/s 7551.554 | - - + +######################################################################################################## +### dtype=Float32 causal=true +######################################################################################################## + +-- FORWARD (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.061 0.064 0.067 | 0.90x 0.95x +small 512 8 | 0.362 0.371 0.489 | 0.74x 0.76x +gpt2-ish 1024 12 | 1.211 1.223 1.233 | 0.98x 0.99x +llama3 L=2k 2048 32 | 3.274 3.298 9.306 | 0.35x 0.35x +llama3 L=4k 4096 32 | 12.472 12.552 35.086 | 0.36x 0.36x +llama3 L=8k 8192 32 | n/s n/s 136.376 | - - + +-- FULL GRADIENT, fwd+bwd (ms) -- +config L H | nnlib lux flash | nnlib/fl lux/fl +tiny 128 4 | 0.335 0.993 0.442 | 0.76x 2.24x +small 512 8 | 2.177 2.552 4.805 | 0.45x 0.53x +gpt2-ish 1024 12 | 6.115 6.351 18.554 | 0.33x 0.34x +llama3 L=2k 2048 32 | 17.189 17.321 226.705 | 0.08x 0.08x +llama3 L=4k 4096 32 | 67.056 66.906 936.796 | 0.07x 0.07x +llama3 L=8k 8192 32 | OOM OOM 3771.119 | - - + +/fl = time relative to flash (>1 → flash faster).