FlashAttention 1 → 2 → 3

Take one attention head of Llama-3-8B at 8192 tokens of context: $d_h = 128$, bf16, no batching, just one head. §3.1 counted the HBM traffic of the textbook implementation; here is the same count summarised, because the whole chapter hangs off it.

The arithmetic is two matmuls: $QK^\top$ and $PV$, each $2T^2 d_h$ FLOPs, so $4 T^2 d_h = 4 \times 8192^2 \times 128 = 34.36$ GFLOP. The traffic is dominated by one object that does not appear in the arithmetic at all — the $T \times T$ score matrix, $8192^2 \times 2 = 134.2$ MB. A textbook implementation writes $S$, reads it back to softmax it, writes $P$, and reads $P$ back for the second matmul: four passes, 536.9 MB. Add $Q$, $K$, $V$ in and $O$ out (8.4 MB) and the head moves 545.3 MB to do 34.36 GFLOP of work.

At 3.35 TB/s that traffic takes at minimum 162.8 µs. The arithmetic, at 989.4 TFLOP/s, takes 34.7 µs. The kernel is 4.7× slower than its own arithmetic requires, and no amount of tuning the matmuls changes that, because the matmuls were never the problem. Every microsecond of the gap is spent moving a matrix that exists only because someone wrote softmax(Q @ K.T / sqrt(d)) @ V as three statements instead of one.

So the optimisation is not "make attention do less work". It is "make attention never write $S$", and you are allowed to spend arithmetic to buy that. This chapter derives how much, walks the three generations of increasingly sophisticated spending, and then reads the exact call that vLLM and SGLang make into the kernel.

Attention is a matmul, a row-wise nonlinearity, and a second matmul. The nonlinearity in the middle is what forces the intermediate to be materialised: softmax over a row needs the whole row before it can normalise any element of it. §3.1 broke that dependency — the online softmax recurrence lets you consume the row in chunks, carrying a running maximum and a running sum, and fix up the partial result as you go. That recurrence is the mathematical precondition for everything below; this chapter is about what you build on top of it.

What you build is a loop that keeps a tile of $Q$, a tile of $K$, a tile of $V$, and a running output accumulator resident in on-chip SRAM and registers, and streams the rest past them. $S$ and $P$ are born inside the tensor-core accumulators, consumed by the softmax, consumed again by the second matmul, and die there. They are never given an address in HBM. The $T \times T$ matrix still conceptually exists; it is just never simultaneously alive.

Two numbers from that figure carry the chapter. The row block $B_r = 128$ is capped by the fp32 output accumulator, which must stay in registers for the entire inner loop. The column block $B_c = 64$ is what the shared-memory budget and the score accumulator jointly allow. §3 derives both.

The intensity identity

Write $T$ for sequence length and $d_h$ for head dimension. For one head, forward only, in a 2-byte dtype:

The naive count is the four passes over the score matrix; the flash count is $Q$, $K$, $V$ in and $O$ out, each $T \times d_h$. Divide, and the $T^2$ and $d_h$ terms collapse into something startlingly clean:

With materialized scores, intensity approaches $d_h/2$ for bf16 as T grows; finite Q/K/V/O traffic makes this an asymptote. The $T/2$ FlashAttention figure counts only compulsory Q/K/V reads and O writes, not actual finite-SRAM tile rereads. Original FlashAttention analyzes HBM access complexity $\Theta(T^2d_h^2/M)$ for SRAM capacity M in its stated regime. Therefore a compulsory-byte crossing at T=590 does not prove the actual kernel becomes compute-bound there. Measure HBM traffic and achieved throughput for the real tile schedule.

What the extra FLOPs actually are

The forward pass pays for the online softmax. After each KV block $j$, the running maximum may have increased, so the accumulator $\tilde O_i$ (shape $B_r \times d_h$) must be rescaled by $\exp(m^{(j-1)} - m^{(j)})$ elementwise. That is $B_r \cdot d_h$ multiplies per (row block, KV block) pair, and there are $(T/B_r)(T/B_c)$ pairs:

