FlashAttention and GPU Data Movement

The one-sentence version

Attention is not slow because the GPU cannot multiply fast enough; it is slow because the standard implementation writes two N × N matrices to HBM and reads them back, and FlashAttention removes that transit by tiling the computation and keeping the softmax state on-chip.

Everything else in this lesson makes that sentence checkable with bytes, a state trace, and an Amdahl calculation on a stated time share.

This is the capstone lesson of Module 1 (Measure and Explain Inference Performance). Near the end of this lesson, you will complete the Module 1 practical exit task: assembling an executive, one-page latency attribution report that isolates a serving bottleneck, applies Amdahl’s Law within strict measurement boundaries, and defends a concrete optimization proposal.

Module 1 prerequisites and context

This lesson builds directly on the previous two lessons in this module:

• Serving Workloads and Performance Metrics (Lesson 1) establishes the measurement boundaries between client TTFT, server queuing, GPU prefill, and token generation speed (TPOT/ITL).
• Where Inference Time and Memory Go (Lesson 2) breaks down weights, KV cache, FLOPs, and HBM bandwidth, showing why prefill is compute-heavy and decode is bandwidth-bound.

You should also be comfortable with scaled dot-product attention (S = QK^T / √d, P = softmax(S), O = PV), causal masking, and basic matrix multiplication.

---

1. Naive attention’s score-matrix materialisation

Start with the version most people write first, in eager PyTorch or in any un-fused kernel chain:

S = (Q @ K.transpose(-2, -1)) / math.sqrt(d)   # kernel 1: GEMM
P = torch.softmax(S, dim=-1)                   # kernel 2: softmax
O = P @ V                                      # kernel 3: GEMM

Call this the explicit materialising baseline: three kernels, with S and P written to HBM in full, in FP16. It is a real implementation pattern, and it is the baseline every FlashAttention byte count is measured against.

Modern PyTorch does not do this by default

In PyTorch 2.15 — the version this lesson was checked against — torch.nn.functional.scaled_dot_product_attention dispatches to a fused kernel — FlashAttention, memory-efficient attention, or cuDNN — whenever layouts, dtype, head dimension, and masks qualify. It falls back to the three-kernel SDPBackend.MATH path only when they do not. That fallback is also not byte-identical to the model below: it upcasts softmax to FP32 and writes FP32 intermediates. So the numbers in this lesson describe an explicit FP16 materialising baseline, not “what PyTorch does” in 2026. The baseline exists to make the mechanism measurable, and §6 shows how to check which backend actually ran.

1.1 The bytes, one head at a time

Fix the parameters: N = 4,096 tokens, head dimension d = 128, FP16 (b = 2 bytes). A single head’s Q, K, V, and O are each:

4,096 × 128 × 2 bytes = 1,048,576 bytes = 1.00 MiB

Each intermediate S and P is N × N:

4,096 × 4,096 × 2 bytes = 33,554,432 bytes = 32.00 MiB

Now follow the bus for one head:

Kernel	Reads HBM	Writes HBM	Traffic
1. GEMM (QK^T)	Q 1.00 MiB + K 1.00 MiB	S 32.00 MiB	34.00 MiB
2. Softmax	S 32.00 MiB	P 32.00 MiB	64.00 MiB
3. GEMM (PV)	P 32.00 MiB + V 1.00 MiB	O 1.00 MiB	34.00 MiB
Total	67.00 MiB (70,254,592 B)	65.00 MiB (68,157,440 B)	132.00 MiB (138,412,032 B)

Two things stand out. First, 128 of the 132 MiB are intermediate round-trips: only 4 MiB is the actual payload of Q, K, V, O. Second, the S and P matrices cross the bus four times — written once each, read once each — and every crossing is a full-width N × N transfer.

That is the memory wall in one number. On an H100 SXM the ridge point is 989 TFLOPS ÷ 3.35 TB/s ≈ 295 FLOP/byte. The softmax kernel’s arithmetic intensity is roughly 1 FLOP/byte — about 4N² FLOPs of elementwise work against 2N²b bytes of S read and P written. That is more than two orders of magnitude below the ridge point. It is not a compute kernel that happens to touch memory; it is a memory kernel that happens to compute.

1.2 Per-head allocation is not the memory you will see

S plus P is 64.00 MiB per head. Turning that into a GPU-memory claim requires saying which lifetime you mean, because the two answers differ by 32×:

Scope	Assumption	Footprint
Inference peak, one layer	32 heads of the active layer run attention in parallel; S and P are both live in FP16	2.0 GiB
Training retention, all layers	No activation checkpointing; every layer’s attention weights are kept for backprop	64.0 GiB

Both are conditional illustrations, not universal peaks. Inference executes layers sequentially, so the buffers for layer l are freed or reused before layer l + 1 runs. If heads are scheduled sequentially, or the caching allocator hands the same block back, the live footprint can fall to 64 MiB or lower. The 64.0 GiB training figure is a cumulative retention sum under an explicit no-checkpointing assumption; a real training budget also carries weights, gradients, and optimizer state, which this sum ignores.

Do not quote 2.0 GiB as the memory attention needs

It is an upper bound on transient scratch for one layer, under a stated scheduling assumption. It is not a memory plan, and it is not what a fused kernel allocates. The fused path allocates tile-sized buffers — on the order of tens of kilobytes — plus the running statistics in §3.3.

---

