Kernel Fusion and the Memory Wall

Consider three operations: a matrix multiply, a bias addition, and ReLU. If each launches a separate kernel, intermediate results can be written to GPU memory and then read back by the next operation. A fused implementation can keep some of those values close to the arithmetic units instead.

Fusion can save both intermediate memory transfers and kernel-launch overhead. Its benefit depends on the operation’s shape and bottleneck.

Compute needs a supply of bytes

A GPU has several storage levels. HBM holds large tensors; caches retain some recently used data; shared memory and registers support work within a streaming multiprocessor. Registers have finite capacity and access costs, and shared memory is explicitly managed rather than simply another automatic cache.

Arithmetic intensity is the workload's arithmetic operations divided by the bytes transferred across the memory boundary being studied. The roofline compares that intensity with peak compute and bandwidth. A small BF16 matrix-vector product may perform roughly two FLOPs per two-byte weight read; a matrix product over many token positions reuses each weight more.

Fusing bias and activation into a matrix-multiply epilogue avoids writing and rereading the intermediate result. The original inputs and weights must still be read, and the final output written.

Softmax needs a global row statistic

For a row of scores $s$, stable softmax computes

An unfused implementation might read the row for its maximum, read it again to write exponentials, read those exponentials for their sum, then read them again to write normalized output. Ignoring the small reduction vectors, that is four reads and two writes of row-sized arrays.

A fused kernel can read and write the row once if it can retain the needed values on chip. If a row is too large, the algorithm needs tiling or additional passes. Normalizing each tile independently would be wrong because the denominator must include the whole row.

FlashAttention: avoid the full score matrix

Dense attention computes $\mathrm{softmax}(QK^\top/\sqrt d)V$. Materializing the score matrix for one head needs $N^2$ elements. At FP16, $N=8192$ means 128 MiB for that matrix alone.

FlashAttention works on tiles and maintains online softmax statistics. It avoids storing the complete score and probability matrices in HBM. Its auxiliary memory is linear in sequence length, but its total HBM traffic is not universally linear: keys and values can be loaded repeatedly for different query tiles. The original paper analyzes an IO bound proportional to $N^2d^2/M$ over its stated SRAM-capacity regime. FlashAttention paper.

For one query tile, retain a running maximum $m$, normalization sum $\ell$, and unnormalized output numerator $a$. When the next score tile has maximum $m_t$, use $m'=\max(m,m_t)$ and rescale the old accumulators by $\alpha=\exp(m-m')$:

After all key tiles, return $a/\ell$. The rescaling makes earlier and later tiles share one softmax denominator.

FlashAttention computes the same attention function, subject to floating-point differences from the order and precision of its operations. Output comparisons therefore need an appropriate numerical tolerance.

A small executable reference

This CPU example demonstrates the online reduction for one unmasked query. It is a correctness sketch, not a GPU implementation or a speed benchmark.

import math
import numpy as np
 
rng = np.random.default_rng(7)
q = rng.normal(size=4)
k = rng.normal(size=(11, 4))
v = rng.normal(size=(11, 3))
 
scores = k @ q / math.sqrt(q.size)
prob = np.exp(scores - scores.max())
prob /= prob.sum()
reference = prob @ v
 
maximum = -np.inf
normalizer = 0.0
numerator = np.zeros(v.shape[1])
for start in range(0, len(k), 3):
    tile_scores = scores[start:start + 3]
    new_maximum = max(maximum, tile_scores.max())
    correction = np.exp(maximum - new_maximum)
    tile_exp = np.exp(tile_scores - new_maximum)
    normalizer = correction * normalizer + tile_exp.sum()
    numerator = correction * numerator + tile_exp @ v[start:start + 3]
    maximum = new_maximum
 
result = numerator / normalizer
np.testing.assert_allclose(result, reference, rtol=1e-12, atol=1e-12)
print(result)

A causal kernel must also apply the correct mask, including partially masked tiles, and safely handle rows with no valid keys. Production implementations add parallel work partitioning, precision choices, and hardware-specific data movement.

FlashAttention-2 improves work partitioning and reduces non-matmul overhead. FlashAttention-3 exploits Hopper's asynchronous execution and low-precision facilities. FlashAttention-2, FlashAttention-3.

The attention operation ends with the weighted value sum; the transformer's learned output projection is a separate operation unless an implementation explicitly combines it. Likewise, an entire gated MLP is not necessarily one kernel. Inspect the actual execution trace before deciding which intermediate tensors were eliminated.

Next, PagedAttention addresses a different memory problem: where growing per-request KV data lives.

Additional reading

• Dao et al., FlashAttention.
• Dao, FlashAttention-2.
• Shah et al., FlashAttention-3.