How FlashAttention Reduces HBM Traffic Exactly

Welcome back. In the previous lesson, you separated weight quantization from KV-cache quantization and used memory and bandwidth budgets to decide which one addresses a serving bottleneck. FlashAttention tackles a related but distinct issue: rather than compressing tensors, it changes the order of computation so the GPU moves far less data between high-bandwidth memory (HBM) and fast on-chip memory.

This lesson focuses on the central claim behind FlashAttention: it computes ordinary softmax attention—without sparsity, low-rank approximation, or altered attention weights—while avoiding the costly materialization of the full $N\times N$ score and probability matrices. You will see why this matters most for long-context training and prefill, how tiling and online softmax make it possible, and why doing more arithmetic can still make the kernel faster.

---

The performance problem is data movement, not the attention formula

For one attention head with sequence length $N$ and head dimension $d$, standard scaled dot-product attention is

$$S=\frac{Q{K}^{\mathsf{T}}}{\sqrt{d}}$$ $$P=\mathrm{softmax}(S)$$ $$O=PV$$

where $Q$, $K$, and $V$ each have shape $N\times d$. A causal mask, padding mask, or other attention mask is normally incorporated before softmax; FlashAttention supports this while preserving the semantics.

The important shape difference is:

Tensor	Shape	Size growth with sequence length
$Q,K,V,O$	$N\times d$	Linear in $N$
Scores $S$	$N\times N$	Quadratic in $N$
Probabilities $P$	$N\times N$	Quadratic in $N$

A conventional decomposition of attention often executes these as separate kernels:

1. Read $Q$ and $K$ from HBM, compute $S$, then write $S$ to HBM.
2. Read $S$, apply masking and softmax, then write $P$ to HBM.
3. Read $P$ and $V$, compute $O$, then write $O$ to HBM.

The score and probability matrices are intermediate results, yet they are enormous. With $N=8192$ and BF16 elements, one $N\times N$ matrix occupies

$$8192^2\times 2=134,217,728\text{ bytes}\approx 128\text{ MiB}$$

per attention head. Writing and later reading both $S$ and $P$ alone costs roughly $512$ MiB of HBM traffic per head, before considering $Q$, $K$, $V$, output, batch size, or multiple heads.

HBM is extremely fast in absolute terms, but on-chip SRAM is much faster and much smaller. Therefore, a GPU kernel can be limited not by its floating-point units, but by the repeated movement of intermediate tensors to and from HBM.

A useful distinction for system design is this:

• The mathematical work remains quadratic. Exact dense attention still evaluates all relevant query-key interactions, requiring $O(N^{2}d)$ arithmetic.
• The HBM traffic is reduced. FlashAttention’s innovation is I/O-aware scheduling, not a reduction in the fundamental dense-attention FLOP count.
• The peak intermediate-memory footprint drops sharply. The full $N\times N$ score and probability matrices never exist in HBM.

That first point prevents a common interview mistake: FlashAttention is not “linear attention,” and it is not an approximation method.

---

FlashAttention’s core move: compute on tiles, not full matrices

The GPU cannot hold full $S$ or $P$ in SRAM. It can, however, hold small blocks of $Q$, $K$, $V$, and temporary results. FlashAttention divides these matrices into tiles and runs the computation for one query tile against one key-value tile while those tiles reside on chip.

Conceptually, the algorithm proceeds as follows:

1. Partition $Q$ into query blocks and $K,V$ into corresponding key and value blocks.
2. Load one $K,V$ block from HBM into SRAM.
3. For each $Q$ block, load that block and its current partial output into SRAM.
4. Compute the block’s scores, masking, softmax contribution, and weighted value contribution entirely on chip.
5. Update the partial output and a small amount of softmax state for that query block.
6. Write only the updated output block and softmax state back to HBM.

The scheduling has a subtle trade-off. A key-value tile can be reused while it sits in SRAM, but each query block may be revisited as the kernel processes subsequent key-value tiles. That causes some re-reading of $Q$ and partial output. This is acceptable because those tensors are $N\times d$, while the avoided matrices are $N\times N$.

In the original FlashAttention analysis, if $M$ denotes available SRAM measured in scalar elements, the conventional implementation has HBM traffic dominated by a quadratic $O(N^2)$ term. FlashAttention’s leading tiled-access term is approximately

$$O\left(\frac{N^2{d}^2}{M}\right)$$

along with unavoidable linear-sized input and output traffic. A rough reduction factor is

$$\frac{M}{d^2}$$

when the SRAM capacity is large relative to one head’s $d^2$ working set.

