GPUs Reward Memory Discipline

Compute grew faster than memory bandwidth, so the bytes you move set the ceiling on every kernel. A modern GPU spends most of its time moving data to and from global memory. Count the bytes a kernel moves before you count its FLOPs.

FlashAttention is the proof. It computes exact attention and never writes the n by n score matrix to global memory. The math stays the same. The traffic drops.

Bytes moved set the ceiling

Progress in language models comes as much from GPU hardware and parallelism as from new model ideas. Language models obey scaling laws, so more compute and more data usually improve results.

For decades that compute arrived as faster single cores. Dennard scaling and Moore's law delivered more transistors, smaller and faster and lower in power.

Then single-thread gains flattened. Transistor counts kept climbing while clock speeds and single-thread performance stopped climbing. Deep learning took its gain from parallelism instead, running more operations at the same time.

The parts of the machine did not scale together.

• Host-device links, PCIe and NVLink, improved slowly.
• Global memory bandwidth improved from GDDR to HBM2E, and the improvement stayed modest.
• Compute, and matmul FLOPs in particular, grew by orders of magnitude.

Older GPUs can be FLOP-bound. Modern GPUs are often memory-bound, and the memory hierarchy is where that shows up.

The hierarchy is strict. Memory near the core is fast and small. Memory far from the core is slow and large.

• Registers are per-thread storage and the fastest memory on the chip.
• Shared memory and L1 sit on the SM, and every thread in a block reads them. They are the right home for tiles and reused data.
• L2 is shared by all SMs, slower than shared memory and faster than DRAM.
• Global memory is off-chip HBM. It is large and much slower than anything on-chip.

If a kernel reads and writes global memory without pause, the SMs wait. Bandwidth decides the result and the arithmetic units sit idle.

Arithmetic intensity is the FLOPs a kernel does for each byte it moves. The roofline model reads it off two axes. The horizontal axis is FLOPs per byte moved. The vertical axis is the achieved FLOPs per second.

Low intensity puts you on the slope, where throughput rises as intensity rises and bandwidth owns the result. High intensity puts you on the flat roof, where peak compute owns it.

The accounting note priced a run before the launch. The roofline asks the same question of a single kernel.

A kernel falls below the roof for four reasons.

• It makes too many reads and writes against global memory.
• Its threads scatter across the address space.
• Its tiling is weak, so reuse stays in global memory.
• Its sizes do not divide cleanly into tiles.

A warp that splits pays for both paths

Manufacturers build a GPU for throughput. It carries many simple compute units with little control logic for each.

An SM is the unit of scheduling and control. It owns registers and shared memory, and it runs many warps at once to hide stalls. When one warp waits on a load, the SM runs another.

An SM holds streaming processors, which are simple arithmetic units. It also holds tensor cores, which are built for matrix multiplication. Tensor core throughput on matmul sits far above the general-purpose FLOP rate.

A model must be matmul-heavy to use the full capability of the chip. An A100 carries more than 100 SMs, each with many streaming processors and tensor cores.

The GPU arranges work into three units.

• A thread is the smallest unit of work, with its own registers and local state.
• A warp is 32 threads running the same instruction on different data at the same time.
• A block is a group of threads scheduled onto one SM, able to share memory and synchronize.

Blocks map to SMs, and warps run inside blocks under SIMT. A warp runs one instruction at a time, so threads that take different branches cannot run at once.

Put a branch inside a warp: if thread_id < 4 do A, else do B.

• Threads 0 to 3 run A while the others stay idle.
• Then threads 4 to 7 run B while the first group stays idle.

The warp runs both paths in sequence and pays the time of both. Every conditional that splits a warp is expensive for that reason. The cost never shows up in a FLOP count.

Four habits recover most of the lost speed

Every habit here does one of two things. It raises arithmetic intensity, or it cuts the bytes moved. Check any new habit against that pair before you spend a week on it.

Lower precision

Fewer bits for each number means fewer bytes moved. Lower precision also often buys more math operations for each cycle.

Take ReLU, x = max(0, x). In float32 the GPU reads 4 bytes and writes 4 bytes, which is 8 bytes for each element. In float16 it reads 2 and writes 2, which is 4 bytes for each element.

The operation count does not change and the traffic halves. Arithmetic intensity doubles, and the tensor cores come into range.

Operator fusion

Naive GPU code launches one kernel for each small operation. Each kernel makes its own round trip to global memory.

Write y = sin(x)² + cos(x)² the naive way.

sin → write s
cos → write c
square s → write s2
square c → write c2
add → write y

Every step in that chain reads global memory and writes it back. The fused kernel loads x once, computes in registers and shared memory, and writes y once.

Global traffic drops and speed rises. torch.compile fuses many chains like this one without help from you.

Recomputation

Backprop stores intermediate activations so the backward pass can reuse them. Those activations live in global memory, and reading them back costs more than making them again.

Three stacked sigmoids show the cost. The naive forward computes s1, s2, and s3, then stores all three and the output. The naive backward reads all three again from global memory.

