Count the HBM bytes FlashAttention reads and writes for one head at N = 8K and d = 128, and compare it with naive attention.

TL;DR: Naive attention moves about 520 MiB for one head at N = 8,192, d = 128 in BF16, almost all of it the 128 MiB score matrix written and read twice. FlashAttention-2's loop re-reads K and V once per 128-row query block, which is 260 MiB by the paper's model, but those 4 MiB of K and V sit in the H100's 50 MB L2, so the HBM traffic is close to the 8 MiB compulsory minimum. At 34.4 GFLOP that is about 4,096 FLOPs per byte: compute-bound, not bandwidth-bound.

How to approach it

Fix the conventions first: BF16 operands, non-causal (halve the score traffic and FLOPs for causal), one head, forward pass only. Then count three numbers, not one: naive traffic, the SRAM-to-HBM model in the FlashAttention paper, and what HBM actually sees once L2 is in the picture. The gap between the last two is what separates a reader of the paper from someone who has profiled the kernel.

A strong answer

The compulsory traffic is reading Q, K and V and writing O, each N by d:

one N x d tensor = 8,192 x 128 x 2 B = 2 MiB
Q + K + V + O    = 4 x 2 MiB         = 8 MiB   (the floor for any kernel)

Naive. The unfused kernel writes $S = QK^\top$, reads it for softmax, writes $P$, and reads $P$ for $PV$:

S (BF16)         = 8,192 x 8,192 x 2 B         = 128 MiB
score traffic    = write S + read S + write P + read P = 4 x 128 = 512 MiB
total            = 512 + 8                     = 520 MiB
at 3.35 TB/s     = 545 x 10^6 B / 3.35e12 B/s  = 163 us

If the scores are kept in FP32, as careful unfused code does, the score traffic doubles.

FlashAttention, paper model. FA2 puts the query loop outermost: each of $T_r = N / B_r$ thread blocks loads one Q tile, sweeps all of K and V through shared memory, and writes its O tile. Counting every K and V load as an HBM access:

Br = 128:  Tr = 8,192 / 128 = 64 query blocks
K and V reads  = 64 x (2 MiB + 2 MiB) = 256 MiB
Q read + O write                       =   4 MiB
total                                  = 260 MiB
(Br = 64 doubles the re-reads: 512 + 4 = 516 MiB)

That is only 2x better than naive, which surprises people who remember "FlashAttention is IO-optimal". The paper's bound is $\Theta(N^2 d^2 / M)$ for SRAM of size $M$; with d = 128 and about 228 KB of shared memory per SM, $d^2/M$ is not small, so the re-read term stays large in this model.

FlashAttention, real HBM. The 64 thread blocks for one head all stream the same 4 MiB of K and V, and they run concurrently: with 132 SMs and blocks scheduled head by head, about two heads are live at once, so about 8 MiB of K and V is in flight. H100 SXM has 50 MB of L2. The first block to touch a K tile pulls it from HBM; the other 63 hit L2.

HBM traffic       ~ 8 MiB (Q, K, V once each, O once) = 8.39 x 10^6 B
FLOPs             = 4 N^2 d = 4 x 8,192^2 x 128   = 34.4 GFLOP  (QK^T and PV)
intensity         = 34.4e9 / 8.39e6               = 4,096 FLOP/B
H100 ridge (BF16) = 989 TFLOP/s / 3.35 TB/s       = 295 FLOP/B
compute time      = 34.4e9 / 989e12               = 34.7 us
HBM time          = 8.39e6 / 3.35e12              = 2.5 us
naive / real HBM  = 520 / 8                       = 65x fewer bytes

So the honest summary: FlashAttention turns a bandwidth-bound op at 163 us into a compute-bound one whose floor is 34.7 us, and the traffic reduction comes from two things the textbook count merges. Fusion removes the score matrix entirely (512 MiB to 0). L2 reuse across concurrent query blocks removes most of the K and V re-reads (256 MiB to about 4 MiB). The first is guaranteed by the algorithm; the second depends on scheduling.

The reversal: the L2 argument breaks when the working set outgrows L2. At N = 128K one head's K and V are 64 MiB, more than L2, and a scheduler that interleaves many heads at once thrashes it. Kernels then order the grid so the blocks sharing a head run together, or use thread-block clusters to multicast K and V tiles across SMs.

What interviewers probe next

"Now it is causal." Fully masked tiles are skipped before loading, so FLOPs halve to 17.2 GFLOP and naive score traffic halves only if the naive kernel also skips, which it does not. The ratio improves for FlashAttention.

"Decode, one query against 8K keys?" Different regime entirely. Traffic is K and V once (4 MiB per KV head) for 4 x 8,192 x 128 FLOPs per query head, about 1 FLOP per byte times the GQA group size. Decode attention is bandwidth-bound regardless of fusion.

"How would you confirm the L2 claim?" Nsight Compute's DRAM bytes and L2 hit rate for the kernel. DRAM bytes near 8 MiB per head confirms reuse; near 260 MiB means the grid order is defeating L2.

How it's scored

Weak: says FlashAttention "reduces memory from N squared to N" without a byte count, or counts naive traffic as one write of S.

Hire: produces 520 MiB naive and the compulsory 8 MiB, and states that fusion removes the score matrix.

Strong Hire: also computes the 260 MiB re-read model, explains why L2 collapses it, compares intensity with the 295 FLOP/B ridge, and names when the L2 argument fails.

Common mistakes

• Quoting the asymptotic bound as if it were a byte count, without plugging in d and SRAM size.
• Forgetting P, so naive traffic is reported as 256 MiB instead of 512 MiB of score traffic.
• Mixing MiB and MB inside one chain: the tensors are binary-sized, the bandwidth is decimal.
• Concluding FlashAttention is bandwidth-bound in prefill because it was in the naive version.