1. Theoretical Motivation & Foundations
Standard multi-head attention computes the full N × N attention matrix S = Q K^T, writes it to High Bandwidth Memory (HBM), computes P = softmax(S), writes P back to HBM, and finally multiplies O = P V. For long sequences (e.g. 32k to 128k tokens), this materialization of intermediate matrices creates prohibitive O(N^2) memory storage and massive IO traffic that throttles GPU execution. FlashAttention resolves this fundamental IO bottleneck by never materializing the N × N attention matrix in global memory. Utilizing online softmax rescaling, it tiles inputs into SRAM-sized blocks, computing intermediate softmax and output projections entirely within fast on-chip register/SRAM caches before writing the final output to HBM. FlashAttention-2 further parallelizes across sequence length blocks and optimizes thread scheduling, while FlashAttention-3 leverages hardware TMA (Tensor Memory Accelerator) and FP8 Tensor Cores on Hopper. At the serving level, Chunked Prefill (Sarathi-Serve) splits massive prefill prompts into budget-constrained chunks that co-run alongside decode steps, preventing prefill requests from starving ongoing generations and eliminating high Inter-Token Latency (ITL) spikes.
2. Mathematical Formulations & Derivations
The governing analytical formulations and proof frameworks for this module:
3. From-Scratch Reference Implementation
Executable, production-tested reference code without magic libraries:
# Online Softmax Incremental Rescaling Verification
import math
def standard_softmax(x):
m = max(x)
exps = [math.exp(v - m) for v in x]
s = sum(exps)
return [v / s for v in exps]
def online_softmax_tiled(chunks):
# Reconstruct softmax without materializing full vector
m_running = float('-inf')
d_running = 0.0
# Pass 1: Streaming accumulator over blocks
for block in chunks:
m_curr = max(block)
m_new = max(m_running, m_curr)
scale_prev = math.exp(m_running - m_new) if m_running != float('-inf') else 0.0
scale_curr = [math.exp(v - m_new) for v in block]
d_running = scale_prev * d_running + sum(scale_curr)
m_running = m_new
# Pass 2: Output normalization
out = []
for block in chunks:
out.extend([math.exp(v - m_running) / d_running for v in block])
return out
data_chunks = [[1.2, 3.4, 0.5], [2.1, 4.8, 1.1], [3.2, 0.9, 2.7]]
flat_data = [x for chunk in data_chunks for x in chunk]
std_probs = standard_softmax(flat_data)
online_probs = online_softmax_tiled(data_chunks)
max_diff = max(abs(a - b) for a, b in zip(std_probs, online_probs))
print('Exact Numerical Equivalence Diff:', max_diff) # < 1e-15
4. Systems Complexity & Memory Footprint
FlashAttention is not an approximation—it computes mathematically exact attention while operating 2x to 4x faster and using 10x to 20x less HBM memory traffic. Combined with chunked prefill, it enables multi-thousand token contexts to run smoothly without inducing tail-latency spikes in concurrent production generation.
5. Canonical Literature & Primary Research
Original research papers and foundational texts recommended for advanced study:
- Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS.
- Dao, T. (2023). FlashAttention-2: Faster Attention with Better Work Partitioning and Parallelism. ICLR.
- Shah, A., et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and FP8 on Hopper. arXiv:2407.08608.
- Agrawal, A., et al. (2024). Sarathi-Serve: Efficient LLM Serving with Chunked-Prefills. OSDI.