Lesson 91: Flash Attention & Memory-Efficient Transformers

1. Flash Attention & Memory-Efficient Transformers

Title

USAAIO · Lesson 91 · Phase 3

Why standard attention is IO-bound, not compute-bound — and four surgical fixes: Flash Attention tiled SRAM kernels, gradient checkpointing, ALiBi position bias, and Multi-Query Attention.

2. By the end of this lesson you can

Objectives

1. State Flash Attention's IO complexity — O(N) HBM reads vs standard O(N²) — and explain why IO is the bottleneck
2. Trace the tiled online-softmax loop (running max m, sum l, output accumulator O) and reproduce the update equations
3. Implement gradient checkpointing for a transformer layer and measure the activation-memory trade-off
4. Compute ALiBi position bias −m·|i−j| for given slopes and explain why it extrapolates beyond training length
5. Build a Multi-Query Attention (MQA) forward pass and quantify the KV-cache reduction vs standard MHA

3. What survived from KV Cache & Attention Efficiency?

Warm-up

4. Standard attention is IO-bound

Section

Part 1 of 4

5. Where the time goes: HBM vs SRAM

Concept

Modern GPUs have fast on-chip SRAM (≈ 20 MB, ~19 TB/s) and slow off-chip HBM (≈ 80 GB, ~2 TB/s). Standard attention is bottlenecked by how often it writes and reads the N×N score matrix to/from HBM — not by multiply-adds.

Arithmetic intensity of attention is low — the bottleneck is data movement, not floating-point ops. Flash Attention eliminates HBM round-trips by never materializing the N×N matrix.

6. Fill in: capacity for Where the time goes: HBM vs SRAM

Comparison

7. Standard attention IO: O(N²) HBM accesses

Concept

Standard attention writes the N×N score matrix S to HBM, reads it back for softmax, writes the probability matrix P, reads it back for P·V. That is four passes over an N² buffer.

\[ \text{HBM accesses}_{\text{standard}} \approx 3Nd + 4N^2 = O(N^2) \]

8. What each one costs: Standard attention IO: O(N²) HBM accesses

Trade off

9. Intuition: tiles fit in SRAM

Intuition

Flash Attention cuts the matrix into tiles small enough to fit in SRAM. It computes a tile of Q·Kᵀ, immediately applies softmax and multiplies by V — all on-chip. The N×N matrix is never written to HBM.

The trick is an online softmax: you can fuse the normalization with the tile pass using a running maximum m and a running sum l. The final output is correct even though you never saw the full row at once.

Backward pass: Flash Attention 2 recomputes attention tiles during the backward instead of storing them — gradient checkpointing at the op level. Net: O(N) HBM accesses, same FLOPS as standard attention.

10. Guess the shape of the answer: Online softmax: tiled update equations

Estimation

11. Online softmax: tiled update equations

Worked example

For tile (i,j), compute the block score S_ij = Q_i · K_j^T · scale. Update the running max, sum, and partial output, then rescale to merge the new tile.

import torch, torch.nn.functional as F
torch.manual_seed(42)
N, d = 8, 4
Q = torch.randn(N, d); K = torch.randn(N, d); V = torch.randn(N, d)
scale = d ** -0.5
# Reference standard attention
O_std = F.softmax((Q @ K.T) * scale, dim=-1) @ V

# Tiled Flash Attention (BLOCK=4)
BLOCK = 4
O = torch.zeros_like(Q)
l = torch.zeros(N)                   # running normalizer
m = torch.full((N,), float('-inf'))  # running max
for j in range(0, N, BLOCK):
    Kj, Vj = K[j:j+BLOCK], V[j:j+BLOCK]
    for i in range(0, N, BLOCK):
        Sij = (Q[i:i+BLOCK] @ Kj.T) * scale   # (BLOCK,BLOCK)
        mij = Sij.max(dim=-1).values
        m_new = torch.maximum(m[i:i+BLOCK], mij)
        e = torch.exp(Sij - m_new.unsqueeze(1))  # re-center
        l_new = torch.exp(m[i:i+BLOCK] - m_new)*l[i:i+BLOCK] + e.sum(dim=-1)
        O[i:i+BLOCK] = (torch.diag(torch.exp(m[i:i+BLOCK]-m_new)) @ O[i:i+BLOCK]
                        + e @ Vj)
        m[i:i+BLOCK] = m_new; l[i:i+BLOCK] = l_new
O = O / l.unsqueeze(1)   # final normalize
print('max abs diff:', (O_std - O).abs().max().item())