At $T = 8192$, $d_h = 128$, $B_c = 64$ that is 134.2 MFLOP — 0.39% more arithmetic, and the compulsory-traffic comparison suggests up to 98.5% less traffic; actual savings depend on finite-SRAM rereads and caching. That is the trade, quantified, for one head. Note that the $T^2$ exponentials are not extra: naive attention computes exactly the same ones. Only the rescale is new.

The backward pass pays much more, and this is where "more FLOPs" stops being a rounding error. A standard attention backward needs four $T\times T\times d_h$ matmuls ($dV$, $dP$, $dQ$, $dK$) and reads the saved $P$. FlashAttention does not save $P$; it recomputes $S = QK^\top$ and re-softmaxes it from the saved per-row logsumexp, adding a fifth matmul — +25% matmul FLOPs — to avoid storing and reloading $O(T^2)$ bytes per head per layer. The original paper measured that trade as a net win end to end (arXiv:2205.14135, §4: 3× end-to-end speedup training GPT-2 at sequence length 1K, and a 15% wall-clock improvement on BERT-large at sequence length 512 over the MLPerf 1.1 record, on A100 40GB). Inference never runs the backward pass, but the logic is the same one and it is the cleanest statement of the book's central bargain: left of the ridge, recomputation is free.

Deriving the tile sizes from an H100

A thread block owns one row block: $B_r$ queries, all $d_h$ channels, for the whole inner loop over KV. What must be resident?

Shared memory. The $Q$ tile is $2 d_h B_r$ bytes. The $K$ and $V$ tiles are $2 d_h B_c$ each, times $s$ pipeline stages for double buffering. So

With $d_h = 128$, $s = 2$, and $B_r = B_c = B$: $256(B + 4B) = 1280 B \le 233472$, so $B \le 182$. Tiles must be multiples of the MMA tile granularity (16 rows, and in practice 64 for the wgmma N dimension), so the largest legal square tile is 128. At $B_r = 128$, $B_c = 64$, $s = 2$ the block uses $256 \cdot 128 + 512 \cdot 2 \cdot 64 = 98{,}304$ bytes — 96 KB of the 228 KB an H100 SM offers (§0.3). Shared memory is not the binding constraint at $d_h = 128$.

Registers. This is what actually binds. An H100 SM has 65,536 32-bit registers (256 KB) and a hard 255-register-per-thread ceiling. Take a block of 8 warps (256 threads, two Hopper warpgroups). The fp32 output accumulator is $4 B_r d_h = 65{,}536$ bytes, i.e. 256 bytes or 64 registers per thread. The fp32 score accumulator is $4 B_r B_c = 32{,}768$ bytes, another 32 registers. That is 96 before a single address, operand descriptor, or softmax temporary. Doubling $B_r$ to 256 would put the output accumulator at 128 registers per thread, leaving fewer registers for scores and temporaries but not exhausting the 255-register per-thread ceiling by itself; doubling $B_c$ to 128 costs another 32 and is done only at $d_h = 64$, where the output accumulator halves. The rule of thumb falls straight out:

At $d_h = 128$, $B_c = 64$: $B_r \le 266$, and after operands and pipeline bookkeeping, 128 is the practical answer. The output accumulator is the reason $B_r$ is 128 and not 512, and the reason it must stay in registers at all is the FA2 loop-order change in §4.

FlashAttention-1: tiling and recomputation

FA1 (arXiv:2205.14135) established the shape. The outer loop runs over KV blocks $j$; the inner loop over row blocks $i$. For each $(i,j)$ the kernel loads $Q_i$ and the current $O_i$, $m_i$, $\ell_i$ from HBM, computes $S_{ij} = Q_i K_j^\top$, applies the online softmax update, and writes $O_i$, $m_i$, $\ell_i$ back to HBM. Parallelism is one thread block per (batch, head) pair.

