4.4 Writing GPU Kernels with Triton¶

There is a moment in every LLM engineer’s life when torch.compile is not enough, when fusing the operation you want by hand in CUDA C++ would take a week, and when the kernel you need does not yet exist in any library. That is the moment Triton was built for. Triton, originally created by Philippe Tillet and developed at OpenAI, is a Python-embedded language and compiler for writing GPU kernels; it now lives in the vendor-neutral triton-lang GitHub organization, with NVIDIA and AMD backends in-tree and an Intel GPU backend maintained alongside, so most of the kernels in this chapter compile and run unmodified on an MI300X as well as on an H100 (what changes across vendors is the tuning, not the source). You write something that looks like NumPy operating on blocks of data; Triton’s compiler turns it into PTX (NVIDIA’s assembly) that, for many bandwidth-bound and moderately compute-bound kernels, lands within a few percent of expertly hand-tuned CUDA — at a tenth of the development cost.

This is the most hands-on chapter in Part IV. We will build, from nothing, four kernels of escalating difficulty: a vector add (to learn the programming model), a fused softmax (to learn reductions and row-wise work), a matmul (to learn tiling, accumulation, and the L2-cache swizzle), and a simplified FlashAttention (to tie it all together with the online softmax). Every kernel here is runnable. By the end you should be able to read the real FlashAttention and fused_moe Triton kernels in vLLM and understand every line.

This chapter assumes you have internalized GPU Architecture & The Memory Hierarchy (SMs, warps, shared memory, HBM vs SRAM) and The Roofline Model & Performance Engineering (arithmetic intensity, memory- vs compute-bound). The softmax and attention kernels here are the implementation of the IO-aware ideas from FlashAttention I: IO-Awareness & The Online Softmax.

Why Triton Exists: The Abstraction Gap¶

To appreciate Triton you have to feel the pain it removes. In raw CUDA, you are responsible for the full hierarchy of parallelism: you decide the grid of thread blocks, the threads within a block, how those threads cooperate through shared memory, how to coalesce global-memory loads so that the 32 threads of a warp touch contiguous addresses, how to avoid shared-memory bank conflicts, and how to feed the Tensor Cores with the exact fragment layout mma instructions demand. A competent CUDA matmul is several hundred lines; a competitive one is thousands.

Triton raises the abstraction by one crucial level: the program operates on blocks (tiles), not scalars. A Triton “program” (the equivalent of a CUDA thread block) loads a block of data, does block-level arithmetic, and stores a block of results. Inside that block, the compiler decides the thread-to-data mapping, the shared-memory staging, the vectorization, and the Tensor Core fragment layout. You think about which tile of the output this program computes; the compiler thinks about how 1024 threads cooperate to compute it.

The trade is control for productivity. You give up the ability to micro-place every byte (which is why the very best vendor kernels — cuBLAS, CUTLASS, cuDNN, and the FP8 FlashAttention-3 kernels of FlashAttention 2 & 3 — are still hand-written), but you gain the ability to write a fused, correct, autotuned kernel for your exact problem in an afternoon. For the long tail of custom operations in LLM training and serving — fused RMSNorm, custom RoPE, MoE routing, dequantize-and-matmul — Triton is the default tool. We compare it to writing raw CUDA in CUDA Programming Essentials and to the compiler-driven path in Kernel Fusion, torch.compile, CUDA Graphs & Compilers (indeed, torch.compile’s GPU backend, TorchInductor, generates Triton).

The Triton Programming Model¶

Five concepts carry the entire language. Learn these and the rest is detail.

1. @triton.jit — the decorator that marks a Python function as a kernel to be JIT-compiled to GPU code. Inside it, you may only use Triton operations and a restricted Python subset (no arbitrary objects, no lists of tensors; control flow and arithmetic on tl values).
2. The launch grid — when you call kernel[grid](...), grid is a tuple (or a function returning one) giving how many independent program instances to launch. This is the CUDA grid.
3. tl.program_id(axis) — inside the kernel, this returns which program instance you are, along a given axis (0, 1, or 2). It is how each program figures out which slice of data it owns. This is the analog of blockIdx.
4. Pointers and tl.arange — Triton works with raw pointers into tensors. You compute a block of pointers (a vector/tensor of addresses) by adding offsets to a base pointer, then tl.load / tl.store that whole block at once.
5. Masks — because tensor dimensions are rarely exact multiples of your block size, you pass a boolean mask to tl.load/tl.store so that out-of-bounds lanes are skipped (and optionally given an other fill value). Masks are how Triton stays correct at the ragged edges.

The Anatomy of a Pointer Computation¶

The single most important skill in Triton is turning “the elements I want” into “a block of pointers to them.” A PyTorch tensor handed to a Triton kernel decays to a pointer to its first element (plus you usually pass its .stride() values explicitly). Suppose x_ptr points at a 1-D tensor and this program is responsible for the contiguous chunk starting at element block_start. Then:

# offsets is a *vector* of BLOCK_SIZE element indices owned by this program
offsets = block_start + tl.arange(0, BLOCK_SIZE)   # e.g. [1024, 1025, ..., 2047]
# ptrs is a *vector* of BLOCK_SIZE addresses (pointer arithmetic broadcasts)
ptrs = x_ptr + offsets
# mask keeps lanes that are still inside the tensor of length n_elements
mask = offsets < n_elements
# one vectorized, coalesced load of up to BLOCK_SIZE elements
block = tl.load(ptrs, mask=mask, other=0.0)

For 2-D tensors you build a 2-D block of pointers by broadcasting two arange vectors against the row and column strides:

# A BLOCK_M x BLOCK_N tile of pointers into a matrix with strides (stride_m, stride_n)
row = tl.arange(0, BLOCK_M)[:, None]   # shape (BLOCK_M, 1)
col = tl.arange(0, BLOCK_N)[None, :]   # shape (1, BLOCK_N)
ptrs = base_ptr + row * stride_m + col * stride_n   # shape (BLOCK_M, BLOCK_N)

This [:, None] / [None, :] broadcasting (identical to NumPy) is the workhorse. Internalize it: rows down, columns across, strides convert logical indices to memory offsets. BLOCK_M, BLOCK_N, and BLOCK_SIZE are compile-time constants (declared tl.constexpr), which lets the compiler unroll loops, size registers, and pick vector widths. Newer Triton also offers tl.make_block_ptr, a structured block-pointer object that tracks shape, strides, and offsets for you; we use explicit pointer math here because it makes the mechanism visible.

Kernel 1: Vector Add — The “Hello World”¶

The whole machine in one screen. We add two length-n vectors. Each program handles BLOCK_SIZE consecutive elements; the grid has ceil(n / BLOCK_SIZE) programs.

import torch
import triton
import triton.language as tl


