Programming an attention kernel in Triton

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.

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.