{"slug": "flashattention-2-from-pytorch-to-triton", "title": "FlashAttention-2 from PyTorch to Triton", "summary": "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.", "body_md": "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:\n\nFuture articles will relax these assumptions to learn these implementation details. I will state upfront whenever I made simplifying assumptions.\n\nAll 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. \n\nThis article is written with the assistance of AI.\n\nThe 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.\n\nThe 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.\n\nThree running quantities per query tile make that possible, initialised on line 6:\n\n`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)\nThe 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.\n\nAt 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.\n\nBefore 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.\n\n``` python\n# flashattention_autograd_function_pytorch.py\nimport math\nimport torch\nimport einops\n\nclass FlashAttentionPytorch(torch.autograd.Function):\n\n    @staticmethod\n    def forward(ctx, Q, K, V, is_causal=False):\n        # Note: Tile size is fixed at 16 as a simplifying assumption\n        tile_size = 16\n        # Split the sequence dimension into tiles: (..., T, B, d)\n        # Leading dims are arbitrary (batch, heads, ...). The sequence axis N is split\n        # into Tq tiles of Bq rows, so N must be a multiple of tile_size.\n        Q_t = einops.rearrange(Q, \"... (Tq Bq) d -> ... Tq Bq d\", Bq=tile_size)\n        K_t = einops.rearrange(K, \"... (Tk Bk) d -> ... Tk Bk d\", Bk=tile_size)\n        V_t = einops.rearrange(V, \"... (Tk Bv) d -> ... Tk Bv d\", Bv=tile_size)\n\n        O = torch.empty_like(Q)\n        L = torch.empty(Q.shape[:-1], device=Q.device, dtype=Q.dtype)\n        scale = 1.0 / math.sqrt(Q.shape[-1])\n\n        for i in range(Q_t.shape[-3]):                      # outer loop: query tiles\n            Q_i = Q_t[..., i, :, :]\n            O_i = torch.zeros_like(Q_i)\n            l_i = torch.zeros(Q_i.shape[:-1] + (1,), device=Q.device, dtype=Q.dtype)\n            m_i = torch.full(Q_i.shape[:-1] + (1,), -torch.inf, device=Q.device, dtype=Q.dtype)\n\n            for j in range(K_t.shape[-3]):                  # inner loop: key tiles\n                K_j = K_t[..., j, :, :]\n                V_j = V_t[..., j, :, :]\n                S_ij = einops.einsum(Q_i, K_j, \"... Bq d, ... Bk d -> ... Bq Bk\") * scale\n\n                m_new = torch.maximum(m_i, S_ij.amax(dim=-1, keepdim=True))\n                P_ij = torch.exp(S_ij - m_new)\n                alpha = torch.exp(m_i - m_new)               # rescale factor for the old state\n                l_i = alpha * l_i + P_ij.sum(dim=-1, keepdim=True)\n                O_i = alpha * O_i + einops.einsum(P_ij, V_j, \"... Bq Bk, ... Bk d -> ... Bq d\")\n                m_i = m_new\n\n            O[..., i * tile_size:(i + 1) * tile_size, :] = O_i / l_i\n            L[..., i * tile_size:(i + 1) * tile_size] = (m_i + torch.log(l_i)).squeeze(-1)\n\n        ctx.save_for_backward(Q, K, V, O, L)\n        ctx.is_causal = is_causal\n        return O, L\n```\n\nA few points to note:\n\n`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.\nThe 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. \n\nBefore I introduce the implementation of FA2 in Triton, I want to take a detour into introducing `tl.make_block_btr`.\n\n`tl.make_block_btr`\nWe will be taking apart `tl.make_block_ptr` in this section, using the `weighted_sum_fwd` kernel as an example. See code below.\n\n```\nx_block_ptr = tl.make_block_ptr(\n    x_ptr,\n    shape=(NUM_ROWS, D),\n    strides=(x_stride_row, x_stride_dim),\n    offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),\n    block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),\n    order=(1, 0),\n)\n\n# then, inside the loop:\nrow = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option=\"zero\")\nx_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE))\n```\n\nThere 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.\n\nA 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.\n\nBefore block pointers, programmers would built the addresses directly, including the official Triton tutorials. Below is the `weighted_sum_fwd` kernel written the classical way.\n\n```\nrow_tile_idx = tl.program_id(0)\n\n# 1. Which rows this program owns\nrows = row_tile_idx * ROWS_TILE_SIZE + tl.arange(0, ROWS_TILE_SIZE)\nrow_mask = rows < NUM_ROWS\n\noutput = tl.zeros((ROWS_TILE_SIZE,), dtype=tl.float32)\nfor i in range(tl.cdiv(D, D_TILE_SIZE)):\n    # 2. Which columns this step covers\n    cols = i * D_TILE_SIZE + tl.arange(0, D_TILE_SIZE)\n    col_mask = cols < D\n\n    # 3. A 2D grid of addresses, built by broadcasting\n    x_ptrs = (x_ptr + rows[:, None] * x_stride_row\n                    + cols[None, :] * x_stride_dim)\n    w_ptrs = weight_ptr + cols * weight_stride_dim\n\n    # 4. Masks for the edges, combined by hand\n    row = tl.load(x_ptrs, mask=row_mask[:, None] & col_mask[None, :], other=0.0)\n    weight = tl.load(w_ptrs, mask=col_mask, other=0.0)\n\n    output += tl.sum(row * weight[None, :], axis=1)\n\ntl.store(output_ptr + rows * output_stride_row, output, mask=row_mask)\n```\n\n`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.\n\nBelow is the kernel as it stands at the end of this post; the two lines marked `FIX` are the ones Step 3 explains.\n\n``` python\n# flashattention_autograd_function_triton.py\nimport math\nimport torch\nimport triton\nimport triton.language as tl\n\n@triton.jit\ndef flash_fwd_kernel(\n    Q_ptr, K_ptr, V_ptr, O_ptr, L_ptr,\n    stride_qb, stride_qq, stride_qd,\n    stride_kb, stride_kk, stride_kd,\n    stride_vb, stride_vk, stride_vd,\n    stride_ob, stride_oq, stride_od,\n    stride_lb, stride_lq,\n    N_QUERIES, N_KEYS,\n    scale,\n    D: tl.constexpr,\n    Q_TILE_SIZE: tl.constexpr,\n    K_TILE_SIZE: tl.constexpr,\n    is_causal: tl.constexpr,\n):\n    query_tile_index = tl.program_id(0)\n    batch_index = tl.program_id(1)\n\n    # Block pointers: a (rows, D) window into each tensor for this batch element.\n    # Q and O windows start at this program's query tile; K and V start at row 0\n    # and are advanced inside the loop.\n    Q_block_ptr = tl.make_block_ptr(\n        Q_ptr + batch_index * stride_qb,\n        shape=(N_QUERIES, D), strides=(stride_qq, stride_qd),\n        offsets=(query_tile_index * Q_TILE_SIZE, 0),\n        block_shape=(Q_TILE_SIZE, D), order=(1, 0),\n    )\n    K_block_ptr = tl.make_block_ptr(\n        K_ptr + batch_index * stride_kb,\n        shape=(N_KEYS, D), strides=(stride_kk, stride_kd),\n        offsets=(0, 0), block_shape=(K_TILE_SIZE, D), order=(1, 0),\n    )\n    V_block_ptr = tl.make_block_ptr(\n        V_ptr + batch_index * stride_vb,\n        shape=(N_KEYS, D), strides=(stride_vk, stride_vd),\n        offsets=(0, 0), block_shape=(K_TILE_SIZE, D), order=(1, 0),\n    )\n    O_block_ptr = tl.make_block_ptr(\n        O_ptr + batch_index * stride_ob,\n        shape=(N_QUERIES, D), strides=(stride_oq, stride_od),\n        offsets=(query_tile_index * Q_TILE_SIZE, 0),\n        block_shape=(Q_TILE_SIZE, D), order=(1, 0),\n    )\n    L_block_ptr = tl.make_block_ptr(\n        L_ptr + batch_index * stride_lb,\n        shape=(N_QUERIES,), strides=(stride_lq,),\n        offsets=(query_tile_index * Q_TILE_SIZE,),\n        block_shape=(Q_TILE_SIZE,), order=(0,),\n    )\n\n    # Running state, kept in fp32 regardless of input dtype.\n    O_acc = tl.zeros((Q_TILE_SIZE, D), dtype=tl.float32)\n    l_acc = tl.zeros((Q_TILE_SIZE, 1), dtype=tl.float32)\n    m_acc = tl.full((Q_TILE_SIZE, 1), value=float(\"-inf\"), dtype=tl.float32)\n\n    Q_i = tl.load(Q_block_ptr, boundary_check=(0, 1), padding_option=\"zero\")\n\n    q_pos = (query_tile_index * Q_TILE_SIZE + tl.arange(0, Q_TILE_SIZE))[:, None]\n\n    for j in range(tl.cdiv(N_KEYS, K_TILE_SIZE)):\n        k_pos = (j * K_TILE_SIZE + tl.arange(0, K_TILE_SIZE))[None, :]\n\n        K_j = tl.load(K_block_ptr, boundary_check=(0, 1), padding_option=\"zero\")\n        V_j = tl.load(V_block_ptr, boundary_check=(0, 1), padding_option=\"zero\")\n\n        S_ij = tl.dot(Q_i, tl.trans(K_j)) * scale          # (Q_TILE, K_TILE), fp32\n\n        # FIX 2: zero-padded keys past N_KEYS score 0, not -inf. Mask them.\n        keep = k_pos < N_KEYS\n        if is_causal:\n            keep = keep & (k_pos <= q_pos)\n        S_ij = tl.where(keep, S_ij, -1e6)\n\n        m_new = tl.maximum(m_acc, tl.max(S_ij, axis=1, keep_dims=True))\n        P_ij = tl.exp(S_ij - m_new)\n        alpha = tl.exp(m_acc - m_new)\n        l_acc = alpha * l_acc + tl.sum(P_ij, axis=1, keep_dims=True)\n\n        # FIX 1: the cast must be assigned. tl.dot needs both operands in the\n        # same dtype; the fp32 accumulator is passed separately via acc=.\n        P_ij = P_ij.to(V_j.dtype)\n        O_acc = alpha * O_acc\n        O_acc = tl.dot(P_ij, V_j, acc=O_acc)\n        m_acc = m_new\n\n        K_block_ptr = K_block_ptr.advance((K_TILE_SIZE, 0))\n        V_block_ptr = V_block_ptr.advance((K_TILE_SIZE, 0))\n\n    O_i = (O_acc / l_acc).to(O_block_ptr.type.element_ty)\n    tl.store(O_block_ptr, O_i, boundary_check=(0, 1))\n\n    L_i = tl.reshape(m_acc + tl.log(l_acc), (Q_TILE_SIZE,))\n    tl.store(L_block_ptr, L_i, boundary_check=(0,))\n\nclass FlashAttentionTriton(torch.autograd.Function):\n    Q_TILE_SIZE = 16\n    K_TILE_SIZE = 16\n\n    @staticmethod\n    def forward(ctx, Q, K, V, is_causal=False):\n        assert Q.ndim == 3, \"expects (batch, seq, head_dim); flatten (B, H, N, D) to (B*H, N, D)\"\n        assert Q.stride(-1) == 1 and K.stride(-1) == 1 and V.stride(-1) == 1\n        B, N_q, D = Q.shape\n        N_k = K.shape[1]\n        assert D in (16, 32, 64, 128), \"block_shape dims must be powers of two\"\n\n        O = torch.empty_like(Q)\n        L = torch.empty((B, N_q), device=Q.device, dtype=torch.float32)\n        grid = (triton.cdiv(N_q, FlashAttentionTriton.Q_TILE_SIZE), B)\n\n        flash_fwd_kernel[grid](\n            Q, K, V, O, L,\n            Q.stride(0), Q.stride(1), Q.stride(2),\n            K.stride(0), K.stride(1), K.stride(2),\n            V.stride(0), V.stride(1), V.stride(2),\n            O.stride(0), O.stride(1), O.stride(2),\n            L.stride(0), L.stride(1),\n            N_q, N_k,\n            1.0 / math.sqrt(D),\n            D=D,\n            Q_TILE_SIZE=FlashAttentionTriton.Q_TILE_SIZE,\n            K_TILE_SIZE=FlashAttentionTriton.K_TILE_SIZE,\n            is_causal=is_causal,\n        )\n        ctx.save_for_backward(Q, K, V, O, L)\n        ctx.is_causal = is_causal\n        return O\n\n    @staticmethod\n    def backward(ctx, dO):\n        raise NotImplementedError(\"tiled backward is part 4 of this series\")\n```\n\nReading it top to bottom:\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\nThe 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.\n\n``` python\n# test_flashattention_triton.py\nimport math\nimport pytest\nimport torch\nimport torch.nn.functional as F\nfrom flashattention_autograd_function_triton import FlashAttentionTriton\n\ndef reference(q, k, v, is_causal):\n    o = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)\n    s = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(q.shape[-1])\n    if is_causal:\n        tril = torch.ones(s.shape[-2:], dtype=torch.bool, device=q.device).tril()\n        s = s.masked_fill(~tril, float(\"-inf\"))\n    return o, torch.logsumexp(s, dim=-1)\n\n# (batch, n_queries, n_keys, head_dim). Deliberately includes shapes that\n# are not multiples of the 16-row tile and non-square attention.\nSHAPES = [\n    (4, 128, 128, 64),     # the assignment's shape\n    (2, 1024, 1024, 64),\n    (2, 100, 100, 64),     # ragged last tile on both axes\n    (1, 37, 53, 32),       # ragged and non-square\n    (1, 200, 96, 128),\n]\n\n@pytest.mark.parametrize(\"dtype\", [torch.float32, torch.bfloat16])\n@pytest.mark.parametrize(\"is_causal\", [False, True])\n@pytest.mark.parametrize(\"B,Nq,Nk,D\", SHAPES)\ndef test_forward_matches_sdpa(B, Nq, Nk, D, is_causal, dtype):\n    if is_causal and Nq != Nk:\n        pytest.skip(\"keep causal tests square so the mask convention is unambiguous\")\n    torch.manual_seed(0)\n    q = torch.randn(B, Nq, D, device=\"cuda\", dtype=dtype, requires_grad=True)\n    k = torch.randn(B, Nk, D, device=\"cuda\", dtype=dtype, requires_grad=True)\n    v = torch.randn(B, Nk, D, device=\"cuda\", dtype=dtype, requires_grad=True)\n\n    o = FlashAttentionTriton.apply(q, k, v, is_causal)\n    # L is not returned; pull it out of the saved tensors, as the assignment test does.\n    L = [t for t in o.grad_fn.saved_tensors if t.shape == (B, Nq)][0]\n\n    o_ref, L_ref = reference(q, k, v, is_causal)\n    # fp32 tl.dot uses TF32 on tensor cores by default; bf16 rounds P before P@V.\n    tol = dict(atol=2e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-2, rtol=1e-2)\n    torch.testing.assert_close(o, o_ref, **tol)\n    torch.testing.assert_close(L, L_ref, **tol)\n```\n\nTwo 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.\n\nThe scaffold kernel had three lines like this:\n\n```\nP_tilde_ij.to(V_j.type.element_ty)   # \"cast P to bf16 before the dot\"\n...\nO_i.to(O_block_ptr.type.element_ty)\n...\nL_i.to(L_block_ptr.type.element_ty)\n```\n\nIn 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.\n\nThe 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.\n\nThe 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.\n\n`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.\n\nAgain 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.\n\nWith 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.\n\nEvery 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.\n\n``` python\n# bench_flashattention.py\nimport torch\nimport torch.nn.functional as F\nimport triton\nfrom flashattention_autograd_function_triton import FlashAttentionTriton\n\ndef attn_flops(batch, n, d, is_causal):\n    # QK^T and PV: two matmuls of n*n*d multiply-adds each, 2 flops per MAC.\n    f = 4 * batch * n * n * d\n    return f / 2 if is_causal else f\n\ndef bench(B, H, N, D, is_causal, dtype=torch.bfloat16):\n    q, k, v = (torch.randn(B * H, N, D, device=\"cuda\", dtype=dtype) for _ in range(3))\n    ms_ours = triton.testing.do_bench(lambda: FlashAttentionTriton.apply(q, k, v, is_causal))\n\n    q4, k4, v4 = (t.view(B, H, N, D) for t in (q, k, v))\n    ms_sdpa = triton.testing.do_bench(\n        lambda: F.scaled_dot_product_attention(q4, k4, v4, is_causal=is_causal))\n\n    flops = attn_flops(B * H, N, D, is_causal)\n    tflops = lambda ms: flops / ms * 1e-9\n    return ms_ours, ms_sdpa, tflops(ms_ours), tflops(ms_sdpa)\n\nif __name__ == \"__main__\":\n    print(f\"{'N':>6} {'causal':>7} {'ours ms':>9} {'sdpa ms':>9} {'ours TF/s':>10} {'sdpa TF/s':>10} {'% sdpa':>7}\")\n    for N in (512, 1024, 2048, 4096):\n        for causal in (False, True):\n            a, b, c, d = bench(B=4, H=8, N=N, D=64, is_causal=causal)\n            print(f\"{N:>6} {str(causal):>7} {a:>9.3f} {b:>9.3f} {c:>10.1f} {d:>10.1f} {100 * b / a:>6.0f}%\")\n```\n\n`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.\n\nResults on the RTX 4070 Super, batch 4, 8 heads, head dim 64, bf16:\n\n| N | causal | ours (ms) | SDPA (ms) | ours TFLOP/s | SDPA TFLOP/s | % of SDPA | \n|---|---|---|---|---|---|---|\n| 512 | no | [fill] | [fill] | [fill] | [fill] | [fill] | \n| 512 | yes | [fill] | [fill] | [fill] | [fill] | [fill] | \n| 1024 | no | [fill] | [fill] | [fill] | [fill] | [fill] | \n| 1024 | yes | [fill] | [fill] | [fill] | [fill] | [fill] | \n| 2048 | no | [fill] | [fill] | [fill] | [fill] | [fill] | \n| 2048 | yes | [fill] | [fill] | [fill] | [fill] | [fill] | \n| 4096 | no | [fill] | [fill] | [fill] | [fill] | [fill] | \n| 4096 | yes | [fill] | [fill] | [fill] | [fill] | [fill] | \n\n[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.]\n\nTo 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.\n\nThe 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.\n\n**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`.\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\nWhat'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.\n\nPart 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.\n\nThe 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.\n\nIf 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.", "url": "https://wpnews.pro/news/flashattention-2-from-pytorch-to-triton", "canonical_source": "https://dev.to/lewis_won/flashattention-2-from-pytorch-to-triton-4ein", "published_at": "2026-09-27 02:27:24+00:00", "updated_at": "2026-09-27 03:01:07.484365+00:00", "lang": "en", "topics": ["machine-learning", "ai-research", "developer-tools"], "entities": ["Triton", "PyTorch", "FlashAttention-2", "Stanford CS336", "RTX 4070 Super", "einops"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/flashattention-2-from-pytorch-to-triton", "markdown": "https://wpnews.pro/news/flashattention-2-from-pytorch-to-triton.md", "text": "https://wpnews.pro/news/flashattention-2-from-pytorch-to-triton.txt", "jsonld": "https://wpnews.pro/news/flashattention-2-from-pytorch-to-triton.jsonld"}}