2. Tiling: keeping blocks in SRAM

FlashAttention never forms S or P in HBM. It partitions the problem so that every intermediate fits in on-chip memory, and streams the inputs through it.

Pick a query block B_r = 128 rows and a key/value block B_c = 64 rows. With d = 128 and FP16:

• Q tile: 128 × 128 × 2 B = 32 KiB
• K or V tile: 64 × 128 × 2 B = 16 KiB
• Score tile S_ij: 128 × 64 × 2 B = 16 KiB

Those fit inside an SM’s shared memory, which is why the block sizes are what they are. Two numbers are worth separating: an A100 SM has 192 KB of combined L1 + shared memory and an H100 SM has 256 KB, but the largest shared-memory carve-out a block can actually request is smaller — 164 KB on A100 and 228 KB on H100. Vendors and papers often quote the total and the carve-out interchangeably; they are not the same limit. The loop inverts the textbook order:

for each Q tile Q_i, i = 1..32 (rows 128(i-1)+1 .. 128i):  <- outer loop, one thread block
    load Q_i from HBM                        (32 KiB, once per tile)
    initialise m = -inf, l = 0, a = 0
    for each KV tile (K_j, V_j):
        j = 1 .. 2i          if causal   <- skip bound, see 2.1
        j = 1 .. 64          if unmasked
        load K_j, V_j from HBM               (32 KiB per tile)
        S_ij = Q_i K_j^T / sqrt(d)           <- registers / SRAM, never HBM
        if causal and the tile straddles the diagonal:
            S_ij[r, c] = -inf  where 64(j-1)+c > 128(i-1)+r
        update m, l, a with online softmax   <- after masking; §3
        discard S_ij
    write O_i = a / l to HBM                 (32 KiB, once per tile)

2.1 Skip bound and causal mask are two different operations

j <= 2i is the skip bound. It is not the mask. FlashAttention-2 §3.1.1 states the two mechanisms separately, and conflating them is the most common way a hand-rolled tiled attention silently attends to the future.

1. Skip any K/V tile whose every column is in the future. With 1-based tiles, K/V tile j covers columns 64(j−1)+1 … 64j and query tile i covers rows 128(i−1)+1 … 128i. The last tile that can contain a permitted column is j = 2i; every tile with j > 2i is entirely masked, so dropping it is exact and saves both the load and the multiply. Because B_r is twice B_c, a query block spans two column tiles.
2. Mask the tiles that straddle the diagonal — here j = 2i−1 and j = 2i, both diagonal-adjacent:
   • j = 2i−1 (columns 128i−127 … 128i−64) is fully permitted for the query block’s lower 64 rows and triangularly masked for its upper 64 rows.
   • j = 2i (columns 128i−63 … 128i) contributes nothing at all to the upper 64 rows — every one of those columns lies in their future — and is triangularly masked for the lower 64 rows.

The distinction matters because a mask changes the result while a skip only changes the cost. Load tiles up to 2i, never mask them, and the upper half of every query block attends to keys that come after it. Nothing crashes: you get a smooth weighted average, so the output looks reasonable and is wrong. The mask must be applied before the running max and sum see the scores, because those two statistics define the softmax:

S_ij^masked[r, c] = -inf        if column 64(j-1)+c > row 128(i-1)+r
m_tile            = rowmax(S_ij^masked)

An all--inf row region (the upper half of tile 2i) contributes zero to both l and a, which is exactly what “no permitted keys here” should mean — but only if your kernel handles exp(-inf − m) as 0 rather than as a NaN. scripts/verify-flashattention-example.py checks a small causal matrix, including the diagonal tiles, against direct masked attention, and also checks that omitting the mask gives a different answer.

2.2 What the loop moves

Count the traffic that loop implies. Q is read once: 32 tiles × 32 KiB = 1.00 MiB. O is written once: 1.00 MiB. K and V are re-read for every query tile, which is the deliberate trade — 32 passes over 2 MiB is 64 MiB of streaming reads, and nothing quadratic is ever allocated.

The full N × N matrices occupy 0 bytes in HBM. Not “less”; zero. S_ij is created, consumed by the state update, and discarded inside the same inner-loop iteration. This is why the win is a data-movement win: the fused kernel does more arithmetic in the forward pass than the baseline — it rescales the accumulator and re-exponentiates on every tile instead of reading a stored P — and it is faster anyway, because it removed 128 MiB of round-trip traffic per head. Backpropagation adds a second, separate cost: it recomputes S and P from Q, K, V, and L instead of loading them (FA-2 §2.3.2). That is a recomputation-for-bandwidth trade, not the forward rescaling, and the two should be stated separately.

That is also the trade to state honestly: FlashAttention trades recomputation and re-reading of K/V for elimination of the quadratic intermediates. At N = 4,096 the trade is usually lopsided in its favour. At N = 256 it may not be — see §6’s first boundary — because the baseline’s intermediates are small enough that the memory system may absorb them, not because there is a threshold below which the trade reverses for every kernel.

---

3. Online softmax: the mathematical core

Tiling has an obvious obstruction. Softmax needs the row maximum and the row sum over the whole row before any probability is final:

m = max_j S_ij
l = Σ_j exp(S_ij - m)
P_ij = exp(S_ij - m) / l

