Attention deep dive (MHA → MQA → GQA → FlashAttention) — Large Language Models

Why this matters

In production, vanilla multi-head attention is almost never what gets shipped. Modern LLMs use Grouped-Query Attention (GQA), FlashAttention-2/3, sliding windows, and during inference a KV cache. If you do not understand these, you cannot reason about cost, latency, or context length. Top inference-engineering interviews live entirely in this territory.

Learning objectives

Implement masked multi-head attention from scratch.
Compare MHA, MQA, and GQA — pick the right one for a given memory budget.
Explain FlashAttention's memory and speed wins.
Reason about the KV cache during autoregressive decoding.
Quantitatively analyse memory and FLOPs for a long-context inference run.

1. Multi-head attention recap (MHA)

For each layer, with hidden size d and h heads:

4 weight matrices: W_Q, W_K, W_V, W_O, each (d, d).
Total params per layer: 4 d^2.
Per token, KV cache stores 2 tensors per head per layer: 2 * h * d_h = 2d values.

For a batch of B, sequence length T, layers L:

KV cache memory ≈ 2 * B * T * L * d * dtype_bytes

For Llama-3 70B (L=80, d=8192), at B=1, T=128k, bf16:

2 * 1 * 128_000 * 80 * 8192 * 2 ≈ 335 GB

That is unacceptable — the model itself only weighs 140 GB in bf16. Hence GQA.

2. MQA — Multi-Query Attention

Idea (Shazeer 2019): keep h query heads but use one shared K and V across all heads.

Params per layer: 2 d^2 + 2 d * d_h (Q and O are full; K and V are tiny).
KV cache: 2 * d_h per token instead of 2 * d. h× memory reduction.
Quality: small loss (~1% on benchmarks). Used by PaLM, Falcon-7B.

The trade-off was too aggressive — a single shared K/V loses too much capacity.

3. GQA — Grouped-Query Attention

Idea (Ainslie 2023): intermediate ground. Use g groups of K/V; each group is shared by h/g query heads.

g = 1 → MQA.
g = h → MHA.
Modern default: g = 4 or g = 8. KV cache reduced h/g-fold.

Llama-3 8B: h = 32, g = 8 → 4× KV memory reduction, ~no quality loss.

python

