Flash Attention — Module 12.3

How This Lesson Fits the Module & Volume

After KV cache management, the next bottleneck is often attention itself: materializing an N×N score matrix thrashs HBM. Flash Attention is an IO-aware algorithm that tiles Q, K, V through on-chip SRAM, never writing the full attention matrix to GPU memory.

It matters for both training and long-context inference (prefill especially). Later lectures on batching assume kernels that can sustain long sequences without OOM on the attention workspace.

Learning Objectives

By the end of this lesson, students should be able to:

Explain why standard attention is memory-bandwidth bound.
Describe tiling / online softmax as the core FlashAttention idea.
Contrast HBM traffic of naive vs fused attention.
Enable FlashAttention-style kernels in PyTorch / Hugging Face.
Relate FlashAttention to causal masks and KV-cache decode paths.
State when FlashAttention helps most (long sequences, large batches).

Definition

Flash Attention is an exact attention algorithm (same math as softmax(QKᵀ/√d)V) that fuses the computation into GPU SRAM-friendly tiles, using an online softmax so intermediate N×N matrices are never fully materialized in HBM. Variants (FlashAttention-2/3, SDPA backends) improve parallelism and hardware mapping.

The IO Problem

Naive attention writes a huge score tensor to high-bandwidth memory (HBM), reads it back for softmax, then multiplies by V. For long sequences, that IO dominates FLOPs. FlashAttention keeps tiles of Q, K, V in fast on-chip memory and streams partial results, cutting HBM traffic dramatically.

Aspect	Standard attention	Flash Attention
Attention matrix in HBM	Yes (N×N)	No (tiled / fused)
Numerics	Exact softmax	Exact (online softmax)
Typical win	—	Speed + lower peak memory
Best regime	Short seq	Long seq / large batch

Online Softmax Intuition

Load tile

Bring Q/K/V blocks to SRAM

Partial scores

Update running max & sum

Rescale

Correct earlier partial outputs

Write O

Only final output to HBM

Using Scaled Dot-Product Attention in PyTorch

Modern PyTorch routes scaled_dot_product_attention to fused backends (FlashAttention / Memory-Efficient / math) when shapes and dtype allow.

import torch
import torch.nn.functional as F

B, H, T, D = 2, 8, 2048, 64
q = torch.randn(B, H, T, D, device="cuda", dtype=torch.float16)
k = torch.randn(B, H, T, D, device="cuda", dtype=torch.float16)
v = torch.randn(B, H, T, D, device="cuda", dtype=torch.float16)

# Prefer fused kernels when available (Flash / mem-efficient)
with torch.backends.cuda.sdp_kernel(
    enable_flash=True, enable_math=False, enable_mem_efficient=True
):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(out.shape)  # (2, 8, 2048, 64)

# Hugging Face: attn_implementation="flash_attention_2" on supported models
# model = AutoModelForCausalLM.from_pretrained(
#     model_id, torch_dtype=torch.float16, attn_implementation="flash_attention_2"
# )

Prefill vs Decode

Prefill

Many tokens attend at once.
FlashAttention shines.
Cuts activation memory.

Decode

Often T_q = 1 vs long KV.
Different kernel shapes.
Still benefits from fused paths.

Training

Backward also IO-aware.
Enables longer context FT.
Pairs with activation checkpointing.

Strengths and Tradeoffs

Strengths

Exact attention with far less HBM traffic.
Lower peak memory for long contexts.
Widely available via SDPA / FA2 packages.

Tradeoffs

Hardware / dtype / head-dim constraints.
Custom masks can fall back to slow math.
Does not shrink the KV cache itself.

Common Misconception

“Flash Attention is an approximation like sparse attention.” It is mathematically equivalent to standard softmax attention (up to floating-point associativity). The innovation is IO scheduling, not dropping tokens.

Knowledge Check

Short Answer: What resource does Flash Attention primarily optimize? Answer: GPU memory bandwidth / HBM traffic (IO).
True/False: Flash Attention changes the mathematical definition of attention. Answer: False—it is exact (fused).
Multiple Choice: The N×N score matrix is: (a) always stored in HBM, (b) avoided via tiling, (c) replaced by FFT. Answer: (b).
Short Answer: Name the technique that updates softmax stats tile by tile. Answer: Online softmax.
True/False: Flash Attention eliminates the need for a KV cache. Answer: False.
Multiple Choice: Biggest wins appear when sequences are: (a) very short, (b) long, (c) empty. Answer: (b).
Short Answer: Which PyTorch API often dispatches to Flash backends? Answer: scaled_dot_product_attention (SDPA).
True/False: Unsupported masks may force a slower math fallback. Answer: True.
Multiple Choice: Flash Attention primarily helps: (a) tokenizer vocab size, (b) attention compute/memory path, (c) SQL indexes. Answer: (b).
Short Answer: What decoding speedup technique comes next in this module? Answer: Speculative decoding.

Key Takeaways

Standard attention is often IO-bound because of the N×N workspace.
Flash Attention tiles QKV through SRAM with online softmax—exact and faster.
Use SDPA / flash_attention_2 when hardware supports it.
Orthogonal to KV cache size; complementary in serving stacks.
Next: Speculative Decoding.

Trainer’s Guide

Hands-on idea: Benchmark SDPA with flash on vs math-only for T=512 vs T=4096; plot speedup vs length.

Discussion prompt: Why might decode (T_q=1) see smaller speedups than prefill?

Recap: Flash Attention makes exact attention IO-aware so long contexts fit and run faster. Continue with Speculative Decoding.