4.3 FlashAttention 2 & 3: Work Partitioning, Warp Specialization & FP8¶

FlashAttention I (covered in FlashAttention I: IO-Awareness & The Online Softmax) solved the memory problem: by tiling the attention computation and keeping the intermediate softmax statistics in on-chip SRAM, it computes exact attention without ever materializing the \(N \times N\) score matrix in high-bandwidth memory (HBM). That single idea — fuse the whole attention into one kernel, never write the scores out — is the foundation everything in this chapter builds on.

But FlashAttention-1 (FA1) left a lot of GPU performance on the table. Measured against the device’s peak matmul throughput, the original kernel ran at roughly 25–40% of the GPU’s FLOP ceiling on an A100, while a well-tuned dense GEMM hits 80–90%. The gap is not memory traffic — FA1 already minimized that. The gap is occupancy and instruction mix: the kernel was leaving compute units idle and burning too many cycles on non-matmul arithmetic. FlashAttention-2 (FA2, Dao 2023) closes most of that gap on Ampere with three changes to how the work is partitioned. FlashAttention-3 (FA3, Shah et al. 2024) then exploits hardware features specific to NVIDIA’s Hopper architecture — asynchronous tensor-core instructions, the Tensor Memory Accelerator (TMA), and FP8 tensor cores — to push utilization to roughly 75% of bf16 peak and into the petaFLOP range with FP8.

This chapter is about the systems engineering of attention kernels. We assume you already understand the online-softmax recurrence; here we focus on work partitioning (which thread does what), warp specialization (producer/consumer pipelines), and low-precision attention (FP8 with incoherent processing). To follow the hardware reasoning, keep GPU Architecture & The Memory Hierarchy and The Roofline Model & Performance Engineering nearby — we lean heavily on the concepts of warps, warp-groups, shared memory, register pressure, and arithmetic intensity.

What FlashAttention-1 left on the table¶

Let us restate the FA1 inner loop precisely, because the FA2 improvements are best understood as edits to it. We tile the query into blocks of \(B_r\) rows and the keys/values into blocks of \(B_c\) columns. For a fixed query block \(Q_i\) we loop over all key blocks \(K_j, V_j\), maintaining a running max \(m\), running denominator \(\ell\), and a running unnormalized output accumulator \(O\):

After the last key block, the final output is \(O / \ell\). This is exact attention — the online recurrence reproduces the global softmax. Three inefficiencies hide in this loop.

1. The rescaling does too much non-matmul work. Every inner iteration multiplies the accumulator \(O\) by the correction factor \(e^{m - m^{\text{new}}}\). On a GPU, FLOPs are not fungible: a tensor-core matmul instruction on an A100 runs at ~312 TFLOP/s in bf16, but the same hardware executes element-wise transcendental and multiply operations on the much slower CUDA cores at perhaps 1/16^th that rate. So the \(\exp\), the rowmax, and the per-iteration rescale — though they are a small fraction of total FLOPs — consume a disproportionate share of time. FA1 rescales \(O\) once per inner iteration; that is more non-matmul work than necessary.

2. The parallelization wastes blocks. FA1 parallelized over the batch and head dimensions and over query blocks, assigning one thread block per \((\text{batch},\text{head},\text{query-block})\). The number of these is batch * heads * (seqlen / B_r). For training large models with long sequences, batch-times-heads can be small (you trade batch for sequence length), so the grid may not have enough thread blocks to fill all the streaming multiprocessors (SMs). An A100 has 108 SMs; if your grid launches only 60 thread blocks, 44% of the chip is idle.

3. Work is split badly across warps inside a block. A thread block on Ampere is typically 4 or 8 warps (128–256 threads). FA1 split the key/value block across warps: each warp computed \(Q_i K_j^\top\) for a slice of columns. But then every warp needs the full \(Q_i\) in registers, and after computing its partial \(P V\) the warps must communicate through shared memory to combine results, because each warp owns a column-slice of the scores but the output reduction is along columns. This shared-memory round trip — the so-called “split-K” pattern — stalls the warps and adds synchronization.

FA2 fixes all three. None of these changes touch the math; they change who computes what, in what order.

FlashAttention-2: three edits to the inner loop¶

Edit 1 — defer the rescaling¶

The first FA2 trick removes the per-iteration accumulator rescale. Instead of keeping \(O\) correctly normalized at every step, we keep an unnormalized accumulator and only divide by \(\ell\) once at the very end.

Look again at the recurrence. The accumulator update is \(O^{\text{new}} = e^{m - m^{\text{new}}} O + P_{ij} V_j\). The factor \(e^{m - m^{\text{new}}}\) must still be applied — it corrects for the change in the running max — so we cannot drop it. But the normalization by \(\ell\) can be deferred. In FA1’s framing the normalization and the max-correction were tangled together; FA2 separates them. The accumulator is held as \(\tilde{O}\) (unnormalized), and we only compute \(O_i = \tilde{O}_i / \ell_i\) after the key loop finishes. This removes one full division per element per inner iteration and replaces it with a single division at the end.

More importantly, FA2 reorganizes which operations are on the critical path so that the bulk of the inner loop is two back-to-back matmuls (\(Q K^\top\) then \(P V\)) with the minimum possible scalar work between them. The rule of thumb: keep the tensor cores fed; do everything else as rarely as possible.

Here is a from-scratch, runnable reference implementation in pure PyTorch that mirrors the FA2 dataflow. It is not fast — it is meant to make the bookkeeping unambiguous so you can map it onto the CUDA kernel.

import torch