If you only hold one key block at a time, you never see the whole row. The online softmax identity (Milakov & Gimelshteyn 2018; Rabe & Staats 2021; adapted to tiled attention by Dao et al. 2022) removes that requirement: keep a running maximum m and a running normalizer l, and rescale the old partial result whenever the maximum moves.

3.1 The recurrence FlashAttention-2 uses

FlashAttention-2 maintains an unnormalized accumulator a — a d-vector per query row — and divides once, at the very end. For each incoming key/value tile with scores S_tile and values V_tile:

m_new = max(m_old, max(S_tile))
alpha = exp(m_old - m_new)
l_new = alpha * l_old + Σ exp(S_tile - m_new)
a_new = alpha * a_old + Σ exp(S_tile - m_new) * V_tile

After the last tile: O = a / l.

alpha is the whole trick. When a new tile raises the maximum, every exponential already accumulated was computed against the old base, so the old sum and the old accumulator must be multiplied by exp(m_old - m_new) to move them onto the new base. The result is mathematically identical to a full-row softmax; it just arrives one tile late and in a different order.

FlashAttention-1 kept the accumulator normalized instead: O_new = (l_old · alpha / l_new) · O_old + (1 / l_new) · Σ …. Equivalent mathematics, more divisions. In FlashAttention-1’s Algorithm 1 each key/value block rescales the whole N × d accumulator, which works out to N × d rescale-divisions per block; that is a derived count, not a number the paper states. FlashAttention-2’s unnormalized form does the division once per output row. The rest of this lesson uses the FA-2 form because it is what modern kernels implement.

3.2 A two-tile trace, worked

One query row, d = 2 for legibility, two key tiles of two tokens each.

Tile 1 arrives. S_1 = [8.0, 6.0], V_1 = [[1.0, 2.0], [3.0, 4.0]].

m_1   = max(8.0, 6.0) = 8.0
w_1   = [exp(0), exp(-2)] = [1.0000, 0.1353]
l_1   = 1.0000 + 0.1353 = 1.1353
a_1   = 1.0000 * [1.0, 2.0] + 0.1353 * [3.0, 4.0]
      = [1.0000, 2.0000] + [0.4060, 0.5413] = [1.4060, 2.5413]

Tile 2 arrives. S_2 = [10.0, 7.0], V_2 = [[5.0, 6.0], [7.0, 8.0]]. The new maximum is larger, so the old state has to be rescaled:

max(S_2) = 10.0
m_2      = max(8.0, 10.0) = 10.0
alpha    = exp(8.0 - 10.0) = exp(-2) = 0.1353
w_2      = [exp(0), exp(-3)] = [1.0000, 0.0498]
l_2      = 0.1353 * 1.1353 + (1.0000 + 0.0498)
         = 0.1536 + 1.0498 = 1.2034
a_2      = 0.1353 * [1.4060, 2.5413] + (1.0000 * [5.0, 6.0] + 0.0498 * [7.0, 8.0])
         = [0.1903, 0.3440] + [5.3485, 6.3983] = [5.5388, 6.7422]

Final normalization, once, after the last tile:

O = a_2 / l_2 = [5.5388 / 1.2034, 6.7422 / 1.2034] = [4.6025, 5.6025]

Check it against the global softmax over all four scores [8, 6, 10, 7]: weights [0.1125, 0.0152, 0.8310, 0.0414], weighted sum [4.6025, 5.6025]. Exact match. The trace never needed the full row, and the only division happened at the end.

Step	State after the step	What changed
Tile 1 local stats	m = 8.0, l = 1.1353, a = [1.4060, 2.5413]	First tile initializes the state
Tile 2 max update	m = 10.0, alpha = 0.1353	Maximum moved, old state must follow
Tile 2 accumulation	l = 1.2034, a = [5.5388, 6.7422]	Old state scaled, new tile added
Final division	O = [4.6025, 5.6025]	One normalization for the row

3.3 What the statistics cost — three different scopes

m and l are FP32 vectors of length N per head, in the shipped kernels: PyTorch’s flash kernels and the reference flash-attention implementation keep logsumexp state in FP32. FP32 is an implementation choice for range and precision during long accumulation, not a mathematical requirement of the recurrence — FP16 or BF16 state would replay the same arithmetic with less headroom.

2 × 4,096 × 4 bytes = 32,768 bytes = 32 KiB

That arithmetic is correct, but 32 KiB is the aggregate over every query row of one head, not the kernel’s live state. Three scopes get confused with each other:

Scope	What it counts	Size at N = 4,096, d = 128, B_r = 128
Conceptual running state, one head	m and l: one FP32 scalar per query row, over the whole head	2 × 4,096 × 4 B = 32 KiB
Live state in one CTA (one Q block)	m, l for B_r rows, plus the unnormalized accumulator a, B_r × d	2 × 128 × 4 B = 1 KiB, plus 128 × 128 × 4 B = 64 KiB (FP32)
Exported logsumexp, if written	L = m + log l, one FP32 scalar per query row, written to HBM	4,096 × 4 B = 16 KiB per head

So “32 KiB replaces 64 MiB of intermediates” is a statement about the whole head, and it uses the smallest of the three numbers. The per-CTA accumulator is the one that competes with the tile buffers for SRAM and registers: 64 KiB in FP32 at B_r = 128, d = 128, which is why real kernels keep a in registers and frequently in a narrower dtype. None of these numbers is the HBM footprint of a fused kernel except the third.