12. Work backwards from the answer: Online softmax: tiled update equations

Reverse engineer

13. Something is wrong here: Flash Attention reduces FLOPS

Anomaly

14. Trap: Flash Attention reduces FLOPS

Trap

15. Gradient checkpointing in a transformer

Section

Part 2 of 4

16. Activation memory in deep transformers

Concept

Backpropagation requires all intermediate activations to be available. For a transformer with L layers and sequence length N, storing every activation during the forward pass consumes O(L·N·d) memory — the dominant cost at scale.

\[ \text{activation memory} = O(L \cdot N \cdot d) \quad \text{(naive)} \]

Gradient checkpointing keeps only a subset of activations (checkpoints). During backward, it re-runs the forward for each segment to recompute the dropped activations. The trade-off: ≈33% extra FLOPs to save O(√L) or O(1) memory per layer.

17. What has to be given first: torch.utils.checkpoint in a transformer layer

Missing information

18. torch.utils.checkpoint in a transformer layer

Worked example

Wrap a transformer layer's forward with checkpoint(). Verify the gradients are numerically identical to the non-checkpointed version.

import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

class ToyTransformerLayer(nn.Module):
    def __init__(self, d=64):
        super().__init__()
        self.Wq = nn.Linear(d, d)
        self.Wk = nn.Linear(d, d)
        self.Wv = nn.Linear(d, d)
        self.ff  = nn.Linear(d, d)
    def forward(self, x):
        q=self.Wq(x); k=self.Wk(x); v=self.Wv(x)
        scale = x.shape[-1] ** -0.5
        a = F.softmax(q @ k.transpose(-2,-1) * scale, dim=-1) @ v
        return F.relu(self.ff(a + x))

torch.manual_seed(0)
N_tok, d_model = 16, 64
layer = ToyTransformerLayer(d_model)

# Without checkpointing
x1 = torch.randn(1, N_tok, d_model, requires_grad=True)
out1 = layer(x1); out1.sum().backward()
print('Normal:  output shape', out1.shape)

# With gradient checkpointing
x2 = x1.detach().clone().requires_grad_(True)
out2 = checkpoint(layer, x2, use_reentrant=False)
out2.sum().backward()
print('Ckpt:    output shape', out2.shape)
print('Grad match:', torch.allclose(x1.grad, x2.grad, atol=1e-5))

19. Work backwards from the answer: torch.utils.checkpoint in a transformer layer

Reverse engineer

20. ALiBi position bias

Section

Part 3 of 4

21. Rotary and learned position encodings: the extrapolation problem

Concept

Learned absolute position embeddings (BERT) and sinusoidal encodings (Lesson 82: standard Transformer) both degrade at inference sequences longer than training length — the model has never seen those positions.

ALiBi (Attention with Linear Biases) sidesteps this by adding no position embedding to the token vectors at all. Instead it modifies the attention logits directly with a bias that penalizes distant tokens.

\[ \text{logit}(i,j) \;=\; \frac{q_i \cdot k_j}{\sqrt{d}} \;-\; m_h \cdot |i - j| \]

22. Guess the shape of the answer: ALiBi slopes and bias matrix

Estimation

23. ALiBi slopes and bias matrix

Worked example

For H heads, slopes are a geometric sequence m_h = 2^(−h·8/H). Each head uses a different slope — sharper (larger m_h) heads focus locally; flatter heads attend globally.

import torch
N, H = 6, 4
slopes = torch.tensor([2**(-h*8/H) for h in range(1, H+1)])
print('slopes:', slopes.numpy().round(4))

i_idx = torch.arange(N).unsqueeze(1)   # (N,1)
j_idx = torch.arange(N).unsqueeze(0)   # (1,N)
dist  = (i_idx - j_idx).abs().float()  # (N,N)

alibi_h0 = -slopes[0] * dist           # (N,N) for head 0
print('head0 bias row 0:', alibi_h0[0].numpy().round(3))
print('head0 bias row 2:', alibi_h0[2].numpy().round(3))

24. Watch it run: ALiBi slopes and bias matrix

Pattern

25. Why ALiBi extrapolates to longer sequences

Concept

ALiBi adds no new trainable parameters and imposes no upper bound on sequence length. At inference with sequence length > training length, the bias formula −m·|i−j| naturally penalizes the new long-distance pairs — behavior that generalizes from the training distribution.

Empirically (Press et al., 2022): a GPT-style model trained on length 1024 with ALiBi generalizes to length 2048 with lower perplexity than sinusoidal or learned encodings trained directly at 2048.