@triton.jit
def add_kernel(
    x_ptr,            # *Pointer* to first input vector
    y_ptr,            # *Pointer* to second input vector
    out_ptr,          # *Pointer* to output vector
    n_elements,       # Size of the vectors (a runtime int)
    BLOCK_SIZE: tl.constexpr,   # Elements per program — compile-time constant
):
    # 1. Which program am I? There is a 1-D grid, so we read axis 0.
    pid = tl.program_id(axis=0)

    # 2. Compute the slice of the output this program owns.
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)   # vector of indices

    # 3. Guard against the tail block (n may not divide BLOCK_SIZE).
    mask = offsets < n_elements

    # 4. Load BLOCK_SIZE elements of x and y from HBM (masked, coalesced).
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)

    # 5. The actual compute — elementwise add, entirely in registers.
    out = x + y

    # 6. Write the result back to HBM.
    tl.store(out_ptr + offsets, out, mask=mask)


def triton_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    assert x.is_cuda and y.is_cuda and out.is_cuda
    n_elements = out.numel()

    # The grid is a function of META so autotuning can change BLOCK_SIZE.
    # triton.cdiv(a, b) == ceil(a / b).
    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)

    # Launch. The [grid] indexing syntax enqueues the kernel on the GPU.
    add_kernel[grid](x, y, out, n_elements, BLOCK_SIZE=1024)
    return out


if __name__ == "__main__":
    torch.manual_seed(0)
    x = torch.randn(98_432, device="cuda")
    y = torch.randn(98_432, device="cuda")
    out_triton = triton_add(x, y)
    out_torch = x + y
    max_err = (out_triton - out_torch).abs().max().item()
    print(f"max abs error vs torch: {max_err:.3e}")   # ~0.0

Note that 98_432 is not a multiple of 1024 (it is 96 * 1024 + 128), so the last of the 97 programs has 896 masked-off lanes. The kernel is still correct because tl.store’s mask prevents those lanes from writing past the end of out. This is bandwidth-bound: we move 3 floats per element (read x, read y, write out) and do one add, an arithmetic intensity of \(1/12\) FLOP/byte in fp32 — far to the left of the roofline ridge, so the kernel runs at HBM bandwidth and there is nothing clever to do. Triton’s whole value here is that it took 30 lines.

Kernel 2: Fused Softmax — Reductions and the Memory Win¶

Softmax over the last dimension of an M x N matrix is the first kernel where fusion pays off. The naive PyTorch path materializes several intermediate M x N tensors in HBM:

A library implementation reads x (to find the max), reads x again (to exponentiate), writes the exponentials, reads them again (to sum), and reads once more to divide — roughly 5 passes over HBM. The max-subtraction is not optional: it is the standard numerically stable softmax that prevents exp from overflowing for large logits (see Numerical Computing, Floating Point & Precision).

The fused Triton kernel assigns one program per row. The program loads the entire row into registers/SRAM once, computes the max, the exponentials, and the sum without ever round-tripping the intermediates to HBM, and writes the row once. That is 1 read + 1 write — a 2.5x reduction in memory traffic, and since softmax is memory-bound, roughly a 2.5x speedup.

import torch
import triton
import triton.language as tl


@triton.jit
def softmax_kernel(
    out_ptr, in_ptr,
    in_row_stride, out_row_stride,   # how many elements to step to go down one row
    n_cols,                          # number of columns (runtime)
    BLOCK_SIZE: tl.constexpr,        # padded power-of-two >= n_cols
):
    # One program per row. program_id(0) is the row index.
    row_idx = tl.program_id(axis=0)

    # Pointer to the start of this row, then a vector covering all columns.
    row_start = in_ptr + row_idx * in_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE)
    in_ptrs = row_start + col_offsets

    # Load the whole row. For padding lanes (col >= n_cols) we load -inf so
    # they never become the max and contribute exp(-inf) = 0 to the sum.
    mask = col_offsets < n_cols
    row = tl.load(in_ptrs, mask=mask, other=-float("inf"))

    # --- Numerically stable softmax, entirely on-chip ---
    row_max = tl.max(row, axis=0)            # block-level reduction -> scalar
    row = row - row_max                      # subtract max for stability
    numerator = tl.exp(row)                  # exp of every (real) element
    denominator = tl.sum(numerator, axis=0)  # reduction -> scalar
    out = numerator / denominator

    # Write the row back once.
    out_row_start = out_ptr + row_idx * out_row_stride
    tl.store(out_row_start + col_offsets, out, mask=mask)


def triton_softmax(x: torch.Tensor) -> torch.Tensor:
    assert x.dim() == 2 and x.is_cuda
    M, N = x.shape
    # BLOCK_SIZE must be a power of two and cover a full row, so the row
    # fits in one program. (This simple version requires N <= ~64K.)
    BLOCK_SIZE = triton.next_power_of_2(N)
    # More warps for wider rows => more parallel reduction throughput.
    num_warps = 4
    if BLOCK_SIZE >= 2048:
        num_warps = 8
    if BLOCK_SIZE >= 4096:
        num_warps = 16

    out = torch.empty_like(x)
    # One program per row => grid is (M,).
    softmax_kernel[(M,)](
        out, x,
        x.stride(0), out.stride(0),
        N,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
    )
    return out


if __name__ == "__main__":
    x = torch.randn(1823, 781, device="cuda")
    ours = triton_softmax(x)
    ref = torch.softmax(x, dim=1)
    print("max abs error:", (ours - ref).abs().max().item())   # ~1e-7

Two subtleties worth dwelling on. First, tl.max and tl.sum with axis=0 are block reductions: the compiler lowers them to a tree reduction across the warps and threads cooperating on the row, including a shared-memory shuffle stage. You write one line; the compiler emits the log-depth reduction. Second, the other=-inf fill is the elegant trick that lets the mask handle ragged column counts and the numerical stability simultaneously — padded lanes are -inf, so they lose the max comparison and their exp is exactly 0, contributing nothing to the denominator.

This “load row, reduce on chip, write row” pattern is the exact same idea that, scaled up to 2-D tiles that don’t fit in SRAM, becomes FlashAttention. Keep it in mind; we return to it in Kernel 4.

Kernel 3: Matrix Multiplication — Tiling, Accumulation, and Swizzling¶

Matmul is where Triton shows its compute-bound chops, because it is the one operation that should live on the right side of the roofline. We compute \(C = A B\) where \(A\) is \(M\times K\), \(B\) is \(K\times N\), and \(C\) is \(M\times N\). The arithmetic is \(2MNK\) FLOPs against, in the tiled scheme, far fewer bytes — high arithmetic intensity, so a good matmul saturates the Tensor Cores.