Does the logsumexp reach HBM in inference? In the published FlashAttention-2 algorithm it does: Algorithm 1 step 15 writes L_i to HBM and step 17 returns L alongside O, because the backward pass consumes it. A forward-only serving kernel may have no consumer for L and may skip the write — but that is an implementation choice you check in the kernel you pinned, not something the algorithm establishes. If yours does write it, add 4N bytes per head to the §4 model: 16 KiB here, roughly 0.01% of the 132 MiB baseline. Small enough to ignore in the ratio, large enough to state rather than assume away.

---

4. Logical traffic comparison

The numbers below come from a simplified logical transfer model: cold cache, exact byte counts with no sector padding, no register spills, fixed B_r = 128 / B_c = 64. They are algorithmic traffic, not dram__bytes_read.sum. Physical DRAM transactions depend on cache-line residency, coalescing, and scheduling, and will differ.

Logical HBM traffic for one head, N = 4,096, d = 128, FP16 (MiB)

Configuration	Reads	Writes	Total	Reduction vs matched baseline
Naive, unmasked (3 kernels)	67.00	65.00	132.00	1.00× (baseline)
FlashAttention-2, unmasked	65.00	1.00	66.00	2.00× vs naive unmasked
Idealized lower-triangular naive, causal	35.00	33.00	68.00	1.00× (causal baseline)
FlashAttention-2, causal	34.00	1.00	35.00	1.94× vs idealized causal naive

Reads and writes are split by direction, so the reduction ratios compare total traffic within the same masking condition. The unmasked and causal rows are separate experiments, not two measurements of one system.

4.1 Unmasked: the like-for-like 2.00× comparison

Same masking condition on both sides, so the ratio is clean.

Naive: 132.00 MiB, itemized in §1.1.

FlashAttention-2: Q read once (1.00 MiB) + O written once (1.00 MiB) + K and V streamed once per query tile. With T_r = 4,096 ÷ 128 = 32 query tiles:

KV reads = 32 × (1.00 + 1.00) MiB = 64.00 MiB
Total    = 1.00 + 1.00 + 64.00   = 66.00 MiB (69,206,016 B)
Ratio    = 132.00 ÷ 66.00        = 2.00×

Exactly 2.00×, and the two terms that produce it are worth naming: the baseline writes 64 MiB of intermediates and reads them back (128 MiB of round-trips), while the fused kernel spends 64 MiB re-reading K/V across query tiles. The fused kernel moves half the bytes, and all of them are streaming reads of a 2 MiB working set rather than quadratic scratch.

4.2 Causal: a separate comparison, with its own baseline

With causal masking, query tile i reads K/V tiles j = 1 … 2i: it skips tiles that are entirely in the future, and it masks the two diagonal tiles per §2.1. Complete passes over K/V are:

Σ_{i=1}^{32} i / 32 = (32 + 1) / 2 = 16.5 passes

T_r = 32 exactly, so 16.5 is exact for these tile sizes rather than an approximation.

Causal KV reads = 16.5 × 2.00 MiB = 33.00 MiB
Total            = 33.00 + 1.00 (Q) + 1.00 (O) = 35.00 MiB (36,700,160 B)

The matching causal baseline is an idealized lower-triangular naive kernel that stores only the lower triangle of S and P (16 MiB each) — the idealization treats each triangle as exactly half of N × N, ignoring the diagonal and ragged tile fill — and writes each once, reads each once:

4 × 16.00 (S, P round-trips) + 3.00 (Q, K, V reads) + 1.00 (O write) = 68.00 MiB
Ratio = 68.00 ÷ 35.00 = 1.94×

An eager PyTorch causal baseline with an explicit mask tensor is a different baseline again — its traffic depends on the mask dtype and on whether masked_fill is fused or a separate kernel — and this lesson does not quote a ratio for it.

A cross-condition ratio is not a result

Dividing the unmasked baseline (132.00 MiB) by the causal fused kernel (35.00 MiB) produces a large, impressive-looking number. It is meaningless: the causal kernel skips roughly half the score tiles, so that division mixes the masking change with the fusion change and cannot be attributed to either. Compare unmasked to unmasked (2.00×) and causal to causal (1.94×), and always name the baseline you used.

4.3 What the model does not include

• No L2 reuse. Real kernels get some K/V hits from the 40–50 MB L2, so measured traffic can come in below 66 MiB.
• No padding or sector granularity. The memory controller moves 32-byte or 128-byte sectors, so odd tile shapes round up.
• No logsumexp write. The totals assume the kernel does not store L; add 16 KiB per head at these dimensions if it does (§3.3).
• No output-logits or mask traffic, and no allocator effects.
• No speedup claim. 2.00× fewer logical bytes is not 2.00× less time — and conversely, a kernel speedup of 3.5× is not evidence that the byte model is exact. §5 is about that gap.

---

5. From kernel speedup to serving latency

Suppose a profiler says attention accounts for 15% of layer time, and a benchmark says the fused attention kernel is 3.5× faster than the baseline kernel on the same shape. How much faster is the layer?

Amdahl’s Law, with the emphasis on time:

speedup = 1 / ((1 - f_time) + f_time / S_kernel)

With f_time = 0.15 and S_kernel = 3.5:

speedup = 1 / (0.85 + 0.15/3.5)
        = 1 / (0.85 + 0.042857)
        = 1 / 0.892857
        = 1.120×