ALiBi is used in production models including MPT and BLOOM. It is a drop-in replacement for any position encoding — just add the bias matrix to the attention logits before softmax (Lesson 82 callback: same QK^T, different additive term).

26. Where does each piece belong: Lesson 91: Flash Attention &…

Sorting

27. Something is wrong here: ALiBi bias added after softmax

Anomaly

28. Trap: ALiBi bias added after softmax

Trap

29. Break it on purpose: ALiBi bias added after softmax

Break the constraint

30. Multi-Query Attention (MQA)

Section

Part 4 of 4

31. MHA KV cache is the inference bottleneck

Concept

During autoregressive decoding, every token generation step re-reads the full KV cache (keys and values for all past tokens). With H heads, the cache holds H×N×d_k×2 values per layer — it dominates memory bandwidth at long sequences.

\[ \text{KV cache (MHA)} = 2 \cdot H \cdot N \cdot d_k \cdot \text{bytes} \quad \text{per layer} \]

32. MQA: one shared K, V for all query heads

Concept

Multi-Query Attention (Shazeer, 2019) keeps H separate Q projections but collapses K and V to a single head each. All Q heads attend to the same K, V. Memory usage for KV cache drops by H× with minimal quality loss on most tasks.

\[ Q \in \mathbb{R}^{H \times N \times d_k}, \quad K,V \in \mathbb{R}^{1 \times N \times d_k} \;\;(\text{shared across heads}) \]

Grouped-Query Attention (GQA) is a middle ground: G groups of query heads share one K, V. MQA = GQA with G=1; MHA = GQA with G=H. Used in LLaMA-2 (G=8) and Mistral.

33. Guess the shape of the answer: MQA forward pass in PyTorch

Estimation

34. MQA forward pass in PyTorch

Worked example

Build MQA with H=4 query heads, d_model=16, d_k=4. One Wk and one Wv projection (not H of each). Verify shapes.

import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(3)
H, d, dk, N = 4, 16, 4, 6
x = torch.randn(1, N, d)
Wq = nn.Linear(d, H*dk, bias=False)  # H query heads
Wk = nn.Linear(d, dk,   bias=False)  # single K head
Wv = nn.Linear(d, dk,   bias=False)  # single V head
Wo = nn.Linear(H*dk, d, bias=False)

Q = Wq(x).view(1, N, H, dk).transpose(1,2)   # (1,H,N,dk)
K = Wk(x).view(1, N, 1, dk).transpose(1,2)   # (1,1,N,dk) shared
V = Wv(x).view(1, N, 1, dk).transpose(1,2)   # (1,1,N,dk) shared
scores = (Q @ K.transpose(-2,-1)) * (dk**-0.5)  # (1,H,N,N) broadcast
attn   = F.softmax(scores, dim=-1) @ V           # (1,H,N,dk)
out    = Wo(attn.transpose(1,2).contiguous().view(1, N, -1))
print(f'Q:{Q.shape} K:{K.shape} V:{V.shape}')
print(f'scores:{scores.shape}  attn:{attn.shape}  out:{out.shape}')

35. Fill in: MHA shape for MQA forward pass in PyTorch

Comparison

36. Without one step: Memory-efficient attention recipe

Constraint

37. Memory-efficient attention recipe

Pattern

1. IO complexity first: standard attention = O(N²) HBM; Flash Attention = O(N) HBM via SRAM tiling — same FLOPs, faster because bandwidth-bound
2. Tiled softmax: maintain running m (max), l (normalizer), O (partial output); rescale each tile with exp(m_old − m_new) before adding new tile
3. Gradient checkpointing: checkpoint(fn, x) discards activations, recomputes them during backward — 33% FLOP overhead, O(1) activation memory per checkpointed block
4. ALiBi: add −m_h·|i−j| to logits before softmax; slopes 2^(−h·8/H) differ per head; no trainable params, extrapolates past training length
5. MQA / GQA: H Q heads share 1 (or G) K,V heads — reduces KV-cache by H/G×; critical for long-context autoregressive inference (LLaMA-2, Mistral)

38. Where does it stop working: Memory-efficient attention recipe

Edge cases

39. Rule out three: Check yourself — Flash Attention IO

Elimination

40. Check yourself — Flash Attention IO

Check

Work it out before clicking.

41. Answer it before you see the options: Check yourself — ALiBi

Prediction

42. Check yourself — ALiBi

Check

Apply the formula.

43. Rule out three: Check yourself — MQA KV cache