The strategy is classic tiling: each program computes one BLOCK_M x BLOCK_N tile of C. To do so it walks the shared K dimension in steps of BLOCK_K, at each step loading a BLOCK_M x BLOCK_K tile of A and a BLOCK_K x BLOCK_N tile of B, multiplying them with tl.dot (which targets the Tensor Cores), and accumulating into an fp32 register tile. Only after the full K loop does it write the tile to HBM. Each element of A and B is thus loaded from HBM once per output tile but reused BLOCK_N and BLOCK_M times respectively from SRAM — that reuse is the entire point.

import torch
import triton
import triton.language as tl


@triton.autotune(
    # Triton benchmarks each config on the first call for a given problem
    # shape and caches the winner. These are reasonable A100/H100 starts.
    configs=[
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64, "GROUP_M": 8},
                      num_stages=3, num_warps=8),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8},
                      num_stages=4, num_warps=4),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8},
                      num_stages=4, num_warps=4),
    ],
    key=["M", "N", "K"],   # re-autotune when these change
)
@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,    # A strides: row, col
    stride_bk, stride_bn,    # B strides: row, col
    stride_cm, stride_cn,    # C strides: row, col
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    # ---- 1. Which output tile do I compute? (with L2-cache swizzle) ----
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    # Group rows of tiles so that programs running concurrently reuse the
    # same B columns / A rows in L2. This "swizzle" is the single biggest
    # perf lever after tiling. Without it, tiles are visited row-major and
    # L2 reuse is poor.
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # ---- 2. Build the initial block pointers for A and B tiles ----
    offs_m = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M   # row indices of C tile
    offs_n = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N   # col indices of C tile
    offs_k = tl.arange(0, BLOCK_K)
    # A tile: (BLOCK_M, BLOCK_K). B tile: (BLOCK_K, BLOCK_N).
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    # ---- 3. The K-loop: load, multiply, accumulate in fp32 ----
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        # Mask the K tail so we don't read past column K.
        k_mask = offs_k[None, :] < K - k * BLOCK_K
        a = tl.load(a_ptrs, mask=k_mask, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
        # tl.dot issues Tensor Core (mma) instructions; accumulate in fp32.
        acc += tl.dot(a, b)
        # Advance the pointers to the next K-tile.
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    c = acc.to(c_ptr.dtype.element_ty)   # cast back to the output dtype

    # ---- 4. Write the tile, masking the M and N edges ----
    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)


def triton_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    assert a.shape[1] == b.shape[0], "incompatible dims"
    M, K = a.shape
    K, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    # 1-D grid of (num_m_tiles * num_n_tiles) programs; swizzle maps pid->(m,n).
    grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)
    matmul_kernel[grid](
        a, b, c, M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
    )
    return c


if __name__ == "__main__":
    torch.manual_seed(0)
    a = torch.randn((512, 768), device="cuda", dtype=torch.float16)
    b = torch.randn((768, 1024), device="cuda", dtype=torch.float16)
    ours = triton_matmul(a, b)
    ref = torch.matmul(a, b)
    # fp16 accumulation differences => compare with a tolerance.
    print("max abs error:", (ours - ref).abs().max().item())
    print("allclose:", torch.allclose(ours, ref, atol=1e-1, rtol=0))

Three pieces deserve emphasis.

tl.dot and fp32 accumulation. tl.dot(a, b) is the line that lights up the Tensor Cores. The inputs are fp16/bf16 tiles; the accumulator acc is fp32. This mixed-precision accumulate — multiply in low precision, sum in high precision — is exactly the pattern from Mixed Precision, bf16 & FP8 Training, and it is non-negotiable: accumulating a 768-long dot product in fp16 would lose catastrophic precision because each partial sum rounds to ~3 decimal digits.

The group swizzle. The block of code computing pid_m/pid_n from a flat pid reorders which tiles run together. The natural row-major order has many concurrently-running programs hammering different B columns, blowing past L2. Grouping GROUP_M tile-rows so that nearby programs share the same A rows and B columns dramatically improves L2 hit rate — frequently a 10–30% win at zero arithmetic cost. This is the kind of memory-traffic reasoning the roofline model trains you to do.

num_stages and software pipelining. The num_stages knob tells Triton how deep to software-pipeline the K-loop: while the Tensor Cores chew on the current BLOCK_K tile, the loads for the next tile are already in flight (issued via the GPU’s async copy units). More stages hide more memory latency at the cost of more shared memory. Autotuning searches over it because the sweet spot depends on the GPU and the tile shape.

A handwritten Triton matmul like this typically reaches a large fraction of cuBLAS on dense GEMMs. You usually would not ship it to replace cuBLAS; you write it so you can fuse something cuBLAS won’t — a dequantize-then-matmul for INT4 weights (Quantization II), a matmul-plus-activation, or the per-expert GEMMs of an MoE (Mixture-of-Experts).

Kernel 4: A Simplified FlashAttention¶

Now we assemble everything. Attention computes, for queries \(Q\), keys \(K\), values \(V\) (each \(N \times d\) for one head):

The naive route materializes the full \(N \times N\) score matrix \(S = QK^\top/\sqrt{d}\) in HBM. For a 16K-token sequence that is \(16384^2 = 268\)M entries per head — quadratic memory that dominates the runtime and caps context length. FlashAttention’s insight (Dao et al.) is that you never need the full \(S\) in HBM: you can stream over the keys/values in blocks, maintaining a running softmax so the output of each query block is built incrementally. This is the online softmax from Kernel 2, generalized to a reduction that does not fit in SRAM. Its derivation lives in FlashAttention I; here we implement it.

The Online Softmax Recurrence¶

Process key/value blocks \(j = 1, 2, \dots\) For each query block we keep three running statistics: \(m\) (running max of scores), \(\ell\) (running sum of exponentials), and the unnormalized output accumulator \(O\). On seeing a new score block \(S^{(j)}\) with block max \(m^{(j)} = \max S^{(j)}\):

\(\alpha\) is the correction factor: because the max changed, every previously accumulated term was exponentiated against the old max and must be rescaled by \(\alpha\). We then update:

After the last block, divide once: \(O \leftarrow O / \ell\). The output is bit-for-bit the same as a full softmax-then-matmul, but no \(N\times N\) tensor ever touched HBM. The kernel below implements the forward pass for one (batch, head) per program row-block.

import torch
import triton
import triton.language as tl