For example, using $M=65,536$ BF16-equivalent elements and $d=64$,

$$\frac{M}{d^2}=\frac{65,536}{64^2}=16$$

This is a useful first-order intuition, not a promise that a specific GPU will be precisely $16$ times faster. Real block sizes are chosen around SRAM capacity, register pressure, tensor-core tile shapes, occupancy, masking requirements, and the details of the GPU architecture.

---

Why softmax does not prevent tiling

Matrix multiplication is naturally tiled, but attention has an apparent obstacle: softmax normalizes each row across all keys. If a query row’s scores are split across key tiles, how can a kernel normalize correctly before seeing every score?

The answer is online softmax. For a score row $x$, the numerically stable softmax is based on:

$$m=\underset{j}{\max }{x}_j$$ $$l=\sum _j{e}^{x_j-m}$$ $$\mathrm{softmax}(x{)}_j=\frac{e^{x_j-m}}{l}$$

Instead of retaining all of $x$, FlashAttention retains only two running statistics for each query row:

• $m$: the largest score seen so far;
• $l$: the corresponding normalized exponential sum.

Suppose a row has already processed previous key tiles with state $(m_{\mathrm{old}},l_{\mathrm{old}},o_{\mathrm{old}})$. For a new key tile, the kernel computes local statistics $(m_{\mathrm{block}},l_{\mathrm{block}},o_{\mathrm{block}})$, where $o_{\mathrm{block}}$ is the value-weighted result normalized within that tile.

The merged maximum is

$$m_{\mathrm{new}}=\max \left(m_{\mathrm{old}},m_{\mathrm{block}}\right)$$

The merged normalization constant is

$$l_{\mathrm{new}}=e^{m_{\mathrm{old}}-m_{\mathrm{new}}}{l}_{\mathrm{old}}+e^{m_{\mathrm{block}}-m_{\mathrm{new}}}{l}_{\mathrm{block}}$$

The merged output is

$$o_{\mathrm{new}}=\frac{e^{m_{\mathrm{old}}-m_{\mathrm{new}}}{l}_{\mathrm{old}}{o}_{\mathrm{old}}+e^{m_{\mathrm{block}}-m_{\mathrm{new}}}{l}_{\mathrm{block}}{o}_{\mathrm{block}}}{l_{\mathrm{new}}}$$

The rescaling factors account exactly for the fact that an unexpectedly large score may appear in a later tile. Once all key blocks have been incorporated, the resulting $o_{\mathrm{new}}$ is the same attention output that a full-row softmax would produce.

In other words, the algorithm does not compute independent softmaxes for each tile and concatenate them. That would be wrong. It computes each local contribution and then merges it with the correct normalization.

---

Kernel fusion: keeping intermediates on chip

Tiling enables the individual pieces of attention to be fused into one coordinated kernel:

• matrix multiplication for block scores;
• scaling and causal or padding masking;
• softmax statistics and exponentials;
• multiplication by the value tile;
• partial-output update.

A conventional implementation makes each operation a boundary at which a large intermediate may be written to and later read from HBM. FlashAttention makes SRAM the workspace for those intermediates. Only the final $O$, plus small per-row softmax statistics needed during the computation or backward pass, must persist outside the kernel’s on-chip working set.

This is a broader GPU optimization pattern worth retaining:

When a workload is memory-bound, recomputing a value can be faster than storing it and fetching it later.

The statement is conditional. Recompute is beneficial only if its additional arithmetic is cheaper than the memory traffic it avoids. In attention’s backward pass, that condition often holds.

---

Training: recomputation replaces saved attention probabilities

During training, backpropagation needs information related to the attention probabilities $P$. A conventional implementation saves $P$, then reads it later during backward propagation. Since $P$ is $N\times N$, this is another major memory and HBM-traffic cost.

FlashAttention instead saves:

• the final attention output $O$, with shape $N\times d$;
• compact softmax statistics, typically one value or a small set of values per query row.

During the backward pass, it reloads relevant $Q$, $K$, and $V$ blocks, recomputes the needed score and probability blocks in SRAM, and uses them to form gradients with respect to $Q$, $K$, and $V$.

This is selective recomputation, closely related in spirit to activation checkpointing. The difference is important:

Approach	Saved state	Later cost	Typical motivation
Store full $P$	Large $N\times N$ probability matrix	HBM reads in backward pass	Minimize recomputation
FlashAttention	Output plus compact softmax statistics	Recompute score/probability tiles	Reduce HBM traffic and memory
General activation checkpointing	Selected activations	Recompute selected forward regions	Fit a larger training workload in memory