1.120× layer speedup, a 10.71% latency reduction — from a kernel that got 3.5× faster. The other 85% of the time did not move, and it caps the result. So “FlashAttention gives 2–4×” is a kernel claim, not an end-to-end one until you pass it through this calculation with a measured time share: at a 3% share the same kernel buys almost nothing, and at a 65% share it buys a great deal. Kernel speedups do reach end-to-end latency; the size of the transfer is exactly what Amdahl computes.

5.1 Why the FLOP share cannot substitute for f_time

Lesson 2’s count puts attention at 3.2% of prefill FLOPs at P = 2,048. Plug 0.032 into Amdahl with S_kernel = 3.5 and you get 1.023× — a number that would make attention optimisation look pointless. It is not a valid projection, because FLOPs and time are not proportional on a GPU:

• Attention GEMMs run on Tensor Cores (989 TFLOPS FP16 dense on H100).
• Softmax, masking, scaling, and reductions run on CUDA cores and SFUs. Ampere’s GA100 SM has 64 FP32 cores but only 16 SFUs, so the exponential issues at roughly one quarter of the FP32 rate per SM — and even the full FP32 path (19.5 TFLOPS on A100) is far below the Tensor Core path (312 TFLOPS FP16 dense). The ratio, not just the aggregate, is why the non-matmul fraction of a kernel can dominate its time.
• Non-matmul work has very low arithmetic intensity (naive softmax ≈ 1 FLOP/byte), so it also waits on memory.

A kernel can therefore be a few percent of FLOPs and a double-digit percentage of wall-clock time. Only a profile resolves it.

The 15% and 3.5× are hypothetical

They are an exercise, not a measurement, and not a benchmark of your GPU. Substitute your own f_time from torch.profiler or Nsight Systems, and your own S_kernel from timing the fused and unfused paths on the same shape, dtype, and clock settings. The shape of the conclusion survives; the numbers do not transfer.

The sensitivity is steep, which is the useful part. Same 3.5× kernel:

f_time (measured attention share)	Layer speedup	Latency reduction
15%	1.120×	10.71%
40%	1.400×	28.6%
65%	1.867×	46.4%

Attention’s share grows with prompt length, so the same kernel investment pays more at 32K context than at 2K. That is a routing decision you can defend with two profiler runs.

---

6. Boundary conditions

FlashAttention is not universally better. Five cases where the advantage shrinks or the mechanism changes:

Scenario	Mechanism	What it means for you
Short sequences	The intermediates are small: at N = 256, d = 128, FP16, one head’s S is 0.125 MiB (131,072 bytes), which is inside a 40–50 MB L2. Fitting is not the same as hitting — the cache is shared with weights, KV, and other CTAs, and a streaming write can still be evicted before the reader runs.	Profile with dram__bytes_read.sum / dram__bytes_write.sum rather than assuming the round-trips vanished. If they are served from L2, the I/O bound weakens and launch/tile overheads can dominate. There is no universal N cutoff.
Cache residency, not traffic ratio	The §4 model counts logical bytes with a cold cache. At small sizes, or with several layers resident, physical DRAM traffic diverges from the model in either direction.	Use the model to rank designs, not to predict counters on a small workload. The hardware counters are the arbiter.
Single-token decode (Q = 1)	The score “matrix” is a 1 × S vector. Query blocking collapses to T_r = 1, so thread blocks = batch × query heads. At batch 1 with 32 query heads that is 32 blocks against 132 SMs on H100.	Standard FA-1/FA-2 leaves most of the GPU idle. Production stacks dispatch Flash-Decoding (split the KV dimension across blocks, then reduce), PagedAttention, or GEMV kernels instead.
Large head dimension	SRAM is finite, so the tile sizes that fit Q, K, V, and the accumulator shrink as d grows. The specific B_r/B_c values are kernel- and dtype-dependent, and vendor kernels may change implementation or refuse the shape.	Re-derive the traffic model for the tiles your kernel actually uses; §4’s B_r = 128, B_c = 64 is a stated assumption, not a constant. Some libraries cap or split the supported head dimension — check the release notes for your version.
Training retention (backward pass)	An implementation that saves both S and P for backward pays the quadratic allocation the forward pass avoided; that is the accounting behind the 64.0 GiB figure for 32 layers × 32 heads.	Inference with no backward pass does not retain them. Read “FlashAttention is memory-efficient” as forward-pass I/O and activation memory, and check what your training framework actually stores.
PyTorch SDPA dispatch	Backend choice depends on dtype, head dimension, mask layout, and GPU architecture (SM80+).	A silent fallback to memory-efficient or math attention changes both bytes and time. Verify, do not assume.

6.1 Decode is a different geometry, not a smaller prefill

In prefill, the N × N intermediate is the problem, and tiling removes it. In decode, Q is a single row, so there is no quadratic intermediate at all: the cost is streaming the KV cache and the weights. FlashAttention-2’s advantage there is not byte reduction — it is that the whole query dimension is one tile.

T_r = 1 is the mechanical cause of the occupancy problem. FlashAttention-1 parallelised over batch and heads; FlashAttention-2 added parallelism over query blocks. With one query token, that second axis contributes nothing. Flash-Decoding (Dao et al. 2023) restores parallelism along the axis that is actually long — the KV sequence — by splitting it across thread blocks, computing a partial softmax and partial output per split, then combining them with a logsumexp reduction. That is why serving engines route decode to Flash-Decoding, PagedAttention, or GEMV kernels rather than standard FA-2.

