A Triton FlashAttention-2 Prefill kernel and a CUDA WMMA FP16 GEMM, built from a naive baseline up through five progressively optimized versions and benchmarked on an RTX 4090.
- HGEMM reaches 91.7% of cuBLAS at 8192³ (84.1% at 4096³) with a hand-written Tensor Core kernel.
- 3.0× throughput jump moving from SIMT FMA to Tensor Cores (8192³: 37.2 → 111.5 TFLOPS).
- Triton FlashAttention-2 Prefill at 0.81–0.89× of PyTorch SDPA-flash, peaking at 125 TFLOPS, with O(N) intermediate memory.
- Five progressive GEMM versions, each isolating a single optimization and verified against cuBLAS independently.
Measured on RTX 4090 D (Ada, sm_89), CUDA 12.8 (nvcc V12.8.93), PyTorch 2.8.0+cu128 / Triton 3.4.0. GEMM baseline is cuBLAS (timed inside the C++ driver); FA2 baseline is PyTorch SDPA restricted to its flash backend.
| M=N=K | v1 | v2 | v3 | v4 | cuBLAS (TFLOPS) |
|---|---|---|---|---|---|
| 1024 | 25.4% | 66.7% | 81.0% | 81.0% | 61.7 |
| 2048 | 22.0% | 59.1% | 70.3% | 70.3% | 117.3 |
| 4096 | 26.0% | 69.2% | 82.8% | 84.1% | 140.8 |
| 8192 | 25.9% | 77.3% | 90.1% | 91.7% | 144.3 |
v4's share of cuBLAS rises with problem size, reaching 91.7% at 8192³.
On a short reduction dimension (M=N=4096, K=64) v4 regresses below v3:
| M=N=4096, K=64 | v2 | v3 | v4 |
|---|---|---|---|
| % of cuBLAS | 80.6% | 86.2% | 53.0% |
With BK=32, K=64 is only 2 K-tiles. The swizzle's ldmatrix index math and the
cp.async pipeline warm-up are per-K-tile fixed costs that a short reduction
dimension cannot amortize, so the net effect turns negative. This is a shape-dependent
trade-off, not a defect — the shape still passes --verify, and the swizzle pays off
once K is large.
batch=4, heads=32/8 (GQA), head_dim=128, causal, fp16:
| seqlen | Triton (ms) | SDPA (ms) | TFLOPS | speedup |
|---|---|---|---|---|
| 512 | 0.127 | 0.114 | 67.6 | 0.895× |
| 1024 | 0.376 | 0.303 | 91.4 | 0.807× |
| 2048 | 1.260 | 1.059 | 109.1 | 0.841× |
| 4096 | 4.599 | 3.995 | 119.5 | 0.869× |
| 8192 | 17.544 | 15.570 | 125.3 | 0.887× |
The kernel reaches 0.81–0.89× of SDPA-flash for N ≥ 1024, D=128, causal, with
throughput climbing to 125 TFLOPS as sequence length grows. Key implementation
points: online softmax keeps intermediate memory at O(N) (the N×N score matrix is
never materialized); the 1/l normalization is deferred to a single pass after the
loop (the difference from FA1); causal masking uses a two-stage loop — a full-speed
unmasked stage below the diagonal, then only the tiles that intersect it; GQA infers
Hkv from k.shape rather than adding an API parameter.
All numbers above were measured with --verify / pytest passing: each GEMM version is
checked against cuBLAS by relative Frobenius norm, and FA2 is checked against SDPA with
a dual-baseline tolerance.
Each version isolates one optimization, so the performance delta can be attributed to that change alone. Percentages are the 4096³ row above.
| Version | Optimization | @4096³ of cuBLAS | vs previous | Takeaway |
|---|---|---|---|---|
| v0 | One output element per thread | — (orders slower, skipped past 2048³) | — | Correctness anchor; no Tensor Cores |
| v1 | Shared-memory tiling + 8×8 register tile | 26.0% | — | The ceiling of SIMT FMA |
| v2 | WMMA three-level tiling (CTA 128×128×32 / warp 64×32 / MMA 16×16×16) | 69.2% | 2.7× | The single biggest step |
| v3 | cp.async double buffering |
82.8% | +13.6 pts | Hides global→shared latency behind compute |
| v4 | XOR swizzle (hand-written ldmatrix + mma.sync) |
84.1% | +1.3 pts | Small on squares; shared memory 37→32 KB |
Requires Linux, Python 3.10–3.12, CUDA Toolkit 12.x, PyTorch ≥ 2.4, Triton ≥ 3.0,
CMake ≥ 3.24, and an NVIDIA GPU with compute capability 8.0 or newer — the kernels
use cp.async, which is an Ampere (sm_80) feature.
# Install
python -m pip install -e ".[dev]"
# Build the GEMM executable (default target sm_89; sm_80 also works)
cmake -S . -B build && cmake --build build -j
# Correctness
build/csrc/gemm/hgemm_bench --kernel v4 --m 4096 --n 4096 --k 4096 --verify
pytest tests/ -v
# Performance (CSV lands in benchmarks/results/, gitignored)
benchmarks/gemm/run_bench.sh
python benchmarks/prefill/bench_prefill.py
# Profiling
profiling/ncu_gemm.sh v4 4096
profiling/ncu_prefill.sh 2048 128M, N and K must be multiples of 16, so every 16×16 MMA tile is wholly in or out of bounds and the epilogue never handles a partial tile.
MIT License.