Both of those choices turned out to be wrong, and FA2 fixed both.

FlashAttention-2: three changes, each with a named bottleneck

(a) Fewer non-matmul FLOPs. FA1 divided by the running sum on every iteration:

FA2 keeps the accumulator unnormalised — $\tilde O_i \leftarrow e^{m^{(j-1)}-m^{(j)}}\tilde O_i + \tilde P_{ij}V_j$ — and divides by $\ell$ exactly once, after the inner loop ends. It also stores only the logsumexp $L = m + \log \ell$ instead of $m$ and $\ell$ separately. Small-looking, but these are non-matmul FLOPs, and the FA2 paper (arXiv:2307.08691, §3.1) prices them at the A100 ratio: 312 TFLOP/s of bf16 matmul against 19.5 TFLOP/s of FP32 non-matmul, so each non-matmul FLOP costs 16 matmul FLOPs of time. At 0.39% of the arithmetic and a 16× price, the rescale is worth about 6% of the matmul time — worth attacking, and worth overlapping entirely, which is what FA3 does.

(b) Parallelise over sequence length. FA1's grid is $\text{batch} \times \text{heads}$. For Llama-3-8B ($h = 32$) at batch 1 that is 32 thread blocks on a 132-SM H100 — 100 SMs idle, 24% of the machine. FA2 swaps the loops so the row-block loop is outermost, and makes it a grid dimension: $\text{batch} \times \text{heads} \times \lceil T/B_r \rceil$. At $T = 8192$, $B_r = 128$ that is $1 \times 32 \times 64 = 2048$ blocks, 15.5 full waves. This is the change that made long-context, small-batch attention actually use the GPU.

(c) Better warp partitioning. Inside a block, FA1 split $K$ and $V$ across the four warps and shared $Q$ ("split-K"). Every warp then held a partial slab of $S$ covering a subset of the key axis for all the rows — which means no warp owns a complete softmax row. To normalise, all four warps must write their partials to shared memory, `__syncthreads`, and reduce. FA2 splits $Q$ across warps and shares $K$ and $V$ ("split-Q"). Each warp now owns a disjoint set of query rows end to end: its own $S$ rows, its own $m$ and $\ell$, its own $O$ rows. No cross-warp softmax reduction is needed for those owned query rows. Shared-memory tile staging and pipeline synchronization can still require barriers.

The cited result: FA2 is roughly 2× FA1, reaching 50–73% of theoretical peak on A100 80GB SXM4 and up to 225 TFLOP/s of model FLOPs utilisation on GPT-style training (arXiv:2307.08691, abstract and §4).

FlashAttention-3: why Hopper changes the problem

FA2 on an H100 leaves most of the machine on the table, and the reason is the softmax. Count it. Exponentials execute on the SM's special-function units, not the tensor cores. Assuming the standard 16 MUFU lanes per SM at 1.755 GHz across 132 SMs, the whole GPU sustains $132 \times 16 \times 1.755 \times 10^9 = 3.71 \times 10^{12}$ exponentials per second — 267× below the 989.4 TFLOP/s bf16 matmul rate. Our head needs $T^2 = 67.1$M exponentials, which is $67.1\times10^6 / 3.71\times10^{12} = 18.1$ µs against 34.7 µs of matmul.

Overlapping them requires hardware that FA2's programming model does not have. Hopper supplies four pieces, and FA3 uses all four:

On top of that machinery FA3 layers three techniques.

Warp specialisation. The thread block is split by role. One producer warpgroup does nothing but issue TMA loads of $K_j$ and $V_j$ into a circular shared-memory buffer and signal an mbarrier. Two consumer warpgroups do nothing but wait on the barrier, issue wgmma, and run the softmax. This is not the same thing as double buffering: the producer runs a different instruction stream, so a stall in the memory path never occupies an issue slot the tensor cores wanted.

