Small GPU kernel lab for LLM inference primitives in Python, Triton, PyTorch, and CUDA. The repo focuses on readable kernels and reproducible validation: row-wise softmax, FP16 GEMM, and a FlashAttention-style fused attention forward kernel using tiled online softmax.
This is a cleaned-up reconstruction of earlier kernel experiments. Benchmark numbers are intentionally generated by the local harness instead of being baked into the README, because GPU model, driver, Triton version, and tensor shapes change the results materially.
row_softmax: one Triton program per row, numerically stable max-subtract softmax, intended for bandwidth-bound reductions where a row fits in SRAM.fp16_matmul: tiled FP16 GEMM with FP32 accumulation, block-leveltl.dot, grouped program ordering for better L2 reuse, and configurableBLOCK_M/BLOCK_N/BLOCK_K.flash_attention_forward: fused attention forward for prefill and simple decode-style shapes, using tiled QK and PV blocks plus the online softmax recurrence to avoid materializing the full attention matrix.
All kernels are validated against a PyTorch fp32 reference; the tables report max
absolute error alongside performance. Reproduce with
python -m triton_llm_kernel_lab.bench.
The one-program-per-row fused softmax is bandwidth-bound and lands close to the hardware limit on large rows:
| shape (rows × cols) | latency | achieved bandwidth | vs torch.softmax |
max err |
|---|---|---|---|---|
| 8192 × 4096 | 0.048 ms | 2.8 TB/s (~83% of HBM3 peak) | 2.78× | 3.8e-6 |
| 4096 × 2048 | 0.024 ms | — | 1.11× | 3.8e-6 |
On small rows (≤1024 cols) the kernel is slower than torch.softmax: the row fits
easily and kernel-launch overhead dominates the actual work. The fused kernel wins
exactly when it should — large, bandwidth-bound reductions.
Every kernel matches the PyTorch reference to within fp16 tolerance:
| kernel | shapes tested | max abs error |
|---|---|---|
| row softmax | 512–4096 cols, non-power-of-two | 3.8e-6 – 1.5e-5 |
| FP16 GEMM | up to 4096×11008×4096, masked edge tiles | 0.0 (bit-exact accum) |
| FlashAttention forward | prefill + decode, causal & non-causal | 6e-5 – 2e-3 |
The FP16 GEMM and FlashAttention-forward kernels are written for readability and correctness first. They are not tuned to beat cuBLAS / cuDNN-flash / PyTorch SDPA, and on H100 they currently trail those vendor paths — see the roadmap below. The gap is dominated by two known, unimplemented optimizations rather than anything fundamental:
- Autotuning. Tile sizes are fixed (
BLOCK_M=32, BLOCK_N=64); H100 wants far larger tiles and more pipeline stages. A Triton autotuning sweep is the first lever. - Flash-decoding (split-KV). The attention kernel does not split the KV dimension
across SMs, so the
q_len=1decode path is heavily underutilized. This is the single largest opportunity and the next thing on the list.
- Triton autotuning configs for GEMM and attention (tile sizes,
num_warps,num_stages) - Flash-decoding / split-KV path for
q_len=1decode attention - Head-to-head vs FlashAttention-2 once the wheel builds in CI
- Nsight Compute traces for the softmax roofline and the GEMM gap
src/triton_llm_kernel_lab/
bench.py # CLI benchmark harness
configs.py # kernel configs and LLM-like benchmark shapes
reference.py # PyTorch reference implementations
runtime.py # CUDA/Triton availability checks
kernels/
attention.py # FlashAttention-style fused attention forward
gemm.py # FP16 GEMM kernel
softmax.py # row-wise fused softmax kernel
tests/
test_references.py # CPU-safe reference and config checks
test_gpu_kernels.py # CUDA/Triton correctness checks, skipped otherwise
docs/
profiling.md # Nsight Compute workflow and metrics
tradeoffs.md # prefill vs decode kernel selection notes
Use a Linux environment with an NVIDIA GPU for the Triton kernels.
python -m venv .venv
source .venv/bin/activate
pip install -U pip
pip install -e ".[gpu,dev]"On a CPU-only machine, install only the test/dev path:
pip install -e ".[dev]"
pytest tests/test_references.pyThe GPU tests compare every custom kernel against a PyTorch reference and report the max absolute error. On a CUDA machine:
pytest tests/test_gpu_kernels.py -qThe test coverage includes:
- softmax rows with non-power-of-two column counts
- FP16 GEMM with masked edge tiles
- causal and non-causal fused attention forward
The harness uses 50 warmup iterations and 200 timed iterations by default. It prints latency, estimated TFLOPS, estimated memory bandwidth, and max error.
python -m triton_llm_kernel_lab.bench --kernel all
python -m triton_llm_kernel_lab.bench --kernel attention --warmup 50 --iters 200 --csv results/attention.csvRepresentative shape groups are defined in configs.py:
- prefill: longer query/key lengths where QK and PV dominate arithmetic
- decode: short query length with long KV cache where memory traffic dominates
- GEMM: common projection and MLP matrix sizes
- softmax: row lengths that stress SRAM fit and reduction behavior
Use Nsight Compute for detailed GPU metrics:
bash scripts/profile_ncu.sh attentionThe profiling notes in docs/profiling.md track the metrics that matter for
this lab: achieved occupancy, memory throughput, L2 hit rate, warp stalls,
tensor core utilization, and DRAM read/write transactions.
The implementation style was informed by the public Triton tutorials and FlashAttention papers, but the code in this repository is written as a compact teaching/lab version rather than a copy of those tutorials.
- Triton fused softmax tutorial: https://triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html
- Triton matrix multiplication tutorial: https://github.com/triton-lang/triton/blob/main/python/tutorials/03-matrix-multiplication.py
- Triton fused attention tutorial: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
- FlashAttention: https://arxiv.org/abs/2205.14135
- FlashAttention-2: https://tridao.me/publications/flash2/flash2.pdf