6.2 Verify which kernel ran

An eligibility check answers “could the fused kernel run?”, not “did it run?”. The documented way to get an execution answer is to constrain dispatch and then observe the kernel that launches:

# Illustrative, not executed for this lesson. Shapes: one head's worth of a
# 32-head, 4,096-token prefill at d=128, FP16, on an SM80+ device.
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.nn.functional import scaled_dot_product_attention

q, k, v = (
    torch.randn(1, 32, 4096, 128, device="cuda", dtype=torch.float16)
    for _ in range(3)
)

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

with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof:
    scaled_dot_product_attention(q, k, v, is_causal=True)
print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=10))

sdpa_kernel(SDPBackend.FLASH_ATTENTION) enables only that backend, so an input it cannot handle fails instead of quietly falling back to the math path. That is what turns dispatch into a test rather than an assumption — but the function is still marked beta in the 2.15 docs, so pin the version you rely on. Lower-level torch.backends.cuda.can_use_flash_attention() and SDPAParams do exist, but SDPAParams has no documented constructor in the pinned docs, so code that hand-builds one depends on an internal signature that has changed between releases.

For evidence, read the profiler table for the attention kernel name, or take an Nsight Systems kernel list and check dram__bytes_read.sum / dram__bytes_write.sum against the §4 model. If the table shows separate bmm, softmax, and bmm kernels instead of one fused attention kernel, you are on the math path and the §4 baseline describes your workload, not the fused one.

---

Key Takeaways

Key takeaways

• Naive attention writes `S` and `P` to HBM: 64 MiB of intermediates per head at N=4,096, d=128, FP16, and 132.00 MiB of total logical traffic per head. The Q/K/V/O payload is only 4 MiB (1 MiB each) of that — 3.0% of the traffic.
• 64 MiB per head is a per-head allocation, not a memory plan. One active layer across 32 heads is 2.0 GiB under a parallel-heads assumption; cumulative training retention across 32 layers without checkpointing is 64.0 GiB.
• FlashAttention tiles the computation so `S` and `P` never reach HBM (0 bytes), reading `Q` once and writing `O` once. The cost is re-reading `K`/`V` per query tile: 64 MiB unmasked at 32 tiles.
• Causal masking changes the shape of the win, not just the constant. FA-2 skips blocks that are entirely future (`i < j`), but diagonal blocks still need elementwise masking, so the honest comparison halves both sides: 68.00 → 35.00 MiB per head is 1.94×. Dividing the unmasked baseline by the causal fused number is not a valid result.
• Online softmax makes tiling exact: keep a running max `m` and normalizer `l`, rescale the unnormalized accumulator `a` by `exp(m_old - m_new)`, and divide once at the end. At N=4,096 the statistics are 32 KiB of FP32, against 64 MiB of intermediates.
• Amdahl with a stated (here hypothetical) time share: 15% attention at 3.5× kernel speedup is 1.120× end-to-end, a 10.71% latency cut. FLOP share is not a substitute for `f_time`, because softmax runs on CUDA cores and SFUs, not Tensor Cores.
• Know the boundaries: cache residency is not measured hit rate, decode at `Q = 1` needs Flash-Decoding or GEMV kernels, large head dimensions force smaller tiles, training-side memory claims are about saved activations, and SDPA backend selection has to be verified with a profiler.

---

Exercises

Try it yourself

1. 01For one head at N=4,096, d=128, FP16, compute the logical HBM traffic for the explicit three-kernel baseline and for FlashAttention-2 unmasked, then state the reduction. Show reads and writes separately.

   Show a hint

   Baseline: GEMM1 reads Q+K (2 MiB) writes S (32 MiB); softmax reads S (32 MiB) writes P (32 MiB); GEMM2 reads P+V (33 MiB) writes O (1 MiB). Reads 67, writes 65, total 132.00 MiB. FA-2: Q 1 + O 1 + 32 query tiles × 2 MiB of KV = 66.00 MiB. Ratio 132/66 = 2.00×, and both sides are unmasked.

2. 02Trace the FlashAttention-2 unnormalized accumulator for one query row across two tiles. Tile 1 scores [8.0, 6.0] with V=[[1,2],[3,4]]; Tile 2 scores [10.0, 7.0] with V=[[5,6],[7,8]]. Give m, alpha, l, a after each tile, then the final output, and check it against the full-row softmax.

   Show a hint

   Tile 1: m=8.0, w=[1, 0.1353], l=1.1353, a=[1.4060, 2.5413]. Tile 2: m=10.0, alpha=e^-2=0.1353, w=[1, 0.0498], l=0.1353×1.1353+1.0498=1.2034, a=0.1353×[1.4060, 2.5413]+[5.3485, 6.3983]=[5.5388, 6.7422]. O=a/l=[4.6025, 5.6025]. Global softmax over [8,6,10,7] gives the same vector, so the single final division is exact.

3. 03A profile shows attention at 22% of prefill time and your fused kernel benchmarks 2.8× faster than the unfused path on the same shape. What layer speedup should you predict, and what would you have predicted if you had used the 3.2% FLOP share instead? Explain which number you would put in a capacity plan.

   Show a hint

   f=0.22, S=2.8: 1/(0.78+0.0786)=1/0.8586≈1.165×, about a 14.1% latency reduction. With f=0.032: 1/(0.968+0.00914)≈1.023×, about 2.3%. Use the 1.165×, and state it as a projection from one profile at one prompt length — FLOP share omits softmax, masking, and memory stalls, so it cannot stand in for time share.

