FlashAttention-style kernel

The loop

Tile Q once, then stream K and V:

q = desc_q.load([qo_offset_y, 0])          # BLOCK_M x HEAD_DIM, resident
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
for start_n in tl.range(lo, hi, BLOCK_N):
    k = desc_k.load([offset_y + start_n, 0]).T
    qk = tl.dot(q, k) * qk_scale
    m_new = tl.maximum(m_i, tl.max(qk, 1))
    alpha = tl.exp2(m_i - m_new)
    p = tl.exp2(qk - m_new[:, None])
    l_i = l_i * alpha + tl.sum(p, 1)
    acc = acc * alpha[:, None] + tl.dot(p.to(tl.float16), desc_v.load([offset_y + start_n, 0]))
    m_i = m_new
acc = acc / l_i[:, None]

This is the tutorial inner loop, with its names. qk_scale is sm_scale * 1.44269504, so the base-two exponential does the work of exp and the conversion is paid once per score.

What is not there

There is no N x N tensor. The score tile is BLOCK_M x BLOCK_N, it is consumed by the two dots and the row sum, and it is dead before the next iteration. What stays live is the Q tile, the K and V tiles and the accumulator, so memory is linear in sequence length where the materialised version is quadratic. The paper calls this IO-awareness: exact attention, with fewer bytes crossing the memory hierarchy. FlashAttention-2 is the same recurrence, partitioned differently across warps and blocks.

Causal masking

The tutorial splits the loop into an off-band part and a diagonal part so fully masked tiles are skipped, and it prunes configurations with BLOCK_M < BLOCK_N for causal kernels. A tile with no real maximum is not harmless: the -1e6 mask sentinel becomes the maximum, and every masked lane then exponentiates to one, which inflates the running sum.

Things that break

Forgetting acc * alpha gives a kernel that is correct while the context fits in one tile and quietly wrong beyond it, which is the failure a short test cannot see. A K descriptor loaded with the wrong block shape transposes without complaint and produces plausible numbers. And a base-two kernel whose alpha is computed with base e mixes two scales, so the output is off by a smooth factor rather than obviously broken. These snippets follow Triton 3.9; check tl.exp2 and the descriptor API against your installed version, since block pointers were removed in 3.9 and the descriptors replaced them.