RugvAI Labs

Attention Efficiency: FlashAttention, GQA, KV Cache

Naive attention has a quadratic memory bill that capped GPT-3 at 2K tokens. Then in 2022 Tri Dao's FlashAttention paper found a way to never write the N×N matrix to slow GPU memory — and quietly, every frontier lab swapped it in. Claude Sonnet 4.6 at 200K context, Gemini 1.5 Pro at 2M tokens, Llama 3.3 at 128K — none of it works without the tricks in this lesson.

Learning Objectives

After this lesson, you will be able to:

Explain why naive attention is O(N²) in both compute and memory and how FlashAttention's IO-aware tiling avoids ever writing the full N×N attention matrix to GPU HBM
Tell Multi-Head, Multi-Query, and Grouped-Query Attention apart by what they share — and pick the right one based on inference latency, memory budget, and quality targets
Use a KV cache to skip recomputing keys and values during autoregressive generation, and recognize when paged attention (vLLM) is the right next step for serving
Diagnose a long-context bottleneck (memory? throughput? latency?) and reach for the right tool — FlashAttention, sliding window, attention sinks, or linear attention — instead of just buying more H100s

Prerequisites

Before You Start

Why

Why This Matters

How FlashAttention turned 2k context into 128k overnight at every major lab -- before FlashAttention (Dao 2022), the quadratic memory cost of attention capped most models at 2k-4k tokens; after, every frontier lab, OpenAI, Anthropic, Google and Meta, quietly swapped it in and immediately scaled context windows by 30-60x without retraining a thing
Why Llama 2/3 use GQA — it cuts inference memory by 8x with minimal quality loss -- Meta's switch from Multi-Head to Grouped-Query Attention in Llama 2 wasn't a research breakthrough; it was the engineering call that made 70B models actually servable on commodity hardware, with quality drop in the noise
How vLLM's PagedAttention 24x'd serving throughput at Anyscale and Anthropic -- the trick was treating the KV cache like virtual memory: chunk it into fixed-size blocks, allocate on demand, share blocks across requests; suddenly one GPU could serve 24x more concurrent users
Build this --> Take a 7B Llama, measure prefill and decode latency at 2k, 8k, 32k context with naive attention vs FlashAttention vs sliding window — you'll watch a curve go from quadratic blow-up to nearly flat in front of your eyes

Don't worry if "FlashAttention" sounds like a paper you have to skim — the core insight is one sentence: never write the full attention matrix to slow memory. Once you see the picture, every optimization in this lesson is a variation on that theme.

#The Quadratic Wall

\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) V \quad \text{where } Q, K, V \in \mathbb{R}^{N \times d}

Anecdote

Historical Note

Vaswani et al. 2017 introduced standard multi-head attention in "Attention Is All You Need" — the canonical transformer paper. Shazeer 2019 ("Fast Transformer Decoding: One Write-Head is All You Need") proposed MQA to slash inference memory. Dao et al. 2022 ("FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", NeurIPS) reframed attention as an IO problem rather than a FLOPs problem — the paper everyone cites. Ainslie et al. 2023 ("GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints", EMNLP) introduced GQA as the middle ground used in Llama 2/3, Mistral, Gemma. Kwon et al. 2023 ("Efficient Memory Management for Large Language Model Serving with PagedAttention", SOSP) introduced paged attention in vLLM — every modern LLM-serving stack now uses it.

#FlashAttention: The IO-Aware Reframe

The attention output is mathematically the same whether you compute it block-by-block or all at once. FlashAttention exploits this by tiling the computation across blocks of Q, K, V that fit in SRAM. The catch is making the softmax work block-by-block — softmax requires the row max and row sum across the entire row, but you only have a block at a time.

The trick: maintain running statistics (running max m_i, running sum ℓ_i) and rescale outputs as new blocks arrive. This is the online softmax algorithm (Milakov & Gimelshein 2018), wrapped around tiled matmul.

\text{Block size } B = \frac{M_{\text{SRAM}}}{4 d_k} \;\Rightarrow\; \text{HBM traffic } O(N^2 d / M_{\text{SRAM}}) \;\ll\; O(N^2 d)

FlashAttention-2 (2023) rebalanced work between thread blocks for better GPU utilization on long sequences, hitting ~1.7x over FlashAttention-1. FlashAttention-3 (2024) added FP8 / asynchronous Tensor Cores for H100s, hitting ~2x over FA-2 on top of that. By 2026, every serious training and inference stack — PyTorch's scaled_dot_product_attention, HuggingFace, vLLM, TensorRT — defaults to a FlashAttention variant.

The KV cache (covered next) stores keys and values from prior tokens for fast autoregressive generation. For a model with H heads, the KV cache is proportional to H. Cut the number of K and V projections, and you cut KV-cache memory linearly.