@triton.jit
def flash_attn_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr, L_ptr,   # L_ptr: per-row logsumexp out, shape (B,H,N_CTX)
    sm_scale,                       # 1/sqrt(d), folded with log2e below
    stride_qb, stride_qh, stride_qm, stride_qd,   # Q strides (B,H,seq,dim)
    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_CTX,
    BLOCK_M: tl.constexpr,          # query block (rows handled per program)
    BLOCK_N: tl.constexpr,          # key/value block (streamed)
    HEAD_DIM: tl.constexpr,         # d, the per-head dimension
    CAUSAL: tl.constexpr,           # apply the causal (lower-triangular) mask
):
    # ---- Which query block, and which (batch, head)? 2-D grid. ----
    start_m = tl.program_id(0)      # query-block index along the sequence
    off_bh = tl.program_id(1)       # flattened (batch * head) index
    off_b = off_bh // H
    off_h = off_bh % H

    # Base pointers into this (batch, head)'s matrices.
    q_base = Q_ptr + off_b * stride_qb + off_h * stride_qh
    k_base = K_ptr + off_b * stride_kb + off_h * stride_kh
    v_base = V_ptr + off_b * stride_vb + off_h * stride_vh
    o_base = O_ptr + off_b * stride_ob + off_h * stride_oh

    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)   # this block's query rows
    offs_d = tl.arange(0, HEAD_DIM)                       # the head dimension

    # Load THIS query block once; it stays resident the whole kernel.
    q_ptrs = q_base + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
    q = tl.load(q_ptrs, mask=offs_m[:, None] < N_CTX, other=0.0)

    # ---- Online-softmax running state ----
    m_i = tl.full((BLOCK_M,), -float("inf"), dtype=tl.float32)   # running max
    l_i = tl.zeros((BLOCK_M,), dtype=tl.float32)                 # running sum
    acc = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)        # running output

    # Fold 1/sqrt(d) and log2(e) so we can use the faster base-2 exp2.
    qk_scale = sm_scale * 1.44269504   # 1/ln(2)

    # ---- Stream over key/value blocks ----
    # Causal: only visit key blocks up to and including this query block's
    # diagonal (assumes BLOCK_M == BLOCK_N alignment). Non-causal: all blocks.
    hi = (start_m + 1) * BLOCK_M if CAUSAL else N_CTX
    for start_n in range(0, hi, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)
        n_mask = offs_n < N_CTX

        # Load a K block (HEAD_DIM x BLOCK_N) and a V block (BLOCK_N x HEAD_DIM).
        k_ptrs = k_base + offs_n[None, :] * stride_kn + offs_d[:, None] * stride_kd
        v_ptrs = v_base + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
        k = tl.load(k_ptrs, mask=n_mask[None, :], other=0.0)
        v = tl.load(v_ptrs, mask=n_mask[:, None], other=0.0)

        # Scores for this block: (BLOCK_M, BLOCK_N). tl.dot -> Tensor Cores.
        s = tl.dot(q, k) * qk_scale
        # Mask out padded keys so they never win the max / contribute to sum.
        s = tl.where(n_mask[None, :], s, -float("inf"))
        if CAUSAL:
            # query i attends only to keys j <= i within this diagonal block
            s = tl.where(offs_m[:, None] >= offs_n[None, :], s, -float("inf"))

        # ----- online softmax update -----
        m_ij = tl.max(s, axis=1)                 # per-row block max
        m_new = tl.maximum(m_i, m_ij)            # new running max
        alpha = tl.exp2(m_i - m_new)             # correction for prior state
        p = tl.exp2(s - m_new[:, None])          # exp of this block's scores
        l_i = l_i * alpha + tl.sum(p, axis=1)    # rescale + add new mass
        acc = acc * alpha[:, None]               # rescale accumulated output
        acc += tl.dot(p.to(v.dtype), v)          # add this block's contribution
        m_i = m_new                              # commit new max

    # ---- Final normalization (one division at the end) ----
    # Guard fully-masked rows (possible only for a padding-only tail block) so
    # we never divide by zero; those rows are masked off at store time anyway.
    l_safe = tl.where(l_i > 0.0, l_i, 1.0)
    acc = acc / l_safe[:, None]

    # Save the per-row logsumexp for the BACKWARD pass, in NATURAL-LOG space.
    # The exp2 trick kept m_i and the scores in base-2 log space; convert back
    # so ch 4.2's backward can reconstruct P = exp(S - L) with a plain exp:
    #   L_nat = m_nat + ln(l_i) = (m_i + log2(l_i)) / log2(e) = (m_i+log2 l_i)*ln2
    L_i = (m_i + tl.log2(l_safe)) * 0.6931471805599453   # 1/log2(e) = ln(2)
    l_ptrs = L_ptr + off_bh * N_CTX + offs_m
    tl.store(l_ptrs, L_i, mask=offs_m < N_CTX)

    # Write the output block.
    o_ptrs = o_base + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
    tl.store(o_ptrs, acc.to(O_ptr.dtype.element_ty), mask=offs_m[:, None] < N_CTX)


def flash_attention(q, k, v, sm_scale=None, causal=False):
    # q,k,v: (B, H, N_CTX, HEAD_DIM); returns (o, L) where L is the per-row logsumexp
    B, H, N_CTX, HEAD_DIM = q.shape
    if sm_scale is None:
        sm_scale = 1.0 / (HEAD_DIM ** 0.5)
    o = torch.empty_like(q)
    L = torch.empty((B, H, N_CTX), device=q.device, dtype=torch.float32)  # for backward
    BLOCK_M, BLOCK_N = 64, 64
    grid = (triton.cdiv(N_CTX, BLOCK_M), B * H)   # (query blocks, batch*head)
    flash_attn_kernel[grid](
        q, k, v, o, L, sm_scale,
        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
        B, H, N_CTX,
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=HEAD_DIM, CAUSAL=causal,
    )
    return o, L


if __name__ == "__main__":
    torch.manual_seed(0)
    B, H, N, D = 2, 4, 256, 64
    q = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    k = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    v = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    for causal in (False, True):
        ours, L = flash_attention(q, k, v, causal=causal)
        scale = 1.0 / (D ** 0.5)
        s = (q.float() @ k.float().transpose(-2, -1)) * scale
        if causal:
            m = torch.tril(torch.ones(N, N, device="cuda", dtype=torch.bool))
            s = s.masked_fill(~m, float("-inf"))
        ref = torch.softmax(s, dim=-1) @ v.float()
        err = (ours.float() - ref).abs().max().item()
        print(f"causal={causal}  max abs error: {err:.3e}")   # ~1e-2 (fp16)

A few notes that connect this to production kernels.

Why exp2 instead of exp. GPUs have a fast hardware approximation for base-2 exponentiation (exp2 / ex2.approx). FlashAttention folds the \(1/\ln 2\) factor into the QK scale (qk_scale = sm_scale * 1.44269...) so that \(e^{x} = 2^{x \log_2 e}\) becomes a single exp2. It is a small constant-factor win that the real kernels all use.

Causal masking, implemented above. The kernel supports causal (decoder) attention via the CAUSAL: tl.constexpr flag: when it’s set, the loop bound hi = (start_m + 1) * BLOCK_M skips key blocks entirely beyond this query block’s diagonal — a roughly 2x FLOP saving — and the diagonal block itself gets a triangular mask, tl.where(offs_m[:, None] >= offs_n[None, :], s, -inf), so query \(i\) only attends to keys \(j \le i\). The l_i > 0 guard before the final divide protects a fully-padding tail block; a valid causal query row always attends to at least itself, so it is never fully masked, but the guard keeps the padding case safe regardless.

