Standard Self-Attention and Its Limits
Vanilla self-attention computes a full S×S attention matrix for sequence length S. At S=2048, that's 4 million entries per head per batch element. At S=8192, it's 67 million. At S=32768, it's over a billion. The quadratic scaling isn't a theoretical concern — it's the thing that stops you from processing long documents, high-resolution images, or full codebases.
Every variant I'll cover here attacks that quadratic cost differently. Some approximate the full attention. Others restructure it entirely. Each trades something — exactness, generality, or implementation complexity — for the ability to handle longer sequences. After working with most of these in production systems, I've got opinions about which tradeoffs are worth making.
Multi-Query and Grouped-Query Attention
The simplest modification to standard attention doesn't change the attention pattern at all — it changes the key/value projection. In multi-query attention (MQA), all attention heads share a single set of keys and values. Each head still has its own query projection, so the attention patterns differ across heads, but the KV cache during inference drops by a factor equal to the number of heads.
For a model with 32 heads and hidden_dim=4096 (head_dim=128), standard multi-head attention needs 32 × 128 = 4096 dimensions for both keys and values. MQA needs just 128 dimensions for each. During autoregressive generation, the KV cache is the memory bottleneck, and MQA cuts it by 32×.
class GroupedQueryAttention(torch.nn.Module):
def __init__(self, hidden_dim, num_heads, num_kv_heads):
super().__init__()
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.head_dim = hidden_dim // num_heads
self.q_proj = torch.nn.Linear(hidden_dim, num_heads * self.head_dim)
self.k_proj = torch.nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
self.v_proj = torch.nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
self.o_proj = torch.nn.Linear(hidden_dim, hidden_dim)
self.kv_repeat = num_heads // num_kv_heads
def forward(self, x):
B, S, _ = x.shape
q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
# Repeat KV heads to match query head count
k = k.repeat_interleave(self.kv_repeat, dim=1)
v = v.repeat_interleave(self.kv_repeat, dim=1)
scale = self.head_dim ** -0.5
attn = (q @ k.transpose(-2, -1)) * scale
attn = torch.softmax(attn, dim=-1)
out = (attn @ v).transpose(1, 2).reshape(B, S, -1)
return self.o_proj(out)
Grouped-query attention (GQA) is the middle ground. Instead of 1 KV head (MQA) or 32 (standard MHA), you use something like 4 or 8 KV groups. Llama 2 70B uses 8 KV heads with 64 query heads. In my benchmarks, GQA with 4 groups retains 98% of MHA's quality while cutting KV cache by 8×.
Linear Attention
What if you could avoid materializing the S×S attention matrix entirely? Linear attention replaces the softmax(QK^T)V computation with a kernel trick. Instead of computing attention weights explicitly, you decompose the operation so that it can be evaluated in O(S) time instead of O(S²).
See also: Building Production Tokenizers: BPE, WordPiece, and Sentence.
The basic idea: replace softmax with a feature map φ such that attention becomes φ(Q)(φ(K)^T V). If you compute φ(K)^T V first (which is D×D for head dimension D), then multiplying by φ(Q) is O(S×D²) instead of O(S²×D). Since D is typically 64-128 while S can be thousands, this is a massive win.
The catch is that the feature map approximation loses the sharp attention patterns that softmax produces. In practice, I've found linear attention works well for tasks where attention is spread relatively evenly (like document summarization) but struggles on tasks requiring precise token-to-token matching (like code generation or copy mechanisms). It's a meaningful limitation.
Sliding Window Attention
Mistral popularized sliding window attention, where each token attends only to its W nearest neighbors. It's local attention with a fixed window size. The memory cost drops from O(S²) to O(S×W), and with a window of 4096, you can process sequences of 32K or 128K tokens without memory issues.
The question is whether local attention is enough. For many tasks, it's. Adjacent tokens in natural language are highly informative about each other. And because information propagates through the layers — token at position 0 influences token at position W at layer 1, which influences token at position 2W at layer 2 — an L-layer model with window W has an effective receptive field of L×W. A 32-layer model with W=4096 covers 131K tokens in theory.
In practice, information doesn't propagate perfectly through layers. I've seen retrieval accuracy drop when the relevant context is more than 3-4 window hops away from the query. For tasks like multi-document QA where evidence might be scattered across very distant positions, pure sliding window attention isn't ideal.
Related reading: Model Monitoring in Production: Detecting Data Drift and Per.
Sparse Attention Patterns
BigBird, Longformer, and similar architectures combine local attention with sparse global attention. Every token attends to its local neighborhood, and a small set of global tokens attend to the entire sequence. It's a practical compromise — you get O(S) complexity for most of the attention computation, plus O(S) for the global token interactions.
The implementation is trickier than it sounds. Standard matrix multiplication hardware expects dense operations. Custom sparse attention kernels exist, but they don't always match the throughput of dense attention for moderate sequence lengths. In my experience, the crossover point where sparse attention becomes faster than dense FlashAttention on an A100 is around S=8192. Below that, FlashAttention with dense attention is actually faster despite doing more FLOPs.
Flash Attention
FlashAttention doesn't change the attention computation — it computes exact standard attention. What it changes is the memory access pattern. By tiling the computation and keeping intermediate results in SRAM (on-chip memory) instead of writing the full S×S matrix to HBM (high-bandwidth memory), it reduces memory IO by orders of magnitude.
The result: 2-4× speedup over naive attention, and peak memory that scales as O(S) instead of O(S²). This is the single most impactful optimization for transformer training in the last three years. Before FlashAttention, training with S=2048 was practical. With FlashAttention, S=8192 is routine and S=32768 is feasible on high-memory GPUs.
# Using FlashAttention via PyTorch 2.0+ SDPA
import torch.nn.functional as F
# This automatically uses FlashAttention when available
output = F.scaled_dot_product_attention(
query, key, value,
attn_mask=None,
dropout_p=0.0,
is_causal=True,
)
The caveat: FlashAttention doesn't return attention weights. If your application needs to inspect or visualize attention patterns (some interpretability methods require this), you can't use FlashAttention for those layers.
This connects to the ideas in Data Pipeline Orchestration for ML: Airflow, Prefect, and Da.
Ring Attention for Very Long Sequences
When a single sequence is too long to fit on one GPU even with FlashAttention, ring attention distributes the computation across devices. Each GPU holds a chunk of the sequence and rotates KV blocks around a ring topology. This enables million-token contexts by distributing both memory and computation.
I've used ring attention for processing entire codebases (500K+ tokens) and long legal documents. The scaling is nearly linear with the number of GPUs — 8 GPUs give you roughly 8× the context length. The bottleneck becomes the ring communication bandwidth, which on NVLink-connected GPUs is fast enough to keep utilization above 85%.
Picking the Right Variant
There's no universal best attention mechanism. For autoregressive generation where KV cache is the bottleneck, GQA with 4-8 groups is the pragmatic choice — it's what Llama, Mistral, and most production models use. For long-context training, FlashAttention with sliding window gives you the best balance of quality and efficiency. Linear attention has a niche in real-time applications where you need O(1) per-token generation cost, but quality suffers on precision-demanding tasks.
My default recommendation: start with standard multi-head attention + FlashAttention. Only introduce more exotic variants when you've measured a specific bottleneck. Premature optimization of attention patterns usually adds complexity without solving the actual problem.