{"slug": "programming-an-attention-kernel-in-triton", "title": "Programming an attention kernel in Triton", "summary": "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.", "body_md": "# Programming an attention kernel in Triton\n\nEvery PyTorch operation you've ever called, `softmax`, `relu`, `matmul`, is,\nunderneath, someone else's compiled GPU kernel. You never see the kernel. You\ncall the function, the tensor comes back, and the actual work of moving\nnumbers in and out of GPU memory, scheduling threads, and managing on-chip\ncaches happens somewhere you're not invited to look.\n\nI wanted to know what was actually happening in there. Not at the level of\n\"attention is a weighted sum,\" which I already understood from\n[a previous post](grok-grok.html), but at the level of: what does the GPU\nactually do, instruction by instruction, when you call `torch.softmax`?\n\nSo 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.\n\n## Why bother: the thing PyTorch hides\n\nMost of what PyTorch calls a single operation is really several.\n`torch.softmax(x)` internally finds the row max, subtracts it, exponentiates,\nsums, and divides, and on a naive execution model each of those steps can\nmean a separate pass over the data: read from GPU memory, compute, write\nback to GPU memory, repeat. For an operation that conceptually happens\n\"once,\" your data can make several unnecessary round trips between the\nGPU's slow global memory and its fast on-chip compute units.\n\nA 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.\n\n## Kernel 1: vector add, or learning to think like a GPU\n\nThe simplest possible Triton kernel doesn't optimize anything, it just adds\ntwo vectors, `out = x + y`. The point of writing this first wasn't the\noperation, it was learning Triton's actual execution model before any\nalgorithm complexity got involved:\n\n- Triton launches a grid of programs, each identified by a `pid` (program\n  ID). Think of each program as one worker handling one chunk of the data.\n- Each program computes its own offsets into the input arrays,\n  `block_start = pid * BLOCK_SIZE` , then`tl.arange(0, BLOCK_SIZE)` for the\n  positions within its chunk.\n- A mask (`offsets < n_elements` ) handles the case where the data size isn't\n  a clean multiple of the block size. Without it, the last block would read\n  and write past the end of the array.\n\n``` python\n@triton.jit\ndef add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):\n    pid = tl.program_id(axis=0)\n    block_start = pid * BLOCK_SIZE\n    offsets = block_start + tl.arange(0, BLOCK_SIZE)\n    mask = offsets < n_elements\n\n    x = tl.load(x_ptr + offsets, mask=mask)\n    y = tl.load(y_ptr + offsets, mask=mask)\n    tl.store(out_ptr + offsets, x + y, mask=mask)\n```\n\nChecked against plain `x + y` in PyTorch: matches exactly. Nothing\ninteresting happens here algorithmically, that's the point. `pid`, offsets,\nand masking are the three ideas every later kernel in this post reuses, so\ngetting them right on the simplest possible operation first meant later bugs\n(and there was one, further down) were never about \"do I understand\nTriton,\" only ever about the specific algorithm.\n\n## Kernel 2: fused ReLU and dropout, and the problem with testing randomness\n\nThe next step up: fuse two operations into one pass instead of one.\nReLUZero out every negative value, leave positive values unchanged.\nand\ndropoutRandomly zero a fraction `p` of values, scaling the survivors by `1 / (1 - p)` to keep the expected sum unchanged.\nare two ops that would normally be two separate kernel calls in a naive\nimplementation. Fusing them means one load, one combined computation, one\nstore.\n\n``` python\n@triton.jit\ndef relu_dropout_kernel(x_ptr, out_ptr, n_elements, p, seed, BLOCK_SIZE: tl.constexpr):\n    pid = tl.program_id(axis=0)\n    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)\n    mask = offsets < n_elements\n\n    x = tl.load(x_ptr + offsets, mask=mask)\n    x = tl.maximum(x, 0.0)\n    randoms = tl.rand(seed, offsets)\n    dropout_mask = randoms > p\n\n    out = tl.where(dropout_mask, x / (1 - p), 0.0)\n    tl.store(out_ptr + offsets, out, mask=mask)\n```\n\nHere's the part worth dwelling on: you can't validate this the way you\nvalidate the vector-add kernel. `allclose` against a reference only works\nwhen the output is deterministic. Dropout is stochastic by design, there is\nno single \"correct\" output to compare against. Eyeballing five printed\nvalues, which is what I did the first time, tells you almost nothing; it's\nthe same mistake as trusting a single run in the\n[birthday-paradox post](birthday-attack.html) instead of the distribution.\n\nThe actual test has to be statistical: run it on enough elements that\nprobability becomes measurable, then check two things, does the empirical\ndrop rate land near `p`, and are the survivors scaled by exactly\n`1 / (1 - p)`?\n\n``` php\nTesting with size=100000, p=0.1  -> empirical rate 0.1007 (expected 0.10 +/- 0.01)  OK\nTesting with size=100000, p=0.5  -> empirical rate 0.4993 (expected 0.50 +/- 0.01)  OK\nTesting with size=100000, p=0.9  -> empirical rate 0.8999 (expected 0.90 +/- 0.01)  OK\n```\n\nAll 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.\n\n## Kernel 3: fused softmax, and the bug that mattered\n\nSoftmax is the first kernel here with real numerical-stability\nconsiderations: naive `exp(x)` overflows for even moderately large `x`, so\nevery real softmax implementation subtracts the row max before\nexponentiating. The kernel does this per row, in one pass:\n\n``` python\n@triton.jit\ndef softmax(output_ptr, input_ptr, input_row_stride, output_row_stride,\n            n_cols, BLOCK_SIZE: tl.constexpr):\n    row_idx = tl.program_id(0)\n    row_start_ptr = input_ptr + row_idx * input_row_stride\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    input_ptrs = row_start_ptr + col_offsets\n\n    mask = col_offsets < n_cols\n    row = tl.load(input_ptrs, mask=mask, other=float(\"-inf\"))\n    row_max = tl.max(row, axis=0)\n    numerator = tl.exp(row - row_max)\n    denominator = tl.sum(numerator, axis=0)\n    softmax_output = numerator / denominator\n\n    output_ptrs = output_ptr + row_idx * output_row_stride + col_offsets\n    tl.store(output_ptrs, softmax_output, mask=mask)\n```\n\nFirst version, tested on a small `(4, 5)` input: matched PyTorch exactly. I\nmoved on, satisfied.\n\nIt was wrong. Here's the actual bug, left in on purpose:\n\n``` python\ndef triton_softmax(x: torch.Tensor):\n    n_rows, n_cols = x.shape\n    BLOCK_SIZE = 1024   # hardcoded\n    ...\n```\n\n`BLOCK_SIZE` was a fixed number, not tied to `n_cols` at all. For a small\ninput this is harmless, `col_offsets` spans more than enough room, and the\nmask correctly zeroes out the unused tail. But nothing about the mask logic\nchecks whether `BLOCK_SIZE` is large enough, it only checks whether each\nposition is within `n_cols`. So the moment a real row is wider than 1024\ncolumns, the kernel simply never loads, computes, or stores anything past\ncolumn 1024. It doesn't crash or warn you, it just returns a softmax over\nthe first 1024 columns and calls it done.\n\nBefore, reproduced explicitly at `n_cols=2000`:\n\n```\nBuggy Softmax Matches PyTorch (n_cols=2000)?  False\nMax abs diff: 0.00978\n```\n\nAfter, the fix is one line: stop guessing a fixed block size and size it to the actual input.\n\n```\nBLOCK_SIZE = triton.next_power_of_2(n_cols)\nCorrected Softmax Matches PyTorch (n_cols=2000)?  True\n```\n\nThe lesson isn't \"remember to make BLOCK_SIZE dynamic,\" it's narrower and\nmore useful than that: a test that only covers the shape you happen to be\nthinking about will pass right over a bug that only shows up at a different\nshape. The `(4, 5)` test I ran first was real, it wasn't fake, it just\nwasn't the test that mattered.\n\n## Kernel 4: bringing it together, a self-attention kernel\n\nEverything up to here, offsets, masks, a fused numerically-stable softmax, combines into one kernel implementing scaled dot-product attention directly:\n\n``` python\n@triton.jit\ndef attention_kernel(\n    q_ptr, k_ptr, v_ptr, out_ptr,\n    row_stride, col_stride, n_rows, n_cols, scale,\n    BLOCK_Q: tl.constexpr, BLOCK_K: tl.constexpr, BLOCK_V: tl.constexpr,\n):\n    row_pid = tl.program_id(0)\n    row_offsets = row_pid * BLOCK_Q + tl.arange(0, BLOCK_Q)\n    col_offsets = tl.arange(0, BLOCK_K)\n    d_offsets = tl.arange(0, BLOCK_V)\n\n    row_mask = row_offsets < n_rows\n    col_mask = col_offsets < n_cols\n\n    q_offsets = tl.expand_dims(row_offsets, 1) * row_stride + tl.expand_dims(d_offsets, 0) * col_stride\n    q = tl.load(q_ptr + q_offsets, mask=tl.expand_dims(row_mask, 1), other=0.0)\n\n    k_offsets = tl.expand_dims(d_offsets, 1) * col_stride + tl.expand_dims(col_offsets, 0) * row_stride\n    k = tl.load(k_ptr + k_offsets, mask=tl.expand_dims(col_mask, 0), other=0.0)\n\n    qk = tl.dot(q, k) * scale\n    qk = tl.where(tl.expand_dims(col_mask, 0), qk, float(\"-inf\"))\n\n    m = tl.max(qk, axis=1)\n    p = tl.exp(qk - tl.expand_dims(m, 1))\n    l = tl.sum(p, axis=1)\n    weights = p / tl.expand_dims(l, 1)\n\n    v_offsets = tl.expand_dims(col_offsets, 1) * row_stride + tl.expand_dims(d_offsets, 0) * col_stride\n    v = tl.load(v_ptr + v_offsets, mask=tl.expand_dims(col_mask, 1), other=0.0)\n\n    out = tl.dot(weights.to(v.dtype), v)\n    out_offsets = tl.expand_dims(row_offsets, 1) * row_stride + tl.expand_dims(d_offsets, 0) * col_stride\n    tl.store(out_ptr + out_offsets, out, mask=tl.expand_dims(row_mask, 1))\n```\n\nValidated against `torch.nn.functional.scaled_dot_product_attention` across\nseveral shapes, including deliberately mismatched Q/K sequence lengths and a\nsingle-query edge case:\n\n``` php\nQ: (5, 32), K: (15, 32), V: (15, 32)   ->  matches: True\nQ: (1, 64), K: (10, 64), V: (10, 64)   ->  matches: True\n```\n\nGiven the softmax lesson above, testing more than one shape here wasn't optional, it's the only reason I can trust this one.\n\n## The wall, named honestly\n\nHere'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.\n\nThis 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.\n\nThe 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.\n\nNaming 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.\"\n\nShort 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.\n\nThe tiled, online-softmax version is the next rung. I haven't climbed it yet.", "url": "https://wpnews.pro/news/programming-an-attention-kernel-in-triton", "canonical_source": "https://sslog.dpdns.org/programming-an-attention-kernel-in-triton.html", "published_at": "2026-09-07 20:31:26.374909+00:00", "updated_at": "2026-09-07 20:31:27.767439+00:00", "lang": "en", "topics": ["developer-tools", "machine-learning", "ai-research"], "entities": ["Triton", "PyTorch", "FlashAttention"], "alternates": {"html": "https://wpnews.pro/news/programming-an-attention-kernel-in-triton", "markdown": "https://wpnews.pro/news/programming-an-attention-kernel-in-triton.md", "text": "https://wpnews.pro/news/programming-an-attention-kernel-in-triton.txt", "jsonld": "https://wpnews.pro/news/programming-an-attention-kernel-in-triton.jsonld"}}