4. 04With `B_r = 128`, `B_c = 64`, `N = 4,096`, and causal masking: which K/V tiles does query tile `i = 5` load, and which of those tiles need elementwise masking rather than a plain load? For one row in the upper half of that query block, explain what goes wrong if you skip the mask but keep the skip bound.

   Show a hint

   Query tile 5 covers rows 513–640, so `j = 1 … 2i = 10`. Tiles 1–8 are entirely permitted and load unmasked. Tile 9 (columns 513–576) is diagonal for the upper rows 513–576 and fully permitted for the lower rows 577–640. Tile 10 (columns 577–640) is entirely future for rows 513–576 and diagonal for rows 577–640 — so it is loaded but contributes nothing to the upper half. Skipping the mask leaves the scores for those columns in the running max and sum, so the row attends to tokens after it; the output is a valid-looking weighted average of the wrong distribution, and nothing raises an error.

---

Module 1 exit task: Latency attribution report

As the capstone for Module 1, your deliverable is a checkable, one-page latency attribution report evaluating an optimization proposal under strict measurement boundaries.

Industry grounding: Concurrency sweeps and serving trade-offs

In LLM Inference Benchmarking: Performance Tuning with TensorRT-LLM, NVIDIA engineers evaluate serving performance by sweeping concurrency and plotting aggregate GPU token throughput against per-user generation speed. Across concurrency levels, higher concurrency increases total tokens per second as execution units are saturated, while per-user output speed decreases and response times lengthen. This lesson’s scenario is an original synthetic illustration; it is not NVIDIA benchmark data.

The scenario: Document QA latency trace

You are evaluating a document QA endpoint running a nominal 8B parameter dense model (32 layers, d = 128, 32 query heads, FP16, untied LM head).

• Workload: Fixed-concurrency C = 8 closed-loop test (P = 4,096 prompt tokens, O = 128 output tokens).
• Service targets: Client p95 TTFT ≤ 350.0 ms; decode generation ≥ 25.0 tokens/s (TPOT ≤ 40.0 ms).

Hypothetical single-request trace (all timings synthetic):
We explicitly assume a simplified, strictly serial, non-overlapping critical path where client TTFT is the sum of network, queue, GPU prefill, and streaming serialization. (In production, overlapping operations, batch scheduling, and unsynchronized clocks mean intervals cannot be simply summed or subtracted without synchronized tracing).

Serial segment	Hypothetical measurement source	Trace duration	Share of TTFT	Notes
1. Client ↔ Server network	Synchronized socket ingress delta	15.0 ms	3.9%	Synthetic transport estimate
2. Server admission queue	Server enqueue → dequeue timestamps	65.0 ms	17.1%	Dwell waiting for batch slot
3. GPU prefill execution	Hypothetical CUDA profiler span	290.0 ms	76.3%	40% attention (116.0 ms), 60% MLP/norms (174.0 ms)
4. Serialization & stream	First SSE chunk serialization	10.0 ms	2.6%	Emitting first streamed chunk
Total Client TTFT (this trace)	Client socket timer	380.0 ms	100.0%	Single trace; service p95 compliance is unknown

For this modeled trace, client TTFT (380.0 ms) exceeds the 350.0 ms target, and decode TPOT (32.0 ms/token) is below the 40.0 ms budget. A single trace establishes neither service-level p95 compliance nor an end-to-end system bottleneck. Inside this trace’s prefill span, attention accounts for 40.0% (116.0 ms) of execution time.

Deliverable: Latency attribution report & worked check

Produce a one-page report structured as follows, verified against the worked values:

1. Workload & Target: Nominal 8B model, FP16, C = 8 closed-loop, P = 4,096, O = 128. Objectives: p95 TTFT ≤ 350 ms, decode ≥ 25 tokens/s.
2. Trace Timings: Network 15.0 ms, queue 65.0 ms, GPU prefill 290.0 ms, write 10.0 ms (TTFT 380.0 ms); decode TPOT 32.0 ms.
3. Observation & Hypotheses: For this trace, TTFT exceeds target while decode TPOT is within budget. Whether the service meets its p95 SLO is unknown without a request distribution. In this trace, GPU prefill represents 76.3% of TTFT.
4. Proposed Change: Enable FlashAttention-2 causal kernel in prefill, with an assumed S_kernel = 3.00× speedup across the entire timed attention portion (116.0 ms) on matching shapes.
5. Projected Impact (Amdahl’s Law within GPU prefill boundary):
   • GPU Prefill boundary: Attention share f_attn = 0.40, non-attention 0.60.

     S_prefill = 1 / (0.60 + 0.40 / 3.00) = 1 / 0.733333 = 1.3636×
     T_prefill_projected = 174.0 ms + (116.0 ms / 3.00) = 212.67 ms (26.67% reduction)

   • Client TTFT boundary (this trace only): Holding serial non-prefill components fixed (15 + 65 + 212.67 + 10):

     T_client_projected = 302.67 ms (S_client = 380.0 / 302.67 ≈ 1.2555×, 20.35% reduction)

     The 302.67 ms result is a conditional projection for this single trace. Service-level p95 compliance remains unverified until a post-change distribution is collected.