def flash_attention_2_reference(Q, K, V, Br=64, Bc=64, causal=False):
    """
    Exact attention via the FlashAttention-2 recurrence.
    Q, K, V: (N, d). Returns O: (N, d).

    Mirrors the FA2 dataflow:
      - outer loop over QUERY blocks (this is the parallel axis on the GPU)
      - inner loop over KEY/VALUE blocks
      - the output accumulator is kept UNNORMALIZED and divided by l only at the end
    """
    N, d = Q.shape
    O = torch.zeros_like(Q)
    # running statistics, one per query row
    for i in range(0, N, Br):
        qi = Q[i:i+Br]                          # (Br, d)
        Oi = torch.zeros(qi.shape[0], d)        # UNNORMALIZED accumulator
        mi = torch.full((qi.shape[0],), -float('inf'))   # running max
        li = torch.zeros(qi.shape[0])           # running denominator (sum of exp)

        for j in range(0, N, Bc):
            if causal and j > i + Br - 1:
                break                            # whole key block is in the future
            kj = K[j:j+Bc]                       # (Bc, d)
            vj = V[j:j+Bc]                       # (Bc, d)

            # ---- matmul 1: scores ----
            Sij = (qi @ kj.T) / (d ** 0.5)       # (Br, Bc)
            if causal:
                # mask future positions inside the diagonal block
                qpos = torch.arange(i, i+qi.shape[0]).unsqueeze(1)
                kpos = torch.arange(j, j+kj.shape[0]).unsqueeze(0)
                Sij = Sij.masked_fill(kpos > qpos, -float('inf'))

            # ---- online softmax statistics ----
            mij = Sij.max(dim=1).values                  # block rowmax
            mi_new = torch.maximum(mi, mij)
            Pij = torch.exp(Sij - mi_new.unsqueeze(1))   # (Br, Bc)
            alpha = torch.exp(mi - mi_new)               # correction factor

            # rescale running denom and accumulator by alpha
            li = alpha * li + Pij.sum(dim=1)

            # ---- matmul 2: weighted values, accumulated UNNORMALIZED ----
            Oi = alpha.unsqueeze(1) * Oi + Pij @ vj      # (Br, d)

            mi = mi_new

        # single normalization at the very end (FA2's deferred division)
        O[i:i+qi.shape[0]] = Oi / li.unsqueeze(1)
    return O


if __name__ == "__main__":
    torch.manual_seed(0)
    N, d = 256, 64
    Q, K, V = torch.randn(N, d), torch.randn(N, d), torch.randn(N, d)
    ref = torch.softmax(Q @ K.T / d**0.5, dim=-1) @ V
    out = flash_attention_2_reference(Q, K, V)
    print("max abs error vs dense softmax:", (ref - out).abs().max().item())
    # prints something like 1e-6 — exact attention up to floating point

The single line Oi = alpha.unsqueeze(1) * Oi + Pij @ vj plus the final Oi / li is the heart of the deferred-normalization scheme. Compare it to FA1, which would divide inside the loop. The savings are not the division per se but the register and instruction budget freed up for the tensor cores.

Edit 2 — parallelize over the sequence dimension¶

FA2’s second change is at the grid level, not inside the kernel. FA1 assigned one thread block per query block but did not use sequence-length parallelism beyond that. FA2 makes the outer loop over query blocks the parallel axis and additionally allows splitting along sequence length so that the grid has enough thread blocks to saturate every SM even when batch and heads are small.

Concretely, the FA2 launch grid is three-dimensional:

gridDim = (ceil(seqlen_q / Br),   num_heads,   batch_size)

Each thread block owns one query block for one (head, batch) and loops over all key blocks itself. Because the query-block axis is now first and independent, long sequences directly produce more thread blocks. For a sequence of 8192 tokens with Br = 64, that is 128 query blocks per head per batch element — easily enough to fill 108 SMs even at batch=1, head=1.

For the backward pass the story is subtler. The backward pass needs to accumulate gradients \(dK\) and \(dV\), which are reductions over query blocks. FA2 parallelizes the backward over key/value blocks (so each thread block owns a \(K_j, V_j\) and loops over query blocks), avoiding atomic adds into \(dK, dV\). The asymmetry — forward parallel over queries, backward parallel over keys — is a recurring theme in attention kernels: pick the parallel axis that makes the reduction local.

Edit 3 — partition warps over queries, not keys¶

The third FA2 change is the most hardware-specific and the one that gives the kernel its “warp partitioning” name. Recall a thread block on Ampere has, say, 4 warps. The question is how to divide the \(B_r \times B_c\) tile of work among them.

FA1 used the split-K layout: each warp computes \(Q_i K_j^\top\) for all \(B_r\) query rows but only a slice of the \(B_c\) key columns. The problem: after the second matmul \(P_{ij} V_j\), the partial outputs from different warps must be summed (they hold partial sums over disjoint key columns), and that reduction goes through shared memory with a synchronization barrier. Every inner iteration pays for a shared-memory write, a __syncthreads(), and a read.

FA2 flips this to the split-Q layout (sometimes called “warp partitioning over rows”): each warp owns a slice of the \(B_r\) query rows and computes the full row of scores and output for those rows. Now warp 0 produces the final output for query rows 0–15, warp 1 for rows 16–31, and so on — no inter-warp reduction is needed, because each warp’s rows are independent. \(K_j\) and \(V_j\) are shared across warps (read from shared memory), but there is no write-back-and-sum.

SPLIT-K (FA1)                           SPLIT-Q (FA2)
each warp: all Br rows, slice of cols    each warp: slice of Br rows, all cols

  Q_i (all rows)                          Q_i rows 0-15  -> warp 0  (full output)
   |                                      Q_i rows 16-31 -> warp 1  (full output)
   v                                      Q_i rows 32-47 -> warp 2  (full output)
  K_j cols 0-15  -> warp 0  (partial)     Q_i rows 48-63 -> warp 3  (full output)
  K_j cols 16-31 -> warp 1  (partial)
   ...                                    K_j / V_j: SHARED across warps (read-only)
  REDUCE across warps via shared mem  <-- no reduction needed
  + __syncthreads()                       no extra __syncthreads in the hot path

Removing that per-iteration synchronization and shared-memory traffic is worth a large fraction of FA2’s speedup. The cost is that each warp needs the \(K_j, V_j\) tile available; FA2 keeps those in shared memory and lets all warps read them, which is exactly what shared memory is good at (broadcast reads).

