Attention: Flash-Style Tiling and Online Softmax

Core Concept

Standard attention is softmax(QK^T/sqrt(d))V. Flash-style algorithms process Q/K/V tiles while maintaining running row maxima, normalizers, and output accumulators, avoiding a full NxN score matrix.

Why It Affects Performance

For long sequences, materializing scores and probabilities causes enormous HBM traffic. IO-aware tiling exchanges extra on-chip state and recomputation for fewer global-memory transfers.

Example or Pseudocode

s = Q_tile @ K_tile^T
m_new = max(m_old, rowmax(s))
l_new = exp(m_old-m_new)*l_old + rowsum(exp(s-m_new))
o_new = exp(m_old-m_new)*o_old + exp(s-m_new) @ V_tile
O = o/l

Profiling and Performance Judgment

Across sequence lengths and head dimensions, compare peak memory, DRAM bytes, Tensor Core utilization, and end-to-end time. Validate causal masks and extreme logits.

Five-Minute Exercise

Does Flash-style attention mainly reduce mathematical FLOPs?