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 include:

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.

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 of exp(S_i - m_i) (line 12) O_i, the running sum of exp(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 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 NxN matrix too. The backward will be covered in a later article. For now L is computed, stored, and tested, but not used.

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.

import math
import torch
import einops

class FlashAttentionPytorch(torch.autograd.Function):

    @staticmethod
    def forward(ctx, Q, K, V, is_causal=False):
        tile_size = 16
        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:

alpha = 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. B_q x d tile on every inner loop.L is stored, not m and l separately. The backward pass recomputes each tile of S from Q and K, then recovers that tile's probabilities as exp(S - L). For context, assuming N = 4096, NxN per head per batch element will have about 16.8 million entries, or about 33MB in bf16. Nx1 per head per batch element contains 4096 entries, or about 16KB. is_causal flag is accepted and ignored. The PyTorch version does not implement causal masking. 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_btr.

tl.make_block_btr 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),
)

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_btr, illustrated in the image below. I will through each one by one. The illustrations are drawn with Claude Opus 5.5.

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.

Before block pointers, programmers would built 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)

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)):
    cols = i * D_TILE_SIZE + tl.arange(0, D_TILE_SIZE)
    col_mask = cols < D

    x_ptrs = (x_ptr + rows[:, None] * x_stride_row
                    + cols[None, :] * x_stride_dim)
    w_ptrs = weight_ptr + cols * weight_stride_dim

    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.

Below is the kernel as it stands at the end of this post; the two lines marked FIX are the ones Step 3 explains.

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)

    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,),
    )

    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

        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)

        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=FlashAttentionTriton.Q_TILE_SIZE,
            K_TILE_SIZE=FlashAttentionTriton.K_TILE_SIZE,
            is_causal=is_causal,
        )
        ctx.save_for_backward(Q, K, V, O, L)
        ctx.is_causal = is_causal
        return O

    @staticmethod
    def backward(ctx, dO):
        raise NotImplementedError("tiled backward is part 4 of this series")

Reading it top to bottom:

The grid. program_id(0) indexes query tiles, program_id(1) indexes batch elements, so the launch has cdiv(N_q, 16) × B programs. Multi-head attention is handled by flattening (B, H, N, D) into (B*H, N, D) before the call; the kernel never knows about heads. At 16 rows per program and, say, 4096 tokens × 32 heads, that is 8192 programs, which is plenty to fill a GPU. The problem, as part 2 will show, is how little work each one does.

Block pointers. make_block_ptr describes a 2D window into a strided tensor: the base pointer, the full logical shape, the strides, where the window starts, and how big it is. order=(1, 0) says the last dimension is the fastest-varying one, i.e. row-major. The payoff is boundary_check and padding_option on load, which handle the ragged last tile for free, and .advance(), which slides the K and V windows down by one tile per iteration without recomputing addresses.

Accumulators in fp32. O_acc, l_acc and m_acc are fp32 whatever the input dtype, and tl.dot returns fp32 for bf16 inputs. This is what makes it safe to run in bf16: the inputs and the P·V product are low precision, but the sum across key tiles is not.

One Q load, many K/V loads. Q_i is loaded once before the loop. K_j and V_j are loaded per iteration, and every program loads the entire K and V for its batch element. That's inherent to the algorithm (it is the trade FlashAttention makes to avoid the N×N matrix), which is why K/V loads dominate the memory traffic and why tile size matters so much.

The inner loop is a line-for-line translation of the PyTorch version: score block, mask, new max, unnormalised probabilities, rescale factor, update l, update O, carry m forward. tl.dot(P_ij, V_j, acc=O_acc) fuses the matrix multiply and the accumulate.

The store. boundary_check on the store discards rows of the last tile that fall past N_QUERIES, so the padded rows computed with zeroed Q are never written.

