IO-Aware Attention: Tiling the Score Matrix Without Materializing It

A naive attention kernel moves four intermediates through global memory for every layer of every batch element: the $n \times n$ score matrix, the $n \times n$ softmax output, and in the backward pass both of those again. For $n = 4096$ and 16 heads in 16-bit, each is $4096^2 \times 16 \times 2$ bytes, about 512 MiB per layer per batch element. The arithmetic that produces them is a handful of matrix multiplies, and per element it is a small constant number of flops against a handful of bytes. On hardware where a global memory read costs far more than a multiply, that ratio fixes the runtime at the bandwidth limit, and the limit sits far below what the multiply units could sustain.

The roofline argument in one paragraph

The score matrix requires $n^2 d_k$ multiply-accumulates to compute and $n^2$ elements to store. Reading it back for the softmax, writing the softmax result, reading it again for the value product: each pass moves the same $n^2$ elements. Four passes over $n^2$ elements at 2 bytes each move $8n^2$ bytes, while the two matrix products perform about $4 n^2 d_k$ flops, so the arithmetic intensity of the whole chain is roughly $d_k/2$ flops per byte, independent of $n$. For $d_k = 64$ that is about 32 flops per byte. A device whose peak throughput is in the hundreds of flops per byte of bandwidth is not limited by its multiply units; it is limited by the memory system, and the only way to improve the situation is to move fewer bytes.

Tiling does exactly that. If a block of scores can be produced, consumed by the softmax, and multiplied by a block of values while it is still in registers or shared memory, the $n^2$ intermediate never reaches global memory. The bytes moved become the queries, keys, and values themselves: $O(n d)$ per head rather than $O(n^2)$. The flops are unchanged. The trade is that the softmax now has to be computed over blocks that are seen one at a time, which means the normalizing constant is not known until the last block arrives.

Softmax over blocks, exactly

The obstacle is the denominator. The softmax of a row needs $\sum_j e^{s_j}$, and a tiled kernel sees the row’s scores in chunks. Two properties make a streaming version possible.

The first is that the maximum, needed for numerical stability, can be tracked incrementally: $\max(a, b)$ is associative, so a running maximum over blocks is the true maximum once all blocks have been seen. The second is that a weighted sum can be rescaled after the fact. If the running state after block $i$ is a partial sum $\ell_i$ and a partial weighted output $o_i$, both computed relative to a running maximum $m_i$, and the next block raises the maximum to $m_{i+1}$, then the earlier partials must be multiplied by $e^{m_i - m_{i+1}}$ before the new block’s contribution is added. That factor is at most 1, because $m_{i+1} \ge m_i$.

The recurrence for one query row, processing key blocks $B_1, B_2, \dots$:

$$m_{i} = \max(m_{i-1}, \max_{j \in B_i} s_j)$$ $$\ell_{i} = e^{m_{i-1} - m_i} \ell_{i-1} + \sum_{j \in B_i} e^{s_j - m_i}$$ $$o_{i} = e^{m_{i-1} - m_i} o_{i-1} + \sum_{j \in B_i} e^{s_j - m_i} v_j$$

At the end, $o_{\text{final}} / \ell_{\text{final}}$ is exactly the softmax-weighted sum. Not approximately: the rescaling factors telescope, and the result equals the value that a full-row softmax would have produced, up to floating-point rounding. The different order of summation changes the rounding, so the last bits can differ from a reference implementation, which is why kernels are compared to a reference with a tolerance rather than bit-exactly.

The cost is a division at the end instead of a division per row, and one extra multiply-add per accumulated element per block. The extra work is proportional to the number of blocks times the state size, which is small relative to the block’s own work. That is the entire price of not materializing the matrix.

import numpy as np

def online_attention(q, k, v, block=64):
    # q: (d_k,), k: (n, d_k), v: (n, d_v) -> exact softmax-weighted sum of v
    n = k.shape[0]
    m, l, o = -np.inf, 0.0, np.zeros(v.shape[1])
    for start in range(0, n, block):
        kb, vb = k[start:start+block], v[start:start+block]
        s = kb @ q / np.sqrt(q.shape[0])        # scores for this block
        m_new = max(m, s.max())
        scale = np.exp(m - m_new)               # <= 1, rescales prior state
        p = np.exp(s - m_new)
        l = l * scale + p.sum()
        o = o * scale + p @ vb
        m = m_new
    return o / l

What has to stay in fast memory

The tiling decision is about which tensors live where. Queries for a block of rows, the current key and value blocks, and the running state (maximum, sum, and output accumulator) must fit in registers or shared memory for the kernel to avoid global traffic. The accumulator is the binding constraint: for a block of $B_q$ query rows and $d_v$ value dimensions, the output tile is $B_q \times d_v$ floats, and in 32-bit that is 4 bytes per element. With $d_v = 64$ and $B_q = 128$, the accumulator alone is 32 KiB, which is a substantial fraction of the shared memory available on a typical accelerator.