Multi-Head Attention (MHA). Every head has its own Q, K, V. Best quality, biggest cache.
Multi-Query Attention (MQA). H query heads, 1 shared K head, 1 shared V head. Cache is H times smaller. Quality drop on training-from-scratch is real but small; on retrofitting, it can be severe.
Grouped-Query Attention (GQA). H query heads, G key/value groups (G < H). Cache is H/G times smaller. With G = H/8, you keep almost all of MHA's quality and most of MQA's savings.

\text{KV cache size:} \; \text{MHA: } 2 \cdot H \cdot N \cdot d_h \quad \text{MQA: } 2 \cdot 1 \cdot N \cdot d_h \quad \text{GQA: } 2 \cdot G \cdot N \cdot d_h

What Do You Think?

You serve a 70B Llama at 100k context to thousands of concurrent users. Which combination?

GQA is the quality/memory sweet spot Llama 2/3 chose. FlashAttention is non-negotiable for the prefill phase at 100k context — naive attention would OOM instantly. Paged KV cache (vLLM-style) is what lets you batch unrelated requests together by sharing GPU memory blocks dynamically. MHA loses the memory war at this scale; full MQA gives up too much quality. Buying H100s helps but doesn't scale.

#The KV Cache: Don't Recompute the Past

Autoregressive generation is "predict next token, append, repeat." Every step the model runs attention over the entire prefix plus the new token. Without a cache, each step does N forward passes — O(N²) total work per generated token. With a KV cache, only the new token's K and V are computed; everything prior is reused.

Step 1:  attend over [tok_1]                 → cache K_1, V_1
Step 2:  attend over [tok_1, tok_2]          → reuse K_1, V_1; compute K_2, V_2
Step 3:  attend over [tok_1, tok_2, tok_3]   → reuse K_1..2, V_1..2; compute K_3, V_3
...

Without cache: O(N²) compute per generated token. With cache: O(N) compute per token plus O(N) memory growth.

pythonplayground.py · Pyodide

# Sliding-window attention — bound how far each token can look back.
# This is the core of Mistral's long-context strategy: each layer has window W,
# but the receptive field grows linearly with depth (W + (L-1)*W).

import torch
import torch.nn.functional as F

def sliding_window_mask(N, W):
  """Causal mask + sliding window of width W."""
  # Causal: i can attend to j only if j &lt;= i (no looking forward)
  # Window: i can attend to j only if j >= i - W + 1
  i = torch.arange(N).unsqueeze(1)
  j = torch.arange(N).unsqueeze(0)
  causal = j &lt;= i
  window = j >= i - W + 1
  mask = causal & window
  return mask  # True where attention is allowed

def windowed_attention(Q, K, V, W):
  """O(N*W) attention instead of O(N^2)."""
  N, d = Q.shape
  mask = sliding_window_mask(N, W).to(Q.device)
  scores = (Q @ K.T) / (d ** 0.5)
  scores = scores.masked_fill(~mask, float('-inf'))
  P = F.softmax(scores, dim=-1)
  return P @ V

# Demo: 1000 tokens, window of 256
torch.manual_seed(0)
N, d = 1000, 64
Q = torch.randn(N, d)
K = torch.randn(N, d)
V = torch.randn(N, d)

out_full = windowed_attention(Q, K, V, W=N)        # full causal attention
out_win  = windowed_attention(Q, K, V, W=256)      # 256-token window

print(f"Full attention output norm:    {out_full.norm():.4f}")
print(f"Window=256 attention norm:     {out_win.norm():.4f}")
print(f"Output difference:             {(out_full - out_win).norm():.4f}")
print()
print("Token at position 500 attends to:")
mask500 = sliding_window_mask(N, 256)[500]
print(f"  Full attention: positions 0-500 ({mask500.sum()} tokens with full)")
mask500_win = sliding_window_mask(N, 256)[500]
print(f"  Window=256:     positions {500-255}-500 ({mask500_win.sum()} tokens)")

# Stack 8 layers: each adds 256 to the receptive field
print("\nReceptive field growth across stacked layers (W=256):")
for layer in range(1, 9):
  rf = 256 * layer
  print(f"  After layer {layer}: receptive field = {rf} tokens")

Tests · Verify the windowed output differs from full attention but produces stable outputs (no NaN), and that receptive field grows linearly with depth.

#PagedAttention: Virtual Memory for Tokens

Naive KV-cache management allocates one contiguous block per request. Result: severe fragmentation when requests have varying lengths — a 90% wasted GPU memory budget on real serving traffic.

