Skip to content

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](/01-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](/01-modeling/01-architecture/03-memory-management/kv-cache/)
and [Pagedattention](/01-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](/01-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

-