Attention cost and long context

Every token is compared with every other token, so attention work grows with the square of the sequence length. Doubling the context roughly quadruples the attention cost. That is why long-context models are expensive to train and serve, and why so many attention variants exist.

Walking through the numbers

Take one layer with batch size 1, 32 heads and head size 128. Compare two sequence lengths.

Score matrix size. Each head has an n × n score matrix, so one layer holds 32 × n × n scores.

• At n = 128: 32 × 128 × 128 = 524,288 scores.
• At n = 8,192: 32 × 8,192 × 8,192 ≈ 2.1 × 10⁹ scores. In FP16, at 2 bytes each, that is about 4 GB for one layer, if you store the full matrix.

Compute for QKᵀ. The matrix product costs 2 × h × n × n × d_k FLOPs.

• At n = 128: 2 × 32 × 128 × 128 × 128 ≈ 1.34 × 10⁸ FLOPs.
• At n = 8,192: the ratio is (8,192 / 128)² = 4,096, so the cost is about 5.5 × 10¹¹ FLOPs for the same one layer.

The multiplication by V costs the same order again. The MLP block, by contrast, grows only linearly with n. For short sequences the MLP often dominates. For long ones, attention does.

Why fused kernels matter

Fused kernels such as FlashAttention compute attention in blocks and avoid writing the full n × n matrix to memory. This saves memory traffic and makes long inputs practical. It does not remove the n² work. The number of multiplications stays the same.

Variants that reduce the cost

These change how keys and values are shared, or which pairs are allowed. The core QKV idea stays the same.

The KV cache grows with n during generation, so MQA and GQA mainly help serving. Sliding windows and linear attention reduce the training and prefill cost as well.

Common mistakes

• Expecting FlashAttention to make attention linear. It saves memory and time per FLOP, but the n² work remains.
• Counting query heads when sizing the cache. The cache depends on the number of key-value heads, not the query heads.
• Forgetting the MLP. For short inputs, the MLP can cost more than attention. Profile before deciding that attention is the bottleneck.

Check yourself

1. If you double the sequence length, how much does the score matrix grow? Four times, because it has n × n entries.
2. Why do MQA and GQA mainly help serving? They shrink the KV cache, which grows with the sequence during generation.

Key takeaways

• Attention cost grows with n², so doubling the context roughly quadruples the attention work.
• At n = 8,192, one layer's score matrix is about 4 GB in FP16 if you store it in full.
• Fused kernels save memory, but the number of FLOPs stays the same.
• MQA, GQA, sliding windows and linear attention trade quality for a smaller cost.

Further learning

The KV cache page explains how the cache is stored during generation, and FlashAttention covers the exact, memory-aware kernel used in production.