Ping-pong scheduling. With two consumer warpgroups, FA3 offsets them by half a stage using named barriers, so that while warpgroup 0 is in its softmax (SFU-bound), warpgroup 1 is in its wgmma (tensor-core-bound), and then they swap. Both units stay busy. Inside a single warpgroup, FA3 additionally pipelines two stages deep: the wgmma for KV block $j+1$ is issued before the softmax of block $j$ has finished, because wgmma is asynchronous and its accumulator is a different register set.

FP8 with in-kernel scaling. The third technique is numerical. FP8 e4m3 (§0.5) doubles tensor-core throughput but has 3 mantissa bits, and attention inputs have outliers. FA3 handles two separate problems. The layout problem: FP8 wgmma requires both operands contiguous along the reduction dimension, which $V$ is not, so the kernel transposes $V$ in shared memory with LDSM/STSM and hides that under the async copies. The accuracy problem: block-level quantisation (a scale per tile rather than per tensor) plus incoherent processing — multiplying $Q$ and $K$ by a random orthogonal matrix, which leaves $QK^\top$ unchanged but spreads outliers across channels. The paper reports 2.6× lower RMS error than a per-tensor FP8 baseline, and up to 1.2 PFLOP/s on H100 SXM5 (arXiv:2407.08608, abstract and §4.3). The paper's per-tile scaling must not be equated with every runtime descale interface: inspect the q_descale/k_descale/v_descale shapes and kernel path, which may use per-sequence/head scales.

Neither engine writes attention kernels here. Both call the same family of CUDA kernels through a thin Python wrapper, and the interesting engineering is in what they pass. As of a556f3f, vLLM has deleted its own PagedAttention CUDA kernel (commit d715b3aa1e) and vendors FlashAttention under vllm/vllm_flash_attn/.

The signature, argument by argument