Base consistency: exp2 and the saved L. Because the kernel folds log2(e) into qk_scale and runs the online softmax in base-2 log space, m_i is a base-2 quantity. The forward therefore converts before storing: L = (m_i + log2(l_i)) * ln(2), so the saved L is a natural-log logsumexp, and Kernel 5’s backward can reconstruct P = exp(S - L) with a plain exp. Combining the two chapters naively — storing m_i + ln(l_i), i.e. mixing bases — would silently produce an inconsistent L and wrong gradients.

The memory win, quantified. The naive path reads/writes the \(N\times N\) score matrix. Our kernel keeps \(Q\), the running \(O\), \(m\), and \(\ell\) in registers/SRAM and streams \(K\),\(V\) tiles — total HBM traffic is \(O(N d)\) per head, not \(O(N^2)\). This is the reason long-context attention is feasible at all. The backward pass needs the same statistics recomputed (FlashAttention recomputes \(S\) in the backward rather than storing it — trading FLOPs for memory), which we implement in Triton as Kernel 5 below (the forward above already saved the per-row logsumexp L that the backward needs).

Sizing the attention kernel: autotuning, SRAM, and head_dim¶

Kernel 4 hardcodes BLOCK_M = BLOCK_N = 64 for readability, but a production kernel autotunes it exactly like the matmul in Kernel 3:

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 32},  num_warps=4, num_stages=2),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64},  num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64},  num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=4),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64},  num_warps=4, num_stages=4),
    ],
    key=["N_CTX", "HEAD_DIM", "CAUSAL"],
)
@triton.jit
def flash_attn_kernel(
    # ... same parameter list and body as the Kernel 4 definition above ...
    Q_ptr, K_ptr, V_ptr, O_ptr, L_ptr, sm_scale,
    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_CTX,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
    HEAD_DIM: tl.constexpr, CAUSAL: tl.constexpr,
):
    pass   # body unchanged from Kernel 4

Because the grid now depends on the autotuned BLOCK_M, the launch must switch to a meta-driven grid — grid = lambda meta: (triton.cdiv(N_CTX, meta["BLOCK_M"]), B * H) — and BLOCK_M/BLOCK_N drop out of the explicit kwargs passed to flash_attn_kernel[grid](...).

A last wrinkle: the kernel needs HEAD_DIM to be a tl.constexpr power of two (it sizes tl.arange(0, HEAD_DIM)). For a non-power-of-two head dimension — d=80 or d=96 are common — set HEAD_DIM = triton.next_power_of_2(d) and mask offs_d < d (with other=0.0) on every Q/K/V load. Zero-padding the head dimension is exact: the padded lanes contribute 0 to Q @ K^T and produce output columns you simply slice off after the kernel returns. Pre-padding the tensors before the launch is the alternative if you’d rather not thread the extra mask through.

Kernel 5: FlashAttention Backward¶

The backward pass reuses every idea from Kernel 4 — tiling, the online recurrence’s cousin, and pointer arithmetic — plus the two primitives introduced above: tl.trans and tl.atomic_add. Recall the three backward quantities from FlashAttention I’s NumPy derivation: \(dV = P^\top dO\), \(dP = dO\,V^\top\), and the softmax-Jacobian collapses to a per-row scalar \(D_i = \text{rowsum}(dO \odot O)\), giving \(dS = P \odot (dP - D_i)\), then \(dQ = dS\,K / \sqrt d\) and \(dK = dS^\top Q / \sqrt d\). The whole trick that makes this an IO-aware backward — no \(N\times N\) tensor touches HBM — is that \(P\) is reconstructed tile-by-tile as \(P = \exp(S - L_i)\) using the natural-log \(L\) that Kernel 4’s forward saved. That’s a plain tl.exp, no exp2: L is already natural-log (see the base-consistency note above), so folding log2(e) in again here would double-apply the base change and silently corrupt every gradient.

Preprocessing: the per-row scalar D¶

@triton.jit
def attn_bwd_preprocess(
    O_ptr, dO_ptr, D_ptr,
    stride_ob, stride_oh, stride_om, stride_od,
    B, H, N_CTX,
    BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr,
):
    start_m = tl.program_id(0)
    off_bh = tl.program_id(1)
    off_b = off_bh // H
    off_h = off_bh % H

    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, HEAD_DIM)
    m_mask = offs_m < N_CTX

    o_base = O_ptr + off_b * stride_ob + off_h * stride_oh
    do_base = dO_ptr + off_b * stride_ob + off_h * stride_oh   # dO shares O's layout
    ptrs = offs_m[:, None] * stride_om + offs_d[None, :] * stride_od

    o = tl.load(o_base + ptrs, mask=m_mask[:, None], other=0.0)
    do = tl.load(do_base + ptrs, mask=m_mask[:, None], other=0.0)

    # D_i = rowsum(dO * O) -- the scalar that collapses the softmax Jacobian.
    d = tl.sum(o.to(tl.float32) * do.to(tl.float32), axis=1)
    d_ptrs = D_ptr + off_bh * N_CTX + offs_m
    tl.store(d_ptrs, d, mask=m_mask)

The kv-parallel dK/dV/dQ kernel¶

Each program owns one key/value block j (the backward-parallelizes-over-keys strategy from FlashAttention 2 & 3), loops over the query blocks that can see it, and accumulates dK_j, dV_j exclusively (no atomics needed) while scattering its contribution to dQ with tl.atomic_add, since every key block touches every query row’s gradient.

