现代 LLM 架构:技术演化视角
LayerNorm→RMSNorm、绝对位置→RoPE、ReLU→SwiGLU、MHA→GQA、KV Cache——按技术演化组织。
对应课件:2-4 BERT&GPT&前沿大语言模型
本章对应课件:《2-4 BERT&GPT&前沿大语言模型》后半部分。核心知识点全部按统一模板展开。这一章就是 LLaMA/Qwen 这一代模型的骨架说明书。
10.1 KV Cache:推理加速的关键 ★
先懂直觉
生成第 100 个 token 时,前 99 个 token 的 Key/Value 其实已经算过了。不缓存的话每次都要重算一遍,白白浪费算力。KV Cache 就是把历史的 K、V 存下来,新 token 只算自己的 K、V。
公式
逐项拆解
| 符号 | 含义 |
|---|---|
| 历史 token 的 K/V(缓存中,不重算) | |
| 新 token 的 K/V(本次计算) | |
| 拼接维度 | 沿序列维拼接:[B, H, t−1, d] ⊕ [B, H, 1, d] → [B, H, t, d] |
矩阵 shape
数字算例(计算量对比)
生成 4 个 token 的序列:
| 步骤 | 无缓存计算量 | 有缓存计算量 |
|---|---|---|
| 第 1 步 | 1 个 token | 1 个 token |
| 第 2 步 | 2 个 token | 1 个 token |
| 第 3 步 | 3 个 token | 1 个 token |
| 第 4 步 | 4 个 token | 1 个 token |
| 总计 | 1+2+3+4 = 10 | 1+1+1+1 = 4 |
上表数的是 K/V 投影的重复计算次数:无缓存时累计 ,有缓存时 ——被消掉的是「历史 token 的投影与完整前向重算」。
不要说「有了 KV Cache,decode 每步变成 O(1)」。准确地说:
- 新 token 的投影/FFN/Norm:每步 (常量);
- 注意力分数:新 Query 仍要与 t 个历史 K 做点积 → 每步 ,随上下文长度线性增长;
- 生成 n 个 token:注意力累计仍是 ,只是免去了对 prefix 的重复完整前向。
KV Cache 省的是「重算历史」;没有让长上下文 decode 变成与上下文无关的常数成本。此外,每步还要从显存读取 t×d×2 个缓存元素——长上下文时 decode 常受显存带宽限制。
显存公式与算例
例:LLaMA-2-7B(L=32, H_kv=32, d_head=128),S=4096,FP16(2 字节):
单条序列就要 2GB!如果 batch=8 就是 17GB——这就是为什么需要 GQA(10.2 节)。
左边是「无缓存」每次都重算全部历史(↻),右边是「有缓存」只算新 token。下面还有显存估算器,改层数/head 数/序列长度看显存变化。
工程里怎么用(简化版 KV Cache 实现)
class CachedAttention(nn.Module):
def forward(self, x, cache=None):
Q, K, V = self.Wq(x), self.Wk(x), self.Wv(x)
if cache is not None:
K = torch.cat([cache["K"], K], dim=1) # 拼接历史 K
V = torch.cat([cache["V"], V], dim=1) # 拼接历史 V
new_cache = {"K": K, "V": V}
scores = Q @ K.transpose(-2, -1) / math.sqrt(Q.size(-1))
out = F.softmax(scores, dim=-1) @ V
return out, new_cache
KV Cache 减少计算,但增加显存。这是最经典的 trade-off:省的是 FLOPs,花的是显存和带宽。长上下文推理时,显存往往才是瓶颈(batch 维度还会再乘上去)。
Q1:KV Cache 为什么能加速推理?代价是什么?
查看答案
无缓存时,生成第 t 个 token 要对整个 prefix(t 个 token)重跑完整前向(历史 token 的 Q/K/V 投影、注意力、FFN 全部重算);有缓存后,只计算新 token 自己的投影/FFN,历史 K/V 直接读缓存。但注意:新 Query 仍需与全部 t 个历史 K 计算注意力分数(每步 O(t·d),随上下文增长)——缓存省掉的是「重算历史」,不是注意力本身。代价是显存:缓存大小随层数、KV head 数、序列长度线性增长(公式见上),长上下文时非常可观。
Q2:为什么只缓存 K/V,不缓存 Q?
查看答案
Q 只用于当前 token 的查询,用完即弃;而 K/V 会被后续所有 token 反复使用(每一步都要和全部历史做注意力)。缓存 Q 没有复用价值。
10.2 MHA → MQA → GQA ★
先懂直觉
KV Cache 太大,是因为每个 Query head 都配了独立的 KV head。那就让多个 Query head 共享 KV——这就是 GQA。共享得越多,KV Cache 越小,但表达能力越受限。
对比公式
其中 是分组大小(每组 个 Q head 共享 1 组 K/V)。
逐项拆解
| 方案 | Q heads | KV heads | KV Cache | 特点 |
|---|---|---|---|---|
| MHA | 32 | 32 | 基准 2.1 GB | 表达力最强,显存最贵 |
| GQA | 32 | 8 | 1/4 → 0.54 GB | 质量与显存的折中(现代主流) |
| MQA | 32 | 1 | 1/32 → 0.07 GB | 显存最省,质量损失最大 |
矩阵 shape
在 MHA / MQA / GQA 之间切换,观察 Q heads 与 K/V heads 的连接方式,以及 KV Cache 的相对大小。
工程里怎么用(HuggingFace 配置)
from transformers import AutoConfig
cfg = AutoConfig.from_pretrained("meta-llama/Llama-3-8b")
print(cfg.num_attention_heads, cfg.num_key_value_heads) # 32 8
# LLaMA-2-7B: 32 Q / 32 KV(MHA)
# LLaMA-2-70B: 64 Q / 8 KV(GQA)
# LLaMA-3-8B: 32 Q / 8 KV(GQA)
- GQA 不减少 Q head 数:只减少 K/V head 数,Q 的表达力不变。
- GQA 的质量损失很小:这是它取代 MQA 成为主流的原因(MQA 压得太狠,质量掉得多)。
Q:为什么大模型推理用 GQA 而不是 MQA?
查看答案
MQA 把 KV head 压到 1,KV Cache 最省,但所有 Q head 共享同一组 K/V,表达力下降明显、质量损失较大;GQA 用分组折中(如 8 组),KV Cache 降到 1/4~1/8,实验表明质量损失很小(loss/评测上接近 MHA,具体随模型与训练配置变化)。因此现代 LLM(LLaMA-2-70B、LLaMA-3、Qwen)普遍采用 GQA。
10.3 RoPE:旋转位置编码 ★
先懂直觉
RoPE 不去「加」位置向量,而是把 Q/K 向量按位置「旋转」一个角度。两个向量各自旋转后做点积,结果只取决于它们的角度差——也就是相对位置。这是现代 LLM 最常用的位置编码。
公式
二维旋转矩阵:
对位置 的向量应用旋转:
逐项拆解(为什么点积只依赖相对位置)
因为旋转矩阵满足 。所以点积只依赖 ,与绝对位置无关。
| 符号 | 含义 |
|---|---|
| 两个 token 的绝对位置 | |
| 旋转频率(每个二维子空间不同) | |
| 相对位置旋转,点积因此只含相对信息 |
矩阵 shape
实际实现中,向量被分成 个二维子空间,每个用不同频率 旋转。
数字算例(θ = 45°,向量 [1, 0])
| q 位置 m | k 位置 n | q 旋转后 | k 旋转后 | 点积 | 相对位置 m−n |
|---|---|---|---|---|---|
| 1 | 0 | [0.707, 0.707] | [1, 0] | 0.707 | 1 |
| 2 | 1 | [0, 1] | [0.707, 0.707] | 0.707 | 1 |
| 2 | 0 | [0, 1] | [1, 0] | 0 | 2 |
关键:前两行相对位置相同(m−n=1),点积也相同(0.707)——无论绝对位置怎么变,只要相对距离一样,注意力分数就一样。
拖动 m、n 滑块改变两个向量的位置,观察它们各自旋转;点击「同时 +1」按钮——两个向量一起旋转,点积保持不变。
工程里怎么用(RoPE 简化实现)
def apply_rope(x, positions, dim, base=10000.0):
half = dim // 2
freqs = 1.0 / (base ** (torch.arange(0, half, device=x.device) / half))
angles = positions[:, None].float() * freqs[None, :] # [S, half]
cos, sin = angles.cos()[None, None], angles.sin()[None, None]
x1, x2 = x[..., :half], x[..., half:]
return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
- RoPE 不是「加到输入上」:它作用在每一层的 Q/K 上(不是 embedding 阶段),而且只作用于 Q 和 K,不作用于 V。
- 旋转的是二维子空间:不是整个向量一起转。
- 相对位置性质是数学结论:由 保证,不是近似。
Q1:RoPE 相比 Sinusoidal PE 的优势?
查看答案
① 内积天然只依赖相对位置,符合 Attention 的需求;② 外推性更好(配合 NTK/YaRN 插值可扩展上下文);③ 无额外参数、实现简单;④ 作用在 Q/K 上而非输入,每层都能感知位置。
Q2:为什么 RoPE 只作用于 Q 和 K,不作用于 V?
查看答案
位置信息只需要影响「匹配分数」(Q·K),V 是被加权取用的内容,不需要参与位置匹配。对 V 加旋转反而会破坏内容表示。
10.4 RMSNorm:更简单的归一化 ★
先懂直觉
LayerNorm 要算均值和方差;RMSNorm 发现「减均值」这一步其实不那么重要,只除以均方根就够了。少一次归约运算,更快更省。
公式
注意 eps 的位置:在根号内(标准实现,与 LLaMA 系及
torch.nn.RMSNorm一致),即 。
逐项拆解
| 符号 | 含义 | 与 LayerNorm 的差别 |
|---|---|---|
| 均方根(只算平方均值开根) | LN 还要减均值 | |
| 可学习缩放 | LN 还有 平移 | |
| 无 、 | — | 省一次归约、少一组参数 |
数字算例(x = [1, 2, 3, 4])
对比 LayerNorm 的 :RMSNorm 不把均值拉到 0,保留了整体的「偏移」信息。
拖动 x 的每个分量,左右并排看两种归一化的完整计算过程和结果差异。
工程里怎么用(PyTorch 手写 RMSNorm)
class RMSNorm(nn.Module):
def __init__(self, d, eps=1e-6):
super().__init__()
self.g = nn.Parameter(torch.ones(d))
self.eps = eps
def forward(self, x):
rms = x.pow(2).mean(-1, keepdim=True).sqrt()
return x / (rms + self.eps) * self.g
- RMSNorm 不是「简化版 LayerNorm 的近似」:它是一个独立的、效果相当甚至更好的归一化方法。
- 没有 β 不是缺陷:实验表明去掉平移项影响很小,但省了参数和计算。
- 仍然是对每个 token 的 hidden 维归一化,不跨 token。
Q:RMSNorm 为什么能省计算?效果为什么不受影响?
查看答案
省掉了「减均值」的归约和 β 参数,只需平方均值开根。研究发现 LayerNorm 的主要收益来自「尺度不变性」(除以范数),而中心化(减均值)贡献很小,因此 RMSNorm 在保持效果的同时更省带宽——对显存带宽瓶颈的大模型推理尤其有价值。
10.5 Pre-Norm 与 Post-Norm ★
先懂直觉
归一化放在子层前面还是后面?原始 Transformer 放在后面(Post-Norm),结果深层训练很不稳定;现代 LLM 放在前面(Pre-Norm),残差路径变成一条「干净的高速公路」,梯度可以直达底层。
公式
Post-Norm:
Pre-Norm:
逐项拆解
| Post-Norm | Pre-Norm | |
|---|---|---|
| 归一化位置 | 残差之后 | 子层之前 |
| 残差路径 | 穿过 Norm | 纯恒等(干净) |
| 梯度 | 每层乘 Norm 的雅可比 | |
| 深层训练 | 难,需要 warmup | 容易 |
| 代表模型 | 原始 Transformer、BERT | GPT-2 之后几乎所有 LLM |
| 额外处理 | — | 末尾加 final norm |
数字算例(梯度传播对比)
假设每层 Norm 的雅可比约 0.7:
| 层数 | Post-Norm: | Pre-Norm: |
|---|---|---|
| 4 | 0.240 | 1.240 |
| 12 | 0.014 | 1.014 |
| 24 | 0.0002 | 1.0002 |
并排看两种 Block 结构,绿色标出残差路径。Pre-Norm 的残差路径上没有 Norm,梯度可以无阻碍回传。
工程里怎么用(两种写法)
# Pre-Norm(现代 LLM 标准)
x = x + self.attn(self.norm1(x))
x = x + self.ffn(self.norm2(x))
x = self.final_norm(x) # 末尾补一次
# Post-Norm(原始 Transformer)
x = self.norm1(x + self.attn(x))
x = self.norm2(x + self.ffn(x))
- Pre-Norm 也需要 final norm:因为表示方差随深度增长,末尾要归一化后再进 LM Head。
- Pre-Norm 不是「更好」的普适结论:它主要是训练稳定性优势;Post-Norm 在浅层、小模型上也可用。
Q:为什么现代 LLM 用 Pre-Norm?
查看答案
Pre-Norm 的残差路径是纯恒等映射(x + F(Norm(x))),梯度含常数 1、可无衰减直达浅层,深层训练稳定,不需要像 Post-Norm 那样依赖精细的 warmup;代价是表示方差随深度增长,需要末尾加 final norm。
10.6 SwiGLU:现代 FFN ★
先懂直觉
ReLU FFN 只有一条路;SwiGLU 有两条路——一条产生「内容」,另一条产生 0~1 的「门」,两者逐元素相乘。门控让网络能动态决定放行多少信息。
公式
其中 SiLU(Swish):
逐项拆解
| 符号 | 作用 | 形状 |
|---|---|---|
| 门控分支(过 SiLU 变成 0~1 附近) | [d, d_ff] | |
| 内容分支 | [d, d_ff] | |
| 逐元素相乘(门控) | [d_ff] | |
| 输出投影 | [d_ff, d] |
矩阵 shape
数字算例(标量)
设 ,,,:
| 步骤 | 计算 | 结果 |
|---|---|---|
| gate | 1.5 × 0.818 = 1.227 | |
| value | 0.8 | |
| hidden | (门控) | 0.982 |
| output | 1.178 |
拖动 x,看 gate = SiLU(xW₁)、value = xW₃、hidden = gate ⊙ value 的逐步计算;图上对比 SiLU 与 ReLU 的形状。
SwiGLU 有 3 个权重矩阵(W₁、W₂、W₃),ReLU FFN 只有 2 个。为了保持总参数量相当:
这就是 LLaMA 的 FFN 中间维度是 而不是 的原因。
- SwiGLU 的 d_ff 不是 4d:用 8/3 d 才能与 ReLU FFN 参数量相当;直接套 4d 会让参数量多 50%。
- 门控是逐元素相乘,不是矩阵乘。
- SiLU 非单调:在负区间先降后升,这是它与 ReLU 的显著区别。
Q:SwiGLU 相比 ReLU FFN 的优势?
查看答案
门控机制让网络能动态控制信息流(乘性交互比单纯的激活更灵活),实验上在同等参数量下效果更好,因此成为 LLaMA/Qwen 等现代 LLM 的标配。代价是多一个投影矩阵,所以中间维度要调成 8/3 d。
10.7 现代 LLM 骨架:一图记住
Tokenizer
↓
Embedding
↓
┌─ 重复 N 层 ────────────────────────┐
│ Pre-RMSNorm │
│ ↓ │
│ GQA / MHA(RoPE 作用于 Q/K) │
│ ↓ │
│ Residual ⊕ │
│ ↓ │
│ Pre-RMSNorm │
│ ↓ │
│ SwiGLU FFN │
│ ↓ │
│ Residual ⊕ │
└────────────────────────────────────┘
↓
Final Norm
↓
LM Head
↓
Logits
| 组件 | 原始 Transformer (2017) | 现代 LLM(LLaMA 类) |
|---|---|---|
| 位置编码 | Sinusoidal(相加) | RoPE(旋转 Q/K) |
| 归一化 | LayerNorm + Post-Norm | RMSNorm + Pre-Norm |
| Attention | MHA | GQA / MHA |
| FFN 激活 | ReLU | SwiGLU |
| 推理优化 | 无 | KV Cache |
10.8 本章总结
| 组件 | 解决的问题 | 一句话 |
|---|---|---|
| KV Cache | 推理重复计算 | 缓存历史 K/V,省掉 prefix 重算(K/V 投影 ;注意力仍随上下文增长) |
| GQA | KV Cache 太大 | 多个 Q head 共享 K/V(如 32→8,降到 1/4) |
| RoPE | 位置信息 | 旋转 Q/K,点积只依赖相对位置 m−n |
| RMSNorm | 归一化成本 | 不减均值,只除均方根 |
| Pre-Norm | 深层训练 | 残差路径干净,梯度直达(1 + F′) |
| SwiGLU | FFN 表达力 | 双路门控,中间维 8/3 d |
显示答案与解析
正确答案:B
KV Cache 缓存历史的 Key/Value,新 token 不再重算历史投影与完整前向(K/V 投影的重复量从 O(n²) 降到 O(n));但注意力分数仍需与历史 K 计算,随上下文增长。代价是显存占用随序列增长(A 说反了)。显示答案与解析
正确答案:B
GQA 让多个 Query head 共享一组 KV head,直接减少 KV Cache 大小(如 32→8 降到 1/4),显著降低长上下文推理的显存和带宽压力。显示答案与解析
正确答案:B
RoPE 按位置旋转 Q/K,由于旋转矩阵性质 R_mᵀR_n = R_{n−m},点积只依赖 m−n(相对位置)。A 是 sinusoidal PE 的做法。显示答案与解析
正确答案:A
RMSNorm = x / √(mean(x²)+ε) · γ,省掉减均值(和 β),计算更简单、更省带宽,效果接近 LayerNorm,被 LLaMA/Qwen 等采用。显示答案与解析
正确答案:B
Pre-Norm 的残差路径是纯恒等映射,∂x_{t+1}/∂x_t 中始终有常数 1;Post-Norm 每层都要穿过 Norm 的雅可比,深层不稳定。注意 Pre-Norm 反而需要额外加 final norm(D 错)。显示答案与解析
正确答案:B
ReLU FFN 有 2 个矩阵(2 × d × 4d),SwiGLU 有 3 个矩阵(3 × d × d_ff)。令 3·d·d_ff ≈ 2·d·4d 得 d_ff ≈ 8/3 d。