Attention Complexity & Cost¶

Overview¶

Attention is quadratic in sequence length: every token attends to every other token. This single fact drives almost all LLM engineering — kernels (FlashAttention), memory systems (KV cache, PagedAttention), architectural variants (GQA, sliding window), and long-context research. Understanding why it's quadratic, and where the costs land (training vs. inference, prompt vs. decode), is the key to reading the rest of this knowledge base.

Compute: O(seq² · d) per layer for QKᵀ and AV
Memory: the N×N attention matrix (and the KV cache) scale with seq² / seq
Decode phase: only the quadratic part shrinks; memory grows linearly per token
Optimization families: sparse, IO-aware, shared-head, cached — see table below

The Quadratic Problem in One Picture¶

N tokens, every token attends to every token:

   token ──────────┬───────────────────────┬───────────────►
          Qᵢ·K₁    Qᵢ·K₂    Qᵢ·K₃   ...    Qᵢ·K_N   → N scores
          ▼        ▼        ▼               ▼
          ┌──────────────────────────────────────────┐
          │        softmax + weighted sum of V       │
          └──────────────────────────────────────────┘
   N tokens × N scores each  =  N² attention weights

N = 8K  →  N² = 67M weights per layer
N = 32K →  N² = 1.07B weights per layer   (16× more than 8K!)

Double the context → quadruple the compute and memory. That's the whole story.

Compute: FLOPs Breakdown¶

Per layer (prompt/training phase — full parallel attention)¶

N = sequence length (e.g., 8,192 tokens for LLaMA-2), d_k = 128

QKᵀ:        N × N × d_k    → 8,192² × 128   ≈ 8.6 × 10⁹ mults
Softmax:    N × N          → 67M exps
A·V:        N × N × d_v    → 8,192² × 128   ≈ 8.6 × 10⁹ mults

Per layer total:           ≈ 17 × 10⁹ operations  (~17 GFLOPs)
LLaMA-2 7B (32 layers):    ≈ 550 × 10⁹ operations  (~550 GFLOPs per prompt)

Compare: the rest of the model is (nearly) linear¶

FFN (per layer):  ~4 · N · d_model²
  N = 8192, d_model = 4096  →  ~4 × 8192 × 4096² ≈ 550 GFLOPs

For LONG sequences, attention dominates:
  seq = 8K:  attention ≈ FFN   (both ~550 GFLOPs total)
  seq = 32K: attention × 16 → ~8.8 TFLOPs vs FFN × 4 → ~2.2 TFLOPs
             → attention is now ~80% of the compute

Why the quadratic term appears twice¶

S = Q·Kᵀ   → (N×d_k) × (d_k×N) = N×N  →  N²d_k  ops
Out = A·V  → (N×N) × (N×d_v) = N×d_v   →  N²d_v  ops

Both terms scale with N² — you can't dodge it by caching one part;
this is why sparse (SWA) and linear-attention ideas exist.

Memory: The N×N Matrix and the KV Cache¶

1. The attention matrix itself (training / long prompts)¶

Attention weights A: (N, N)

  N = 8,192   → 67M entries × 4 bytes (fp32) ≈ 268 MB per head-layer
  N = 32,768  → 1B entries → 4+ GB  ← explodes quickly!

Per layer with 32 heads: ×32.
  FlashAttention's core trick: never materialize A on GPU memory
  → recompute pieces on the fly, keep it in SRAM
  (see [Flash Attention](/02-llm-modeling/01-architecture/02-attention-optimization/flash-attention/))

2. The KV cache (inference, autoregressive decoding)¶

During decode, every generated token needs to attend to ALL past tokens.
Instead of recomputing their K/V every step, cache them:

Cache size = N_tokens × n_layers × n_heads × head_dim × bytes

Concrete (LLaMA-2 7B, fp16, 8K context):
  8,192 × 32 × 32 × 128 × 2 bytes = 2 × 10⁹ bytes ≈ 2 GB per sequence!

This is why GQA/MQA (fewer KV heads) and PagedAttention (fragmentation
control) exist — see [Kv Cache](/02-llm-modeling/01-architecture/03-memory-management/kv-cache/)
and [Pagedattention](/02-llm-modeling/01-architecture/03-memory-management/pagedattention/).