@triton.jit
def attn_bwd_dkdv_dq(
    Q_ptr, K_ptr, V_ptr, L_ptr, D_ptr, dO_ptr,
    dQ_ptr, dK_ptr, dV_ptr,
    sm_scale,
    stride_qb, stride_qh, stride_qm, stride_qd,
    stride_kb, stride_kh, stride_kn, stride_kd,
    stride_vb, stride_vh, stride_vn, stride_vd,
    B, H, N_CTX,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr,
    CAUSAL: tl.constexpr,
):
    start_n = tl.program_id(0)      # this program's KV block index
    off_bh = tl.program_id(1)
    off_b = off_bh // H
    off_h = off_bh % H

    q_base = Q_ptr + off_b * stride_qb + off_h * stride_qh
    k_base = K_ptr + off_b * stride_kb + off_h * stride_kh
    v_base = V_ptr + off_b * stride_vb + off_h * stride_vh
    do_base = dO_ptr + off_b * stride_qb + off_h * stride_qh   # dO shares Q's layout
    dq_base = dQ_ptr + off_b * stride_qb + off_h * stride_qh

    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)   # this program's key rows
    offs_d = tl.arange(0, HEAD_DIM)
    n_mask = offs_n < N_CTX

    k_ptrs = k_base + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
    v_ptrs = v_base + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
    k = tl.load(k_ptrs, mask=n_mask[:, None], other=0.0)   # (BLOCK_N, HEAD_DIM)
    v = tl.load(v_ptrs, mask=n_mask[:, None], other=0.0)   # (BLOCK_N, HEAD_DIM)

    dk_acc = tl.zeros((BLOCK_N, HEAD_DIM), dtype=tl.float32)
    dv_acc = tl.zeros((BLOCK_N, HEAD_DIM), dtype=tl.float32)

    # Causal: key block j only receives gradient from query blocks i >= j
    # (the same diagonal that let the forward skip blocks the other way).
    lo = (start_n * BLOCK_N // BLOCK_M) * BLOCK_M if CAUSAL else 0

    for start_m in range(lo, N_CTX, BLOCK_M):
        offs_m = start_m + tl.arange(0, BLOCK_M)
        m_mask = offs_m < N_CTX

        q_ptrs = q_base + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
        do_ptrs = do_base + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
        q = tl.load(q_ptrs, mask=m_mask[:, None], other=0.0)     # (BLOCK_M, HEAD_DIM)
        do = tl.load(do_ptrs, mask=m_mask[:, None], other=0.0)   # (BLOCK_M, HEAD_DIM)

        l_i = tl.load(L_ptr + off_bh * N_CTX + offs_m, mask=m_mask, other=0.0)
        d_i = tl.load(D_ptr + off_bh * N_CTX + offs_m, mask=m_mask, other=0.0)

        # ---- recompute the score/probability tile (no exp2: L is natural-log) ----
        s = tl.dot(q, tl.trans(k)) * sm_scale             # (BLOCK_M, BLOCK_N)
        if CAUSAL:
            s = tl.where(offs_m[:, None] >= offs_n[None, :], s, -float("inf"))
        p = tl.exp(s - l_i[:, None])                       # plain exp; P = exp(S - L)

        # ---- dV_j += P^T dO_i ----
        dv_acc += tl.dot(tl.trans(p).to(do.dtype), do)

        # ---- dP = dO_i V_j^T ; dS = P * (dP - D_i) * sm_scale ----
        dp = tl.dot(do, tl.trans(v))                        # (BLOCK_M, BLOCK_N)
        ds = p * (dp - d_i[:, None]) * sm_scale             # scale folded once here

        # ---- dK_j += dS^T Q_i ----
        dk_acc += tl.dot(tl.trans(ds).to(q.dtype), q)

        # ---- dQ_i += dS K_j  (scattered: every key block hits every query row) ----
        dq_part = tl.dot(ds.to(k.dtype), k)                  # (BLOCK_M, HEAD_DIM)
        dq_ptrs = dq_base + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
        tl.atomic_add(dq_ptrs, dq_part, mask=m_mask[:, None])

    dk_ptrs = dK_ptr + off_b * stride_kb + off_h * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
    dv_ptrs = dV_ptr + off_b * stride_vb + off_h * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
    tl.store(dk_ptrs, dk_acc.to(dK_ptr.dtype.element_ty), mask=n_mask[:, None])
    tl.store(dv_ptrs, dv_acc.to(dV_ptr.dtype.element_ty), mask=n_mask[:, None])

Note the dtype discipline: tl.dot’s operands are cast down to the working dtype (do.dtype, q.dtype, k.dtype) right before the call, while dk_acc, dv_acc, and the dQ buffer stay fp32 throughout — the same “multiply low, accumulate high” pattern as the forward and the Kernel 3 matmul.

Driver¶

def flash_attn_backward(q, k, v, o, L, do, sm_scale, causal=False):
    B, H, N_CTX, HEAD_DIM = q.shape
    BLOCK_M, BLOCK_N = 64, 64

    D = torch.empty((B, H, N_CTX), device=q.device, dtype=torch.float32)
    grid_pre = (triton.cdiv(N_CTX, BLOCK_M), B * H)
    attn_bwd_preprocess[grid_pre](
        o, do, D,
        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
        B, H, N_CTX,
        BLOCK_M=BLOCK_M, HEAD_DIM=HEAD_DIM,
    )

    dq = torch.zeros_like(q, dtype=torch.float32)   # fp32: atomic_add target
    dk = torch.empty_like(k)
    dv = torch.empty_like(v)
    grid_bwd = (triton.cdiv(N_CTX, BLOCK_N), B * H)
    attn_bwd_dkdv_dq[grid_bwd](
        q, k, v, L, D, do,
        dq, dk, dv,
        sm_scale,
        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
        B, H, N_CTX,
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=HEAD_DIM, CAUSAL=causal,
    )
    return dq.to(q.dtype), dk, dv

Wiring it into torch.autograd¶

class FlashAttentionFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q, k, v, causal, sm_scale):
        o, L = flash_attention(q, k, v, sm_scale=sm_scale, causal=causal)
        ctx.save_for_backward(q, k, v, o, L)
        ctx.causal = causal
        ctx.sm_scale = sm_scale
        return o

    @staticmethod
    def backward(ctx, do):
        q, k, v, o, L = ctx.saved_tensors
        dq, dk, dv = flash_attn_backward(q, k, v, o, L, do, ctx.sm_scale, ctx.causal)
        return dq, dk, dv, None, None   # None, None for the non-tensor causal/sm_scale args


def flash_attn_func(q, k, v, causal=False, sm_scale=None):
    if sm_scale is None:
        sm_scale = 1.0 / (q.shape[-1] ** 0.5)
    return FlashAttentionFn.apply(q, k, v, causal, sm_scale)

ctx.save_for_backward(q, k, v, o, L) is the standard autograd.Function mechanism for stashing tensors needed later without keeping the whole autograd graph alive. We save L, not the \(N\times N\) probability matrix P — that’s the entire point of the recompute-don’t-store strategy from FlashAttention I: L is \(O(N)\) per head, P is \(O(N^2)\), and the backward regenerates P tile-by-tile from L at the cost of one extra tl.dot and one tl.exp.

Correctness test¶

if __name__ == "__main__":
    torch.manual_seed(0)
    B, H, N, D = 2, 4, 256, 64
    for causal in (False, True):
        q = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16, requires_grad=True)
        k = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16, requires_grad=True)
        v = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16, requires_grad=True)
        do = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)

        o = flash_attn_func(q, k, v, causal=causal)
        o.backward(do)
        dq, dk, dv = q.grad.clone(), k.grad.clone(), v.grad.clone()

        # fp32 PyTorch-autograd reference (not gradcheck: gradcheck needs
        # float64, and tl.dot's Tensor Core path does not support float64).
        qf = q.detach().float().requires_grad_()
        kf = k.detach().float().requires_grad_()
        vf = v.detach().float().requires_grad_()
        scale = 1.0 / (D ** 0.5)
        s = (qf @ kf.transpose(-2, -1)) * scale
        if causal:
            mask = torch.tril(torch.ones(N, N, device="cuda", dtype=torch.bool))
            s = s.masked_fill(~mask, float("-inf"))
        ref_o = torch.softmax(s, dim=-1) @ vf
        ref_o.backward(do.float())

        print(f"causal={causal}")
        print("  max dQ error:", (dq.float() - qf.grad).abs().max().item())   # ~2e-2 (fp16 kernel vs fp32 ref)
        print("  max dK error:", (dk.float() - kf.grad).abs().max().item())   # ~2e-2
        print("  max dV error:", (dv.float() - vf.grad).abs().max().item())   # ~2e-2