The CS336 test for this kernel runs one shape, (4, 128, 128, 64), in fp32, at a tolerance of 1e-2. My kernel passed it. It was still wrong in two ways, and both were found by the test file below, which is the first thing I'd tell anyone starting from the scaffold to write.

import math
import pytest
import torch
import torch.nn.functional as F
from flashattention_autograd_function_triton import FlashAttentionTriton

def reference(q, k, v, is_causal):
    o = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
    s = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(q.shape[-1])
    if is_causal:
        tril = torch.ones(s.shape[-2:], dtype=torch.bool, device=q.device).tril()
        s = s.masked_fill(~tril, float("-inf"))
    return o, torch.logsumexp(s, dim=-1)

SHAPES = [
    (4, 128, 128, 64),     # the assignment's shape
    (2, 1024, 1024, 64),
    (2, 100, 100, 64),     # ragged last tile on both axes
    (1, 37, 53, 32),       # ragged and non-square
    (1, 200, 96, 128),
]

@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("is_causal", [False, True])
@pytest.mark.parametrize("B,Nq,Nk,D", SHAPES)
def test_forward_matches_sdpa(B, Nq, Nk, D, is_causal, dtype):
    if is_causal and Nq != Nk:
        pytest.skip("keep causal tests square so the mask convention is unambiguous")
    torch.manual_seed(0)
    q = torch.randn(B, Nq, D, device="cuda", dtype=dtype, requires_grad=True)
    k = torch.randn(B, Nk, D, device="cuda", dtype=dtype, requires_grad=True)
    v = torch.randn(B, Nk, D, device="cuda", dtype=dtype, requires_grad=True)

    o = FlashAttentionTriton.apply(q, k, v, is_causal)
    L = [t for t in o.grad_fn.saved_tensors if t.shape == (B, Nq)][0]

    o_ref, L_ref = reference(q, k, v, is_causal)
    tol = dict(atol=2e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-2, rtol=1e-2)
    torch.testing.assert_close(o, o_ref, **tol)
    torch.testing.assert_close(L, L_ref, **tol)

Two details in the harness. L isn't part of the function's return value, so the test digs it out of o.grad_fn.saved_tensors, which is the same trick the assignment's test uses and requires the inputs to have requires_grad=True. And the fp32 tolerance is looser than you might expect because Triton's tl.dot on fp32 inputs uses TF32 on Ampere and later unless you pass input_precision="ieee"; the reference is true fp32. Worth knowing before you spend an evening chasing a 1e-3 discrepancy that isn't a bug.

The scaffold kernel had three lines like this:

P_tilde_ij.to(V_j.type.element_ty)   # "cast P to bf16 before the dot"
...
O_i.to(O_block_ptr.type.element_ty)
...
L_i.to(L_block_ptr.type.element_ty)

In Triton, as in PyTorch, .to() returns a new tensor. It doesn't mutate. All three lines were no-ops, and the comment above the first one described a cast that wasn't taking place. The kernel still passed the assignment test because the test runs in fp32, where casting fp32 to fp32 changes nothing.

The bf16 test found it immediately, at compile time rather than as a numerical error: tl.dot requires both operands to share a dtype, and after the no-op P_tilde_ij was still fp32 while V_j was bf16. The fix is one character of intent: P_ij = P_ij.to(V_j.dtype). The O and L casts are fixed the same way, folded into the store lines.

The embarrassing part isn't the bug. It's that I'd written a comment claiming to have done the one thing the line didn't do. The lesson I'm taking from it is to check the dtype of every tl.dot operand when a kernel is first written, not when it first fails. Part 3 measures what the bf16 cast of P actually costs in accuracy.

tl.load with boundary_check and padding_option="zero" fills rows past the end of the tensor with zeros. For Q that's harmless: padded query rows produce garbage that the store then discards. For K it is not harmless. A zero key row gives a score of exactly 0 for every query, and 0 is a perfectly ordinary score, so the padded keys enter the softmax with weight exp(0 - m) each. When N_KEYS isn't a multiple of the key tile, the output is a weighted average over real keys plus some phantom ones.