6. Boundary Errors to Avoid:
   • Kernel vs. end-to-end: Applying 3.00× to TTFT predicts 380.0 / 3.00 = 126.7 ms, ignoring non-attention compute and queue dwell.
   • FLOP share vs. time share: Attention is ~6% of FLOPs at P = 4,096 (Lesson 2). FLOP share cannot substitute for observed execution-time share (40%), as it omits non-matmul operations and execution stalls.
   • Unmeasured components: Holding queue dwell at 65.0 ms is a static assumption; queuing dynamics under load must be empirically measured.
7. Next Measurement Plan:
   • Confirm fused kernel dispatch via with sdpa_kernel(SDPBackend.FLASH_ATTENTION):.
   • Validate GPU prefill execution time against the 212.67 ms projection using matched profiler spans (CUDA events / torch.profiler).
   • Monitor server histograms (request_prefill_time_seconds, request_queue_time_seconds) for server-side elapsed-time distributions.
   • Collect client-observed request samples to accept or reject the service p95 TTFT SLO (≤ 350 ms).
   • Evaluate physical memory traffic independently with Nsight Compute (dram__bytes_read.sum, dram__bytes_write.sum); the §4 logical model is an algorithmic count, not a physical hardware guarantee.

---

Sources and scope

Claim	Source	Status / limitation
Standard attention needs Θ(Nd + N²) HBM accesses; tiling reduces this to Θ(N²d²/M) for SRAM size M	Dao et al., FlashAttention, arXiv:2205.14135, §1, §3.1, Theorem 2	Establishes the I/O complexity bound; specific byte counts depend on tile sizes, dtype, and cache behaviour.
Online softmax updates a running maximum and normalizer without global row synchronisation	Milakov & Gimelshteyn (2018); Rabe & Staats (2021); Dao et al. 2022, Algorithm 1	The identity predates FlashAttention; FlashAttention contributes the SRAM-tiled formulation.
FlashAttention-2 inverts the loop (Q outer, KV inner), keeps an unnormalized accumulator, and divides once	Dao, FlashAttention-2, arXiv:2307.08691, §2.1–2.2, Algorithm 1	Applies to the FA-2 reference implementation; vendor kernels may differ in tiling and scheduling.
Causal attention skips blocks entirely in the future but still masks blocks that straddle the diagonal	Dao 2023, §3.1.1 (causal masking); Algorithm 1, steps 15 and 17	The paper states the two mechanisms separately; skip bounds depend on the B_r/B_c ratio and are not universal.
FA-2 forward attention kernel reaches 50–73% of peak FLOP/s on A100 (up to 230 TFLOP/s)	Dao 2023, §4.1	Kernel-level compute utilization, not end-to-end training MFU (up to ~72%) and not SM occupancy.
Non-matmul operations run at lower throughput than Tensor Core GEMMs, so FLOP share misleads	Dao 2023, §1–2; NVIDIA Ampere GA100 whitepaper (SM diagram: 64 FP32 cores / 16 SFUs); A100 datasheet (19.5 TFLOPS FP32, 312 TFLOPS FP16 dense)	Qualitative claim. The 64-FP32-core / 16-SFU ratio and the 19.5-vs-312 TFLOPS split are Ampere GA100 figures and do not transfer unchanged to other architectures.
Single-token decode at Q = 1 underutilizes SMs in low-batch regimes; Flash-Decoding splits the KV dimension	Dao et al., Flash-Decoding, 2023	Scoped to low-batch, long-context decode; optimal kernel choice is configuration-dependent and should be profiled.
SDPA dispatches to fused backends on SM80+ under supported dtype/shape/mask; sdpa_kernel forces or rejects a backend	PyTorch 2.15 SDPA docs and torch.nn.attention.sdpa_kernel	Version-pinned documentation. SDPAParams is listed without a documented constructor, so hand-built eligibility checks are version-fragile; confirm selection with a profiler.
H100 SXM: 989 TFLOPS FP16 dense, 3.35 TB/s HBM3, ridge ≈ 295 FLOP/byte; A100: 312 TFLOPS, 2.039 TB/s, ≈ 153 FLOP/byte	NVIDIA H100 and A100 product pages; Dao et al. 2022 §2.1	Peak, dense, not sustained. Ridge point is arithmetic on peaks, not a property of any kernel.
Traffic model: cold cache, exact bytes, no sector padding, no spills, B_r = 128, B_c = 64	Internal: research/flashattention-gpu-data-movement.md §6; checked by scripts/verify-flashattention-example.py	Logical model only. dram__bytes_read.sum / dram__bytes_write.sum in Nsight Compute report physical transactions, which will differ.
Concurrency sweeps evaluate serving trade-offs by plotting aggregate GPU throughput against per-user generation speed	NVIDIA: LLM Inference Benchmarking	Explains concurrency sweeps and the throughput versus per-user speed trade-off using TensorRT-LLM (trtllm-bench). Exit task scenario is an original synthetic illustration, not NVIDIA data.

---

Related lessons

• Inference Optimisation & AI Infrastructure · Measure and Explain Inference PerformanceServing Workloads and Performance Metrics12 min
• Inference Optimisation & AI Infrastructure · Measure and Explain Inference PerformanceWhere Inference Time and Memory Go13 min
• LLM Internals & Architecture · Attention mechanisms from scratchSelf-Attention: Queries, Keys, and Values14 min