E8.1 · LESSON 1 OF 6

IO-aware attention

Compare kernel, parallelism and disaggregated-serving choices using memory, communication and latency constraints.

THE IDEA

Attention is limited by memory traffic

Standard attention builds the whole n × n score matrix, writes it to GPU memory, reads it back for the softmax and again for the weighted sum. The arithmetic is cheap compared with all that memory movement. IO-aware attention (FlashAttention) restructures the computation into tiles that stay in fast on-chip memory, never materialising the full matrix in slow memory.

It is the same attention, not a shortcut: an online softmax tracks a running maximum and sum so the result can be built block by block, and equals the standard result up to normal floating-point differences.

sequence length n4,096tile size128

full score matrix per head: 4,096² × 2 bytes = 32 MiB (written to and read back from slow memory) one 128×128 tile: 32,768 bytes = 32 KiB (kept in fast on-chip memory) key/value blocks per query row: 32

The trick is an online softmax: the scores [1, 3, 2, 5, 4, 0, 6, 2] are seen in blocks of 4. The kernel keeps only a running maximum m, a running sum l and a running weighted value sum, and fixes them up when the maximum grows.

Block 1 of 2

after block 1: scores seen [1, 3, 2, 5] running max m = 5, running sum l = Σ exp(score − m) = 1.2034 when a later block has a larger maximum, the old l is multiplied by exp(old m − new m)

Change n and the tile size, then step through the online softmax.

Predict first. Is FlashAttention simply dropping attention pairs as an approximation?

Chapter context

Scaling serving changes how work is placed and moved. An operation can be mathematically unchanged while its memory traffic changes dramatically. Device count alone does not describe the performance of the resulting system.

TRY IT YOURSELF

Part 1. Do it hands-on

Size the score matrix and a tile by hand, then implement block-by-block attention with an online softmax and verify it equals the standard result.

sequence length n4,096tile size128

Block 1 of 2

after block 1: scores seen [1, 3, 2, 5] running max m = 5, running sum l = Σ exp(score − m) = 1.2034 when a later block has a larger maximum, the old l is multiplied by exp(old m − new m)

MiB for a 4096×4096 score matrix (2 bytes) bytes in a 128×128 tile blocks of 128 in 4096

To do: Standard attention writes the whole score matrix (n × n numbers per head) to slow GPU memory. For n = 4096 and 2 bytes per number, how many MiB per head?
To do: An IO-aware kernel processes the scores in small tiles that stay in fast on-chip memory. A 128 × 128 tile of 2-byte numbers is how many bytes?
To do: Keys and values are visited in blocks of 128 positions. For 4096 positions, how many blocks per query row?

Part 2. Do it in code

standard_attention(q, K, V): softmax(K @ q / sqrt(d)) @ V for one query, with the full score vector materialised.
online_attention(q, K, V, block): the same result computed over blocks of `block` keys, never holding more than one block of scores. Keep a running maximum m, a running sum l of exp(score − m) and a running weighted value sum acc. For each block: new_m = max(m, block max); rescale the old l and acc by exp(m − new_m); add the block's exp(score − new_m) terms; at the end return acc / l. block must be at least 1 (ValueError).
score_bytes(n, bytes_per_value=2): the size of a full n × n score matrix.
tile_bytes(t, bytes_per_value=2): the size of a t × t tile.
num_blocks(n, block): ceil(n / block).

import numpy as np

def standard_attention(q, K, V):
    pass

def online_attention(q, K, V, block):
    pass

def score_bytes(n, bytes_per_value=2):
    pass

def tile_bytes(t, bytes_per_value=2):
    pass

def num_blocks(n, block):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Is FlashAttention simply dropping attention pairs as an approximation?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

A weight-traffic lower-bound estimate is

t_{step}\gtrsim M_{read}/BW.

This assumes the stated bytes really traverse the bottleneck bandwidth. KV transfer adds roughly bytes_to_move/network_bandwidth before protocol and synchronization overhead. Batching and caching can change the bytes per produced token.

Worked example

A fictional 14-GB dense weight scan across a sustained 700-GB/s memory path takes at least 14/700=0.02 seconds, giving a 50-step/s bandwidth ceiling under the one-scan-per-step assumption. Real execution includes other memory traffic and work. A 512-MiB KV transfer over 32 GiB/s has a payload-only lower bound of 1/64 second.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import isclose
assert isclose(14/700,.02)
assert isclose(1/(14/700),50)
transfer=(512*1024**2)/(32*1024**3)
assert isclose(transfer,1/64)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Peak device bandwidth is not guaranteed sustained application bandwidth. Exact attention can be computed with different memory layouts. Additional GPUs can add communication bottlenecks.