class GQA(nn.Module):
    def __init__(self, d, h_q, h_kv):
        super().__init__()
        assert d % h_q == 0 and h_q % h_kv == 0
        self.h_q, self.h_kv = h_q, h_kv
        self.dk = d // h_q
        self.W_Q  = nn.Linear(d, h_q  * self.dk, bias=False)
        self.W_KV = nn.Linear(d, 2 * h_kv * self.dk, bias=False)
        self.W_O  = nn.Linear(d, d, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        Q  = self.W_Q(x).view(B, T, self.h_q, self.dk).transpose(1, 2)         # (B, h_q, T, dk)
        kv = self.W_KV(x).view(B, T, 2, self.h_kv, self.dk).permute(2, 0, 3, 1, 4)
        K, V = kv[0], kv[1]                                                    # (B, h_kv, T, dk)
        # repeat K, V to match Q's heads
        rep = self.h_q // self.h_kv
        K = K.repeat_interleave(rep, dim=1)
        V = V.repeat_interleave(rep, dim=1)
        scores = (Q @ K.transpose(-2, -1)) / (self.dk ** 0.5)
        # ... causal mask ...
        out = scores.softmax(-1) @ V
        return self.W_O(out.transpose(1, 2).reshape(B, T, -1))

In practice you do not repeat_interleave — you broadcast inside the attention kernel (Flash supports it natively).

4. FlashAttention — the kernel that ate the world

Vanilla attention computes the full (T, T) score matrix in HBM (high-bandwidth memory). For T=128k, that matrix is 128k × 128k × 4 bytes ≈ 64 GB per head — impossible.

FlashAttention (Dao 2022, v2 2023, v3 2024): reorder the computation so:

Tiles of Q and K are loaded into the GPU's tiny but fast SRAM (~100KB).
Softmax is computed online (running max/normaliser).
The full T×T matrix is never materialised in HBM.

Result:

~5-10× speedup on attention.
Memory O(T) instead of O(T^2).
Numerically identical to standard attention (within fp16 noise).

In PyTorch 2.x you get FlashAttention-style execution by calling:

python

import torch.nn.functional as F
out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

Set torch.backends.cuda.enable_flash_sdp(True) and you get the kernel for free on Ampere/Hopper/Blackwell. Always use SDPA in production code.

FlashAttention-3 (2024) further optimises for Hopper FP8 and async copies — used by vLLM and TGI internally.

5. The KV cache — autoregressive decoding's secret

During training, you process all T tokens at once. During generation, you produce one new token at a time. If you naively re-ran the full forward each step, you would do O(T^2) work to generate T tokens.

KV cache stores K and V tensors per layer and appends to them as each new token is generated. Then attention at step t only needs:

Q_t  : (B, h_q,  1, d_h)        # just for the new token
K_all: (B, h_kv, t, d_h)        # cached + new
V_all: (B, h_kv, t, d_h)

So per generated token: O(t) work. To generate N tokens: O(N^2) total — same as a single forward pass on a N-token sequence. That is why long generations work.

python

# Conceptual cache
cache = {l: {"K": [], "V": []} for l in range(L)}

def step(x_new):
    for l, block in enumerate(blocks):
        Q, K, V = block.qkv(x_new)
        cache[l]["K"].append(K)
        cache[l]["V"].append(V)
        K_all = torch.cat(cache[l]["K"], dim=2)
        V_all = torch.cat(cache[l]["V"], dim=2)
        x_new = attn(Q, K_all, V_all) + x_new
        x_new = block.mlp(block.norm2(x_new)) + x_new
    return x_new

In production engines (vLLM, TGI, SGLang) the KV cache lives in:

Paged Attention (vLLM): KV stored in fixed-size pages → no fragmentation, batches with different lengths share memory like an OS page table.
Prefix caching: if multiple requests share a system prompt, store its KV once and reuse.
Quantized KV cache (FP8 / INT8): halve KV memory at small quality loss.

These are Phase 5 topics; today install the concept.

6. Sliding-window attention

Mistral-7B (and Phi, Gemma) use a sliding window: each token only attends to the last W tokens (e.g., W=4096). Combined with causal mask:

python

mask[i, j] = (j <= i) AND (i - j < W)

Reduces attention from O(T^2) to O(T*W), allowing long contexts at lower cost. Multiple stacked windowed layers can still propagate information further (similar to dilated convolutions).

Modern Mixtral / Llama-3 do not use sliding window by default — they rely on FlashAttention efficiency at full attention.

7. The math you should be able to reproduce in interviews

For a decoder-only forward pass on T tokens, with L layers, d model:

Attention FLOPs:  4 * L * T^2 * d            (dominates at long T)
FFN FLOPs:        2 * L * T * d * d_ffn     (dominates at short T)
Total ≈           6 * params * T              (rule of thumb for fwd+bwd)

For inference (single new token, KV cache):

Per-step FLOPs ≈ 2 * params + (linear-in-T attention to KV)

→ Inference is memory-bandwidth bound for small batch sizes (you have to read every weight once per generated token). Increasing batch size → more compute per byte loaded → higher throughput. This is why vLLM aggressively batches.

Hands-on lab (4 hours)

attention_deep.ipynb:

Implement MultiHeadAttention (full MHA) with is_causal switch. Verify against F.scaled_dot_product_attention.
Implement GQA (32 query heads, 8 KV heads) and verify shape compatibility.
Benchmark MHA vs F.scaled_dot_product_attention on (B=4, T=2048, d=1024, h=16) with bf16 on GPU. Report tokens/s.
Implement a manual KV-cache for your MHA module. Generate 256 tokens. Compare wall-time with no-cache (re-running forward each step).
For Llama-3 8B parameters (L=32, d=4096, h=32, h_kv=8), compute KV cache size at T=32k, B=1, bf16. Confirm against documented numbers.
Bonus: implement sliding-window mask (W=128). Show that perplexity stays close to full attention on a small task while attention compute drops.

Common pitfalls

Forgetting is_causal=True when using scaled_dot_product_attention — silently leaks future info during training.
Mismatched (h_q, h_kv) divisibility — h_q must be a multiple of h_kv.
Confusing T (current input length) and T_kv (cache length) when implementing a KV cache.
Using a Python list for the cache and forgetting to batch concatenate efficiently — slow.
Treating FlashAttention as a different model — it is the same math, only the implementation differs.

Self-check

Memory cost of the KV cache as a function of B, T, L, d, h_kv?
Why is GQA preferred over MQA in modern models?
What does FlashAttention avoid materialising?
Why is inference memory-bandwidth bound?
When would you choose sliding-window attention over full attention?

References

Vaswani et al. (2017), "Attention Is All You Need."
Shazeer (2019), "Fast Transformer Decoding: One Write-Head Is All You Need" (MQA).
Ainslie et al. (2023), "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints."
Dao et al. (2022, 2023, 2024), FlashAttention v1, v2, v3 papers.
vLLM blog, "PagedAttention."
HuggingFace scaled_dot_product_attention docs.