Programming an attention kernel in Triton

A developer documented writing GPU kernels in Triton to understand PyTorch operations, starting with vector add and fused ReLU/dropout, highlighting the benefit of fused kernels in reducing memory round trips. The post details the execution model of Triton and sets up for a discussion of FlashAttention.

Programming an attention kernel in Triton Every PyTorch operation you've ever called, softmax , relu , matmul , is, underneath, someone else's compiled GPU kernel. You never see the kernel. You call the function, the tensor comes back, and the actual work of moving numbers in and out of GPU memory, scheduling threads, and managing on-chip caches happens somewhere you're not invited to look. I wanted to know what was actually happening in there. Not at the level of "attention is a weighted sum," which I already understood from a previous post grok-grok.html , but at the level of: what does the GPU actually do, instruction by instruction, when you call torch.softmax ? So I wrote the kernels myself, in Triton, starting from the simplest thing that could possibly be called a kernel and working up. This is what that ladder looked like, including the one real bug I found and the wall I hit at the end, named honestly, not smoothed over. Why bother: the thing PyTorch hides Most of what PyTorch calls a single operation is really several. torch.softmax x internally finds the row max, subtracts it, exponentiates, sums, and divides, and on a naive execution model each of those steps can mean a separate pass over the data: read from GPU memory, compute, write back to GPU memory, repeat. For an operation that conceptually happens "once," your data can make several unnecessary round trips between the GPU's slow global memory and its fast on-chip compute units. A fused kernel does all of those steps in one pass: load the data into fast on-chip memory once, do the whole computation, write the result back once. That's the entire premise behind Triton, and it's the same underlying idea that motivates FlashAttention, which I'll come back to at the end: fewer memory round trips, not fewer FLOPs, is often the real lever for speed on a GPU. Kernel 1: vector add, or learning to think like a GPU The simplest possible Triton kernel doesn't optimize anything, it just adds two vectors, out = x + y . The point of writing this first wasn't the operation, it was learning Triton's actual execution model before any algorithm complexity got involved: - Triton launches a grid of programs, each identified by a pid program ID . Think of each program as one worker handling one chunk of the data. - Each program computes its own offsets into the input arrays, block start = pid BLOCK SIZE , then tl.arange 0, BLOCK SIZE for the positions within its chunk. - A mask offsets < n elements handles the case where the data size isn't a clean multiple of the block size. Without it, the last block would read and write past the end of the array. python @triton.jit def add kernel x ptr, y ptr, out ptr, n elements, BLOCK SIZE: tl.constexpr : pid = tl.program id axis=0 block start = pid BLOCK SIZE offsets = block start + tl.arange 0, BLOCK SIZE mask = offsets < n elements x = tl.load x ptr + offsets, mask=mask y = tl.load y ptr + offsets, mask=mask tl.store out ptr + offsets, x + y, mask=mask Checked against plain x + y in PyTorch: matches exactly. Nothing interesting happens here algorithmically, that's the point. pid , offsets, and masking are the three ideas every later kernel in this post reuses, so getting them right on the simplest possible operation first meant later bugs and there was one, further down were never about "do I understand Triton," only ever about the specific algorithm. Kernel 2: fused ReLU and dropout, and the problem with testing randomness The next step up: fuse two operations into one pass instead of one. ReLUZero out every negative value, leave positive values unchanged. and dropoutRandomly zero a fraction p of values, scaling the survivors by 1 / 1 - p to keep the expected sum unchanged. are two ops that would normally be two separate kernel calls in a naive implementation. Fusing them means one load, one combined computation, one store. python @triton.jit def relu dropout kernel x ptr, out ptr, n elements, p, seed, BLOCK SIZE: tl.constexpr : pid = tl.program id axis=0 offsets = pid BLOCK SIZE + tl.arange 0, BLOCK SIZE mask = offsets < n elements x = tl.load x ptr + offsets, mask=mask x = tl.maximum x, 0.0 randoms = tl.rand seed, offsets dropout mask = randoms p out = tl.where dropout mask, x / 1 - p , 0.0 tl.store out ptr + offsets, out, mask=mask Here's the part worth dwelling on: you can't validate this the way you validate the vector-add kernel. allclose against a reference only works when the output is deterministic. Dropout is stochastic by design, there is no single "correct" output to compare against. Eyeballing five printed values, which is what I did the first time, tells you almost nothing; it's the same mistake as trusting a single run in the birthday-paradox post birthday-attack.html instead of the distribution. The actual test has to be statistical: run it on enough elements that probability becomes measurable, then check two things, does the empirical drop rate land near p , and are the survivors scaled by exactly 1 / 1 - p ? php Testing with size=100000, p=0.1 - empirical rate 0.1007 expected 0.10 +/- 0.01 OK Testing with size=100000, p=0.5 - empirical rate 0.4993 expected 0.50 +/- 0.01 OK Testing with size=100000, p=0.9 - empirical rate 0.8999 expected 0.90 +/- 0.01 OK All three land within a hundredth of a percent of the target. That's the right way to trust a kernel whose output is supposed to be random: proving the distribution is correct, not any single sample. Kernel 3: fused softmax, and the bug that mattered Softmax is the first kernel here with real numerical-stability considerations: naive exp x overflows for even moderately large x , so every real softmax implementation subtracts the row max before exponentiating. The kernel does this per row, in one pass: python @triton.jit def softmax output ptr, input ptr, input row stride, output row stride, n cols, BLOCK SIZE: tl.constexpr : row idx = tl.program id 0 row start ptr = input ptr + row idx input row stride col offsets = tl.arange 0, BLOCK SIZE input ptrs = row start ptr + col offsets mask = col offsets < n cols row = tl.load input ptrs, mask=mask, other=float "-inf" row max = tl.max row, axis=0 numerator = tl.exp row - row max denominator = tl.sum numerator, axis=0 softmax output = numerator / denominator output ptrs = output ptr + row idx output row stride + col offsets tl.store output ptrs, softmax output, mask=mask First version, tested on a small 4, 5 input: matched PyTorch exactly. I moved on, satisfied. It was wrong. Here's the actual bug, left in on purpose: python def triton softmax x: torch.Tensor : n rows, n cols = x.shape BLOCK SIZE = 1024 hardcoded ... BLOCK SIZE was a fixed number, not tied to n cols at all. For a small input this is harmless, col offsets spans more than enough room, and the mask correctly zeroes out the unused tail. But nothing about the mask logic checks whether BLOCK SIZE is large enough, it only checks whether each position is within n cols . So the moment a real row is wider than 1024 columns, the kernel simply never loads, computes, or stores anything past column 1024. It doesn't crash or warn you, it just returns a softmax over the first 1024 columns and calls it done. Before, reproduced explicitly at n cols=2000 : Buggy Softmax Matches PyTorch n cols=2000 ? False Max abs diff: 0.00978 After, the fix is one line: stop guessing a fixed block size and size it to the actual input. BLOCK SIZE = triton.next power of 2 n cols Corrected Softmax Matches PyTorch n cols=2000 ? True The lesson isn't "remember to make BLOCK SIZE dynamic," it's narrower and more useful than that: a test that only covers the shape you happen to be thinking about will pass right over a bug that only shows up at a different shape. The 4, 5 test I ran first was real, it wasn't fake, it just wasn't the test that mattered. Kernel 4: bringing it together, a self-attention kernel Everything up to here, offsets, masks, a fused numerically-stable softmax, combines into one kernel implementing scaled dot-product attention directly: python @triton.jit def attention kernel q ptr, k ptr, v ptr, out ptr, row stride, col stride, n rows, n cols, scale, BLOCK Q: tl.constexpr, BLOCK K: tl.constexpr, BLOCK V: tl.constexpr, : row pid = tl.program id 0 row offsets = row pid BLOCK Q + tl.arange 0, BLOCK Q col offsets = tl.arange 0, BLOCK K d offsets = tl.arange 0, BLOCK V row mask = row offsets < n rows col mask = col offsets < n cols q offsets = tl.expand dims row offsets, 1 row stride + tl.expand dims d offsets, 0 col stride q = tl.load q ptr + q offsets, mask=tl.expand dims row mask, 1 , other=0.0 k offsets = tl.expand dims d offsets, 1 col stride + tl.expand dims col offsets, 0 row stride k = tl.load k ptr + k offsets, mask=tl.expand dims col mask, 0 , other=0.0 qk = tl.dot q, k scale qk = tl.where tl.expand dims col mask, 0 , qk, float "-inf" m = tl.max qk, axis=1 p = tl.exp qk - tl.expand dims m, 1 l = tl.sum p, axis=1 weights = p / tl.expand dims l, 1 v offsets = tl.expand dims col offsets, 1 row stride + tl.expand dims d offsets, 0 col stride v = tl.load v ptr + v offsets, mask=tl.expand dims col mask, 1 , other=0.0 out = tl.dot weights.to v.dtype , v out offsets = tl.expand dims row offsets, 1 row stride + tl.expand dims d offsets, 0 col stride tl.store out ptr + out offsets, out, mask=tl.expand dims row mask, 1 Validated against torch.nn.functional.scaled dot product attention across several shapes, including deliberately mismatched Q/K sequence lengths and a single-query edge case: php Q: 5, 32 , K: 15, 32 , V: 15, 32 - matches: True Q: 1, 64 , K: 10, 64 , V: 10, 64 - matches: True Given the softmax lesson above, testing more than one shape here wasn't optional, it's the only reason I can trust this one. The wall, named honestly Here's what this kernel is not: it is not FlashAttention, and it's worth being precise about why, rather than letting the name imply more than the code does. This kernel loads the entire K and V into fast on-chip memory SRAMThe GPU's small, extremely fast on-chip memory, as opposed to its much larger but far slower off-chip global memory. in one shot before doing anything else. That's fine at the sequence lengths tested here, but SRAM is small tens of kilobytes per streaming multiprocessor, not gigabytes , and K/V grow linearly with sequence length. At some sequence length, "all of K and V" simply stops fitting, and this kernel breaks, not gracefully, it just runs out of room. The real FlashAttention trick is to never need all of K/V in SRAM at once: process it in chunks, and keep a running max and running sum as you go, correcting the accumulated output every time a new chunk reveals a bigger max than anything seen so far. That's a genuinely different, harder algorithm than anything in this post, the online-softmax rescaling has no analogue in kernels 1 through 4, and it's the specific thing I haven't built yet. Naming that clearly here, instead of leaving it unsaid, is the whole point: this post is "I climbed four rungs of a real ladder, correctly, and found a real bug along the way," not "I built FlashAttention." Short and unforced: the useful thing wasn't the final kernel, it was the softmax bug. It's the one moment on this ladder where "it passed my test" and "it's actually correct" turned out to be different claims, the same gap that showed up when I mixed up the median and the mean in the birthday-paradox post. Different domain, same shape of mistake, same fix: test the case you didn't think to test. The tiled, online-softmax version is the next rung. I haven't climbed it yet.