# FlashAttention-2 from PyTorch to Triton

> Source: <https://dev.to/lewis_won/flashattention-2-from-pytorch-to-triton-4ein>
> Published: 2026-09-27 02:27:24+00:00

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.