Elimination

44. Check yourself — MQA KV cache

Check

Use the formula.

45. Your turn: implement Flash + MQA

Section

Project

46. Project: memory-efficient attention from scratch

Concept

Three milestones: tiled Flash Attention matching F.scaled_dot_product_attention, a checkpointed transformer layer, and a complete MQA forward pass with KV cache simulation.

Build rules: type every line, use fixed seeds (torch.manual_seed(42)), and print shapes at every intermediate step.

47. Break it if you can: Project: memory-efficient attention from scratch

Counterexample

48. Milestone 1 — tiled Flash Attention

Worked example

Your turn: implement the tiled online-softmax loop. Predict: will the output match F.scaled_dot_product_attention to < 1e-6 max abs error?

Hint: for each (i,j) tile: m_new = max(m_old, tile_max), rescale O and l by exp(m_old − m_new), accumulate exp(Sij − m_new) @ Vj.

import torch, torch.nn.functional as F
torch.manual_seed(42)
N, d, BLOCK = 8, 4, 4
Q = torch.randn(N, d); K = torch.randn(N, d); V = torch.randn(N, d)
scale = d**-0.5
O_ref = F.scaled_dot_product_attention(
    Q.unsqueeze(0).unsqueeze(0),
    K.unsqueeze(0).unsqueeze(0),
    V.unsqueeze(0).unsqueeze(0)).squeeze()

O = torch.zeros_like(Q)
l = torch.zeros(N)
m = torch.full((N,), float('-inf'))
for j in range(0, N, BLOCK):
    Kj, Vj = K[j:j+BLOCK], V[j:j+BLOCK]
    for i in range(0, N, BLOCK):
        Sij = (Q[i:i+BLOCK] @ Kj.T) * scale
        mij = Sij.max(dim=-1).values
        m_new = torch.maximum(m[i:i+BLOCK], mij)
        e = torch.exp(Sij - m_new.unsqueeze(1))
        l_new = torch.exp(m[i:i+BLOCK]-m_new)*l[i:i+BLOCK]+e.sum(dim=-1)
        O[i:i+BLOCK] = torch.diag(torch.exp(m[i:i+BLOCK]-m_new))@O[i:i+BLOCK]+e@Vj
        m[i:i+BLOCK]=m_new; l[i:i+BLOCK]=l_new
O = O / l.unsqueeze(1)
print('max abs diff:', (O_ref - O).abs().max().item())

49. Milestone 2 — gradient checkpointing

Worked example

Your turn: wrap your transformer layer with checkpoint(). Predict: will the loss and gradients match the non-checkpointed forward?

Hint: from torch.utils.checkpoint import checkpoint; out = checkpoint(layer, x, use_reentrant=False). Same five backward steps as Lesson 40.

import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

class ToyLayer(nn.Module):
    def __init__(self, d=32):
        super().__init__()
        self.ff = nn.Linear(d, d)
    def forward(self, x):
        return F.relu(self.ff(x))

torch.manual_seed(0)
d = 32; N = 16
layer = ToyLayer(d)
x1 = torch.randn(1, N, d, requires_grad=True)
out1 = layer(x1); out1.sum().backward()
print('Normal grad norm:', x1.grad.norm().item())

x2 = x1.detach().clone().requires_grad_(True)
out2 = checkpoint(layer, x2, use_reentrant=False)
out2.sum().backward()
print('Ckpt   grad norm:', x2.grad.norm().item())

50. Milestone 3 — MQA with KV cache sizing

Worked example

Your turn: build the MQA forward. Before running, predict the shape of K after view and transpose, and predict the KV cache size at N=2048.

Hint: K = Wk(x).view(1, N, 1, dk).transpose(1,2) — shape (1, 1, N, dk). It broadcasts against (1, H, N, dk) Q during matmul.

import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(3)
H, d, dk, N = 4, 16, 4, 6
x = torch.randn(1, N, d)
Wq=nn.Linear(d,H*dk,bias=False); Wk=nn.Linear(d,dk,bias=False)
Wv=nn.Linear(d,dk,bias=False);  Wo=nn.Linear(H*dk,d,bias=False)
Q=Wq(x).view(1,N,H,dk).transpose(1,2)
K=Wk(x).view(1,N,1,dk).transpose(1,2)
V=Wv(x).view(1,N,1,dk).transpose(1,2)
scores=(Q@K.transpose(-2,-1))*(dk**-0.5)
attn=F.softmax(scores,dim=-1)@V
out=Wo(attn.transpose(1,2).contiguous().view(1,N,-1))
print(f'Q:{Q.shape} K:{K.shape} V:{V.shape}')
kv_mha = 2*H*2048*dk*4;  kv_mqa = 2*1*2048*dk*4
print(f'KV MHA N=2048: {kv_mha/1024:.1f} KB  MQA: {kv_mqa/1024:.1f} KB  ratio: {kv_mha//kv_mqa}x')