Putting the three edits together, FA2 roughly doubles FA1’s throughput on A100, reaching on the order of 50–73% of theoretical bf16 matmul peak depending on head dimension and sequence length. In practice you never write this yourself — you call the library:

import torch
import torch.nn.functional as F

# PyTorch's scaled_dot_product_attention dispatches to a FlashAttention-2
# backend on Ampere+/CUDA when the inputs qualify (fp16/bf16, head_dim <= 256,
# contiguous, etc.). This is the FA2 kernel in production form.
q = torch.randn(2, 16, 8192, 128, device="cuda", dtype=torch.bfloat16)  # (B,H,N,d)
k = torch.randn_like(q)
v = torch.randn_like(q)

with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

print(out.shape)  # (2, 16, 8192, 128)

The backward pass: recompute, and the one identity that makes it tile-able¶

The forward is only half the kernel — training spends roughly twice as long in the backward, and a from-scratch implementation has to get both right (this is precisely the CS336 systems assignment, written in Triton; see Writing GPU Kernels with Triton). The backward keeps the same IO discipline: it never materializes \(S\) or \(P\) in HBM. That is possible because the forward saves just two things per query row: the output \(O_i\) and the logsumexp \(L_i = m_i + \log \ell_i\), one fp32 number per row (this is the softmax_lse tensor flash_attn optionally returns). From \(L\) alone the probabilities are recoverable exactly, with no running max and no second online pass, since \(L_i\) is already the finished global normalizer for that row:

Given the incoming gradient \(dO\), and writing \(\circ\) for the elementwise product and \(\tau = 1/\sqrt{d}\), the whole backward is:

Two points carry the design. First, \(dS = P \circ (dP - D)\) is the row-softmax Jacobian in closed form — you never build the per-row \(N \times N\) Jacobian. Second, and this is the identity that makes the backward tile-able at all: \(D_i\) looks like a reduction over the entire key axis (so a tile working on key block \(j\) could not compute it), but it collapses to the row dot product of two tensors you already have,

So \(D\) is computed once in a tiny elementwise preprocessing kernel over \((dO, O)\) and every tile reads it locally. Without it, a tiled backward would need an extra full pass over the keys.

That is five matmuls (recompute \(S\), then \(dV\), \(dP\), \(dQ\), \(dK\)) against two in the forward, which is why the backward runs at roughly \(2.5\times\) the forward FLOPs. The parallelization asymmetry mentioned above also falls out here: give each thread block a \(K_j, V_j\) tile and \(dK_j, dV_j\) accumulate locally in registers, while \(dQ\) — a reduction over key blocks — becomes the one gradient that must be combined across blocks, via fp32 atomic adds or a separate pass. Those atomics are also why attention backward is non-deterministic by default; flash_attn’s deterministic=True flag (and PyTorch’s use_deterministic_algorithms) selects a slower, reproducible path.

import torch

def flash_attention_2_backward_reference(Q, K, V, O, L, dO, causal=False):
    """
    FA2-style backward for one head. The only things carried over from the
    forward are O and the per-row logsumexp L -- S and P are RECOMPUTED
    (here all at once for clarity; the real kernel does it tile by tile).
    Q, K, V, O, dO: (N, d).  L: (N,), L_i = logsumexp_j S_ij.
    """
    N, d = Q.shape
    scale = d ** -0.5
    S = (Q @ K.T) * scale                        # matmul 1: recompute scores
    if causal:
        pos = torch.arange(N)
        S = S.masked_fill(pos.unsqueeze(0) > pos.unsqueeze(1), -float('inf'))
    P = torch.exp(S - L.unsqueeze(1))            # exact probabilities from L alone

    dV = P.T @ dO                                # matmul 2
    dP = dO @ V.T                                # matmul 3
    D = (dO * O).sum(dim=1)                      # rowsum(P .* dP) == rowdot(dO, O)
    dS = P * (dP - D.unsqueeze(1))               # softmax Jacobian, closed form
    dQ = (dS @ K) * scale                        # matmul 4  (reduces over key blocks)
    dK = (dS.T @ Q) * scale                      # matmul 5
    return dQ, dK, dV


if __name__ == "__main__":
    torch.manual_seed(0)
    N, d = 128, 32
    Q, K, V = (torch.randn(N, d, requires_grad=True) for _ in range(3))

    # reference forward with autograd, saving ONLY O and L (as FA2 does)
    S = (Q @ K.T) * d ** -0.5
    pos = torch.arange(N)
    S = S.masked_fill(pos.unsqueeze(0) > pos.unsqueeze(1), -float('inf'))
    L = torch.logsumexp(S, dim=1)
    O = torch.softmax(S, dim=1) @ V
    dO = torch.randn_like(O)
    O.backward(dO)

    dQ, dK, dV = flash_attention_2_backward_reference(
        Q.detach(), K.detach(), V.detach(), O.detach(), L.detach(), dO, causal=True)
    for name, mine, ref in [("dQ", dQ, Q.grad), ("dK", dK, K.grad), ("dV", dV, V.grad)]:
        print(name, "max abs err:", (mine - ref).abs().max().item())
    # all ~1e-6 -- the recomputed backward is exact, not an approximation

Note how this differs from ordinary gradient checkpointing (Memory-Efficient Training: Checkpointing, Offloading & LoRA Math): checkpointing re-runs a region of the graph to save activations; FlashAttention’s backward recomputes \(S\) and \(P\) inside the kernel, in SRAM, so the recomputed tensor never reaches HBM at all. The activation memory of attention drops from \(O(N^2)\) to \(O(N d)\) for \(O\) plus \(O(N)\) for \(L\), which is what makes 8k–128k contexts trainable.

Calling the real kernel: flash-attn, SDPA, and packed sequences¶

Everything above is for understanding the kernel; in production you call one of three entry points, and they are not drop-in interchangeable.