Again the assignment test couldn't see it, because 128 is a multiple of 16. The (2, 100, 100, 64) case fails at once. The fix is the keep = k_pos < N_KEYS mask, applied whether or not attention is causal. It costs one compare and one select per score, which is nothing next to the two matrix multiplies, and it means the kernel is now correct for any sequence length rather than any sequence length divisible by 16.

With both fixes, all 20 cases in the file pass. Together they are the reason this post exists as a baseline rather than as an answer: the assignment's test proves the algorithm is right, and nothing else.

Every later part of this series reports one number: the kernel's speed as a percentage of PyTorch's scaled_dot_product_attention on the same inputs, same card. SDPA is the right yardstick because at head dim 64 in bf16 it dispatches to the FlashAttention-2 kernel that ships inside PyTorch, so "100% of SDPA" means "as fast as FA2". The harness is short and it never changes after this post.

import torch
import torch.nn.functional as F
import triton
from flashattention_autograd_function_triton import FlashAttentionTriton

def attn_flops(batch, n, d, is_causal):
    f = 4 * batch * n * n * d
    return f / 2 if is_causal else f

def bench(B, H, N, D, is_causal, dtype=torch.bfloat16):
    q, k, v = (torch.randn(B * H, N, D, device="cuda", dtype=dtype) for _ in range(3))
    ms_ours = triton.testing.do_bench(lambda: FlashAttentionTriton.apply(q, k, v, is_causal))

    q4, k4, v4 = (t.view(B, H, N, D) for t in (q, k, v))
    ms_sdpa = triton.testing.do_bench(
        lambda: F.scaled_dot_product_attention(q4, k4, v4, is_causal=is_causal))

    flops = attn_flops(B * H, N, D, is_causal)
    tflops = lambda ms: flops / ms * 1e-9
    return ms_ours, ms_sdpa, tflops(ms_ours), tflops(ms_sdpa)

if __name__ == "__main__":
    print(f"{'N':>6} {'causal':>7} {'ours ms':>9} {'sdpa ms':>9} {'ours TF/s':>10} {'sdpa TF/s':>10} {'% sdpa':>7}")
    for N in (512, 1024, 2048, 4096):
        for causal in (False, True):
            a, b, c, d = bench(B=4, H=8, N=N, D=64, is_causal=causal)
            print(f"{N:>6} {str(causal):>7} {a:>9.3f} {b:>9.3f} {c:>10.1f} {d:>10.1f} {100 * b / a:>6.0f}%")

triton.testing.do_bench handles warm-up, flushes L2 between runs, and reports the median, so a single call is enough. The FLOP count is the standard one for attention: two matmuls of N×N×d multiply-adds, halved for causal because the kernel is only supposed to do the lower triangle. Note the word supposed. This kernel computes every tile and masks half of them, so its causal FLOP rate is flattered by the accounting. That's deliberate: the accounting is what FA2 gets credit for, and the gap it opens is exactly the thing part 3 closes.

Results on the RTX 4070 Super, batch 4, 8 heads, head dim 64, bf16:

[Replace this paragraph with two or three sentences on what the table shows: the headline % of SDPA at N=4096 non-causal, whether the gap widens or narrows with N, and how much worse causal is than non-causal relative to SDPA. Then move the N=4096 non-causal percentage into the opening paragraph of the post.]

To check which SDPA backend you're actually racing, wrap the reference call in torch.nn.attention.sdpa_kernel([SDPBackend.FLASH_ATTENTION]): if it errors, you're not comparing against FA2 and the percentages mean something else.

The kernel is correct and it is slow. Here is why, in the order I plan to fix it. Each item names the mechanism, not just the symptom, because the mechanism is what the next post has to demonstrate.