FlashAttention may perform more FLOPs than a conventional attention backward pass. Yet its wall-clock time can be lower because the avoided HBM transfers dominate the extra arithmetic. This directly applies the roofline reasoning from earlier in the course: fewer FLOPs does not necessarily mean lower latency.

---

Exact attention does not mean bitwise-identical output

“Exact” in FlashAttention means it evaluates the same dense softmax-attention operation, rather than changing the mathematical operator through approximation.

It does not necessarily mean a bit-for-bit identical tensor compared with every other implementation. Floating-point arithmetic is not associative:

$$(a+b)+c\ne a+(b+c)$$

in general under finite precision. A tiled implementation changes reduction and accumulation order, and different kernels may use different internal precisions. Small numerical differences can therefore occur, just as they can across two ordinary matrix-multiply implementations.

The relevant distinction is:

Method	Does every relevant query-key interaction contribute?	Changes attention mathematics?
FlashAttention	Yes	No
Sliding-window attention	No, distant positions excluded	Yes
Block-sparse attention	No, selected blocks excluded	Yes
Linear-attention approximation	Not evaluated as dense softmax attention	Yes
Low-rank approximation	Approximates the score or probability structure	Yes

This is why FlashAttention is broadly deployable: it improves the standard attention primitive without requiring a model to be trained around a different attention mechanism.

---

Where FlashAttention helps in LLM systems

FlashAttention matters most when many query tokens are processed together:

• Training: forward and backward passes often process long sequences, so avoiding the full attention matrix substantially lowers memory use and HBM traffic.
• Prefill: a serving system processes all prompt tokens to establish the KV cache. Long prompts produce a large query-by-key attention region, making FlashAttention highly valuable for TTFT and memory feasibility.
• Decode: each active sequence typically has only one new query token at a step. There is no full $N\times N$ intermediate to eliminate in the same way, although specialized attention kernels still matter. Decode remains strongly affected by reading model weights and growing KV caches.

This distinction connects FlashAttention to the serving metrics from earlier lessons:

Serving phase	Main FlashAttention contribution
Prompt prefill	Reduces intermediate-memory pressure and HBM traffic for dense prompt attention; can improve TTFT
Autoregressive decode	Less direct benefit from avoiding an $N\times N$ matrix; KV-cache and weight bandwidth remain primary concerns
Long-context workloads	Particularly valuable because conventional attention intermediates grow quadratically with prompt length

Do not frame FlashAttention as a replacement for KV-cache quantization. They attack different traffic sources:

• FlashAttention avoids materializing ephemeral attention-score and probability matrices.
• KV-cache quantization shrinks the persistent per-request $K,V$ state that decoding repeatedly reads.
• Weight quantization reduces fixed model-residency memory and repeated weight movement.

A production serving stack commonly uses all three, subject to hardware and kernel support.

---

A concise architecture-review answer

For an interview or design review, a strong explanation can follow this structure:

Standard attention typically writes the $N\times N$ score matrix and the $N\times N$ probability matrix to HBM between separate operations. For long sequences, that HBM traffic and intermediate memory dominate runtime. FlashAttention tiles $Q$, $K$, and $V$, loads small blocks into fast SRAM, and fuses score computation, masking, softmax, and value aggregation. Online softmax maintains running maxima and normalization sums, so the output can be accumulated tile by tile with the same dense softmax-attention result. In training, it saves compact normalization statistics and recomputes attention blocks during backward propagation rather than storing the full probability matrix. It retains quadratic dense-attention computation, but substantially reduces HBM traffic and peak memory.

---

Key takeaways

FlashAttention is an I/O-aware exact-attention algorithm.

• Standard dense attention’s key bottleneck is often the HBM read/write traffic caused by materializing $S$ and $P$, each of shape $N\times N$.
• FlashAttention tiles $Q$, $K$, and $V$, performs attention subcomputations in SRAM, and avoids storing the full score and probability matrices in HBM.
• Online softmax keeps only running maxima and normalization sums, allowing blockwise computation while preserving correct global softmax normalization.
• The algorithm still performs dense $O(N^{2}d)$ attention arithmetic; it reduces memory traffic and peak intermediate storage rather than changing the attention operator.
• In the backward pass, recomputation trades extra FLOPs for substantially lower HBM traffic, which is often faster for a memory-bound workload.
• The strongest serving benefit is generally during long-prompt prefill; decode has different dominant costs, especially model-weight and KV-cache reads.

Next, you will analyze speculative decoding and calculate when a draft model’s accepted tokens save enough verification work to improve generation latency and throughput.