• torch.nn.functional.scaled_dot_product_attention (SDPA) — in-tree, no extra dependency, dispatches among a FlashAttention backend, the xFormers-derived memory-efficient backend, a cuDNN backend, and a math fallback. Layout (B, H, N, d). This is the right default.
• The flash-attn package (Dao-AILab/flash-attention) — the reference implementation and usually the first place new features land (variable-length batching, GQA-aware kernels, sliding windows, logit softcapping, and the FA3 Hopper kernels, which ship as a separate flash_attn_interface module built from the repo’s hopper/ directory). Its layout is (B, N, H, d) — heads third, not second — and passing an SDPA-shaped tensor to it is a classic silently-wrong-results bug.
• Serving engines (vLLM, SGLang, TensorRT-LLM) wrap their own paged variants of these kernels; see PagedAttention & KV-Cache Memory Management.

# pip install flash-attn --no-build-isolation    # compiles CUDA kernels; needs nvcc
import torch
from flash_attn import flash_attn_func, flash_attn_varlen_func

B, N, H, Hkv, d = 2, 8192, 16, 4, 128    # GQA: 16 query heads share 4 KV heads
q = torch.randn(B, N, H,   d, device="cuda", dtype=torch.bfloat16)   # (B, N, H, d)!
k = torch.randn(B, N, Hkv, d, device="cuda", dtype=torch.bfloat16)
v = torch.randn_like(k)

# GQA is native: nheads_k need only divide nheads_q -- no KV replication in HBM.
out = flash_attn_func(q, k, v, causal=True)                     # (B, N, H, d)

# Sliding-window attention: window_size=(left, right) counts keys strictly to
# each side, so (4095, 0) with causal=True is a 4096-token window including self
# -- the same mask you implement by hand in Exercise 5, but block-skipped in HW.
out_swa = flash_attn_func(q, k, v, causal=True, window_size=(4095, 0))

Packed variable-length batches. Pretraining concatenates many short documents into one fixed-length row to avoid padding waste. Run plain causal attention on such a row and tokens attend across document boundaries — a quiet data-contamination bug. flash_attn_varlen_func is the fix: flatten the batch to (total_tokens, H, d) and hand the kernel a cumulative-length index, and its block scheduler simply refuses to schedule key blocks outside the current document (the same block-skipping logic as the causal aside above, applied per segment, so you also save the masked FLOPs rather than multiplying by zero).

total, H, Hkv, d = 2048, 16, 4, 128
# Three documents of length 300, 1000, 748 packed into one 2048-token row.
lengths = torch.tensor([300, 1000, 748], device="cuda")
cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device="cuda"),
                        lengths.cumsum(0).to(torch.int32)])   # [0, 300, 1300, 2048]

qf = torch.randn(total, H,   d, device="cuda", dtype=torch.bfloat16)
kf = torch.randn(total, Hkv, d, device="cuda", dtype=torch.bfloat16)
vf = torch.randn_like(kf)

out_packed = flash_attn_varlen_func(
    qf, kf, vf,
    cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
    max_seqlen_q=int(lengths.max()), max_seqlen_k=int(lengths.max()),
    causal=True)                     # (total, H, d); no cross-document attention

In Hugging Face transformers you normally select the backend at load time rather than calling either API yourself, and DataCollatorWithFlattening produces the flattened batches that route a model to the varlen path:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "HuggingFaceTB/SmolLM2-135M", dtype=torch.bfloat16, device_map="cuda",
    attn_implementation="flash_attention_2",   # or "sdpa" (default) / "eager"
)
print(model.config._attn_implementation)       # confirm what you actually got

(Recent transformers releases also expose an FA3/Hopper option under a separate implementation name; check your version’s docs rather than assuming, and always verify with the line above.) Stack-100M trains on exactly this pattern — 2048-token packed windows, document-boundary-aware attention — in The Pretraining Run: A Complete Single-GPU Training Loop, fed by the packing described in Data: Sourcing, Filtering, Dedup, Tokenize & Pack ~20B Tokens.

Finally, the cheapest way to check that you are on the fast path is to demand it: restricting SDPA to a single backend makes it raise instead of silently falling back.

import torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend

q_bhnd = q.transpose(1, 2)   # SDPA wants (B, H, N, d); flash-attn wants (B, N, H, d)
k_bhnd, v_bhnd = k.transpose(1, 2), v.transpose(1, 2)

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):    # raises if FA cannot run these inputs
    out = F.scaled_dot_product_attention(q_bhnd, k_bhnd, v_bhnd,
                                         enable_gqa=True, is_causal=True)

If it raises, one of the fast-path conditions failed — unsupported dtype (fp32 never qualifies), head dimension out of range, a materialized additive float mask instead of is_causal=True, or awkward strides. That single line turns “I think FlashAttention is on” into a fact.

FlashAttention-3: exploiting Hopper¶

FA2 is close to optimal for Ampere’s programming model, where the tensor-core instruction (mma) is synchronous: a warp issues it and waits for the result. NVIDIA’s Hopper architecture (H100) added new hardware that breaks that synchronous assumption, and FA3 is essentially “FA2 rewritten to use everything Hopper offers.” Three Hopper features matter.

Warpgroup-wide asynchronous matmul (WGMMA). Hopper’s wgmma instruction is issued by a warpgroup (4 warps = 128 threads acting together) and runs asynchronously: the warpgroup issues it, then continues executing other instructions while the tensor cores grind, and later waits on the result. This is the key that lets us overlap matmul with softmax.

The Tensor Memory Accelerator (TMA). TMA is a dedicated hardware unit that performs bulk asynchronous copies between global memory and shared memory, with the address generation done in hardware. Instead of every thread computing addresses and issuing loads, a single thread kicks off a TMA copy of a whole tile and the unit streams it in the background. This frees the warps from address arithmetic and overlaps data movement with compute.

FP8 tensor cores. Hopper tensor cores can do matmuls in 8-bit floating point (e4m3/e5m2), roughly doubling throughput over bf16 — on the order of ~2 PFLOP/s for FP8 on an H100 versus ~1 PFLOP/s bf16 (dense). For attention, where the matmuls dominate, FP8 is a huge prize if you can manage the accuracy.

FA3’s design is built around two ideas: warp specialization (a producer/consumer pipeline) and overlapping the two matmuls with the softmax.

