Learning kernels, starting at softmax

I want to learn to write fast GPU kernels by hand. Specifically, Triton, with each line shaped on purpose. Two reasons.

One. I’ve spent time at the modeling layer and time at the system layer and almost none at the layer where the two meet. The kernel is where math hits hardware. If you can’t read or write at that layer, you cede the most important performance decisions to whoever shipped the library.

Two. The economics of every model I want to train are dominated by attention and a few normalizations. Shave a constant off softmax or attention in Triton and that constant shows up in every forward pass for the next year.

Starting points.

Softmax. The canonical first kernel. Three passes naive: max, exp plus sum, divide. Two passes with online softmax that fuses max and sum into a single streaming reduction. One pass when the row fits in SRAM and you can normalize on the fly. Writing all three in Triton and benchmarking them is the way to feel the cost of memory traffic.

Flash attention. After softmax. The point is the tiling pattern: blocks of Q against blocks of K and V, keeping the running softmax state across blocks, never materializing the full attention matrix. It is the most cited modern kernel for a reason. If I can implement it from scratch and match the reference within a constant factor, I’ll trust myself to write the next one.

Goal. By end of summer, a small library of kernels I wrote and understand. One published note per kernel, covering the math, the tile shape, the bottleneck I hit, and the speedup over a naive PyTorch baseline.