Run the same test with q, k, v in fp32 and you should see errors tighten to roughly 1e-3 — the fp16 case’s larger tolerance is accumulation-order noise, not a bug. Two failure signatures worth memorizing so you recognize them instantly: dropping the \(D_i\) subtraction (using dS = P * dP * sm_scale instead of P * (dP - D_i) * sm_scale) makes dQ wrong by a per-row constant offset, because the softmax Jacobian’s rank-1 correction term vanishes; dropping the sm_scale fold (forgetting to multiply dS by sm_scale, or applying it twice) makes every gradient off by a factor of \(\sqrt d\).

Full-spectrum hardware notes. On a laptop or CPU-only box, set TRITON_INTERPRET=1 and shrink the problem (B=1, H=1, N=64, D=16) to check correctness only — the interpreter is very slow, so this is a debugging mode, not a benchmark. On a single A100, H100, or B200, use realistic sizes (N = 1024 to 8192) for real timing; the backward does roughly 2x the forward’s FLOPs (it recomputes S/P and additionally produces dQ, dK, dV), so expect the backward to cost about 2–2.5x the forward’s latency at matched shapes. On AMD (ROCm) the same source compiles through Triton’s AMD backend, but the tuning constants do not transfer: a CDNA wavefront is 64 lanes wide rather than NVIDIA’s 32, so num_warps=4 is 256 threads there and 128 here — re-run @triton.autotune per vendor instead of shipping one config table. Multi-GPU is orthogonal here — this is a per-(batch, head) kernel; sharding across devices is handled by the training framework (data/tensor/context parallelism), not by anything inside the kernel.

Library mapping. The official Triton tutorial 06-fused-attention.py implements exactly this structure: the forward stores M (the running max) and L-equivalent statistics, _attn_bwd_preprocess computes delta = sum(o * do), and _attn_bwd_dkdv / _attn_bwd_dq split the backward the atomics-free way described above. Dao-AILab/flash-attention is the reference implementation to read once this kernel makes sense.

Autotuning, Debugging, and Performance Practice¶

Writing a correct kernel is half the job; making it fast and trusting it is the other half. A few tools and habits.

Autotuning. As shown in the matmul, @triton.autotune takes a list of triton.Configs (each a set of constexpr meta-parameters plus num_warps and num_stages) and a key of argument names. On the first launch for each distinct key, Triton benchmarks every config and caches the winner. The knobs that matter most:

Debugging. Three indispensable tools:

# 1. Run the kernel on the CPU in pure Python, element by element, so you can
#    use print() and a debugger. Slow, but exact — your first line of defense.
import os
os.environ["TRITON_INTERPRET"] = "1"   # set BEFORE importing triton

# 2. Inside an interpreted kernel you can even print tiles:
#    tl.device_print("scores", s)

# 3. Sanity-check shapes/strides — most bugs are pointer-arithmetic bugs.
#    Always diff against an eager-PyTorch reference with torch.allclose.

The single most valuable habit: write the eager-PyTorch reference first, then make the kernel match it under torch.allclose with an appropriate tolerance (looser for fp16/bf16 because of accumulation order differences). Almost every Triton bug is a pointer/stride mistake or a missing mask, and a reference test localizes it immediately.

Reading what the compiler produced. Triton lowers your kernel through a stack of MLIR dialects: Triton IR (your block program, hardware-agnostic) → TritonGPU IR (where tile layouts, the num_warps thread mapping, shared-memory allocation, and the num_stages pipeline are actually decided) → LLVM IR → PTX → cubin. Launching a kernel hands back the compiled object, and every stage is readable on it:

compiled = add_kernel[(1,)](x, y, out, n_elements, BLOCK_SIZE=1024)
print(compiled.asm.keys())                 # e.g. 'ttir', 'ttgir', 'llir', 'ptx', 'cubin'
print(compiled.n_regs, compiled.n_spills)  # registers per thread, and register SPILLS

(The exact key names track the backend and Triton version.) n_spills > 0 is the most actionable red flag in Triton: your tiles no longer fit the register file, so the compiler spills to “local” memory — which physically lives in HBM — and your supposedly on-chip kernel is quietly doing HBM round-trips inside the inner loop. The fix is a smaller BLOCK_M/BLOCK_N, fewer num_stages, or more num_warps (which splits the tile across more registers); this is a large part of why the autotuner’s biggest config so often loses. Compiled binaries are cached on disk under TRITON_CACHE_DIR (~/.triton/cache by default) — clear it if you suspect a stale build, and set TRITON_PRINT_AUTOTUNING=1 to see which config the autotuner actually picked.

Benchmarking. Use triton.testing.do_bench, which handles GPU warmup, CUDA-stream synchronization, and the L2-cache flush between runs that naive time.time() benchmarks get wrong:

import triton

ms = triton.testing.do_bench(lambda: triton_matmul(a, b))
flops = 2 * M * N * K
print(f"{ms:.3f} ms,  {flops / (ms * 1e-3) / 1e12:.1f} TFLOP/s")

Where Triton fits in the stack. TorchInductor — the backend of torch.compile — generates Triton code for fused pointwise and reduction kernels automatically. So even if you never write a @triton.jit function, you are running Triton when you torch.compile a model on an NVIDIA GPU. Writing kernels by hand is for the cases the compiler can’t fuse well: novel attention variants, quantized matmuls, MoE dispatch, custom losses. See Kernel Fusion, torch.compile, CUDA Graphs & Compilers.

Making your kernel torch.compile-safe. A @triton.jit launch is opaque Python as far as TorchDynamo is concerned, so dropping one into a torch.compiled model risks a graph break — the compiled region splits in two around your kernel and you lose the surrounding fusions. The supported fix (PyTorch ≥ 2.6) is to register the launch as a real custom operator with torch.library.triton_op, marking the launch itself with torch.library.wrap_triton:

import torch
import triton
from torch.library import triton_op, wrap_triton


# Registers a genuine PyTorch operator, "stackbook::add", implemented by a
# Triton kernel. Dynamo now sees a known op instead of untraceable Python.
@triton_op("stackbook::add", mutates_args={})
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    n_elements = out.numel()
    # wrap_triton makes the launch *traceable*: under torch.compile it is
    # captured into the graph rather than executed as a black box.
    wrap_triton(add_kernel)[(triton.cdiv(n_elements, 1024),)](
        x, y, out, n_elements, BLOCK_SIZE=1024
    )
    return out


