摘要
- 注意力机制是Transformer的核心：占LLM 60%+计算量与显存，优化注意力=优化整个模型
- 5代注意力演进：MHA→MQA→GQA→MLA→Flash Attention，每代都在计算效率与模型质量间寻找平衡
- MLA(DeepSeek独创)是2026年最激进的创新：KV Cache压缩95%+，推理成本降低10×
- Flash Attention 3是硬件级优化：H100上实现75%理论峰值利用率，比FA2快2×
- 本文提供5代注意力的数学推导、代码实现与生产部署对比
目录
注意力在LLM中的资源占比
| 组件 |
计算量占比 |
显存占比 |
优化价值 |
| Self-Attention |
40-50% |
50-60% |
最高 |
| FFN/MLP |
40-50% |
25-30% |
高 |
| Embedding |
2-5% |
5-10% |
低 |
| LayerNorm |
<1% |
<1% |
最低 |
5代注意力演进路线
| 代际 |
时间 |
代表 |
KV Cache |
计算效率 |
模型质量 |
| 第1代 MHA |
2017 |
Transformer |
基线 |
基线 |
最好 |
| 第2代 MQA |
2019 |
PaLM |
-75% |
+30% |
略降 |
| 第2代 GQA |
2023 |
LLaMA-2 |
-50% |
+20% |
接近MHA |
| 第3代 MLA |
2024 |
DeepSeek-V2 |
-95% |
+40% |
接近MHA |
| 第4代 FA3 |
2024 |
FlashAttn3 |
不变 |
+100% |
不变 |
| 第5代 稀疏 |
2025-2026 |
MoBA |
-80% |
+50% |
中 |
第1代：MHA多头注意力
MHA数学公式
标准MHA计算流程：
Input X ∈ R^(B×S×D)
Q = XW_Q, K = XW_K, V = XW_V
分头：
Q_i = Q[:, :, i*d_h:(i+1)*d_h] for i = 0, ..., h-1
K_i = K[:, :, i*d_h:(i+1)*d_h]
V_i = V[:, :, i*d_h:(i+1)*d_h]
注意力计算：
Attn_i = softmax(Q_i K_i^T / √d_h) V_i
合并：
Output = Concat(Attn_0, ..., Attn_{h-1}) W_O
KV Cache大小 = 2 × S × h × d_h × sizeof(dtype)
= 2 × S × D × sizeof(dtype)
对于70B模型, S=4096, D=8192, FP16:
KV Cache = 2 × 4096 × 8192 × 2 = 128MB/层
64层总计 = 8GB
MHA实现
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开销
| 模型 |
层数 |
隐层维度 |
头数 |
KV Cache/层 |
总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 |
第2代：MQA与GQA分组注意力
MQA(Multi-Query Attention)
MHA vs MQA:
MHA: 每个头独立的K和V
Q: [h × d_h] K: [h × d_h] V: [h × d_h]
KV Cache = 2 × h × d_h × S
MQA: 所有头共享一组K和V
Q: [h × d_h] K: [1 × d_h] V: [1 × d_h]
KV Cache = 2 × d_h × S (减少h倍)
例如 h=64:
MQA的KV Cache = MHA的 1/64 ≈ 减少98.4%
GQA(Grouped-Query Attention)
GQA: g组头共享K和V
Q: [h × d_h] K: [g × d_h] V: [g × d_h]
KV Cache = 2 × g × d_h × S
当 g=1: 退化为MQA
当 g=h: 退化为MHA
当 g=8: KV Cache减少8倍
LLaMA-2 70B使用 g=8:
KV Cache = MHA的 8/64 = 1/8
GQA实现
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对比
| 方案 |
KV组数 |
KV Cache |
模型质量 |
代表模型 |
| MHA |
h=64 |
100% |
最好 |
GPT-3 |
| GQA-8 |
g=8 |
12.5% |
接近MHA |
LLaMA-2/3 |
| GQA-4 |
g=4 |
6.25% |
略降 |
Mistral |
| MQA |
g=1 |
1.56% |
降5-8% |
PaLM |
第3代：MLA多头潜在注意力
MLA核心思想
┌──────────────────────────────────────────────────────────────┐
│ MLA核心创新 │
│ │
│ 传统MHA/GQA: │
│ K, V直接存储 → KV Cache大 │
│ │
│ MLA: │
│ 1. 投影到低维潜在空间: c_kv = Compress(X) │
│ c_kv维度 << K,V维度 → Cache压缩95%+ │
│ │
│ 2. 推理时从潜在空间恢复: │
│ K = W_k_up × c_kv │
│ V = W_v_up × c_kv │
│ │
│ 3. 吸收技术(Absorption): │
│ Q × K^T = Q × (W_k_up × c_kv)^T │
│ = (Q × W_k_up^T) × c_kv^T │
│ = Q' × c_kv^T │
│ 避免显式恢复K，直接在低维空间计算注意力 │
└──────────────────────────────────────────────────────────────┘
MLA实现
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对比
| 方案 |
KV维度/头 |
总KV维度 |
KV Cache |
压缩比 |
| MHA(h=128) |
128 |
16384 |
100% |
1× |
| GQA(g=8) |
128 |
1024 |
6.25% |
16× |
| MLA(d_c=512) |
- |
512 |
3.1% |
32× |
| MLA(d_c=256) |
- |
256 |
1.56% |
64× |
第4代：Flash Attention硬件级优化
Flash Attention原理
┌──────────────────────────────────────────────────────────────┐
│ Flash Attention核心思想 │
│ │
│ 标准Attention: │
│ Q,K,V → S=QK^T → P=softmax(S) → O=PV │
│ 问题: S和P是S×S矩阵，显存O(S²)，且多次读写HBM │
│ │
│ Flash Attention: │
│ 1. 分块计算(Tiling): 将Q,K,V分成小块 │
│ 2. 在SRAM中完成softmax(在线softmax) │
│ 3. 只写回最终结果O到HBM │
│ │
│ 显存: O(S) vs O(S²) │
│ HBM读写: O(S²d/N) vs O(S²d) → 减少N倍 │
│ N = SRAM大小/块大小 │
└──────────────────────────────────────────────────────────────┘
Flash Attention 3特性
| 特性 |
FA1 |
FA2 |
FA3 |
| GPU支持 |
A100 |
A100/H100 |
H100+ |
| 数据类型 |
FP16/BF16 |
FP16/BF16 |
FP16/BF16/FP8 |
| 异步 |
否 |
否 |
是(WGMMA) |
| 流水线 |
否 |
否 |
是 |
| 理论利用率 |
50% |
62% |
75% |
| 相对速度 |
1× |
2× |
4× |
Flash Attention使用
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性能实测
| 序列长度 |
标准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 |
第5代：稀疏注意力与混合架构
稀疏注意力类型
| 类型 |
稀疏模式 |
计算量 |
适用场景 |
| 局部窗口 |
固定窗口 |
O(S×W) |
长文档 |
| 全局+局部 |
少量全局token |
O(S×(W+G)) |
通用 |
| 膨胀注意力 |
膨胀卷积模式 |
O(S×logS) |
层次结构 |
| MoBA |
混合分块 |
O(S×√S) |
通用 |
| 线性注意力 |
核方法 |
O(S×D²) |
超长序列 |
MoBA(Mixture of Block Attention)实现
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)
总结与引流
关键要点回顾
- MHA是基础：质量最好但KV Cache最大，适合小模型
- GQA是平衡点：KV Cache减少8×，质量接近MHA，LLaMA-2/3标配
- MLA最激进：KV Cache压缩95%+，DeepSeek独创，推理成本降低10×
- Flash Attention是硬件优化：不改变算法，4×加速，与所有方案兼容
- 稀疏注意力是未来：MoBA等方案适合超长上下文
注意力方案推荐
| 场景 |
推荐方案 |
原因 |
| 通用训练 |
GQA + Flash Attention |
质量与效率平衡 |
| 推理服务 |
MLA + Flash Attention |
KV Cache最小 |
| 超长上下文 |
MoBA + Flash Attention |
计算量可控 |
| 小模型 |
MHA + Flash Attention |
质量优先 |
需要格式化注意力权重的JSON数据？试试我们的JSON格式化工具和cURL转代码。
延伸阅读