FlashAttention-2 from PyTorch to Triton

A developer implemented the FlashAttention-2 forward pass twice — first as a deliberately slow PyTorch reference with a Python double loop over tiles, then as a Triton kernel — to teach kernel optimization. The Triton version keeps three running quantities per query tile (running max m_i, running sum l_i, and unnormalized output O_i) and applies online softmax tile by tile, avoiding materialization of the N×N score matrix. It also stores the per-row logsumexp L for use in a future backward pass, with all numbers measured on a single RTX 4070 Super.

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 https://github.com/stanford-cs336/assignment2-systems/blob/main/cs336 assignment2 systems.pdf 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 https://dev.to/lewis won/online-softmax-by-hand-4h13 , 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. python flashattention autograd function pytorch.py import math import torch import einops class FlashAttentionPytorch torch.autograd.Function : @staticmethod def forward ctx, Q, K, V, is causal=False : Note: Tile size is fixed at 16 as a simplifying assumption tile size = 16 Split the sequence dimension into tiles: ..., T, B, d Leading dims are arbitrary batch, heads, ... . The sequence axis N is split into Tq tiles of Bq rows, so N must be a multiple of tile size. Q t = einops.rearrange Q, "... Tq Bq d - ... Tq Bq d", Bq=tile size K t = einops.rearrange K, "... Tk Bk d - ... Tk Bk d", Bk=tile size V t = einops.rearrange V, "... Tk Bv d - ... Tk Bv d", Bv=tile size O = torch.empty like Q L = torch.empty Q.shape :-1 , device=Q.device, dtype=Q.dtype scale = 1.0 / math.sqrt Q.shape -1 for i in range Q t.shape -3 : outer loop: query tiles Q i = Q t ..., i, :, : O i = torch.zeros like Q i l i = torch.zeros Q i.shape :-1 + 1, , device=Q.device, dtype=Q.dtype m i = torch.full Q i.shape :-1 + 1, , -torch.inf, device=Q.device, dtype=Q.dtype for j in range K t.shape -3 : inner loop: key tiles K j = K t ..., j, :, : V j = V t ..., j, :, : S ij = einops.einsum Q i, K j, "... Bq d, ... Bk d - ... Bq Bk" scale m new = torch.maximum m i, S ij.amax dim=-1, keepdim=True P ij = torch.exp S ij - m new alpha = torch.exp m i - m new rescale factor for the old state l i = alpha l i + P ij.sum dim=-1, keepdim=True O i = alpha O i + einops.einsum P ij, V j, "... Bq Bk, ... Bk d - ... Bq d" m i = m new O ..., i tile size: i + 1 tile size, : = O i / l i L ..., i tile size: i + 1 tile size = m i + torch.log l i .squeeze -1 ctx.save for backward Q, K, V, O, L ctx.is causal = is causal return O, L A few points to note: 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 , then, inside the loop: row = tl.load x block ptr, boundary check= 0, 1 , padding option="zero" x block ptr = x block ptr.advance 0, D TILE SIZE There are a total of 6 arguments in tl.make block 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 1. Which rows this program owns rows = row tile idx ROWS TILE SIZE + tl.arange 0, ROWS TILE SIZE row mask = rows < NUM ROWS output = tl.zeros ROWS TILE SIZE, , dtype=tl.float32 for i in range tl.cdiv D, D TILE SIZE : 2. Which columns this step covers cols = i D TILE SIZE + tl.arange 0, D TILE SIZE col mask = cols < D 3. A 2D grid of addresses, built by broadcasting x ptrs = x ptr + rows :, None x stride row + cols None, : x stride dim w ptrs = weight ptr + cols weight stride dim 4. Masks for the edges, combined by hand row = tl.load x ptrs, mask=row mask :, None & col mask None, : , other=0.0 weight = tl.load w ptrs, mask=col mask, other=0.0 output += tl.sum row weight None, : , axis=1 tl.store output ptr + rows output stride row, output, mask=row mask tl.load takes a whole grid of addresses, one for each element in the tile. The kernel makes that grid from two index vectors: rows goes down and cols goes across. Writing :, None turns rows into a column, and None, : turns cols into a row. Adding them fills in the grid by broadcasting. The mask is built the same way. Below is the kernel as it stands at the end of this post; the two lines marked FIX are the ones Step 3 explains. python flashattention autograd function triton.py import math import torch import triton import triton.language as tl @triton.jit def flash fwd kernel Q ptr, K ptr, V ptr, O ptr, L ptr, stride qb, stride qq, stride qd, stride kb, stride kk, stride kd, stride vb, stride vk, stride vd, stride ob, stride oq, stride od, stride lb, stride lq, N QUERIES, N KEYS, scale, D: tl.constexpr, Q TILE SIZE: tl.constexpr, K TILE SIZE: tl.constexpr, is causal: tl.constexpr, : query tile index = tl.program id 0 batch index = tl.program id 1 Block pointers: a rows, D window into each tensor for this batch element. Q and O windows start at this program's query tile; K and V start at row 0 and are advanced inside the loop. Q block ptr = tl.make block ptr Q ptr + batch index stride qb, shape= N QUERIES, D , strides= stride qq, stride qd , offsets= query tile index Q TILE SIZE, 0 , block shape= Q TILE SIZE, D , order= 1, 0 , K block ptr = tl.make block ptr K ptr + batch index stride kb, shape= N KEYS, D , strides= stride kk, stride kd , offsets= 0, 0 , block shape= K TILE SIZE, D , order= 1, 0 , V block ptr = tl.make block ptr V ptr + batch index stride vb, shape= N KEYS, D , strides= stride vk, stride vd , offsets= 0, 0 , block shape= K TILE SIZE, D , order= 1, 0 , O block ptr = tl.make block ptr O ptr + batch index stride ob, shape= N QUERIES, D , strides= stride oq, stride od , offsets= query tile index Q TILE SIZE, 0 , block shape= Q TILE SIZE, D , order= 1, 0 , L block ptr = tl.make block ptr L ptr + batch index stride lb, shape= N QUERIES, , strides= stride lq, , offsets= query tile index Q TILE SIZE, , block shape= Q TILE SIZE, , order= 0, , Running state, kept in fp32 regardless of input dtype. O acc = tl.zeros Q TILE SIZE, D , dtype=tl.float32 l acc = tl.zeros Q TILE SIZE, 1 , dtype=tl.float32 m acc = tl.full Q TILE SIZE, 1 , value=float "-inf" , dtype=tl.float32 Q i = tl.load Q block ptr, boundary check= 0, 1 , padding option="zero" q pos = query tile index Q TILE SIZE + tl.arange 0, Q TILE SIZE :, None for j in range tl.cdiv N KEYS, K TILE SIZE : k pos = j K TILE SIZE + tl.arange 0, K TILE SIZE None, : K j = tl.load K block ptr, boundary check= 0, 1 , padding option="zero" V j = tl.load V block ptr, boundary check= 0, 1 , padding option="zero" S ij = tl.dot Q i, tl.trans K j scale Q TILE, K TILE , fp32 FIX 2: zero-padded keys past N KEYS score 0, not -inf. Mask them. keep = k pos < N KEYS if is causal: keep = keep & k pos <= q pos S ij = tl.where keep, S ij, -1e6 m new = tl.maximum m acc, tl.max S ij, axis=1, keep dims=True P ij = tl.exp S ij - m new alpha = tl.exp m acc - m new l acc = alpha l acc + tl.sum P ij, axis=1, keep dims=True FIX 1: the cast must be assigned. tl.dot needs both operands in the same dtype; the fp32 accumulator is passed separately via acc=. P ij = P ij.to V j.dtype O acc = alpha O acc O acc = tl.dot P ij, V j, acc=O acc m acc = m new K block ptr = K block ptr.advance K TILE SIZE, 0 V block ptr = V block ptr.advance K TILE SIZE, 0 O i = O acc / l acc .to O block ptr.type.element ty tl.store O block ptr, O i, boundary check= 0, 1 L i = tl.reshape m acc + tl.log l acc , Q TILE SIZE, tl.store L block ptr, L i, boundary check= 0, class FlashAttentionTriton torch.autograd.Function : Q TILE SIZE = 16 K TILE SIZE = 16 @staticmethod def forward ctx, Q, K, V, is causal=False : assert Q.ndim == 3, "expects batch, seq, head dim ; flatten B, H, N, D to B H, N, D " assert Q.stride -1 == 1 and K.stride -1 == 1 and V.stride -1 == 1 B, N q, D = Q.shape N k = K.shape 1 assert D in 16, 32, 64, 128 , "block shape dims must be powers of two" O = torch.empty like Q L = torch.empty B, N q , device=Q.device, dtype=torch.float32 grid = triton.cdiv N q, FlashAttentionTriton.Q TILE SIZE , B flash fwd kernel grid Q, K, V, O, L, Q.stride 0 , Q.stride 1 , Q.stride 2 , K.stride 0 , K.stride 1 , K.stride 2 , V.stride 0 , V.stride 1 , V.stride 2 , O.stride 0 , O.stride 1 , O.stride 2 , L.stride 0 , L.stride 1 , N q, N k, 1.0 / math.sqrt D , D=D, Q TILE SIZE=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. python test flashattention triton.py 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 batch, n queries, n keys, head dim . Deliberately includes shapes that are not multiples of the 16-row tile and non-square attention. 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 is not returned; pull it out of the saved tensors, as the assignment test does. 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 fp32 tl.dot uses TF32 on tensor cores by default; bf16 rounds P before P@V. 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. python bench flashattention.py 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 : QK^T and PV: two matmuls of n n d multiply-adds each, 2 flops per MAC. 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: | 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. 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 https://github.com/searlion/cs336-2026-assignment2-systems/tree/main/flash forward ; the earlier posts on online softmax https://dev.to/lewis won/online-softmax-by-hand-4h13 and FlashAttention by hand https://dev.to/lewis won/flashattention-by-hand-34im cover the algorithm at the level of individual numbers if the recap above went by too fast.