# Composes with the rest of the graph — no break, and Inductor can still fuse
# the surrounding pointwise ops.
compiled = torch.compile(lambda a, b: torch.sin(add(a, b)))

Because triton_op lets the tracer walk the wrapper body (that is exactly what wrap_triton enables), you generally do not need to hand-write the fake/meta implementation that the lower-level torch.library.custom_op requires for shape inference. Attach a gradient with torch.library.register_autograd("stackbook::add", backward_fn, setup_context=...) and the op is differentiable everywhere — the same job FlashAttentionFn did above, but additionally visible to torch.compile and torch.export. Rule of thumb: autograd.Function is fine for eager-only research code; register a custom op the moment the kernel has to live inside a compiled or exported model.

Don’t write what already exists. Before writing a kernel, check whether someone shipped it. Liger-Kernel provides Triton LigerRMSNorm, LigerSwiGLUMLP, LigerRopeFunction, and — the big one — LigerFusedLinearCrossEntropy, which fuses the LM head with the loss so the \((B{\cdot}T, V)\) logit tensor never lands in HBM; HuggingFace Trainer and TRL enable the whole set with one flag (use_liger_kernel=True). vLLM and SGLang ship Triton kernels for fused MoE routing and paged attention; unslothai/unsloth rewrites the fine-tuning backward in Triton. This is concrete for our own model: the fused cross-entropy is what keeps the loss head off the memory budget in The Pretraining Run: A Complete Single-GPU Training Loop, and the fused RMSNorm/SwiGLU swap-ins are called out in The Stack-100M Architecture. Hand-written Triton is for the op that isn’t in a library yet.

Key Takeaways¶

Further reading¶

• Philippe Tillet, H. T. Kung, David Cox, Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations (2019) — the original paper.
• The official OpenAI Triton tutorials (vector-add, fused softmax, matrix multiplication, and the FlashAttention example) — the canonical, maintained reference implementations these kernels follow.
• Tri Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022), and FlashAttention-2 (2023) — the online-softmax and work-partitioning ideas behind Kernel 4.
• NVIDIA CUTLASS documentation — for the layout/tiling/pipelining concepts that Triton automates, seen from the hand-written CUDA side.
• The vLLM and Unsloth repositories — production Triton kernels (paged/flash attention, fused MoE, RMSNorm, RoPE) worth reading once you can follow this chapter’s code.

Exercises¶

1. In Kernel 1 (vector add) the driver launches on x = torch.randn(98_432) with BLOCK_SIZE = 1024. Exactly how many programs does the grid contain, and how many lanes of the last program are masked off? Then explain what would go wrong at runtime if you deleted the mask=mask argument from the tl.store on line 101 (keeping everything else the same).

2. The chapter states the naive multi-op softmax makes about 5 passes over HBM (read to find max, read to exponentiate, write the exponentials, read them to sum, read once more to divide) while the fused kernel makes 2 (one read, one write). Take an input of shape M = 2048, N = 4096 stored in fp16 (2 bytes/element), on a GPU with 2.0 TB/s of HBM bandwidth. Compute the HBM bytes moved by each version, the resulting memory-bound time for each, and the speedup. Why can no fused softmax kernel beat this speedup by much?

3. Consider the matmul of Kernel 3 with a different autotune config: BLOCK_M = 64, BLOCK_N = 128, BLOCK_K = 64, num_stages = 4, fp16 inputs (2 bytes) and an fp32 accumulator (4 bytes). Compute (a) the shared-memory footprint of the pipelined A and B tiles across all stages, and (b) the size of the fp32 accumulator tile. Does the shared-memory footprint fit on an A100 (164 KiB) and an H100 (228 KiB)?

4. Implement a fused RMSNorm kernel in Triton in the exact “one program per row” style of Kernel 2 (fused softmax). RMSNorm over the last dimension is

where w is a length-N learnable gain vector and eps is a small constant. Write both the @triton.jit kernel and its Python driver, and explain why loading padded (out-of-bounds) lanes with other=0.0 gives the correct mean-of-squares.

import torch
import triton
import triton.language as tl


@triton.jit
def rmsnorm_kernel(
    out_ptr, in_ptr, w_ptr,
    in_row_stride, out_row_stride,
    n_cols, eps,
    BLOCK_SIZE: tl.constexpr,        # padded power-of-two >= n_cols
):
    row_idx = tl.program_id(axis=0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols

    # Load this row and the gain vector. Padding lanes read 0.0.
    row_start = in_ptr + row_idx * in_row_stride
    x = tl.load(row_start + col_offsets, mask=mask, other=0.0)
    w = tl.load(w_ptr + col_offsets, mask=mask, other=0.0)

    # Mean of squares over the REAL columns (divide by n_cols, not BLOCK_SIZE).
    x = x.to(tl.float32)
    mean_sq = tl.sum(x * x, axis=0) / n_cols     # block reduction -> scalar
    rstd = 1.0 / tl.sqrt(mean_sq + eps)

    out = (x * rstd) * w.to(tl.float32)

    out_row_start = out_ptr + row_idx * out_row_stride
    tl.store(out_row_start + col_offsets, out.to(out_ptr.dtype.element_ty), mask=mask)


def triton_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    assert x.dim() == 2 and x.is_cuda and w.is_cuda
    M, N = x.shape
    assert w.shape == (N,)
    BLOCK_SIZE = triton.next_power_of_2(N)
    num_warps = 4
    if BLOCK_SIZE >= 2048:
        num_warps = 8
    if BLOCK_SIZE >= 4096:
        num_warps = 16

    out = torch.empty_like(x)
    rmsnorm_kernel[(M,)](
        out, x, w,
        x.stride(0), out.stride(0),
        N, eps,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
    )
    return out


if __name__ == "__main__":
    torch.manual_seed(0)
    x = torch.randn(1823, 781, device="cuda")
    w = torch.randn(781, device="cuda")
    ours = triton_rmsnorm(x, w)
    ref = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * w
    print("max abs error:", (ours - ref).abs().max().item())   # ~1e-6

5. Work the online-softmax recurrence of Kernel 4 by hand for a single query row streamed as two key blocks of one key each. The (already scaled) scores are \(S^{(1)} = 2\) and \(S^{(2)} = 4\); the corresponding values (with HEAD_DIM = 1) are \(v_1 = 1\) and \(v_2 = 3\). Starting from \(m = -\infty\), \(\ell = 0\), \(O = 0\), apply the recurrence block by block: report the correction factor \(\alpha\) used when block 2 arrives, the final \(\ell\) and \(O\) before normalization, and the normalized output \(O/\ell\). Verify it equals the plain softmax-weighted combination of \(v_1, v_2\).