vLLM's PagedAttention treats the KV cache like an OS treats RAM. Memory is split into fixed-size blocks (typically 16 tokens). Each request keeps a block table mapping logical positions to physical blocks. Free blocks return to a pool. Adjacent requests can share blocks (e.g., shared prefixes for parallel sampling) via copy-on-write.

\text{Block table: } \text{logical}_i \;\longmapsto\; \text{physical}_{j} \quad\text{(}j \text{ may not be contiguous)}

#Sliding Window, Attention Sinks, and Beyond

For very long sequences (1M+ tokens), even FlashAttention's O(N²) compute is impractical. Several approximations dominate:

Sliding-window attention (Mistral). Each token attends only to the last W tokens; receptive field grows with depth.
Attention sinks (StreamingLLM, Xiao 2023). The first few tokens accumulate all the attention from later positions; preserve them and you can stream indefinitely without quality collapse.
Linear attention (Performer, Linformer, Longformer). Kernel approximations that bring attention to O(N) at the cost of some accuracy.
State Space Models (Mamba, RWKV). Covered in the Mamba lesson; RNN-like recurrence with parallel training.

#2025 frontier: NSA, MoBA, and hybrid stacks

Sliding-window and attention-sink tricks were the 2023 answer. The 2025 wave is a generation more sophisticated: instead of one fixed sparsity pattern, modern long-context models route each query through several attention pathways at once and let the model learn which pathway matters where. Four designs dominate the frontier.

NSA (Native Sparse Attention, DeepSeek 2025) ships in DeepSeek-V3 and is the most influential of the new wave. NSA runs three parallel attention pathways for every query: a compression branch that attends to coarse block summaries (one key/value per 64-token block, giving a low-resolution view of the entire context), a selection branch where a small router scores blocks and the query attends only to the top-k highest-scoring blocks at full resolution, and a sliding-window branch that always attends to the most-recent W tokens for local fidelity. The three branches' outputs are fused with learned gates. The whole pattern is trained end-to-end and stays hardware-friendly because each branch is itself a dense, kernel-friendly matmul.

MoBA (Mixture of Block Attention, Moonshot 2025) powers Kimi K1.5's million-token context. MoBA chunks the KV-cache into blocks, scores each block against the query through a lightweight gating network, and lets the query attend only to its top-k blocks. Unlike NSA's hand-designed three pathways, MoBA's routing is fully learned and looks closer to a Mixture-of-Experts router operating over context blocks rather than over MLP experts. The gating is differentiable enough to train from scratch, which is what distinguishes MoBA from earlier retrieval-style "select-then-attend" approaches.

Hybrid local + global stacks in Gemma 2 and Llama 4 alternate layer types: roughly every fourth layer uses dense global attention while the rest use sliding-window-only attention. The intuition is that you do not need every layer to see every token — a few global layers per stack are enough to mix long-range information, and the rest can stay cheap. This is a tiny architectural change that buys most of the cost benefit of sparse attention without changing the kernel.

Mamba2-Hybrid in Jamba 1.5 (AI21) goes further by mixing layer types: some layers are state-space (Mamba) blocks with O(N) recurrence, others are full attention, and a sprinkle of MoE FFN layers sits between them. The attention layers handle the few global lookups; the SSM layers carry the long-distance state cheaply. Jamba 1.5 Mini fits 256K context on a single 80 GB GPU because most of its layers no longer pay the O(N²) bill.

Quick check

You are choosing a long-context attention strategy for a 1M-token coding assistant that runs on a single H100. Which trade-off correctly maps each technique to its strongest use case?

#Key Takeaways

Attention's wall is memory bandwidth, not FLOPs. Naive attention writes O(N²) bytes to HBM. FlashAttention's IO-aware tiling drops HBM traffic to O(N), giving 5-10x wall-clock speedup with no quality change.
GQA is the modern default. Between Multi-Head (full quality, full cache) and Multi-Query (smallest cache, quality risk), Grouped-Query Attention is the sweet spot every frontier 70B+ model has chosen since Llama 2.
The KV cache is what makes inference cheap. Without it, generating each token costs O(N²) compute; with it, O(N). It also dominates inference memory, which is why GQA matters so much.
PagedAttention is the OS analogy that 24x'd serving. Chunking the KV cache into fixed blocks and tracking them with a block table eliminates fragmentation and enables prefix sharing.
Long context (>32k) requires a stack, not a single trick. FlashAttention + GQA + sliding window + attention sinks together; or switch to alternative architectures (Mamba, hybrid).

#Quick Check

Quick Check1 / 4

What is the core insight behind FlashAttention?

Attention's quadratic wall is the central engineering challenge of the LLM era. Next we look at the inference-side optimizations — quantization, speculative decoding, and the serving stacks that turn a 70B model into a 200ms-per-token API.