Warp specialization: producer/consumer with TMA¶

In FA2, every warp does the same thing: load, matmul, softmax, matmul. FA3 instead specializes warps into roles, turning the thread block into a software pipeline.

• Producer warps do almost no compute. Their job is to issue TMA loads that stream the next \(K_j, V_j\) tiles from HBM into a circular buffer in shared memory, staying one or two stages ahead of consumption.
• Consumer warpgroups do the math. They wait until a tile is present in the shared-memory buffer (signaled by the producer via an asynchronous barrier), then run wgmma for \(QK^\top\), the softmax, and wgmma for \(PV\).

This is the classic producer/consumer decoupling. Because the producer uses TMA (which is asynchronous and doesn’t occupy the tensor cores or the math pipes), data movement for iteration \(j+1\) happens while the consumer computes iteration \(j\). The shared-memory buffer is a small ring (often 2–3 stages deep); synchronization uses Hopper’s mbarrier asynchronous transaction barriers rather than __syncthreads().

                         shared-memory ring buffer (stages 0,1,2)
   PRODUCER warps   --TMA-->  [ K0 V0 | K1 V1 | K2 V2 ]  --read-->  CONSUMER warpgroups
   (issue async copies,                                            (wgmma QK^T, softmax,
    stay ahead)         <-- "buffer empty" signal (mbarrier) --     wgmma PV, advance)

   time --->
   prod:  load0  load1  load2  load3  load4 ...
   cons:         gemm0  gemm1  gemm2  gemm3 ...   (always one step behind, never starved)

The payoff: the tensor cores rarely stall waiting for data, and the warps doing math never spend cycles on address generation. Register usage is also better balanced — Hopper lets you reallocate registers between warpgroups (setmaxnreg), giving the compute warpgroups more registers and the producer warps fewer.

Overlapping the two matmuls with softmax (ping-pong & pipelined)¶

The deeper FA3 trick attacks the worked-example problem above: the softmax (an \(\exp\)-heavy non-matmul op running on the slow special-function units) sits between the two matmuls and stalls the tensor cores. FA3 overlaps them in two complementary ways.

Inter-warpgroup “ping-pong” scheduling. With two consumer warpgroups, FA3 staggers them so that while warpgroup 1 is doing its softmax (on the special-function units / CUDA cores), warpgroup 2 is doing its wgmma (on the tensor cores), and then they swap. Because softmax uses different execution units than the matmul, the two warpgroups interleave to keep both unit types busy. The scheduler uses synchronization barriers to force this alternation, so the tensor cores are almost never idle waiting on an \(\exp\).

PING-PONG across two consumer warpgroups (WG):

WG1:  GEMM(QK)  | softmax  | GEMM(PV) | GEMM(QK)  | softmax  | ...
WG2:  ........  | GEMM(QK) | softmax  | GEMM(PV)  | GEMM(QK) | ...
            ^ while WG1 does softmax (SFU), WG2 does GEMM (tensor cores)
              -> tensor cores and SFUs both stay busy

Intra-warpgroup pipelining. Within a single warpgroup, FA3 also software-pipelines across inner-loop iterations: the \(QK^\top\) matmul of iteration \(j+1\) is issued (asynchronously, via wgmma) before the softmax of iteration \(j\) finishes, so the score matmul for the next block overlaps the current block’s exponentiation. Because wgmma is asynchronous, you issue it and move on; the dependency is resolved later by a wait. This is exactly the kind of overlap that is impossible on Ampere, where mma is synchronous.

The combination — TMA-fed producer/consumer plus matmul/softmax overlap — is why FA3 reaches roughly 75% of H100 bf16 peak (on the order of 700+ TFLOP/s) where FA2 on the same hardware sits closer to 35–45%. FA2 simply cannot express the overlap; it was designed for a synchronous instruction set.

A from-scratch model of the producer/consumer pipeline¶

You cannot write wgmma in Python, but you can simulate the dataflow to internalize the dependency structure. The following models the ring buffer and the producer/consumer roles, computing exact attention while showing how a consumer is “always one step behind” the producer.

import torch
from collections import deque

def fa3_pipeline_model(Q, K, V, Bc=64, n_stages=2):
    """
    A *dataflow model* of FA3's producer/consumer pipeline for one query block.
    It does NOT use real async hardware; it shows the ordering:
      - producer 'prefetches' up to n_stages KV tiles into a ring buffer
      - consumer pops a tile only after it is present, then does the GEMMs/softmax
    Result is exact attention for the single query block Q.
    """
    N, d = K.shape
    ring = deque()                      # the shared-memory ring buffer of (Kj, Vj)
    Oi = torch.zeros(Q.shape[0], d)     # unnormalized accumulator
    mi = torch.full((Q.shape[0],), -float('inf'))
    li = torch.zeros(Q.shape[0])

    key_blocks = [(j, K[j:j+Bc], V[j:j+Bc]) for j in range(0, N, Bc)]
    it = iter(key_blocks)

    def produce():                      # PRODUCER: issue one "TMA load"
        try:
            ring.append(next(it))       # copy a KV tile into the ring
            return True
        except StopIteration:
            return False

    # prime the pipeline: producer runs ahead by n_stages
    for _ in range(n_stages):
        produce()

    while ring:                         # CONSUMER drains the ring
        j, kj, vj = ring.popleft()      # wait-for-tile then take it
        produce()                       # producer refills behind the consumer

        Sij = (Q @ kj.T) / (d ** 0.5)          # wgmma #1 (async on real HW)
        mij = Sij.max(dim=1).values
        mi_new = torch.maximum(mi, mij)
        Pij = torch.exp(Sij - mi_new.unsqueeze(1))   # softmax on SFU (overlapped)
        alpha = torch.exp(mi - mi_new)
        li = alpha * li + Pij.sum(dim=1)
        Oi = alpha.unsqueeze(1) * Oi + Pij @ vj      # wgmma #2 (async on real HW)
        mi = mi_new

    return Oi / li.unsqueeze(1)