def flash_attn_varlen_func(
    q,
    k,
    v,
    max_seqlen_q,
    cu_seqlens_q,
    max_seqlen_k,
    cu_seqlens_k=None,  # only used for non-paged prefill
    seqused_k=None,
    q_v=None,
    dropout_p=0.0,
    softmax_scale=None,
    causal=False,
    window_size: list[int] | None = None,
    softcap=0.0,  # 0.0 means deactivated
    alibi_slopes=None,
    deterministic=False,
    return_attn_probs=False,
    block_table=None,
    return_softmax_lse=False,
    out=None,
    # FA3 Only
    scheduler_metadata=None,
    q_descale=None,
    k_descale=None,
    v_descale=None,
    num_splits: int = 0,
    # FA4 Only
    output_scale=None,
    # Version selector
    fa_version: int = DEFAULT_FA_VERSION,
    s_aux=None,

Map that onto §3, one argument at a time.

The causal alignment detail is worth pausing on, because getting it wrong is a silent correctness bug in any prefix-caching engine. From the same docstring:

    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
        1 1 1 1 0
        1 1 1 1 1
    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
        0 0
        0 0
        0 0
        1 0
        1 1
    If the row of the mask is all zero, the output will be zero.

Chunked prefill and prefix caching both produce exactly the first case: $q$ new tokens against $C + q$ cached-plus-new keys. Bottom-right alignment gives each new token its full cached context plus the new tokens up to itself. Top-left alignment would silently truncate every request that ever hit the prefix cache.

vLLM's call site

                flash_attn_varlen_func(
                    q=query[:num_actual_tokens],
                    k=key_cache,
                    v=value_cache,
                    out=output[:num_actual_tokens],
                    cu_seqlens_q=cu_seqlens_q,
                    max_seqlen_q=max_seqlen_q,
                    seqused_k=seqused_k,
                    max_seqlen_k=max_seqlen_k,
                    softmax_scale=self.scale,
                    causal=causal,
                    alibi_slopes=self.alibi_slopes,
                    window_size=sliding_window_size,
                    block_table=block_table,
                    softcap=self.logits_soft_cap,
                    scheduler_metadata=scheduler_metadata,
                    fa_version=self.vllm_flash_attn_version,
                    q_descale=q_descale,
                    k_descale=k_descale,
                    v_descale=v_descale,
                    dynamic_causal=dynamic_causal,
                    num_splits=attn_metadata.max_num_splits,
                    s_aux=self.sinks,

Note what is absent: cu_seqlens_k. vLLM always passes the paged cache, so it supplies seqused_k instead — and the wrapper enforces exactly that pairing at vllm/vllm_flash_attn/flash_attn_interface.py:L276-L278: assert block_table is None or seqused_k is not None.

Every one of those arguments is read out of a metadata struct built once per forward pass:

    num_actual_tokens: int  # Number of tokens excluding padding.
    max_query_len: int
    query_start_loc: torch.Tensor
    max_seq_len: int
    seq_lens: torch.Tensor
    block_table: torch.Tensor
    slot_mapping: torch.Tensor

query_start_loc becomes cu_seqlens_q; seq_lens becomes seqused_k. And self.scale traces back to the model definition — vllm/model_executor/models/llama.py:L159 is self.scaling = self.head_dim**-0.5, which for Llama-3-8B is $128^{-1/2} = 0.0884$.

Version selection, and where the engines diverge

        # 1. default version depending on platform
        if device_capability.major == 9 and is_fa_version_supported(3):
            # Hopper (SM90): prefer FA3
            fa_version = 3
        elif device_capability.major == 10 and is_fa_version_supported(4):
            # Blackwell (SM100+, restrict to SM100 for now): prefer FA4
            fa_version = 4
        else:
            # Fallback to FA2
            fa_version = 2

FA3 is gated hard to SM90 — vllm/vllm_flash_attn/flash_attn_interface.py:L62-L69 returns "FA3 is only supported on devices with compute capability 9.x". Everything in §4 about warp specialisation, TMA, and ping-pong exists only on Hopper, and on an A100 vLLM silently runs FA2. The rest of get_flash_attn_version is a list of demotions: ALiBi forces FA2, batch-invariant mode can demote FA4 to FA2 but does not universally demote FA3, head sizes over 128 on Blackwell force FA2.

SGLang reaches the same kernels but with a structurally different policy. Its FlashAttentionBackend has no FA2 path at all:

        # Select version
        self.fa_impl_ver = fa_impl_ver
        device_capability = get_device_capability()
        if self.fa_impl_ver == 3:
            from sgl_kernel.flash_attn import (
                flash_attn_varlen_func,
                flash_attn_with_kvcache,
                get_scheduler_metadata,
            )
# ...
        elif self.fa_impl_ver == 4:
# ...
        else:
            raise ValueError(f"Invalid version: {self.fa_impl_ver=}")

The backend is registered twice, as "fa3" and "fa4" (python/sglang/srt/layers/attention/attention_registry.py:L209-L241), and the default chooser picks it only on Hopper:

        if not use_mla_backend:
            # MHA architecture

            if is_hopper_with_cuda_12_3() and is_no_spec_infer_or_topk_one(
                resolved_view(self)
            ):
                # Note: flashinfer 0.6.1 caused performance regression on Hopper attention kernel
                # Before the kernel is fixed, we choose fa3 as the default backend on Hopper MHA
                # ref: https://github.com/sgl-project/sglang/issues/17411
                return "fa3"

They also use different entry points. vLLM funnels everything — prefill, decode, chunked prefill — through flash_attn_varlen_func with a block_table. SGLang calls flash_attn_varlen_func only when the KV is not paged (python/sglang/srt/layers/attention/flashattention_backend.py:L1479-L1494), and otherwise calls flash_attn_with_kvcache with page_table= and cache_seqlens=:

                result = flash_attn_with_kvcache(
                    q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
                    k_cache=key_cache,
                    v_cache=value_cache,
                    page_table=page_table,
                    cache_seqlens=cache_seqlens,
                    cu_seqlens_q=cu_seqlens_q,
                    cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
                    max_seqlen_q=max_seqlen_q,
                    softmax_scale=layer.scaling,
                    causal=False if use_cascade_attn else causal,
                    window_size=window_size,
                    softcap=layer.logit_cap,
                    return_softmax_lse=use_cascade_attn,
                    num_splits=self.num_splits,
                    out=_fa_out,
                    ver=self.fa_impl_ver,
                    **kwargs,
                )

page_table/cache_seqlens are the same two tensors as vLLM's block_table/seqused_k under different names, and ver= is the same selector as fa_version=. The kernels underneath are the same Dao-AILab CUDA.

One Llama-3-8B request, 8192 prompt tokens, no prefix cache hit, on an H100, layer 0, through vLLM.

1. FlashAttentionMetadataBuilder.build() (vllm/v1/attention/backends/flash_attn.py:L538-L740) runs once per forward pass. It sets query_start_loc = [0, 8192], seq_lens = [8192], max_query_len = 8192, and block_table to the 8192/16 = 512 physical pages the block manager allocated.
2. Because get_flash_attn_version() == 3 on H100, self.aot_schedule is True (vllm/v1/attention/backends/flash_attn.py:L447), so build() calls _get_scheduler_metadata(), which calls into torch.ops._vllm_fa3_C.get_scheduler_metadata (vllm/vllm_flash_attn/flash_attn_interface.py:L147-L172) with the batch shape, head counts, head dim, page size, and causal flag. FA3's persistent tile scheduler is thereby computed on the host and handed to the kernel as a tensor, instead of being derived per-block on device — which is what makes the kernel CUDA-graph capturable.
3. FlashAttentionImpl.forward() unpacks the KV cache. The cache is stored as (num_blocks, num_kv_heads, block_size, 2 * head_size) (vllm/v1/attention/backends/flash_attn.py:L141-L151) — K and V packed into one contiguous content dimension — and split at flash_attn.py:L983-L984: key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1).
4. Descale tensors are built with .expand() rather than materialised: descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) at flash_attn.py:L1016. For bf16 KV they are the scalar 1.0 broadcast to $(1, 8)$; for FP8 KV they are the per-sequence per-head scales that §4's in-kernel scaling consumes.
5. flash_attn_varlen_func(..., fa_version=3) dispatches to torch.ops._vllm_fa3_C.fwd (vllm/vllm_flash_attn/flash_attn_interface.py:L349-L385). Note the positional None, # pack_gqa — FA3 can pack the 4 query heads that share each KV head into one row block, filling $B_r$ with (position, head) pairs rather than positions alone; passing None lets the kernel heuristic decide.
6. On device: the logical query tiling has $1\times32\times\lceil8192/128\rceil=2048$ tiles; persistent launch scheduling and GQA packing can change the actual grid. Each loads its 32 KB $Q$ tile, then runs 128 iterations of the inner loop over 64-key blocks — halved by causal=True for the average row block, so about 65 iterations for 128-row query tiles against 64-key tiles (ends cover 2,4,...,128 key tiles). The coarse half-of-128 estimate is 64, not 32. Each iteration: producer TMA-loads $K_j, V_j$; consumers wgmma $Q_i K_j^\top$, exp2, rescale the register-resident $\tilde O_i$, wgmma $\tilde P_{ij} V_j$. Nothing of size $T \times T$ is ever addressed in HBM.
7. After the loop, divide by $\ell$ once, write $O_i$ (32 KB) back through out=, which vLLM pre-allocated so the kernel writes in place and the attention layer does no extra copy.

