Evolution of AI Attention Mechanisms: A Complete Guide from MHA to MLA and Flash Attention
Summary
- Attention mechanism is the core of Transformer: accounting for 60%+ of LLM computation and memory, optimizing attention = optimizing the entire model
- 5 generations of attention evolution: MHA -> MQA -> GQA -> MLA -> Flash Attention, each seeking balance between computational efficiency and model quality
- MLA (DeepSeek's original innovation) is the most radical innovation of 2026: KV Cache compression 95%+, inference cost reduced 10x
- Flash Attention 3 is hardware-level optimization: achieving 75% theoretical peak utilization on H100, 2x faster than FA2
- This article provides mathematical derivation, code implementation, and production deployment comparison for all 5 generations of attention
Table of Contents
- Attention Mechanism: The Heart of Transformer
- Generation 1: MHA Multi-Head Attention
- Generation 2: MQA and GQA Grouped Attention
- Generation 3: MLA Multi-Head Latent Attention
- Generation 4: Flash Attention Hardware-Level Optimization
- Generation 5: Sparse Attention and Hybrid Architecture
- Summary and Further Reading
Attention Mechanism: The Heart of Transformer
Attention Resource Proportion in LLM
| Component | Compute Proportion | Memory Proportion | Optimization Value |
|---|---|---|---|
| Self-Attention | 40-50% | 50-60% | Highest |
| FFN/MLP | 40-50% | 25-30% | High |
| Embedding | 2-5% | 5-10% | Low |
| LayerNorm | <1% | <1% | Lowest |
5 Generations of Attention Evolution
| Generation | Time | Representative | KV Cache | Compute Efficiency | Model Quality |
|---|---|---|---|---|---|
| Gen 1 MHA | 2017 | Transformer | Baseline | Baseline | Best |
| Gen 2 MQA | 2019 | PaLM | -75% | +30% | Slightly lower |
| Gen 2 GQA | 2023 | LLaMA-2 | -50% | +20% | Close to MHA |
| Gen 3 MLA | 2024 | DeepSeek-V2 | -95% | +40% | Close to MHA |
| Gen 4 FA3 | 2024 | FlashAttn3 | Unchanged | +100% | Unchanged |
| Gen 5 Sparse | 2025-2026 | MoBA | -80% | +50% | Medium |
Generation 1: MHA Multi-Head Attention
MHA Mathematical Formulation
` Standard MHA computation flow:
Input X ∈ R^(B×S×D)
Q = XW_Q, K = XW_K, V = XW_V
Split heads: Q_i = Q[:, :, i*d_h:(i+1)d_h] for i = 0, ..., h-1 K_i = K[:, :, id_h:(i+1)d_h] V_i = V[:, :, id_h:(i+1)*d_h]
Attention computation: Attn_i = softmax(Q_i K_i^T / √d_h) V_i
Merge: Output = Concat(Attn_0, ..., Attn_{h-1}) W_O
KV Cache size = 2 × S × h × d_h × sizeof(dtype) = 2 × S × D × sizeof(dtype)
For 70B model, S=4096, D=8192, FP16: KV Cache = 2 × 4096 × 8192 × 2 = 128MB/layer 64 layers total = 8GB `
MHA Implementation
`python import torch import torch.nn as nn import math
class MultiHeadAttention(nn.Module): def init(self, hidden_size=8192, num_heads=64): super().init() self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.scale = self.head_dim ** -0.5
self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
def forward(self, x, attention_mask=None, past_kv=None):
B, S, D = 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_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
if past_kv is not None:
past_k, past_v = past_kv
k = torch.cat([past_k, k], dim=2)
v = torch.cat([past_v, v], dim=2)
attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scale
if attention_mask is not None:
attn_weights = attn_weights.masked_fill(attention_mask == 0, float('-inf'))
attn_weights = torch.softmax(attn_weights, dim=-1)
attn_output = torch.matmul(attn_weights, v)
attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, D)
return self.o_proj(attn_output), (k, v)
`
MHA KV Cache Overhead
| Model | Layers | Hidden Size | Heads | KV Cache/Layer | Total KV Cache |
|---|---|---|---|---|---|
| 7B | 32 | 4096 | 32 | 2MB | 64MB |
| 14B | 40 | 5120 | 40 | 2.5MB | 100MB |
| 70B | 64 | 8192 | 64 | 8MB | 512MB |
| 405B | 80 | 16384 | 128 | 32MB | 2.56GB |
Generation 2: MQA and GQA Grouped Attention
MQA (Multi-Query Attention)
` MHA vs MQA:
MHA: Each head has independent K and V Q: [h × d_h] K: [h × d_h] V: [h × d_h] KV Cache = 2 × h × d_h × S
MQA: All heads share one set of K and V Q: [h × d_h] K: [1 × d_h] V: [1 × d_h] KV Cache = 2 × d_h × S (reduced by h times)
For example h=64: MQA KV Cache = MHA's 1/64 ≈ 98.4% reduction `
GQA (Grouped-Query Attention)
` GQA: g groups of heads share K and V
Q: [h × d_h] K: [g × d_h] V: [g × d_h] KV Cache = 2 × g × d_h × S
When g=1: degenerates to MQA When g=h: degenerates to MHA When g=8: KV Cache reduced 8x
LLaMA-2 70B uses g=8: KV Cache = MHA's 8/64 = 1/8 `
GQA Implementation
`python class GroupedQueryAttention(nn.Module): def init(self, hidden_size=8192, num_heads=64, num_kv_heads=8): super().init() self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = hidden_size // num_heads self.kv_dim = self.num_kv_heads * self.head_dim self.scale = self.head_dim ** -0.5 self.n_rep = num_heads // num_kv_heads
self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.k_proj = nn.Linear(hidden_size, self.kv_dim, bias=False)
self.v_proj = nn.Linear(hidden_size, self.kv_dim, bias=False)
self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
def _repeat_kv(self, x):
if self.n_rep == 1:
return x
B, g, S, d = x.shape
return (
x[:, :, None, :, :]
.expand(B, g, self.n_rep, S, d)
.reshape(B, self.num_heads, S, d)
)
def forward(self, x, attention_mask=None, past_kv=None):
B, S, D = 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)
if past_kv is not None:
past_k, past_v = past_kv
k = torch.cat([past_k, k], dim=2)
v = torch.cat([past_v, v], dim=2)
k_expanded = self._repeat_kv(k)
v_expanded = self._repeat_kv(v)
attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) * self.scale
if attention_mask is not None:
attn_weights = attn_weights.masked_fill(attention_mask == 0, float('-inf'))
attn_weights = torch.softmax(attn_weights, dim=-1)
attn_output = torch.matmul(attn_weights, v_expanded)
attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, D)
return self.o_proj(attn_output), (k, v)
`
MHA/MQA/GQA Comparison
| Method | KV Groups | KV Cache | Model Quality | Representative Model |
|---|---|---|---|---|
| MHA | h=64 | 100% | Best | GPT-3 |
| GQA-8 | g=8 | 12.5% | Close to MHA | LLaMA-2/3 |
| GQA-4 | g=4 | 6.25% | Slightly lower | Mistral |
| MQA | g=1 | 1.56% | 5-8% lower | PaLM |
Generation 3: MLA Multi-Head Latent Attention
MLA Core Idea
┌──────────────────────────────────────────────────────────────┐ │ MLA Core Innovation │ │ │ │ Traditional MHA/GQA: │ │ K, V stored directly → Large KV Cache │ │ │ │ MLA: │ │ 1. Project to low-dimensional latent space: c_kv = Compress(X)│ │ c_kv dimension << K,V dimension → Cache compression 95%+ │ │ │ │ 2. Recover from latent space during inference: │ │ K = W_k_up × c_kv │ │ V = W_v_up × c_kv │ │ │ │ 3. Absorption technique: │ │ Q × K^T = Q × (W_k_up × c_kv)^T │ │ = (Q × W_k_up^T) × c_kv^T │ │ = Q' × c_kv^T │ │ Avoid explicit K recovery, compute attention directly │ │ in low-dimensional space │ └──────────────────────────────────────────────────────────────┘
MLA Implementation
`python class MultiHeadLatentAttention(nn.Module): def init( self, hidden_size=4096, num_heads=32, kv_latent_dim=512, ): super().init() self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.kv_latent_dim = kv_latent_dim
self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.kv_compress = nn.Linear(hidden_size, kv_latent_dim, bias=False)
self.k_up_proj = nn.Linear(kv_latent_dim, hidden_size, bias=False)
self.v_up_proj = nn.Linear(kv_latent_dim, hidden_size, bias=False)
self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.scale = self.head_dim ** -0.5
def forward(self, x, attention_mask=None, past_c_kv=None):
B, S, D = x.shape
q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
c_kv = self.kv_compress(x)
if past_c_kv is not None:
c_kv_full = torch.cat([past_c_kv, c_kv], dim=1)
else:
c_kv_full = c_kv
q_absorbed = torch.matmul(
q.reshape(B * self.num_heads, S, self.head_dim),
self.k_up_proj.weight.T,
).view(B, self.num_heads, S, self.kv_latent_dim)
attn_weights = torch.matmul(
q_absorbed, c_kv_full.transpose(-2, -1)
) * (self.kv_latent_dim ** -0.5)
if attention_mask is not None:
attn_weights = attn_weights.masked_fill(attention_mask == 0, float('-inf'))
attn_weights = torch.softmax(attn_weights, dim=-1)
attn_to_c = torch.matmul(attn_weights, c_kv_full)
v = self.v_up_proj(attn_to_c.reshape(B, S, self.kv_latent_dim))
v = v.view(B, S, D)
return self.o_proj(v), c_kv
`
MLA KV Cache Comparison
| Method | KV Dim/Head | Total KV Dim | KV Cache | Compression Ratio |
|---|---|---|---|---|
| MHA(h=128) | 128 | 16384 | 100% | 1x |
| GQA(g=8) | 128 | 1024 | 6.25% | 16x |
| MLA(d_c=512) | - | 512 | 3.1% | 32x |
| MLA(d_c=256) | - | 256 | 1.56% | 64x |
Generation 4: Flash Attention Hardware-Level Optimization
Flash Attention Principle
┌──────────────────────────────────────────────────────────────┐ │ Flash Attention Core Idea │ │ │ │ Standard Attention: │ │ Q,K,V → S=QK^T → P=softmax(S) → O=PV │ │ Problem: S and P are S×S matrices, memory O(S²), and │ │ multiple reads/writes to HBM │ │ │ │ Flash Attention: │ │ 1. Tiled computation: split Q,K,V into small blocks │ │ 2. Compute softmax in SRAM (online softmax) │ │ 3. Only write final result O back to HBM │ │ │ │ Memory: O(S) vs O(S²) │ │ HBM reads/writes: O(S²d/N) vs O(S²d) → reduced by N times │ │ N = SRAM size / block size │ └──────────────────────────────────────────────────────────────┘
Flash Attention 3 Features
| Feature | FA1 | FA2 | FA3 |
|---|---|---|---|
| GPU Support | A100 | A100/H100 | H100+ |
| Data Type | FP16/BF16 | FP16/BF16 | FP16/BF16/FP8 |
| Async | No | No | Yes (WGMMA) |
| Pipeline | No | No | Yes |
| Theoretical Utilization | 50% | 62% | 75% |
| Relative Speed | 1x | 2x | 4x |
Flash Attention Usage
`python from flash_attn import flash_attn_func
def flash_attention_forward(q, k, v, causal=True): output = flash_attn_func( q, k, v, causal=causal, softmax_scale=None, ) return output
import torch.nn as nn
class FlashAttentionLayer(nn.Module): def init(self, hidden_size=8192, num_heads=64): super().init() self.num_heads = num_heads self.head_dim = hidden_size // num_heads
self.qkv_proj = nn.Linear(hidden_size, 3 * hidden_size, bias=False)
self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
def forward(self, x):
B, S, D = x.shape
qkv = self.qkv_proj(x)
q, k, v = qkv.chunk(3, dim=-1)
q = q.view(B, S, self.num_heads, self.head_dim)
k = k.view(B, S, self.num_heads, self.head_dim)
v = v.view(B, S, self.num_heads, self.head_dim)
output = flash_attn_func(q, k, v, causal=True)
output = output.reshape(B, S, D)
return self.o_proj(output)
`
Flash Attention Performance Benchmarks
| Sequence Length | Standard Attn | FA1 | FA2 | FA3 (H100) |
|---|---|---|---|---|
| 1K | 2.1ms | 0.8ms | 0.5ms | 0.3ms |
| 4K | 28ms | 3.2ms | 1.8ms | 0.9ms |
| 16K | OOM | 14ms | 7ms | 3.5ms |
| 32K | OOM | 32ms | 15ms | 7ms |
| 128K | OOM | OOM | 68ms | 32ms |
Generation 5: Sparse Attention and Hybrid Architecture
Sparse Attention Types
| Type | Sparse Pattern | Computation | Use Case |
|---|---|---|---|
| Local Window | Fixed window | O(S×W) | Long documents |
| Global+Local | Few global tokens | O(S×(W+G)) | General |
| Dilated Attention | Dilated convolution pattern | O(S×logS) | Hierarchical structure |
| MoBA | Hybrid block | O(S×√S) | General |
| Linear Attention | Kernel method | O(S×D²) | Ultra-long sequences |
MoBA (Mixture of Block Attention) Implementation
`python class MoBABlock: def init(self, block_size=256): self.block_size = block_size
def compute_block_importance(self, q, k_blocks):
scores = []
for k_block in k_blocks:
score = torch.matmul(
q.mean(dim=1),
k_block.mean(dim=1).T,
).max()
scores.append(score)
return torch.tensor(scores)
class MoBAAttention(nn.Module): def init(self, hidden_size, num_heads, block_size=256, top_k_blocks=4): super().init() self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.block_size = block_size self.top_k_blocks = top_k_blocks
self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
def forward(self, x):
B, S, D = 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_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
num_blocks = (S + self.block_size - 1) // self.block_size
k_blocks = k.chunk(num_blocks, dim=2)
v_blocks = v.chunk(num_blocks, dim=2)
block_importance = self._score_blocks(q, k_blocks)
top_blocks = torch.topk(block_importance, min(self.top_k_blocks, num_blocks))
selected_k = torch.cat([k_blocks[i] for i in top_blocks.indices], dim=2)
selected_v = torch.cat([v_blocks[i] for i in top_blocks.indices], dim=2)
attn_weights = torch.matmul(q, selected_k.transpose(-2, -1)) / (self.head_dim ** 0.5)
attn_weights = torch.softmax(attn_weights, dim=-1)
output = torch.matmul(attn_weights, selected_v)
output = output.transpose(1, 2).contiguous().view(B, S, D)
return self.o_proj(output)
def _score_blocks(self, q, k_blocks):
scores = []
for k_block in k_blocks:
score = torch.matmul(q.mean(dim=-1), k_block.mean(dim=-1).transpose(-2, -1)).mean()
scores.append(score)
return torch.stack(scores)
`
Summary and Further Reading
Key Takeaways
- MHA is the foundation: Best quality but largest KV Cache, suitable for small models
- GQA is the balance point: KV Cache reduced 8x, quality close to MHA, standard for LLaMA-2/3
- MLA is the most radical: KV Cache compression 95%+, DeepSeek's original innovation, inference cost reduced 10x
- Flash Attention is hardware optimization: Does not change the algorithm, 4x speedup, compatible with all methods
- Sparse attention is the future: MoBA and similar methods are suitable for ultra-long contexts
Attention Method Recommendations
| Scenario | Recommended Method | Reason |
|---|---|---|
| General Training | GQA + Flash Attention | Balance of quality and efficiency |
| Inference Service | MLA + Flash Attention | Minimal KV Cache |
| Ultra-Long Context | MoBA + Flash Attention | Controllable computation |
| Small Models | MHA + Flash Attention | Quality first |
Need to format attention weights as JSON data? Try our JSON Formatter and cURL to Code.
Further Reading
- LLM Inference Acceleration: Three-Engine Showdown — Inference acceleration solutions
- LLM Long Context Optimization — Long context attention
- AI Chip HBM Memory Bottleneck — Memory optimization
- GQA: Training Generalized Multi-Query Transformer — GQA paper
- DeepSeek-V2: MLA Technical Report — MLA explained
- Flash Attention 3 — FA3 paper
Try these browser-local tools — no sign-up required →