Recomputation drops the internal sigmoids in the forward pass. The backward pass rebuilds them from x inside the kernel, then computes the gradients.

You spend compute to buy traffic. Compute is cheap on a modern GPU and bandwidth is scarce, so the trade usually pays.

This is checkpointing under another name, and speed is the reason to reach for it here. The same trade returns in the parallelism note, where activation memory is the last bottleneck standing.

Coalesced access

Global memory moves data in bursts. Fetching element 0 pulls a full aligned chunk, and the start of that transfer is the expensive part. Bytes next to it are cheap once the burst runs.

Threads in a warp that read nearby addresses let the hardware fold those reads into a few bursts. Threads that read scattered addresses start many bursts and waste the bandwidth they paid for.

Coalesced access holds the effective bandwidth high. Non-coalesced access collapses it.

Check the layout before you check the loop. A thread that walks across a row can still be strided in memory, and the warp loses its bursts.

Tiling moves the reuse onto the chip

Matmul is the workload the GPU is built for. A naive kernel reads the same values of A and B from global memory many times over.

Tiling moves that reuse onto the chip.

• Split A and B into tiles.
• Load one A tile and one B tile into shared memory.
• Accumulate partial results for a tile of C.
• Repeat for every tile pair that feeds that C tile.

Each global value now loads far fewer times. Inside a tile the reuse happens in shared memory, and global reads drop by about the tile size factor. The work moves onto fast on-chip memory.

Four things limit the tile you can pick.

• Shared memory size caps how much of A and B a block holds.
• The warp structure fixes how threads cover the tile.
• Coalescing decides whether the tile loads ride full bursts.
• Divisibility decides whether the last tile is full or ragged.

Miss any one of the four and the tile stops paying.

Matmul speed looks wavy because sizes must divide cleanly

Throughput generally rises as the matrix grows. The curve still carries dips and waves as it climbs.

Sizes that are multiples of the tile, warp, and burst sizes keep the hardware busy. Sizes that are not give you partial tiles, idle threads, and extra memory transactions. An awkward size runs much slower than a clean size beside it.

Tile count against SM count adds its own dip. Each tile maps to a block and then to an SM. A matrix that needs 98 tiles on a GPU with 108 SMs runs in one wave.

A matrix that needs 120 tiles runs 108 tiles first, then 12 more in a second wave at low occupancy. Throughput drops at exactly those sizes.

Burst boundaries add another dip. Tile widths that line up with DRAM bursts let a row fit into a few of them. Add one column, the rows cross a boundary, and the number of bursts doubles.

These effects combine, and the measured curve waves. Round your dimensions to the hardware before you blame the kernel.

FlashAttention keeps the math exact and refuses to write the n by n matrix

Standard attention runs three steps.

For sequence length n, S is n by n. Storing S or the softmax weights costs O(n²) memory. Moving that matrix in and out of global memory is the real bill.

The math stays O(n²) for general attention. The traffic does not have to. FlashAttention keeps HBM access far below the naive version and returns the same numbers, so nothing downstream changes.

Softmax is row-global, which looks like a barrier to tiling. Online softmax removes the barrier.

Stable softmax takes four steps.

• Take the max of the row.
• Compute exp of x minus that max.
• Sum the exps.
• Divide by the sum.

Online softmax walks the row in chunks instead. It carries a running max m and a running sum d over the prefix it has seen. Each new chunk updates both, and a larger max rescales d.

The kernel can now stream over the score tiles. The normalization updates while the stream runs, so the full row never lands in global memory.

The forward pass combines tiled matmuls, online softmax, and on-chip storage.

• Partition Q, K, and V along the sequence dimension.
• For each query tile, load the Q tile and one K tile into shared memory.
• Compute the score tile, then update the running max and running sum for each row.
• Move to the next K tile.

After the last K tile the normalization is known. The kernel forms the output with the intermediates still on chip.

The backward pass is harder, because naive gradients need the softmax outputs, and those are n by n. FlashAttention recomputes them instead.

• Loop over the tiles again.
• Rebuild the local scores and softmax values from Q, K, and V with the forward formulas.
• Compute the gradients for that tile.
• Discard the intermediates.

The pass spends more compute and cuts both traffic and storage.

Tiling, online softmax, recomputation, and matmul hardware turn attention from a memory problem into a workload that fits the chip.

The Builder Test

Pick the operation you want faster. Divide its FLOPs by the bytes it moves.

Name the habit that moves that number: precision, fusion, recomputation, coalescing, or a tile that fits the hardware. Write the name down before you write any code.

If no habit raises the intensity, the kernel already sits at the roof. Look for the time somewhere else in the step.

What Carries

A hardware-efficient algorithm is a memory-efficient algorithm. The roofline tells you which half of the machine you are fighting.

A GPU runs many threads at once and stays fragile in two places. A split warp pays for both paths, and a scattered access pattern wastes the bursts it starts.

You know the habits now. The next question is which of them is worth a kernel you write yourself.