51. What each one costs: Milestone 3 — MQA with KV cache sizing

Trade off

52. The full program

Concept

import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

# --- Flash Attention (tiled) ---
torch.manual_seed(42)
N, d, BLOCK = 8, 4, 4
Q=torch.randn(N,d); K=torch.randn(N,d); V=torch.randn(N,d); scale=d**-0.5
O=torch.zeros_like(Q); l=torch.zeros(N); m=torch.full((N,),float('-inf'))
for j in range(0,N,BLOCK):
    Kj,Vj=K[j:j+BLOCK],V[j:j+BLOCK]
    for i in range(0,N,BLOCK):
        Sij=(Q[i:i+BLOCK]@Kj.T)*scale; mij=Sij.max(dim=-1).values
        m_new=torch.maximum(m[i:i+BLOCK],mij)
        e=torch.exp(Sij-m_new.unsqueeze(1))
        l_new=torch.exp(m[i:i+BLOCK]-m_new)*l[i:i+BLOCK]+e.sum(dim=-1)
        O[i:i+BLOCK]=torch.diag(torch.exp(m[i:i+BLOCK]-m_new))@O[i:i+BLOCK]+e@Vj
        m[i:i+BLOCK]=m_new; l[i:i+BLOCK]=l_new
O=O/l.unsqueeze(1)
ref=F.scaled_dot_product_attention(Q.view(1,1,N,d),K.view(1,1,N,d),V.view(1,1,N,d)).squeeze()
print('flash diff:', (O-ref).abs().max().item())

# --- MQA ---
torch.manual_seed(3)
H,d,dk,Ns=4,16,4,6; x=torch.randn(1,Ns,d)
Wq=nn.Linear(d,H*dk,bias=False); Wk=nn.Linear(d,dk,bias=False)
Wv=nn.Linear(d,dk,bias=False);   Wo=nn.Linear(H*dk,d,bias=False)
Q_=Wq(x).view(1,Ns,H,dk).transpose(1,2)
K_=Wk(x).view(1,Ns,1,dk).transpose(1,2)
V_=Wv(x).view(1,Ns,1,dk).transpose(1,2)
out_=Wo((F.softmax((Q_@K_.transpose(-2,-1))*(dk**-0.5),dim=-1)@V_).transpose(1,2).contiguous().view(1,Ns,-1))
print('MQA out:', out_.shape)

All four techniques target the same bottleneck: memory bandwidth. Sequence lengths double when you halve the HBM footprint — that is why GPT-4, LLaMA-2, and Gemini all use Flash Attention + MQA/GQA.

53. Fill in: notes for The full program

Comparison

54. Show it off

Concept

Out loud, slides closed: (1) explain why standard attention is IO-bound and how Flash Attention's tile loop eliminates HBM round-trips; (2) state the three variables maintained in the online softmax and their update rule; (3) contrast ALiBi with sinusoidal position encoding on the extrapolation question; (4) compute the KV-cache size for MHA vs MQA for your choice of H, d_k, N.

Stretch (homework): implement F.scaled_dot_product_attention profiling against your manual standard attention for N in [512, 1024, 2048] and measure wall-clock time; add ALiBi bias to the toy attention block; implement Grouped-Query Attention (GQA) with G=2 heads per group and verify it lies between MHA and MQA in KV cache size. Next up: Lesson 92 — sparse attention, linear attention, and long-context efficiency.

55. Connect it up: Lesson 91: Flash Attention & Memory-Efficient Transformers

Connect it up

56. What you can do now

Recap

• Derive Flash Attention's IO complexity: O(N) HBM via SRAM tiling vs standard O(N²) — FLOPS unchanged
• Trace the online-softmax loop: running max m, sum l, partial output O; rescale with exp(m_old−m_new) per tile
• Apply torch.utils.checkpoint to a transformer layer — 33% FLOP overhead saves full activation memory
• Add ALiBi bias −m_h·|i−j| before softmax; slopes per head; generalizes to unseen sequence lengths
• Build MQA (one K,V shared by all H Q heads) — KV-cache shrinks by H×