Passing both cu_seqlens_k and seqused_k. The wrapper asserts: "cu_seqlens_k and seqused_k cannot be provided at the same time" (vllm/vllm_flash_attn/flash_attn_interface.py:L269-L275). Writing a new backend, the natural instinct is to fill in everything you know. The two arguments describe mutually exclusive KV layouts and the kernel picks its addressing path from which one is present.

Silently getting FA2 when you expected FA3. get_flash_attn_version demotes on ALiBi, on applicable version-specific batch-invariant restrictions, and on unsupported head sizes. The only signal is a logger.info_once("Using FlashAttention version %s", ...) at vllm/v1/attention/backends/flash_attn.py:L868-L871. Grep the log for that line before profiling anything.

Degenerate strides break TMA. A real one, and a good illustration of how §4's hardware leaks into Python. From vllm/v1/attention/backends/flash_attn.py:L985-L987: "Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP). FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment." When tensor parallelism reduces a GQA model to one KV head per rank, PyTorch reports an arbitrary stride for that size-1 dimension, and TMA's descriptor validation rejects it. vLLM canonicalises the strides before the call.

FP8 KV requires FA3 or FA4, and never e5m2. flash_attn_supports_kv_cache_dtype (vllm/v1/attention/backends/fa_utils.py:L223-L243) returns False immediately for fp8_e5m2, and otherwise requires FA3-on-SM90 or FA4-on-SM100. Setting --kv-cache-dtype fp8_e5m2 on Hopper does not get you a faster kernel; it gets you a different backend.

