从 PyTorch 到 Triton 实现 FlashAttention-2

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: Fixing the batch size to 8 to avoid the need for a multihead attention. Hard coding to 16x16 tiles. Future articles will relax these assumptions to learn these implementation details. I will state upfront whenever I made simplifying assumptions. All numbers in this series are from a single RTX 4070 Super (Ada, sm_89, 12 GB, roughly 100 KB of shared memory per SM) under Linux. This article is written with the assistance of AI. What we are building The FlashAttention-2 (FA2) forward pass: for each query tile Q_i, iterate over key and value tiles K_j and V_j, rescaling the running output O_i on every step and dividing by l only once at the end (FA1 re-normalised at every step). The outer loop goes through the query tiles where on the GPU, each query tile gets its own program, and all the programs run in parallel. The kernel computes standard scaled-dot-product attention without materialising the N×N score matrix. For a query tile Q_i of B_q rows and a key tile K_j of B_k rows, it computes the B_q × B_k score block, folds it into a running softmax, and moves to the next key tile. When all key tiles have been visited, the output rows for Q_i are complete and are written once. Three running quantities per query tile make that possible, initialised on line 6: m_i: the running maximum score for each row (line 10) l_i, the running row sum 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. Step 1: the tiled forward in PyTorch Before diving into Triton, I coded FA2 with PyTorch to familiarise with the algorithm before introducing Triton concepts. The PyTorch implementation is slow by design because it has a Python double loop over tiles, and launches a handful of small CUDA ops per iteration. The goal here is to be correct, not fast. # flashattention_autograd_function_pytorch.py import math import torch import einops class FlashAttentionPytorch(torch.autograd.Function): @staticmethod def forward(ctx, Q, K, V, is_causal=False): # Note: Tile size is fixed at 16 as a simplifying assumption tile_size = 16 # Split the sequence dimension into tiles: (..., T, B, d) # Leading dims are arbitrary (batch, heads, ...). The sequence axis N is split # into Tq tiles of Bq rows, so N must be a multiple of tile_size. Q_t = einops.rearrange(Q, "... (Tq Bq) d -> ... Tq Bq d", Bq=tile_size) K_t = einops.rearrange(K, "... (Tk Bk) d -> ... Tk Bk d", Bk=tile_size) V_t = einops.rearrange(V, "... (Tk Bv) d -> ... Tk Bv d", Bv=tile_size) O = torch.empty_like(Q) L = torch.empty(Q.shape[:-1], device=Q.device, dtype=Q.dtype) scale = 1.0 / math.sqrt(Q.shape[-1]) for i in range(Q_t.shape[-3]): # outer loop: query tiles Q_i = Q_t[..., i, :, :] O_i = torch.zeros_like(Q_i) l_i = torch.zeros(Q_i.shape[:-1] + (1,), device=Q.device, dtype=Q.dtype) m_i = torch.full(Q_i.shape[:-1] + (1,), -torch.inf, device=Q.device, dtype=Q.dtype) for j in range(K_t.shape[-3]): # inner loop: key tiles K_j = K_t[..., j, :, :] V_j = V_t[..., j, :, :] S_ij = einops.einsum(Q_i, K_j, "... Bq d, ... Bk d -> ... Bq Bk") * scale m_new = torch.maximum(m_i, S_ij.amax(dim=-1, keepdim=True)) P_ij = torch.exp(S_ij - m_new) alpha = torch.exp(m_i - m_new) # rescale factor for the old state l_i = alpha * l_i + P_ij.sum(dim=-1, keepdim=True) O_i = alpha * O_i + einops.einsum(P_ij, V_j, "... Bq Bk, ... Bk d -> ... Bq d") m_i = m_new O[..., i * tile_size:(i + 1) * tile_size, :] = O_i / l_i L[..., i * tile_size:(i + 1) * tile_size] = (m_i + torch.log(l_i)).squeeze(-1) ctx.save_for_backward(Q, K, V, O, L) ctx.is_causal = is_causal return O, L Enter fullscreen mode Exit fullscreen mode A few points to note: m_i starts at negative infinity so that the first tile's row maximum "wins" in line 10 of algorithm 1. m_i starting at negative infinity also ensures on the first step 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. O_i is the weighted sum of value rows, and the softmax needs that sum divided by the total weight l_i. There are two ways to book-keep: (i) keep the sum and divide once at the end (FA2); or (ii) keep the average and re-devide every time a new tile arrives (FA1). FA2's approach saves two elementwise passess over a 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, S with dimension NxN per head per batch element will have about 16.8 million entries, or about 33MB in bf16. L with dimension Nx1 per head per batch element contains 4096 entries, or about 16KB. The is_causal flag is accepted and ignored. The PyTorch version does not implement causal masking. Step 2: the Triton forward kernel The Triton version has the same shape as the PyTorch one, with the outer loop over query tiles replaced by the launch grid. Each program (Triton's name for a thread block) gets one query tile and one batch element, loads Q_i once, and loops over key tiles on its own. Before I introduce the implementation of FA2 in Triton, I want to take a detour into introducing tl.make_block_btr. A quick intro to 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), ) # then, inside the loop: row = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero") x_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE)) Enter fullscreen mode Exit fullscreen mode 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. Why use a block pointer A Triton program works on a small tile of data at a time. To read a tile, it needs the memory address of every element in the tile. It also must not read past the edges of the tensor. There are two ways to do this. The classic way with pointer math Before block pointers, programmers would 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) # 1. Which rows this program owns rows = row_tile_idx * ROWS_TILE_SIZE + tl.arange(0, ROWS_TILE_SIZE) row_mask = rows 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}%") Enter fullscreen mode Exit fullscreen mode 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: N causal ours (ms) SDPA (ms) ours TFLOP/s SDPA TFLOP/s % of SDPA 512 no [fill] [fill] [fill] [fill] [fill] 512 yes [fill] [fill] [fill] [fill] [fill] 1024 no [fill] [fill] [fill] [fill] [fill] 1024 yes [fill] [fill] [fill] [fill] [fill] 2048 no [fill] [fill] [fill] [fill] [fill] 2048 yes [fill] [fill] [fill] [fill] [fill] 4096 no [fill] [fill] [fill] [fill] [fill] 4096 yes [fill] [fill] [fill] [fill] [fill] [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. What's wrong with this kernel 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. What's next 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.