if __name__ == "__main__":
    torch.manual_seed(1)
    N, d = 512, 64
    Q = torch.randn(128, d); K = torch.randn(N, d); V = torch.randn(N, d)
    ref = torch.softmax(Q @ K.T / d**0.5, dim=-1) @ V
    out = fa3_pipeline_model(Q, K, V)
    print("max abs error:", (ref - out).abs().max().item())  # ~1e-6

The takeaway is the ordering: produce() is called before the consumer touches the tile, and again after it pops one, so the buffer stays full and the consumer never waits. On real Hopper hardware the produce() is a TMA instruction that the GPU executes concurrently with the consumer’s wgmmas.

FP8 attention and incoherent processing¶

The biggest single lever in FA3 is FP8. Hopper’s FP8 tensor cores roughly double matmul throughput, and since attention is matmul-bound, FP8 attention can nearly double the kernel’s speed. But naive FP8 attention is inaccurate, and FA3 introduces a specific technique — borrowed from quantization literature — to make it usable.

Why FP8 attention is hard¶

FP8 e4m3 has 1 sign bit, 4 exponent bits, 3 mantissa bits — about 3–4 bits of precision, representing a small dynamic range with coarse steps. Two problems arise.

Outliers blow up the scale. To quantize a tensor to FP8 you pick a scale \(s\) so that values map into the representable range. If a few entries of \(Q\) or \(K\) are large outliers (and attention activations are notoriously outlier-prone — see Quantization I: Post-Training Quantization (GPTQ, AWQ, SmoothQuant)), the scale must stretch to cover them, crushing all the typical values into just a handful of FP8 levels. The quantization error on the bulk of the data explodes.

Layout mismatch for the second matmul. FP8 wgmma on Hopper requires specific operand layouts; the probability matrix \(P\) produced by softmax and the value matrix \(V\) may not be in the right layout, forcing in-kernel transposes/shuffles. FA3 handles this with byte-permute tricks and careful layout choices, but it is real engineering, not free.

Incoherent processing: rotate away the outliers¶

FA3’s accuracy fix for the outlier problem is incoherent processing, an idea from the QuIP quantization work (Chee et al.). The insight: multiply \(Q\) and \(K\) by a random orthogonal matrix \(M\) before quantizing. Because \(M\) is orthogonal, \(M^\top M = I\), so

the attention scores are unchanged. But the rotation spreads each outlier across many coordinates: a single huge entry becomes many moderate entries, so no coordinate dominates and the FP8 scale no longer has to stretch to cover an outlier. The quantization becomes far more uniform and the error drops.

In practice FA3 uses a structured orthogonal matrix that is cheap to apply — a random-sign Hadamard transform, \(M = \frac{1}{\sqrt{d}} H D\), where \(H\) is the Hadamard matrix (applied in \(O(d \log d)\) via the fast Hadamard transform) and \(D\) is a random diagonal of \(\pm 1\). The cost is a small pre-rotation of \(Q\) and \(K\); the benefit is that FP8 attention error approaches that of a properly-scaled quantization without outlier contamination.

import torch

def hadamard(n):
    """Build the n x n (Sylvester) Hadamard matrix; n must be a power of 2."""
    assert n & (n - 1) == 0, "n must be a power of two"
    H = torch.ones(1, 1)
    while H.shape[0] < n:
        H = torch.cat([torch.cat([H,  H], dim=1),
                       torch.cat([H, -H], dim=1)], dim=0)
    return H / (H.shape[0] ** 0.5)            # orthonormal: H^T H = I

def incoherent_rotation(d, seed=0):
    g = torch.Generator().manual_seed(seed)
    H = hadamard(d)
    D = torch.where(torch.rand(d, generator=g) > 0.5, 1.0, -1.0)  # random +-1 signs
    return H @ torch.diag(D)                   # M = H D, orthogonal

def fake_fp8_quant(x):
    """Crude e4m3-like quantization: per-tensor scale, ~3 mantissa bits."""
    s = x.abs().max() / 448.0                  # 448 = max finite e4m3 magnitude
    if s == 0:
        s = torch.tensor(1.0)
    xs = x / s                                 # bring values into FP8 range
    # simulate ~3 mantissa bits by rounding to 2^-3 relative steps
    # (illustrative, not a bit-exact e4m3 emulator)
    mant = torch.round(xs * 8) / 8
    return mant * s

if __name__ == "__main__":
    torch.manual_seed(0)
    d = 128
    Q = torch.randn(64, d); K = torch.randn(64, d)
    Q[:, 7] *= 25.0; K[:, 7] *= 25.0           # inject an outlier feature
    ref = Q @ K.T

    # --- FP8 WITHOUT incoherent processing ---
    err_plain = (fake_fp8_quant(Q) @ fake_fp8_quant(K).T - ref).abs().mean()

    # --- FP8 WITH incoherent processing (rotate, quantize, scores invariant) ---
    M = incoherent_rotation(d)
    Qr, Kr = Q @ M, K @ M                       # (QM)(KM)^T == QK^T exactly
    err_rot = (fake_fp8_quant(Qr) @ fake_fp8_quant(Kr).T - ref).abs().mean()

    print(f"mean score error  plain FP8 : {err_plain:.4f}")
    print(f"mean score error  rotated   : {err_rot:.4f}")
    # the rotated version has substantially lower error because the
    # outlier in feature 7 is spread across all d coordinates

Run it and the rotated version shows markedly lower error: the outlier in feature 7, which forced a huge scale and wrecked the plain FP8 matmul, is smeared across all \(d\) coordinates after the Hadamard rotation. This is the same principle that powers Hadamard-based quantization schemes elsewhere in the stack (QuaRot, SpinQuant); see Quantization II: INT4/INT8/FP8, GGUF, bitsandbytes & QAT and the FP8 training discussion in Mixed Precision, bf16 & FP8 Training.

What stays in higher precision¶