Expecting FA2's parallelism at decode. Covered in §4 — if you benchmark decode-phase attention and find the GPU 24% occupied, that is not a bug, that is $T_q = 1$. §3.3.

Confirm which version your deployment actually runs, then force the other one and compare:

# What version did vLLM pick? Look for "Using FlashAttention version N".
VLLM_LOGGING_LEVEL=INFO vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --max-model-len 8192 2>&1 | grep -i "FlashAttention version"

# Request FA2 explicitly; confirm acceptance and selected version in the log.
# The config field read at vllm/v1/attention/backends/fa_utils.py:L102-L110 is attention_config.flash_attn_version.
vllm bench latency --model meta-llama/Llama-3.1-8B-Instruct \
  --input-len 8192 --output-len 1 --batch-size 1 \
  --attention-config '{"flash_attn_version": 2}'

# SGLang: name the backend explicitly rather than relying on the Hopper default.
python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct \
  --attention-backend fa3

Read vllm/v1/attention/backends/fa_utils.py top to bottom while you do it — 305 lines, and the single most useful file in the repo for understanding why your kernel is the kernel it is. The --input-len 8192 --output-len 1 shape isolates prefill, the regime this chapter describes; flip to --input-len 128 --output-len 512 and you are measuring §3.3 instead.

• Dao, Fu, Ermon, Rudra, Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv:2205.14135). §3 has the tiling proof and the $\Theta(N^2 d^2 M^{-1})$ IO complexity bound; §4 has the A100 40GB measurements.
• Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (arXiv:2307.08691). §3.1 is the non-matmul FLOP argument and the 16× figure; §3.2 is the parallelism change; §3.3 is split-Q.
• Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arXiv:2407.08608). §3 is warp specialisation and ping-pong; §4 is FP8 with block quantisation and incoherent processing.
• Dao-AILab/flash-attention — hopper/flash_fwd_kernel_sm90.h is the warp-specialised kernel described in §4, and hopper/tile_scheduler.hpp is what get_scheduler_metadata feeds.
• vllm-project/vllm#47361 — the PR that removed vLLM's bespoke PagedAttention CUDA kernel, leaving FlashAttention/FlashInfer/Triton as the attention path. Context for why this chapter's call site is the only one.
• sgl-project/sglang#17411 — the issue cited in server_args.py for why FA3, not FlashInfer, is SGLang's Hopper MHA default at this SHA.
• Colfax Research's FlashAttention-3 write-up — the clearest available walkthrough of Hopper's TMA/wgmma/mbarrier programming model, by the people who co-authored the paper.