Writing FlashAttention 2 in Triton — a step-by-step build

This is the load-bearing project of Phase 3. When you finish, you own the algorithm at the kernel level. No amount of paper reading substitutes.

Reference implementations to consult (in order):

  1. Official Triton tutorial 06 — https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html

  2. triton/python/tutorials/06-fused-attention.py in the repo

  3. Dao-AILab flash-attention csrc/flash_attn/kernels/ (CUDA reference)

  4. Umar Jamil’s walkthrough of FA2 in Triton on YouTube (search “Umar Jamil FlashAttention Triton”)

Prerequisites (do not skip):

  • Online softmax derived on paper (01_online_softmax.md)

  • FA1 paper read end-to-end (arXiv <phone_number_or_numberic_id_or_random_id_54>)

  • Triton fused softmax written from memory

  • Triton matmul from tutorial 03 understood, autotuned, and profiled

Step 0 — the scaffolding

import torch
import triton
import triton.language as tl

@triton.jit
def _fa_fwd_kernel(
    Q, K, V, O, L,               # tensors (fp16/bf16) and log-sum-exp buffer (fp32)
    stride_qb, stride_qh, stride_qm, stride_qd,
    stride_kb, stride_kh, stride_kn, stride_kd,
    stride_vb, stride_vh, stride_vn, stride_vd,
    stride_ob, stride_oh, stride_om, stride_od,
    B, H, N_Q, N_KV,             # runtime shapes
    softmax_scale,               # 1 / sqrt(d)
    IS_CAUSAL: tl.constexpr,
    BLOCK_M: tl.constexpr,       # Q tile rows
    BLOCK_N: tl.constexpr,       # KV tile rows
    HEAD_DIM: tl.constexpr,      # power-of-two, 64 or 128
):
    ...

Notes:

  • Store the log-sum-exp L per Q-row for the backward pass. Even if you’re not writing backward, do this — it’s free correctness insurance and gives you a debugging channel.

  • HEAD_DIM MUST be a constexpr. Runtime head dim in a Triton attention kernel is a footgun.

Step 1 — grid layout: one program per (batch, head, Q-tile)

grid = (triton.cdiv(N_Q, BLOCK_M), B * H, 1)   # (M_tiles, B*H, 1)

Inside the kernel:

pid_m = tl.program_id(0)              # which Q-tile
pid_bh = tl.program_id(1)             # which (batch, head)
off_b = pid_bh // H
off_h = pid_bh %  H

Why this layout: FA2’s contribution vs FA1 is exactly this — parallelize along the M (Q sequence) dimension. Program pid_m owns Q rows [pid_m*BLOCK_M, (pid_m+1)*BLOCK_M), iterates over all KV tiles.

Step 2 — compute Q, K, V base pointers

q_ptr = Q + off_b * stride_qb + off_h * stride_qh
k_ptr = K + off_b * stride_kb + off_h * stride_kh
v_ptr = V + off_b * stride_vb + off_h * stride_vh
o_ptr = O + off_b * stride_ob + off_h * stride_oh

Use tl.make_block_ptr for cleaner slicing on Hopper+ (it maps to TMA descriptors when possible):

Q_block_ptr = tl.make_block_ptr(
    base=q_ptr,
    shape=(N_Q, HEAD_DIM),
    strides=(stride_qm, stride_qd),
    offsets=(pid_m * BLOCK_M, 0),
    block_shape=(BLOCK_M, HEAD_DIM),
    order=(1, 0),
)

Repeat for K (as (HEAD_DIM, N_KV) for the QK^T contraction) and V (as (N_KV, HEAD_DIM)).

Step 3 — load Q once, initialize state

q = tl.load(Q_block_ptr)                                 # (BLOCK_M, HEAD_DIM), bf16
q = (q * softmax_scale).to(tl.bfloat16)                  # premultiply by scale

m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")   # running max
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)                  # running denominator
o_i = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)        # running output

Critical: m_i, l_i, o_i are fp32. Even though inputs/outputs are bf16, all softmax state accumulates in fp32. This is non-negotiable.

Step 4 — the main loop over KV tiles