A crucial detail: FA3 does not do everything in FP8. The matmul inputs (\(Q\), \(K\), \(P\), \(V\)) are FP8, but the tensor cores accumulate in FP32, and the softmax statistics (\(m\), \(\ell\)) and the running output accumulator are kept in higher precision. You quantize the operands going into the matmul, not the running state. This mixed strategy — low precision for the bandwidth- and throughput-bound multiply, high precision for the sensitive accumulation and the softmax — is the same pattern as bf16 mixed-precision training and is what keeps end-to-end error tolerable.

Comparison, trade-offs, and when each wins¶

It is worth stepping back to place these kernels against each other and against alternatives.

A few important clarifications:

Each generation is a rewrite, not a patch. The through-line of the table is that FlashAttention gets re-derived whenever the hardware changes what “asynchronous” means. Blackwell (B200) continues the pattern: its fifth-generation tensor cores are driven by the tcgen05 MMA instructions, whose accumulators live in a dedicated on-chip Tensor Memory (TMEM) rather than in each warp’s registers, and two SMs in a cluster can cooperate on a single MMA. Because the accumulator no longer sits in the same register file as the softmax, FA3’s warpgroup ping-pong is no longer the right shape; and since tensor-core throughput grew faster than the special-function unit’s exp2 rate, FA4 reportedly goes as far as emulating the exponential in the regular math pipes instead of paying for the SFU (see the SoTA box). The lesson to carry forward: the algorithm — tiling plus online softmax — has been stable since 2022; the kernel is rewritten every hardware generation, which is precisely why you should call a maintained library rather than freeze your own.

xFormers and the lineage of memory_efficient_attention. xFormers is Meta’s library of composable attention/kernel building blocks, and its memory_efficient_attention was one of the first widely-used fused, tiled attention kernels: it never materializes the full \(N \times N\) score matrix, using the same online-softmax-plus-tiling idea as FlashAttention. Its lineage traces to Rabe & Staats’s “Self-Attention Does Not Need \(O(N^2)\) Memory” (2021) combined with FlashAttention-style IO-aware tiling, and its fast path is a set of CUTLASS-based CUDA C++ kernels (with a fallback Triton/flash path). The API differs from SDPA’s layout — q/k/v are shaped (batch, seqlen, num_heads, head_dim) rather than (B, H, N, d), and causal masking is passed as a bias object rather than a boolean flag:

from xformers.ops import memory_efficient_attention, LowerTriangularMask

# q, k, v: (batch, seqlen, num_heads, head_dim)
out = memory_efficient_attention(q, k, v, attn_bias=LowerTriangularMask())

The library-mapping point that matters most: PyTorch SDPA’s EFFICIENT_ATTENTION backend (torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION) descends directly from these xFormers CUTLASS kernels, so when SDPA does not select the FLASH_ATTENTION backend, this is often the path it falls back to. Today memory_efficient_attention is mostly worth reaching for directly in two cases: on older GPUs where prebuilt flash-attn wheels are unavailable or unsupported — e.g., Turing/Volta (sm < 80) — since it has broader architecture/head-dim coverage, and in diffusion stacks, where Hugging Face diffusers exposes pipe.enable_xformers_memory_efficient_attention() and it long predated SDPA becoming the default. On modern Ampere/Hopper training and LLM inference, prefer scaled_dot_product_attention or flash-attn directly, which dispatch to the FA2/FA3 kernels documented above.

Forward vs backward. All the FA2/FA3 ideas apply to the forward pass cleanly. The backward pass is harder: it has more matmuls (five vs two), more intermediate quantities, and a different optimal parallelization (over keys, as noted). FA3’s FP8 backward in particular is delicate; many production setups run the forward in FP8 but keep the backward in bf16.

Decode is a different regime. During autoregressive inference decode the query is a single token (\(B_r = 1\)), so the \(QK^\top\) matmul is a matrix-vector product — memory-bound, not compute-bound. FA2’s seqlen parallelism over queries does nothing when there is one query. The right tool is FlashDecoding, which parallelizes over the KV sequence (split the long KV cache across SMs, each computes a partial softmax, then combine) to keep all SMs busy. This is the same online-softmax combine logic, applied along a different axis. See The Anatomy of LLM Inference: Prefill, Decode & The KV Cache and PagedAttention & KV-Cache Memory Management.

Head dimension matters. These kernels are tuned for head dimensions up to 128 or 256. Very large head dims, or unusual ones, may fall off the fast path and dispatch to a slower fallback. Variants like GQA/MQA/MLA change the effective shapes (Multi-Head Attention, MQA, GQA & MLA); modern kernels have GQA-aware variants that avoid replicating KV.

Alternatives to attention. None of this helps the asymptotic \(O(N^2)\) scaling — FlashAttention makes the constant factor and memory tiny, but the work still grows quadratically in sequence length. Sub-quadratic architectures (state-space models, linear attention) attack the asymptotics directly and are covered in Beyond Attention: SSMs, Mamba, RWKV & Linear Attention. In practice, FlashAttention-style exact attention remains dominant because its constant is so small that quadratic-but-fast beats subquadratic-but-clunky up to quite long contexts.

Further reading¶

• Tri Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (2023) — the definitive source for the deferred-rescale, sequence-parallel, and split-Q warp-partitioning ideas.
• Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (2024) — warp specialization, TMA producer/consumer, wgmma overlap, and FP8 with incoherent processing.
• Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) — the original; read alongside FlashAttention I: IO-Awareness & The Online Softmax.
• Tri Dao et al., Flash-Decoding for long-context inference (2023) — the split-KV technique for the decode regime.
• Jerry Chee, Yaohui Cai, Volodymyr Kuleshov, Christopher De Sa, QuIP: 2-Bit Quantization of Large Language Models With Guarantees (2023) — origin of incoherent processing via random orthogonal/Hadamard rotations.
• NVIDIA, Hopper Architecture Whitepaper and the CUTLASS library / CUDA programming guide — authoritative references for wgmma, TMA, asynchronous barriers, and FP8 tensor cores.
• The Dao-AILab/flash-attention repository — production kernels and the GQA/causal/varlen variants you will actually call.

Exercises¶

