Save memory traffic while retaining dense arithmetic

For one noncausal self-attention head with N tokens and equal query, key, and value width d, multiplying Q by K transpose produces N² scores. Multiplying the resulting probability matrix by V has the same leading matrix-multiply cost. Under the conventional count of two FLOPs per multiply-add, the two products total about four N²d FLOPs, excluding softmax, masking, and projections. Tiling changes their execution order rather than removing those pairs.

Materializing both score and probability matrices stores two N-by-N arrays if they are kept separately. In-place reuse can reduce that constant, but one quadratic intermediate remains. Online tiled execution instead keeps row statistics and output accumulators plus bounded working tiles. Its saved forward state can grow linearly with N for fixed d, even though the dense pair count remains quadratic. Report the exact memory boundary, including inputs and outputs, instead of calling every byte in the program auxiliary attention memory.

A tile schedule loads blocks of Q, K, and V into the device’s faster local storage, computes scores, updates online summaries, and reuses values before eviction. The fast-memory capacity constrains tile shapes, output accumulators, and temporary fragments. Bigger tiles can improve reuse but consume registers or shared memory, reducing concurrent resident work. Consequently, choosing the largest legal tile is not automatically the fastest schedule on a real accelerator.

The original FlashAttention paper analyzes a two-level memory model with SRAM capacity M measured in scalar elements and HBM transfers counted in elements. Its dense algorithm has the stated asymptotic access count within the specified capacity range. That model is not a literal measurement of cache misses, transaction sizes, synchronization overhead, or every later kernel variant. It explains why reuse matters; actual latency still requires architecture- and shape-specific measurement.

FlashAttention-2 improves work partitioning, reduces non-matmul operations, and changes how warps share work, while preserving the dense operator. These changes can improve utilization without altering the leading quadratic pair arithmetic. An asymptotic analysis would miss a speedup caused by less synchronization or better distribution of short-sequence work. Conversely, removing a quadratic materialization does not guarantee a speedup for tiny shapes whose overheads dominate.

When comparing kernels, distinguish logical bytes saved from measured memory traffic and actual runtime. The same tile may reread input blocks several times while still avoiding expensive full intermediate writes. Also distinguish prefill from incremental decoding: a single query attending to a long key cache has a rectangular score matrix and a different reuse problem from N-by-N training attention. Use the appropriate pair count and shape when reporting effective throughput.