# For causal attention: only iterate over KV tiles ≤ this Q tile
if IS_CAUSAL:
    hi = (pid_m + 1) * BLOCK_M
else:
    hi = N_KV

for start_n in range(0, hi, BLOCK_N):
    start_n = tl.multiple_of(start_n, BLOCK_N)

    k = tl.load(K_block_ptr)                            # (HEAD_DIM, BLOCK_N)
    v = tl.load(V_block_ptr)                            # (BLOCK_N, HEAD_DIM)

    # 4a) scores
    s = tl.dot(q, k)                                    # (BLOCK_M, BLOCK_N), accumulated fp32

    # 4b) mask
    if IS_CAUSAL:
        offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_n = start_n + tl.arange(0, BLOCK_N)
        s = tl.where(offs_m[:, None] >= offs_n[None, :], s, float("-inf"))

    # 4c) online softmax update
    m_new = tl.maximum(m_i, tl.max(s, axis=1))          # (BLOCK_M,)
    alpha = tl.exp(m_i - m_new)                         # rescale factor for old state
    p     = tl.exp(s - m_new[:, None])                  # unnormalized weights, (BLOCK_M, BLOCK_N)

    l_i = l_i * alpha + tl.sum(p, axis=1)
    o_i = o_i * alpha[:, None] + tl.dot(p.to(v.dtype), v)  # p in bf16 for tensor-core matmul

    m_i = m_new

    # advance block pointers
    K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
    V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))

Read this loop line by line and compare to §5 of 01_online_softmax.md. It is the exact recurrence, with (a) causal masking, (b) fp32 accumulate.

Step 5 — final normalization and store

o_i = o_i / l_i[:, None]                                # deferred divide (FA2 optimization)
L_i = m_i + tl.log(l_i)                                 # log-sum-exp for backward

O_block_ptr = tl.make_block_ptr(base=o_ptr, ...)
tl.store(O_block_ptr, o_i.to(O.dtype.element_ty))
tl.store(L_ptr + offs_m, L_i)

FA2 optimization: the divide by l_i happens ONCE at the end, not per tile. In FA1 it was per tile. This is the “reduce non-matmul FLOPs” refinement.

Step 6 — Python wrapper + autotune

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64},  num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64},  num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 128}, num_warps=8, num_stages=3),
    ],
    key=["N_Q", "N_KV", "HEAD_DIM", "IS_CAUSAL"],
)
@triton.jit
def _fa_fwd_kernel(...): ...

def flash_attention_forward(q, k, v, causal=False):
    B, H, N_Q, D = q.shape
    _, _, N_KV, _ = k.shape
    o = torch.empty_like(q)
    L = torch.empty((B, H, N_Q), dtype=torch.float32, device=q.device)
    softmax_scale = 1.0 / (D ** 0.5)
    grid = lambda META: (triton.cdiv(N_Q, META["BLOCK_M"]), B * H, 1)
    _fa_fwd_kernel[grid](q, k, v, o, L, ...strides..., B, H, N_Q, N_KV,
                         softmax_scale, causal, HEAD_DIM=D)
    return o, L

The bug catalog — every mistake I know about

