FlashAttention-2 from PyTorch to Triton
FlashAttention-2 from PyTorch to Triton
This article is for readers who are familiar with PyTorch, and want to learn kernel optimisation. I use Triton as it is my gateway domain-specific language (DSL) into the world of kernel optimisation, and the Stanford CS336 assignment 2 scaffold which makes some assumptions to simplify the implementation of FlashAttention and focus on key concepts in Triton.
Assumptions:
- Fixing the batch size to 8 to avoid the need for multihead attention
- Hard coding to 16×16 tiles
- Future articles will relax these assumptions to learn these implementation details
I will state upfront whenever I made simplifying assumptions. All numbers in this series are from a single RTX 4070 Super (Ada, sm_89, 12 GB, roughly 100 KB of shared memory per SM) under Linux. This article is written with the assistance of AI.
What We Are Building
The FlashAttention-2 (FA2) forward pass: for each query tile Q_i, iterate over key and value tiles K_j and V_j, rescaling the running output O_i on every step and dividing by l only once at the end (FA1 re-normalised at every step). The outer loop goes through the query tiles, where on the GPU, each query tile gets its own program, and all the programs run in parallel.
The kernel computes standard scaled-dot-product attention without materialising the N×N score matrix. For a query tile Q_i of B_q rows and a key tile K_j of B_k rows, it computes the B_q × B_k score block, folds it into a running softmax, and moves to the next key tile. When all key tiles have been visited, the output rows for Q_i are complete and are written once.
Three running quantities per query tile make that possible, initialised on line 6:
m_i- the running maximum score for each row (line 10)l_i- the running row sum ofexp(S_i - m_i)(line 12)O_i- the running sum ofexp(S_i - m_i)times the value rows, kept unnormalised (line 13)
The superscript (j) marks the value after the j-th key tile. On every step, the previous l_i^(j-1) and O_i^(j-1) are rescaled by exp(m_i^(j-1) - m_i^(j)), which is 1 when the maximum is unchanged, before the new tile's contribution is added. This is the online softmax, which is applied one tile at a time.
At the end of each query tile, the kernel also writes a second output L, which is computed on line 16 and stored on line 18. L is the per-row logsumexp of the scaled scores: L_i = m_i + log(l_i), which equals log Σ_k exp(S_ik) over the whole row.
The backward pass will use it to recompute the softmax probabilities tile by tile as P = exp(S - L), since exp(S - m) / l = exp(S - m - log(l)) = exp(S - L). Storing L instead of P is what lets FA's backward avoid the N×N matrix too.
The backward will be covered in a later article. For now L is computed, stored, and tested, but not used.
Step 1: The Tiled Forward in PyTorch
Before diving into Triton, I coded FA2 with PyTorch to familiarise with the algorithm before introducing Triton concepts. The PyTorch implementation is slow by design because it has a Python double loop over tiles, and launches a handful of small CUDA ops per iteration. The goal here is to be correct, not fast.
# flashattention_autograd_function_pytorch.py
import math
import torch
import einops
class FlashAttentionPytorch(torch.autograd.Function):
@staticmethod
def forward(ctx, Q, K, V, is_causal=False):
# Note: Tile size is fixed at 16 as a simplifying assumption
tile_size = 16
# Split the sequence dimension into tiles: (..., T, B, d)
# Leading dims are arbitrary (batch, heads, ...). The sequence axis N is split
# into Tq tiles of Bq rows, so N must be a multiple of tile_size.
Q_t = einops.rearrange(Q, "... (Tq Bq) d -> ... Tq Bq d", Bq=tile_size)
K_t = einops.rearrange(K, "... (Tk Bk) d -> ... Tk Bk d", Bk=tile_size)
V_t = einops.rearrange(V, "... (Tk Bv) d -> ... Tk Bv d", Bv=tile_size)
O = torch.empty_like(Q)
L = torch.empty(Q.shape[:-1], device=Q.device, dtype=Q.dtype)
scale = 1.0 / math.sqrt(Q.shape[-1])
for i in range(Q_t.shape[-3]): # outer loop: query tiles
Q_i = Q_t[..., i, :, :]
O_i = torch.zeros_like(Q_i)
l_i = torch.zeros(Q_i.shape[:-1] + (1,), device=Q.device, dtype=Q.dtype)
m_i = torch.full(Q_i.shape[:-1] + (1,), -torch.inf, device=Q.device, dtype=Q.dtype)
for j in range(K_t.shape[-3]): # inner loop: key tiles
K_j = K_t[..., j, :, :]
V_j = V_t[..., j, :, :]
S_ij = einops.einsum(Q_i, K_j, "... Bq d, ... Bk d -> ... Bq Bk") * scale
m_new = torch.maximum(m_i, S_ij.amax(dim=-1, keepdim=True))
P_ij = torch.exp(S_ij - m_new)
alpha = torch.exp(m_i - m_new) # rescale factor for the old state
l_i = alpha * l_i + P_ij.sum(dim=-1, keepdim=True)
O_i = alpha * O_i + einops.einsum(P_ij, V_j, "... Bq Bk, ... Bk d -> ... Bq d")
m_i = m_new
O[..., i * tile_size:(i + 1) * tile_size, :] = O_i / l_i
L[..., i * tile_size:(i + 1) * tile_size] = (m_i + torch.log(l_i)).squeeze(-1)
ctx.save_for_backward(Q, K, V, O, L)
ctx.is_causal = is_causal
return O, L
A few points to note:
m_istarts at negative infinity so that the first tile's row maximum "wins" in line 10 of algorithm 1.m_istarting at negative infinity also ensures on the first stepalpha = exp(-inf - m_new)is exactly 0, not NaN, so that the rescale in lines 12 and 13 runs unchanged and there is no special case to handle for the first iteration.O_iis the weighted sum of value rows, and the softmax needs that sum divided by the total weightl_i. There are two ways to book-keep: (i) keep the sum and divide once at the end (FA2); or (ii) keep the average and re-divide every time a new tile arrives (FA1). FA2's approach saves two elementwise passes over aB_q × dtile on every inner loop.Lis stored, notmandlseparately. The backward pass recomputes each tile ofSfromQandK, then recovers that tile's probabilities asexp(S - L).
For context, assuming N = 4096, S with dimension N×N per head per batch element will have about 16.8 million entries, or about 33 MB in bf16. L with dimension N×1 per head per batch element contains 4096 entries, or about 16 KB.
The is_causal flag is accepted and ignored. The PyTorch version does not implement causal masking.
Step 2: The Triton Forward Kernel
The Triton version has the same shape as the PyTorch one, with the outer loop over query tiles replaced by the launch grid. Each program (Triton's name for a thread block) gets one query tile and one batch element, loads Q_i once, and loops over key tiles on its own.
Before I introduce the implementation of FA2 in Triton, I want to take a detour into introducing tl.make_block_ptr.
A Quick Intro to tl.make_block_ptr
We will be taking apart tl.make_block_ptr in this section, using the weighted_sum_fwd kernel as an example. See code below.
x_block_ptr = tl.make_block_ptr(
x_ptr,
shape=(NUM_ROWS, D),
strides=(x_stride_row, x_stride_dim),
offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),
block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
order=(1, 0),
)
# then, inside the loop:
row = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")
x_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE))
There are a total of 6 arguments in tl.make_block_ptr, illustrated in the image below. I will go through each one by one. The illustrations are drawn with Claude Opus 5.5.
Why Use a Block Pointer?
A Triton program works on a small tile of data at a time. To read a tile, it needs the memory address of every element in the tile. It also must not read past the edges of the tensor. There are two ways to do this.
The Classic Way with Pointer Math
Before block pointers, programmers would build the addresses directly, including the official Triton tutorials. Below is the weighted_sum_fwd kernel written the classical way.
row_tile_idx = tl.program_id(0) # 1. Which rows this program owns
rows = row_tile_idx * ROWS_TILE_SIZE + tl.arange(0, ROWS_TILE_SIZE)
row_mask = rows < NUM_ROWS
output = tl.zeros((ROWS_TILE_SIZE,), dtype=tl.float32)
for i in range(tl.cdiv(D, D_TILE_SIZE)):
# 2. Which columns this step covers
cols = i * D_TILE_SIZE + tl.arange(0, D_TILE_SIZE)
col_mask = cols < D
# 3. A 2D grid of addresses, built by broadcasting
x_ptrs = (x_ptr + rows[:, None] * x_stride_row + cols[None, :] * x_stride_dim)
w_ptrs = weight_ptr + cols * weight_stride_dim
# 4. Masks for the edges, combined by hand
row = tl.load(x_ptrs, mask=row_mask[:, None] & col_mask[None, :], other=0.0)
weight = tl.load(w_ptrs, mask=col_mask, other=0.0)
output += tl.sum(row * weight[None, :], axis=1)
tl.store(output_ptr + rows * output_stride_row, output, mask=row_mask)
tl.load takes a whole grid of addresses, one for each element in the tile. The kernel makes that grid from two index vectors: rows goes down and cols goes across. Writing [:, None] turns rows into a column, and [None, :] turns cols into a row. Adding them fills in the grid by broadcasting. The mask is built the same way.
Back to Triton FA2
Below is the kernel as it stands at the end of this post; the two lines marked FIX are the ones Step 3 explains.
# flashattention_autograd_function_triton.py
import math
import torch
import triton
import triton.language as tl
@triton.jit
def flash_fwd_kernel(
Q_ptr, K_ptr, V_ptr, O_ptr, L_ptr,
stride_qb, stride_qq, stride_qd,
stride_kb, stride_kk, stride_kd,
stride_vb, stride_vk, stride_vd,
stride_ob, stride_oq, stride_od,
stride_lb, stride_lq,
N_QUERIES, N_KEYS, scale,
D: tl.constexpr,
Q_TILE_SIZE: tl.constexpr,
K_TILE_SIZE: tl.constexpr,
is_causal: tl.constexpr,
):
query_tile_index = tl.program_id(0)
batch_index = tl.program_id(1)
# Block pointers: a (rows, D) window into each tensor for this batch element.
# Q and O windows start at this program's query tile; K and V start at row 0
# and are advanced inside the loop.
Q_block_ptr = tl.make_block_ptr(
Q_ptr + batch_index * stride_qb,
shape=(N_QUERIES, D),
strides=(stride_qq, stride_qd),
offsets=(query_tile_index * Q_TILE_SIZE, 0),
block_shape=(Q_TILE_SIZE, D),
order=(1, 0),
)
K_block_ptr = tl.make_block_ptr(
K_ptr + batch_index * stride_kb,
shape=(N_KEYS, D),
strides=(stride_kk, stride_kd),
offsets=(0, 0),
block_shape=(K_TILE_SIZE, D),
order=(1, 0),
)
V_block_ptr = tl.make_block_ptr(
V_ptr + batch_index * stride_vb,
shape=(N_KEYS, D),
strides=(stride_vk, stride_vd),
offsets=(0, 0),
block_shape=(K_TILE_SIZE, D),
order=(1, 0),
)
O_block_ptr = tl.make_block_ptr(
O_ptr + batch_index * stride_ob,
shape=(N_QUERIES, D),
strides=(stride_oq, stride_od),
offsets=(query_tile_index * Q_TILE_SIZE, 0),
block_shape=(Q_TILE_SIZE, D),
order=(1, 0),
)
L_block_ptr = tl.make_block_ptr(
L_ptr + batch_index * stride_lb,
shape=(N_QUERIES,),
strides=(stride_lq,),
offsets=(query_tile_index * Q_TILE_SIZE,),
block_shape=(Q_TILE_SIZE,),
order=(0,),
)
# Running state, kept in fp32 regardless of input dtype.
O_acc = tl.zeros((Q_TILE_SIZE, D), dtype=tl.float32)
l_acc = tl.zeros((Q_TILE_SIZE, 1), dtype=tl.float32)
m_acc = tl.full((Q_TILE_SIZE, 1), value=float("-inf"), dtype=tl.float32)
Q_i = tl.load(Q_block_ptr, boundary_check=(0, 1), padding_option="zero")
q_pos = (query_tile_index * Q_TILE_SIZE + tl.arange(0, Q_TILE_SIZE))[:, None]
for j in range(tl.cdiv(N_KEYS, K_TILE_SIZE)):
k_pos = (j * K_TILE_SIZE + tl.arange(0, K_TILE_SIZE))[None, :]
K_j = tl.load(K_block_ptr, boundary_check=(0, 1), padding_option="zero")
V_j = tl.load(V_block_ptr, boundary_check=(0, 1), padding_option="zero")
S_ij = tl.dot(Q_i, tl.trans(K_j)) * scale # (Q_TILE, K_TILE), fp32
# FIX 2: zero-padded keys past N_KEYS score 0, not -inf. Mask them.
keep = k_pos < N_KEYS
if is_causal:
keep = keep & (k_pos <= q_pos)
S_ij = tl.where(keep, S_ij, -1e6)
m_new = tl.maximum(m_acc, tl.max(S_ij, axis=1, keep_dims=True))
P_ij = tl.exp(S_ij - m_new)
alpha = tl.exp(m_acc - m_new)
l_acc = alpha * l_acc + tl.sum(P_ij, axis=1, keep_dims=True)
# FIX 1: the cast must be assigned. tl.dot needs both operands in the
# same dtype; the fp32 accumulator is passed separately via acc=.
P_ij = P_ij.to(V_j.dtype)
O_acc = alpha * O_acc
O_acc = tl.dot(P_ij, V_j, acc=O_acc)
m_acc = m_new
K_block_ptr = K_block_ptr.advance((K_TILE_SIZE, 0))
V_block_ptr = V_block_ptr.advance((K_TILE_SIZE, 0))
O_i = (O_acc / l_acc).to(O_block_ptr.type.element_ty)
tl.store(O_block_ptr, O_i, boundary_check=(0, 1))
L_i = tl.reshape(m_acc + tl.log(l_acc), (Q_TILE_SIZE,))
tl.store(L_block_ptr, L_i, boundary_check=(0,))
class FlashAttentionTriton(torch.autograd.Function):
Q_TILE_SIZE = 16
K_TILE_SIZE = 16
@staticmethod
def forward(ctx, Q, K, V, is_causal=False):
assert Q.ndim == 3, "expects (batch, seq, head_dim); flatten (B, H, N, D) to (B*H, N, D)"
assert Q.stride(-1) == 1 and K.stride(-1) == 1 and V.stride(-1) == 1
B, N_q, D = Q.shape
N_k = K.shape[1]
assert D in (16, 32, 64, 128), "block_shape dims must be powers of two"
O = torch.empty_like(Q)
L = torch.empty((B, N_q), device=Q.device, dtype=torch.float32)
grid = (triton.cdiv(N_q, FlashAttentionTriton.Q_TILE_SIZE), B)
flash_fwd_kernel[grid](
Q, K, V, O, L,
Q.stride(0), Q.stride(1), Q.stride(2),
K.stride(0), K.stride(1), K.stride(2),
V.stride(0), V.stride(1), V.stride(2),
O.stride(0), O.stride(1), O.stride(2),
L.stride(0), L.stride(1),
N_q, N_k,
1.0 / math.sqrt(D),
D=D, Q_TILE_SIZE=Q_TILE_SIZE,
The article text ends here, with the kernel launch call appearing to be cut off mid-line.
Comments
No comments yet. Start the discussion.