Programming an attention kernel in Triton

September 08, 2026 gpu programming, triton, machine learning source ↗

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, 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.

@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.

@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 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)?

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:

@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:

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.

n_cols 512

BLOCK_SIZE 1024 columns dropped 0 matches PyTorch? True

Drag n_cols past 1024 in "hardcoded" mode and the bar stops filling exactly where the bug does: everything past the line is real data the kernel never reads, writes, or warns about. Switch to "dynamic" and the same drag never breaks anything, because BLOCK_SIZE grows with n_cols instead of sitting fixed.

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:

@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:

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.