The block sizes are therefore chosen so that the accumulators fit and the arithmetic per byte stays high. A larger $B_q$ reuses each key and value block across more queries, which reduces the key/value traffic per query; a larger $B_k$ amortizes the query block’s residency. Both are bounded by the memory available for the accumulators and by the register file, and the optimum depends on the head width: a head of width 128 has twice the accumulator of width 64 at the same block size, so the block size halves.

The backward pass adds a second constraint. It needs the softmax output, which is not stored, so it is recomputed from the saved queries, keys, values, and the per-row normalization constants. That recomputation is one extra forward pass over the score blocks, paid in flops instead of in bytes: the same trade, applied again. The saved state for backward is $O(n d)$ per head plus the $O(n)$ row constants, so the activation memory drops from quadratic to linear in sequence length, which is what makes long-context training feasible at all rather than merely faster.

The result is that attention’s activation footprint becomes comparable to the feed-forward layers’ rather than dominant over them. A model whose attention activations were $n^2$-sized could not be trained at long context regardless of compute, because the memory would exceed the device before the batch did.

Parallelism and the cost of the split

A tiled kernel has two independent axes to parallelize: blocks of query rows, and heads with batch elements. The number of query blocks is $n / B_q$, which at $n = 4096$ and $B_q = 128$ is 32, enough to fill a device only when multiplied by heads and batch. The earlier formulation of the algorithm parallelized over query blocks and heads but kept the key loop sequential inside each block, which leaves the per-block work as a long serial chain.

Splitting the key dimension across blocks as well gives more parallelism at the cost of a reduction: each of $P$ key partitions produces a partial output and a partial normalizer, and they must be combined with the same rescaling recurrence applied across partitions. The combination is a small amount of work over $P$ partials, but it requires a synchronization point and either a second kernel or a cooperative reduction. The benefit is that each block’s serial key loop is $P$ times shorter, which matters when the number of query blocks alone does not saturate the device, which is the case of short sequences or small batch.

The trade is the same shape as any split-reduce: more parallelism, more traffic to combine. At long sequence lengths the query-block axis alone is usually sufficient, and the split is unnecessary overhead. At short sequences it is the difference between using the device and leaving most of it idle. Implementations therefore choose the partitioning strategy by shape rather than fixing it.

A second-order effect concerns the causal case. With a causal mask, half the score blocks are entirely below the diagonal and need no mask, half are entirely above and can be skipped, and only the diagonal blocks need element-wise masking. A kernel that processes blocks in order can simply stop the key loop at the diagonal for each query block, which halves the work and, more importantly, halves the bytes read for keys and values. The block-skipping arithmetic is covered in causal masking and mask placement; the point here is that the saving is free only when the block structure aligns with the mask structure.

Where the approach stops helping

Three regimes are worth naming.

At very short sequences, the quadratic term is small enough that the tiling overhead can exceed the saving: the extra rescaling multiplies, the loop bookkeeping, and the reduced occupancy from large accumulator tiles all count against it. For $n$ in the low hundreds, a straightforward kernel that materializes the matrix may be competitive, and the crossover depends on the hardware’s cache sizes and the head width. Tiled kernels are usually still faster, but the margin narrows and the tuning becomes shape-dependent.

At decode time with a single query row, there is nothing to tile over in the query dimension, and the operation becomes a bandwidth-bound streaming reduction over the cached keys and values. The blockwise-softmax recurrence still applies and is still useful, because it avoids materializing a $1 \times n$ score vector, a small saving, but the dominant cost is reading the cache, which tiling cannot reduce. This is the regime where cache size, not kernel design, sets throughput, as described in KV cache arithmetic.

At very long sequences the $O(n)$ memory is no longer negligible: at 128,000 positions and $d_{\text{model}} = 4096$ in 16-bit, the key and value tensors for one layer are about 2 GiB, and the linear term starts to matter as much as the quadratic one did at shorter lengths. The asymptotic improvement is real but the constant is large, and the practical context limit is set by the cache rather than by the score matrix.

The final constraint is numerical. The rescaling factors $e^{m_{i-1} - m_i}$ are computed in whatever precision the accumulators use, and a long chain of them compounds rounding. In 16-bit accumulation the error is visible; implementations therefore accumulate the output in 32-bit even when the inputs are 16-bit, and the rescaling is done in the wider format. That mixed-precision choice is not optional, and a kernel that accumulates in 16-bit will produce results that differ from the reference by more than the tolerance that training expects.