IO-Aware Algorithms

Definition

An IO-aware algorithm is designed around minimizing traffic across the memory hierarchy — chiefly HBM reads/writes — rather than around FLOP count, treating the hierarchy as a first-class co-design constraint. The governing insight of modern accelerators: past the memory wall, FLOPs are cheap and bytes are expensive, so you win by moving fewer bytes even at the cost of more arithmetic. It is the unifying principle behind tiling, fusion, recomputation, and FlashAttention.

How it works

The roofline says a kernel is memory-bound whenever its arithmetic intensity $I=\text{FLOPs}/\text{bytes}$ sits left of the ridge $I^{\ast }=\text{peak}/\text{bandwidth}$ (H100 BF16 $I^{\ast }\approx 296$). Most non-GEMM work — elementwise, softmax, norms, attention as launched — lives there, so runtime $\approx$ bytes moved. IO-aware design therefore optimizes the denominator. The concrete tactics, all “do more work per byte pulled from the slow tier”:

• Tiling — stage a chunk in SMEM/registers once, reuse it many times (GEMM traffic $2MNK\to 2MNK/S$).
• Fusion — keep intermediates on-chip so op chains read/write HBM once, not $k$ times.
• Recomputation — recompute cheap activations instead of storing/reloading them (spend idle FLOPs to save scarce bandwidth).
• Streaming reformulation — e.g. online softmax, which lets a global reduction be computed incrementally so nothing large is materialized.

FlashAttention combines all four: it does more arithmetic (recomputes in backward) yet runs far faster, because it replaces $\Theta (N^2)$ score-matrix traffic with $\Theta (Nd)$ activation traffic. The complementary framing is the roofline-per-tier lens: a kernel can be HBM-compute-bound yet SMEM-bandwidth-bound — register blocking is IO-awareness applied between SMEM and registers.

Co-design loop. The binding resource shifts with hardware: FA-1 targeted HBM, FA-2 occupancy, FA-3 async stalls (Hopper TMA/wgmma), FA-4 SFU exp throughput and SMEM bandwidth (Blackwell). IO-awareness is the discipline of asking “what resource is starved, and how do I restructure the algorithm to feed it?” rather than “how do I do fewer FLOPs?”

Why it matters

On the far side of the memory wall, performance engineering is data-movement engineering. Every high-leverage kernel in an LLM stack — attention, fused epilogues, MoE routing, quantized matmuls — is an exercise in IO-awareness. It reframes optimization from counting math to counting bytes and matching computation to the units that produce them.

Taught in

• 06-memory-wall-and-flashattention — §9 algorithm–hardware co-design as the through-line.
• 03-performance-modeling-roofline — the roofline lens.
• 05-optimizing-gemm — tiling/register blocking as per-tier IO-awareness.

See also

• flash-attention — the flagship IO-aware kernel
• memory-wall — the condition that makes IO dominate
• roofline-model · arithmetic-intensity
• shared-memory-tiling · kernel-fusion · recomputation · online-softmax
• gpu-memory-hierarchy
• gpu-systems-for-llms