Correctness bugs (found by comparing to torch.nn.functional.scaled_dot_product_attention):

  1. Forgot to accumulate in fp32. Output looks right at small N; drifts at N > 4k. Fix: m_i, l_i, o_i are fp32 regardless of input dtype.

  2. Applied softmax scale after tl.dot instead of premultiplying Q. Numerically equivalent in fp32, but at bf16 you lose bits. Premultiply Q by 1/sqrt(d) before the loop.

  3. Masked with 0 instead of -inf. After exp(0 - m_new) you get 1, not 0. Correct mask value is float("-inf").

  4. Wrong causal comparison direction. offs_m >= offs_n means Q-row m attends to K-col n if m >= n. Get it wrong and you either look ahead (bad) or exclude the diagonal (very bad).

  5. Update order in the recurrence. You must compute m_new first, then alpha = exp(m_i - m_new), then apply to l_i and o_i, THEN assign m_i = m_new. Overwrite m_i too early and everything is silently wrong.

  6. p.to(v.dtype) missing. tl.dot on Hopper wants both operands in the same low-precision dtype for tensor cores. If p stays fp32 you fall off tensor cores and lose ~10× perf.

  7. Non-power-of-2 head_dim. Triton’s dot requires HEAD_DIM to be a compile-time power of two. Pad K/V to the nearest power of two or route non-standard head dims to SDPA.

  8. Fully-masked row → NaN. If a causal row has zero valid keys (shouldn’t happen with causal + Q index ≥ 0, but happens with sliding window at position 0), l_i = 0 → divide by zero. Guard: o_i = o_i / tl.maximum(l_i, 1e-6)[:, None].

  9. Stride confusion between Q/K when N_Q ≠ N_KV. Cross-attention has different Q and KV sequence lengths. Two separate strides, two separate offset variables.

  10. Autotune cache poisoning. After changing kernel logic during development, wipe ~/.triton/cache/. Otherwise you’re running an old compiled version and going insane.

Performance bugs:

  1. num_warps too high. 8 is aggressive; 4 is often faster at small BLOCK_M. Autotune.

  2. num_stages too high on small SMEM. More pipeline stages = more SMEM pressure. On Ampere 3 is a good start; Hopper handles 4–5 with TMA.

  3. BLOCK_N = 128 at HEAD_DIM = 128 may spill because S = BLOCK_M × BLOCK_N = 128×128 fp32 = 64 KB. Watch for spill in Nsight (stall_reason.long_scoreboard on local memory).

  4. Not using tl.multiple_of. Tells the compiler about alignment; drops guard instructions inside the loop.

Numerics verification harness

def test_flash_attention():
    torch.manual_seed(0)
    B, H, N, D = 2, 16, 4096, 64
    q = torch.randn(B, H, N, D, dtype=torch.bfloat16, device="cuda")
    k = torch.randn(B, H, N, D, dtype=torch.bfloat16, device="cuda")
    v = torch.randn(B, H, N, D, dtype=torch.bfloat16, device="cuda")

    o_ref = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
    o_ours, _ = flash_attention_forward(q, k, v, causal=True)

    torch.testing.assert_close(o_ours, o_ref, atol=1e-2, rtol=1e-2)
    print("OK")

    # Run across shape sweep too
    for shape in [(1,8,512,64), (1,8,2048,64), (4,32,4096,128),
                  (1,8,8192,128), (2,16,16384,64)]:
        ...

Tolerances: bf16 with fp32 accumulate → atol=1e-2, rtol=1e-2. If you can’t hit that, you have a bug. If you’re comfortably inside 1e-3 for a specific shape, your accumulator or scale ordering is probably fine.

Profiling checklist (Nsight Compute, --set full)

  • SM Utilization should be ≥ 70% at (B=2, H=16, N=4096, D=64). If < 40%, launch config is wrong or loop is memory-bound.

  • Tensor Core Utilization should be ≥ 50%. If 0%, you didn’t cast p back to bf16 before tl.dot.

  • Long Scoreboard Stalls should be low (< 20%). If high, K/V loads are not overlapped — check num_stages.

  • Shared Memory usage — if spilling to local memory, reduce BLOCK_M × BLOCK_N.

  • Compare to torch.nn.functional.scaled_dot_product_attention(..., is_causal=True) timings. Target: within 2× of SDPA’s flash backend on H100 for causal bf16 at N=4096. If you land within 1.5×, you’ve written a real FA2.

What to do after your FA2 works

  1. Backward pass. Use the stored L (log-sum-exp) tensor. Backward has its own tile schedule and its own bugs. Do this only if time permits — a working forward is the priority.

  2. Read FA3. The Hopper primitives (warp-spec + wgmma + pingpong) will now feel concrete because you know what you’re accelerating.

  3. Sliding-window mask. Add a WINDOW_SIZE constexpr, change the causal mask formula. Half a day.

  4. ALiBi bias. One-line change in the score computation. Understand where in the loop you’d add it.

  5. Contribute to KernelBench. Your FA2 forward is a legitimate submission for the attention task category.

References