1. The tiles are far too small. 16×16 tiles mean each tl.dot is a 16×64 by 64×16 product. Tensor cores want bigger operands than that to reach their throughput, and each program does so little work per K/V load that the kernel spends its time moving data rather than multiplying it. Every program also reloads all of K and V from L2 or HBM, and with a 16-row Q tile the ratio of loads to useful FLOPs is about eight times worse than with a 128-row tile. The right sizes depend on the card: Ada has roughly 100 KB of shared memory per SM against Hopper's 228 KB, so the 128×64 or 128×128 configurations in the FA2 paper aren't automatically right here. Part 2 does the shared-memory budget arithmetic, sweeps tile sizes together with num_warps and num_stages, and wraps the winner in @triton.autotune.

2. Causal attention does all the work and throws half away. The mask is applied to every tile, including tiles that lie entirely above the diagonal, where every element is masked. Those tiles contribute nothing to the output and the kernel still loads K and V for them, runs both matmuls and the exponentials, then discards it all. The fix is to stop the key loop at the diagonal, which roughly halves causal runtime, and to only apply the mask on the one tile that actually straddles it. Part 3.

3. exp instead of exp2. The GPU's fast exponential unit computes 2^x; tl.exp is exp2(x · log₂e), an extra multiply on every score. The standard trick is to fold scale · log₂e into the scaling of Q once per tile, then use tl.math.exp2 directly. It's a small win but it's in the inner loop. Part 3.

4. The -1e6 mask sentinel. Masking with a large negative number rather than -inf sidesteps a real hazard: a row whose every score is -inf produces exp(-inf - (-inf)) = NaN. But it also means a fully masked tile still contributes exp(-1e6 - m) per element, which is zero in fp32 only because the exponent underflows. That's fine at this scale and worth understanding properly before the mask logic changes shape in part 3, where the rows-with-no-valid-key case gets handled explicitly.

5. No backward pass. backward raises NotImplementedError, and the PyTorch reference backward materialises the full N×N score matrix, so it isn't FlashAttention's backward at all. The tiled backward needs the L we've been carrying, a precomputed D = rowsum(dO ∘ O), and a kernel that iterates over query tiles for each key tile to build dK and dV, with dQ accumulated either by atomics or by a second pass. It's the hardest part of the algorithm and it's part 4.

6. The mask is computed unconditionally. The k_pos < N_KEYS compare and the tl.where run on every tile, including the full interior ones where nothing is masked. It's cheap, but in a kernel this small everything in the inner loop counts. Once tile-skipping is in place the mask can be restricted to boundary tiles. Part 3.

7. No profile. Everything above is reasoning from first principles. None of it has been confirmed with Nsight Compute, and reasoning about GPU performance without a profiler is how people end up optimising the wrong thing. Part 5 profiles the finished forward and backward, plots them against the card's roofline, and names what's left between this kernel and FA2.

What's deliberately not on the list: multi-query and grouped-query attention, sliding windows, variable-length batches, dropout, and anything that needs Hopper or Blackwell features. Those are branches, not steps, and they come after the trunk is fast.

Part 2 takes the kernel above and changes nothing except the tile sizes, num_warps and num_stages. It works through the shared-memory budget for an Ada SM, sweeps the configuration space, and shows the heatmap, then wraps the winner in @triton.autotune and measures what that costs in compile time. The benchmark harness and the test file stay exactly as they are here, so the number in part 2 is directly comparable to the one above.

The rest of the track, in order: part 3 goes inside the inner loop for causal tile-skipping, exp2 and the P cast; part 4 is the tiled backward; part 5 is the profiler and the roofline. After that the same kernel gets rewritten in TileLang and CuTe DSL, and then moved to Blackwell, where the interesting question is which of these optimisations turn out to have been hardware-independent.

If you spot something wrong in the kernel above, I'd rather hear it now than after part 4 has built on it. The code is in the repo; the earlier posts on online softmax and FlashAttention by hand cover the algorithm at the level of individual numbers if the recap above went by too fast.