1. FlashAttention-1 already minimized HBM traffic to the theoretical minimum for exact attention. Yet FA2 is roughly \(2\times\) faster on the same A100. Explain in one or two sentences why the two facts are not contradictory, and name the bottleneck FA2 actually attacks.

2. In the FA2 reference recurrence, the per-iteration correction factor \(\alpha = e^{m - m^{\text{new}}}\) is applied to both the running denominator \(\ell\) and the unnormalized accumulator \(\tilde{O}\), but the division by \(\ell\) is deferred to the end. Suppose a well-meaning “optimizer” also tries to defer the \(\alpha\) multiplication on \(\tilde{O}\) (i.e., drops the alpha.unsqueeze(1) * Oi term and just does Oi = Oi + Pij @ vj). Does the kernel still compute exact attention? Explain precisely what breaks.

3. (Quantitative.) Reproduce and extend the chapter’s “FLOPs are not fungible” calculation for a smaller head. Take one attention head, full (non-causal) attention, sequence length \(N = 4096\), head dimension \(d = 64\). Assume the two matmuls cost \(4N^2 d\) FLOPs total, the non-matmul softmax work is \(\approx 5N^2\) FLOPs, tensor cores run at \(312\) TFLOP/s, and the special-function unit doing exp/max runs at \(20\) TFLOP/s. Compute (a) the matmul time, (b) the non-matmul time, and © the non-matmul share of total runtime assuming no overlap. Then (d) state what the number becomes if FA3’s ping-pong scheduling perfectly overlaps the softmax with matmul.

4. In FA2’s split-Q warp layout each warp owns a slice of the \(B_r\) query rows and needs no inter-warp reduction, whereas FA1’s split-K layout forced a shared-memory reduction plus a __syncthreads() every inner iteration. (a) Why does split-Q eliminate the reduction — what property of the query-row partition makes the warps independent? (b) What does split-Q now require of the \(K_j, V_j\) tiles, and why is shared memory well suited to satisfying it? © The chapter says the backward pass instead parallelizes over key blocks. Give the one-sentence reason the forward and backward pick opposite axes.

5. (Implementation.) The reference flash_attention_2_reference in the chapter handles a full and a causal mask. Extend it to support a sliding-window (local) attention mask of width \(w\): query position \(q\) may attend only to key positions \(k\) with \(q - w < k \le q\) (a causal window of the last \(w\) keys, including itself). Your implementation must (a) still skip entire key blocks that fall completely outside the window (the load-balancing point from the chapter’s causal aside), and (b) mask the partial-overlap blocks correctly. Verify against a dense reference.

import torch

def flash_attention_2_sliding_window(Q, K, V, w, Br=64, Bc=64):
    """
    Exact FA2 recurrence with a sliding-window causal mask of width w:
    query q attends to keys k with (q - w) < k <= q.
    Skips key blocks fully outside the window; masks partial-overlap blocks.
    """
    N, d = Q.shape
    O = torch.zeros_like(Q)
    for i in range(0, N, Br):
        qi = Q[i:i+Br]
        br = qi.shape[0]
        Oi = torch.zeros(br, d)
        mi = torch.full((br,), -float('inf'))
        li = torch.zeros(br)

        q_lo, q_hi = i, i + br - 1                     # query row range in this block
        for j in range(0, N, Bc):
            kj = K[j:j+Bc]; vj = V[j:j+Bc]
            bc = kj.shape[0]
            k_lo, k_hi = j, j + bc - 1                 # key range in this block

            # ---- block-level skip (load balancing) ----
            # entirely in the future: even the last query can't reach k_lo
            if k_lo > q_hi:
                break                                 # all later blocks are future too
            # entirely past the window: even the oldest query (q_lo) can't reach k_hi
            # q_lo attends to k > q_lo - w, so block is dead if k_hi <= q_lo - w
            if k_hi <= q_lo - w:
                continue

            Sij = (qi @ kj.T) / (d ** 0.5)            # (br, bc)
            qpos = torch.arange(q_lo, q_lo + br).unsqueeze(1)   # (br,1)
            kpos = torch.arange(k_lo, k_lo + bc).unsqueeze(0)   # (1,bc)
            # causal upper bound AND window lower bound
            mask = (kpos > qpos) | (kpos <= qpos - w)
            Sij = Sij.masked_fill(mask, -float('inf'))

            mij = Sij.max(dim=1).values
            mi_new = torch.maximum(mi, mij)
            # A block can survive the block-level skip yet still leave SOME
            # query rows fully masked (e.g. the newest rows of this query
            # block vs. an old key block). Those rows have mi_new = -inf and
            # exp(-inf - (-inf)) = nan. Guard them: they contribute nothing,
            # so force Pij -> 0 and alpha -> 1 (their Oi, li stay 0).
            dead = torch.isinf(mi_new)
            safe = torch.where(dead, torch.zeros_like(mi_new), mi_new)
            Pij = torch.exp(Sij - safe.unsqueeze(1)).masked_fill(dead.unsqueeze(1), 0.0)
            alpha = torch.where(dead, torch.ones_like(mi), torch.exp(mi - safe))
            li = alpha * li + Pij.sum(dim=1)
            Oi = alpha.unsqueeze(1) * Oi + Pij @ vj
            mi = mi_new

        O[i:i+br] = Oi / li.unsqueeze(1)
    return O


if __name__ == "__main__":
    torch.manual_seed(0)
    N, d, w = 256, 64, 48
    Q, K, V = torch.randn(N, d), torch.randn(N, d), torch.randn(N, d)

    # dense reference with the same sliding-window mask
    S = Q @ K.T / d**0.5
    q = torch.arange(N).unsqueeze(1); k = torch.arange(N).unsqueeze(0)
    allowed = (k <= q) & (k > q - w)
    S = S.masked_fill(~allowed, -float('inf'))
    ref = torch.softmax(S, dim=-1) @ V

    out = flash_attention_2_sliding_window(Q, K, V, w)
    print("max abs error:", (ref - out).abs().max().item())   # ~1e-6