Training vs. Inference: Where the Quadratic Cost Hides¶

┌─────────────────────┬──────────────────────────────────┬───────────────────────────────┐
│ Phase               │ Compute pattern                  │ Quadratic term present?       │
├─────────────────────┼──────────────────────────────────┼───────────────────────────────┤
│ Training            │ full N×N attention every step    │ ✅ fully — N² on every batch  │
│ Inference: prompt   │ process N input tokens at once   │ ✅ fully — N² per request     │
│ Inference: decode   │ generate 1 token; attend to N    │ ❌ score is 1×N (linear)      │
│                     │ past tokens                      │   but KV cache memory grows   │
└─────────────────────┴──────────────────────────────────┴───────────────────────────────┘

Decode step in detail:
  Q_new: (1, d_k)   vs   K_cache: (N, d_k)
  scores: 1 × N  →  linear in N ✅   (this is why decode is "cheap" FLOPs-wise)

  BUT: cache grows by 1 token per step → memory grows with N
  and the memory access pattern (read all K/V every step) becomes
  the real bottleneck — FlashAttention-style IO optimization helps here too

Scaling Rules (Memorize These)¶

Context length:  N  →  2N
────────────────────────────────────────────
Attention compute:        ×4   (quadratic)
KV cache memory:          ×2   (linear in tokens)
FFN compute:              ×2   (linear in tokens)
End-to-end latency (long): ≈×2–4 depending on phase

The Optimization Families (Attack the Quadratic)¶

Family	Technique	What it fixes	Complexity
IO-aware	Flash Attention (v1/v2)	memory bandwidth, N×N materialization	O(N²) but ~10× faster
Sparse	Sliding Window Attention	quadratic compute → local window W	O(N·W)
Shared	[Gqa](/02-llm-modeling/01-architecture/01-core-designs/multi-query-attention-(mqa)-&-grouped-query-attention-(gqa)	mqa/gqa/)	KV cache size (÷heads factor)
Cached	Kv Cache	decode recomputation	linear memory per token
Managed	Pagedattention	fragmentation, sharing	linear memory, better utilization
Kernel	Kernel Fusion	launch overhead, intermediate writes	—

Decision guide¶

Problem                           → Technique
Training too slow                 → FlashAttention, kernel fusion
Long context OOM (compute)        → Sliding window / sparse attention
Long context OOM (KV cache)       → GQA/MQA, KV quantization
Decode latency high               → KV cache, continuous batching
Many concurrent long requests     → PagedAttention (vLLM)

Worked Budget Example¶

Serve LLaMA-2 7B (32 layers, 32 heads, head_dim 128) at 8K context, fp16:

Attention FLOPs per token (prompt):
  ≈ 4 · N² · d_k per layer  ≈ 4 × 67M × 128 ≈ 34 GFLOPs
  × 32 layers ≈ 1.1 TFLOPs (prompt processing)

KV cache per token:
  32 layers × 32 heads × 128 dim × 2 bytes = 256 KB
  8K tokens → 2 GB per concurrent sequence
  A server with 80 GB VRAM can host ~40 sequences' caches
  (plus weights ~14 GB fp16, activations, etc.)

Moral: cache capacity, not FLOPs, often limits long-context serving —
which is why GQA (÷8 cache) and PagedAttention are deployed in practice.

Key Takeaways¶

⏳ Attention is O(N²) — double context ⇒ 4× cost
🧮 Two quadratic terms: QKᵀ and AV both scale with N²
💾 Two memory costs: the N×N matrix (training) + KV cache (inference)
🔀 Phases differ: prompt is quadratic; decode is linear FLOPs but memory-bound
🧰 Every major technique is a fix for this one problem: FlashAttention, SWA, GQA, KV cache, PagedAttention
📏 Scale rules: N→2N ⇒ compute ×4, cache ×2

00 Attention Mechanisms — chapter overview
02 Scaled Dot Product Attention — the formula with the quadratic cost
Flash Attention — IO-aware O(N²) with usable memory
Sliding Window Attention — O(N·W) sparse alternative
Mqa & Gqa — KV-cache reduction
Kv Cache and Pagedattention — memory systems for decoding