04 - 推理优化轨道 - 让模型飞起来 🚀
你已经学会了如何训练和对齐模型,现在是时候让它跑得更快了!
推理优化关注的是:训练后的模型如何高效部署?
🎯 本轨道内容
| 算法 | 核心问题 | 解决方案 | 难度 |
|---|---|---|---|
| Flash Attention | O(N²)内存爆炸 | 分块计算,在线softmax | ⭐⭐⭐⭐ |
| KV-Cache | 生成时重复计算 | 缓存Key/Value | ⭐⭐⭐ |
| Quantization | 模型太大 | INT8/INT4压缩 | ⭐⭐⭐ |
| SSM/Mamba | Attention太慢 | 线性复杂度状态空间 | ⭐⭐⭐⭐ |
| RoPE | 绝对位置编码不灵活 | 旋转位置嵌入 | ⭐⭐⭐ |
| GQA | KV-Cache内存大 | 分组查询减少KV头 | ⭐⭐⭐ |
| Speculative Decoding | 生成慢 | 小模型推测+大模型验证 | ⭐⭐⭐ |
| PagedAttention | KV-Cache碎片化 | 虚拟内存分页管理 | ⭐⭐⭐⭐ |
| Beam Search | 贪心解码质量差 | 维护多条候选路径 | ⭐⭐ |
| Discretization | SSM的连续/离散转换 | Euler/Tustin方法 | ⭐⭐⭐⭐ |
| Complex SSM | SSM表达力不足 | 复数状态空间 | ⭐⭐⭐⭐ |
学习路径: Flash Attention → KV-Cache → Quantization → SSM/Mamba → 其他优化
💡 算法一: Flash Attention - 内存高效的注意力
什么是Flash Attention?
Flash Attention (Dao et al., 2022) 解决了Transformer的核心瓶颈:O(N²)内存。
标准注意力的问题:
# 标准注意力
S = Q @ K.T # [N, N] 巨大矩阵!
P = softmax(S)
O = P @ V当N=8192时,S矩阵需要 8192² ≈ 67M 个float(256MB) — GPU显存爆炸!
Flash Attention方案:
- 分块计算:把Q,K,V切成小块,一次只处理一个tile
- 在线Softmax:不存完整S矩阵,边算边归一化
- 结果: 内存从O(N²) → O(N),数学上等价
对话1: 为什么标准Attention内存是O(N²)?
🤔 学生: 不是最后输出O只有[N,d]吗,为什么说内存O(N²)?
💡 导师: 中间过程才是内存杀手!
标准Attention计算流程:
步骤1: S = Q @ K.T 内存: [N,d] @ [d,N] = [N,N] ← 这里!
步骤2: P = softmax(S) 内存: [N,N] (可以in-place)
步骤3: O = P @ V 内存: [N,N] @ [N,d] = [N,d]
具体数字:
| 序列长度N | S矩阵大小 | 内存(FP16) |
|---|---|---|
| 512 | 512² = 262K | 0.5 MB |
| 2048 | 4M | 8 MB |
| 8192 | 67M | 256 MB |
| 32768 | 1B | 2 GB |
| 131072 | 17B | 32 GB |
为什么不能直接释放S?
- 反向传播需要S来计算梯度
- 即使推理时(没有反向传播),softmax也需要完整的S
✨ 关键洞察: Attention的”全局信息”优势来自N×N矩阵——但这也是它的致命弱点。Flash Attention证明:可以不存矩阵也能算出正确结果。
对话2: 分块计算会改变结果吗?
🤔 学生: 分块计算softmax,难道不会算出不同的值吗?
💡 导师: 这是Flash Attention的核心魔法:在线Softmax算法。
问题场景:
# 标准softmax(需要看到所有值)
scores = [2.0, 1.0, 3.0, 0.5] # 必须知道全部才能归一化
max_score = 3.0
exp_sum = exp(2-3) + exp(1-3) + exp(3-3) + exp(0.5-3)
probs = [exp(s-3)/exp_sum for s in scores]分块怎么办? 假设分两块:[2.0, 1.0] 和 [3.0, 0.5]
块1:
- max1=2.0, sum1=exp(2-2)+exp(1-2) = 1.0 + 0.368 = 1.368
- 局部probs1 = [1.0/1.368, 0.368/1.368]
块2:
- max2=3.0, sum2=exp(3-3)+exp(0.5-3) = 1.0 + 0.082 = 1.082
如何合并? (关键算法)
# 步骤1: 更新全局max
global_max = max(max1, max2) = 3.0
# 步骤2: 重新缩放块1的sum (因为max变了)
sum1_rescaled = sum1 * exp(max1 - global_max)
= 1.368 * exp(2 - 3) = 1.368 * 0.368 = 0.503
# 步骤3: 重新缩放块2的sum
sum2_rescaled = sum2 * exp(max2 - global_max)
= 1.082 * exp(3 - 3) = 1.082
# 步骤4: 总和
global_sum = sum1_rescaled + sum2_rescaled = 0.503 + 1.082 = 1.585
# 步骤5: 重新缩放之前的输出
# (如果之前有输出O1 = probs1 @ V1, 现在要调整)
O1_rescaled = O1 * (sum1_rescaled / global_sum)数学证明(为什么等价):
原始: exp(s_i - max_all) / sum_all
分块: [exp(s_i - max_block) * exp(max_block - max_all)] /
[sum_block * exp(max_block - max_all)]
= exp(s_i - max_all) / sum_block_rescaled
通过动态维护running_max和running_sum,每个新块到来时:
- 更新max → 重新缩放旧sum
- 加上新sum
- 重新缩放旧输出
✨ 关键洞察: Softmax的归一化性质允许”延迟归一化”——先算局部,后全局调整。Flash Attention把这个数学技巧变成了工程现实。
对话3: Flash Attention在GPU上为什么快?
🤔 学生: Python实现Flash Attention反而更慢,为什么论文说快10倍?
💡 导师: 因为GPU内存层次结构!
GPU内存两层:
| 层级 | 容量 | 速度 | 用途 |
|---|---|---|---|
| HBM(显存) | 40 GB | 慢(600 GB/s) | 存模型参数,大矩阵 |
| SRAM(L2缓存) | 20 MB | 快(19 TB/s) | GPU核心运算 |
速度差异:30倍!
标准Attention的内存访问:
1. 从HBM读Q,K → SRAM计算S → 写S回HBM
2. 从HBM读S → SRAM计算softmax → 写P回HBM
3. 从HBM读P,V → SRAM计算O → 写O回HBM
总HBM访问:3次读 + 3次写 (每次都是巨大的N²矩阵)
Flash Attention的内存访问:
对每个tile:
1. 从HBM读Q_tile, K_tile, V_tile → SRAM
2. 在SRAM内: 算S_tile → softmax → 算O_tile (全程SRAM!)
3. 更新running统计(极小内存)
4. 写最终O回HBM
总HBM访问:1次读(Q,K,V) + 1次写(O) — N²矩阵从不碰HBM!
具体加速:
- N=8192, d=64, block=256
- 标准Attention: 读写 2×(8192²×2字节) ≈ 256 MB × 6次 = 1.5 GB
- Flash Attention: 读写 2×(8192×64×2字节) ≈ 2 MB × 2次 = 4 MB
带宽节省: 1.5 GB / 4 MB ≈ 375倍!
✨ 关键洞察: Flash Attention不是”算得更少”,而是”搬运数据更少”。现代AI是内存绑定(memory-bound),不是计算绑定(compute-bound)。
代码片段: 在线Softmax核心
def flash_attention_simplified(Q, K, V, block_size=64):
"""
Flash Attention简化版(教学用)
Q,K,V: [seq_len, d_model]
"""
N, d = Q.shape
O = np.zeros((N, d)) # 输出
row_max = np.full(N, -np.inf) # 每行的running max
row_sum = np.zeros(N) # 每行的running sum
# 外循环: 遍历Q的块
for q_start in range(0, N, block_size):
q_end = min(q_start + block_size, N)
Q_block = Q[q_start:q_end] # [block, d]
# 内循环: 遍历K,V的块
for k_start in range(0, N, block_size):
k_end = min(k_start + block_size, N)
K_block = K[k_start:k_end] # [block, d]
V_block = V[k_start:k_end]
# 计算局部scores
S_tile = Q_block @ K_block.T / np.sqrt(d) # [Bq, Bk]
# 在线softmax更新(核心!)
for i in range(len(Q_block)):
global_i = q_start + i
# 当前块的max
m_new = np.max(S_tile[i])
m_old = row_max[global_i]
m = max(m_old, m_new)
# 重新缩放旧sum
sum_old_rescaled = row_sum[global_i] * np.exp(m_old - m)
# 新块的sum
sum_new = np.sum(np.exp(S_tile[i] - m))
# 更新全局sum
sum_total = sum_old_rescaled + sum_new
# 重新缩放旧输出
O[global_i] *= (sum_old_rescaled / sum_total) * np.exp(m_old - m)
# 加上新贡献
weights = np.exp(S_tile[i] - m) / sum_total
O[global_i] += weights @ V_block
# 更新running统计
row_max[global_i] = m
row_sum[global_i] = sum_total
return O关键点:
- 双层循环:外层Q块,内层K/V块
- running统计:每行维护max和sum
- 指数重缩放:
exp(m_old - m_new)调整旧值 - 增量更新:每个块贡献累加到O
⚡ 算法二: KV-Cache - 消除重复计算
什么是KV-Cache?
问题:生成文本时,每生成一个token就要重算整个序列的Key和Value!
例子:
已生成: "The cat sat"
要生成: "on"
标准做法(浪费!):
1. 计算 "The" 的 K,V
2. 计算 "cat" 的 K,V
3. 计算 "sat" 的 K,V
4. 计算 "on" 的 K,V
5. Attention(Q_on, [K_all], [V_all])
KV-Cache做法(聪明!):
1. 从缓存读 K_The, V_The, K_cat, V_cat, K_sat, V_sat
2. 只计算 K_on, V_on (新token)
3. Attention(Q_on, [K_cached + K_on], [V_cached + V_on])
加速:生成100 tokens,从100次全量计算 → 100次增量计算,节省99%计算!
对话4: KV-Cache为什么有效?
🤔 学生: 为什么可以缓存K,V,但不能缓存Q?
💡 导师: 因为因果注意力的非对称性!
Attention公式回顾:
Q = X @ W_Q # [seq, d]
K = X @ W_K
V = X @ W_V
Scores = Q @ K.T # [seq, seq]
O = softmax(Scores) @ V
关键观察:
-
Q是查询:当前token”on”的Q只用来查询别人
- Q_on 只在计算第4步的Attention时用
- 下一步生成”the”时,Q_on没用了(因为”the”有自己的Q_the)
-
K,V是被查询的:历史token的K,V会被未来token重复使用
- K_The 在生成 “cat”, “sat”, “on”, “the”… 时都会被查
- V_The 也是,用来加权输出
- 因果掩码保证:未来token可以看过去,过去token的K,V永远有效
数学上:
生成token t时:
Q_t 只与 K_0..t 交互一次,然后丢弃
K_t, V_t 会被 Q_{t+1}, Q_{t+2}, ... Q_T 使用多次
缓存收益:
K_t, V_t 被复用 (T-t) 次
T=100, t=10 → 被复用90次!
✨ 关键洞察: KV-Cache利用了自回归生成的”左看右不看”性质——历史是只读的,所以可以缓存。
对话5: KV-Cache的内存代价是多少?
🤔 学生: 缓存听起来很好,但内存会爆炸吗?
💡 导师: 是的,长序列会吃光显存!
内存分析:
每层缓存大小:
K: [seq_len, d_model]
V: [seq_len, d_model]
总: 2 × seq_len × d_model × sizeof(float16)
具体数字(GPT-3规模):
- d_model=12288, n_layers=96, fp16 (2字节)
- 每个token需要: 2 × 96 × 12288 × 2 = 4.7 MB
| 序列长度 | KV-Cache大小 |
|---|---|
| 512 | 2.4 GB |
| 2048 | 9.6 GB |
| 8192 | 38.4 GB |
| 32768 | 153 GB (爆!) |
批处理更糟:
- batch_size=64, seq_len=2048
- 总内存: 9.6 GB × 64 = 614 GB (单卡放不下!)
解决方案:
- PagedAttention: 像OS分页内存一样管理KV-Cache
- Multi-Query Attention (MQA): 所有头共享K,V → 减少n_heads倍
- Grouped-Query Attention (GQA): MQA和多头的折中
- 量化: KV-Cache用INT8/INT4存储
- Offloading: 把不常用的K,V放CPU内存
工业界实践:
- vLLM: PagedAttention,显存利用率提升4倍
- TensorRT-LLM: KV-Cache量化到INT8
- llama.cpp: CPU推理,256k上下文(靠CPU内存)
✨ 关键洞察: KV-Cache是”时间换空间”的反面——“空间换时间”。长序列推理的瓶颈从计算变成了内存容量。
代码片段: KV-Cache实现
class GPTWithKVCache:
def __init__(self, model):
self.model = model
self.kv_cache = None # 会在generate时初始化
def generate(self, prompt_ids, max_new_tokens=50):
"""
自回归生成(带KV-Cache)
"""
# 初始化cache: [n_layers, 2, batch=1, n_heads, seq_len, head_dim]
self.kv_cache = self._init_cache()
# 第一步: prefill阶段(处理prompt)
# 这一步需要计算所有prompt token的K,V
input_ids = prompt_ids # [1, prompt_len]
position_offset = 0
logits = self.forward_with_cache(
input_ids,
position_offset=position_offset,
use_cache=True
) # 内部填充cache
next_token = torch.argmax(logits[:, -1], dim=-1) # 取最后一个token的预测
generated = [next_token.item()]
# 后续步骤: decode阶段(每次只处理1个new token)
for _ in range(max_new_tokens - 1):
position_offset += len(input_ids[0])
input_ids = next_token.unsqueeze(0) # [1, 1] 只有新token!
logits = self.forward_with_cache(
input_ids,
position_offset=position_offset,
use_cache=True # 使用并更新cache
)
next_token = torch.argmax(logits[:, -1], dim=-1)
generated.append(next_token.item())
if next_token.item() == EOS_TOKEN:
break
return prompt_ids + generated
def forward_with_cache(self, input_ids, position_offset=0, use_cache=False):
"""
前向传播(支持KV-Cache)
Args:
input_ids: [batch, seq_len] 当前输入(可能只有1个token)
position_offset: 位置偏移(已生成token数)
use_cache: 是否使用cache
"""
batch, seq_len = input_ids.shape
# Embedding + Position
x = self.model.embed(input_ids) + self.model.pos_embed(
torch.arange(position_offset, position_offset + seq_len)
)
# Transformer层
for layer_idx, layer in enumerate(self.model.layers):
# Self-Attention with KV-Cache
if use_cache:
x, new_k, new_v = layer.attn_with_cache(
x,
kv_cache=self.kv_cache[layer_idx], # 传入这层的cache
position_offset=position_offset
)
# 更新cache (append新的K,V)
self.kv_cache[layer_idx][0] = torch.cat([
self.kv_cache[layer_idx][0], new_k
], dim=2) # dim=2 是seq_len维度
self.kv_cache[layer_idx][1] = torch.cat([
self.kv_cache[layer_idx][1], new_v
], dim=2)
else:
x = layer.attn(x)
# MLP
x = layer.mlp(x)
# LM Head
logits = self.model.lm_head(x)
return logits
def _init_cache(self):
"""初始化KV-Cache为空列表"""
# [n_layers, 2(K和V), batch, n_heads, 0(初始seq=0), head_dim]
return [
[
torch.empty(1, self.model.n_heads, 0, self.model.head_dim), # K
torch.empty(1, self.model.n_heads, 0, self.model.head_dim), # V
]
for _ in range(self.model.n_layers)
]关键细节:
- Prefill vs Decode: 第一步处理整个prompt(并行),后续每次1个token(串行)
- Cache更新: 每步append新K,V到cache
- Position offset: 新token的位置 = 已生成长度
🗜️ 算法三: Quantization - 模型压缩
什么是量化?
量化(Quantization): 把浮点数权重压缩成低比特整数,减少内存和加速推理。
例子:
原始权重(FP32): 0.123456789 (32 bit)
量化到INT8: 31 (8 bit, 缩小4倍!)
核心公式:
量化: q = round( (w - zero_point) / scale )
反量化: w ≈ scale * q + zero_point
对话6: 量化为什么不会严重损失精度?
🤔 学生: 从32bit降到8bit,丢掉24bit信息,模型不会变傻吗?
💡 导师: 因为神经网络权重的分布特性!
观察: LLM权重的99%值在一个很小的范围内
例子(真实GPT-2权重统计):
均值: μ = 0.001
标准差: σ = 0.08
99%的权重在 [-3σ, +3σ] = [-0.24, +0.24]
量化策略(Symmetric):
假设权重范围 [-0.24, +0.24]
INT8范围: [-128, 127]
scale = 0.24 / 127 ≈ 0.00189
映射:
-0.24 → -127
0.00 → 0
+0.24 → +127
精度: 每个整数代表 0.00189 的变化
为什么精度够用?
1. 神经网络鲁棒性:
- 训练时已经有dropout/噪声
- 网络学到的是”模糊模式”,不是精确数值
- 0.123 vs 0.125 对最终预测几乎无影响
2. 重要权重保留:
- 大权重(|w| > 0.1)量化误差 < 1%
- 小权重(|w| < 0.01)本身对输出贡献就小
- 异常值可以用混合精度(INT8主体 + FP16异常值)
实验数据(BERT量化):
| 精度 | 模型大小 | 准确率 | 相对损失 |
|---|---|---|---|
| FP32 | 440 MB | 84.5% | - |
| INT8 | 110 MB | 84.2% | -0.3% |
| INT4 | 55 MB | 83.1% | -1.4% |
| INT2 | 28 MB | 76.8% | -7.7% |
✨ 关键洞察: 量化是”精度换效率”——损失1%准确率换4倍压缩,大多数场景非常值得。
对话7: INT8 vs INT4量化有什么区别?
🤔 学生: INT4比INT8再小一半,为什么精度损失更大?
💡 导师: 因为量化位数决定了离散化粒度。
INT8量化:
范围: [-128, 127] (256个离散值)
权重 -0.24 到 +0.24
step = 0.48 / 256 ≈ 0.0019
误差界: ±0.0019 / 2 ≈ ±0.001 (最大量化误差)
相对误差: 0.001 / 0.08 ≈ 1.25%
INT4量化:
范围: [-8, 7] (16个离散值)
step = 0.48 / 16 = 0.03
误差界: ±0.015
相对误差: 0.015 / 0.08 ≈ 18.75%
对比:
| 量化 | 离散值个数 | 步长 | 相对误差 |
|---|---|---|---|
| INT8 | 256 | 0.0019 | ~1% |
| INT4 | 16 | 0.03 | ~19% |
INT4的问题:
- 粗粒度:小权重容易被”四舍五入”归零
- 异常值敏感:一个大权重会挤压整个量化范围
- 层间累积:每层1%误差,100层后指数放大
INT4可行的条件(QLoRA的秘密):
- NF4(NormalFloat4):专门为正态分布设计的非均匀量化
- 在0附近分配更多量化点
- 权重集中区域精度高
- Double Quantization:量化scale本身也量化(进一步省内存)
- 分块量化:每64个权重一组,独立量化(避免全局异常值影响)
✨ 关键洞察: 量化不只是”换整数”,还要设计量化点分布。NF4针对神经网络权重分布优化,比均匀INT4好10倍。
代码片段: 对称量化实现
def quantize_weights_int8(weights):
"""
对称量化(symmetric quantization)到INT8
Args:
weights: np.array, 原始FP32权重
Returns:
q_weights: int8量化权重
scale: float, 反量化缩放因子
"""
# 1. 计算权重绝对值最大值
abs_max = np.max(np.abs(weights))
# 2. 计算scale (映射到[-127, 127])
# 留一个值给padding,用127不用128
scale = abs_max / 127.0
# 3. 量化
q_weights = np.round(weights / scale).astype(np.int8)
return q_weights, scale
def dequantize_weights_int8(q_weights, scale):
"""反量化"""
return q_weights.astype(np.float32) * scale
def quantized_linear(x, W_q, scale, bias=None):
"""
量化线性层推理
Args:
x: 输入激活(FP32)
W_q: 量化权重(INT8)
scale: 缩放因子
bias: 偏置(FP32)
"""
# 方法1: 反量化后计算(简单但慢)
# W = dequantize_weights_int8(W_q, scale)
# return x @ W.T + bias
# 方法2: INT8矩阵乘(快!需要硬件支持)
# 核心技巧: (x @ W) = (x @ (scale * W_q)) = scale * (x @ W_q)
# 先算INT8乘法,最后乘scale
# x转INT8 (激活量化)
x_scale = np.max(np.abs(x)) / 127.0
x_q = np.round(x / x_scale).astype(np.int8)
# INT8矩阵乘(这里用int32累加,防止溢出)
result_q = x_q.astype(np.int32) @ W_q.T.astype(np.int32)
# 反量化
result = result_q.astype(np.float32) * (x_scale * scale)
if bias is not None:
result += bias
return result关键点:
- 对称量化:zero_point=0,简化计算
- 激活也量化:动态量化,每个batch计算scale
- INT8矩阵乘:现代CPU/GPU都有专用指令(VNNI, DP4A)
🌊 算法四: SSM/Mamba - 线性复杂度的序列建模
什么是SSM?
状态空间模型(State Space Model): 用状态压缩历史信息,实现O(N)复杂度。
核心思想:
- Transformer: 每个token看所有历史 → O(N²)
- RNN: 用固定大小隐状态 → 梯度消失
- SSM: 可学习的状态更新规则 → 线性时间 + 长程记忆
连续系统:
dx/dt = A·x + B·u (状态方程)
y = C·x (输出方程)
x: 隐状态 [N维]
u: 输入 (1个token)
y: 输出
A,B,C: 可学习参数
离散化(处理token序列):
x_k = A_bar · x_{k-1} + B_bar · u_k
y_k = C · x_k
A_bar = exp(Δ·A) ≈ I + Δ·A (Euler近似)
B_bar = Δ·B
对话8: SSM vs Attention,谁更强?
🤔 学生: SSM既然线性复杂度,为什么GPT还用Attention?
💡 导师: 因为表达力 vs 效率的权衡!
Attention优势(表达力):
- 全局连接:任意token pair可直接交互
- 内容寻址:动态决定关注谁(Query-Key匹配)
- 并行训练:所有位置同时算
SSM优势(效率):
- 线性时间:O(N) vs O(N²)
- 常数内存:状态大小固定,不随N增长
- 流式推理:不需要缓存历史K,V
任务差异:
| 任务 | 需要什么 | 谁更好 |
|---|---|---|
| 检索(找文中某个名字) | 全局搜索 | Attention |
| 复制(重复输入) | 精确记忆 | Attention |
| 语言建模(预测下个词) | 局部语法 + 长程主题 | 都可以 |
| 超长文档(100k tokens) | 效率 | SSM |
| 时间序列预测 | 连续动态 | SSM |
Mamba的突破(2023):
- 选择性机制: A,B,C变成输入的函数 → 内容感知!
# 传统SSM: A,B,C是固定参数
x_k = A · x_{k-1} + B · u_k
# Mamba: A,B,C依赖输入
B_k = Linear(u_k) # 不同输入,不同过滤
A_k = Softplus(Linear(u_k))
x_k = A_k · x_{k-1} + B_k · u_k # 动态状态转移!现在A,B,C随内容变化 → SSM也能”看重点”了!
最新进展(2024-2026):
- Mamba-2: 半分离SSM,训练速度提升5倍
- Hybrid架构: Attention层(每N层) + SSM层(主力)
- 例: Jamba模型 (AI21 Labs)
- Mamba-3: 复数SSM + 硬件优化,逼近Transformer精度
✨ 关键洞察: SSM vs Attention不是”谁替代谁”,而是”在不同约束下的最优解”。Attention是”全知全能但昂贵”,SSM是”压缩记忆但高效”。
对话9: SSM的状态x到底记住了什么?
🤔 学生: SSM的状态维度才几十,怎么可能记住成千上万个token的信息?
💡 导师: 因为状态x是压缩表示,不是逐字记忆!
类比: 人类记忆
- 你看完一本书,不会记住每个字(那需要GB内存)
- 你记住的是主题、情节、情感(几百KB?)
- 需要时,从压缩表示”重建”细节
SSM的状态也是:
输入序列: "The cat sat on the mat"
Attention做法(显式存储):
K,V_cache = 存每个token的表示 [6 × 768 = 4.6k floats]
SSM做法(压缩):
x = [x1, x2, ..., x_64] (64维状态)
x 编码了什么?
- x[0]: 可能是"是否提到动物"
- x[5]: 可能是"语法主语-动词-宾语结构"
- x[30]: 可能是"当前句子长度"
- x[60]: 可能是"情感倾向"
但这些都是连续值,不是离散标签!
状态更新过程:
初始: x = [0, 0, ..., 0]
读"The": x变成 [0.1, 0.02, ..., 0.5] (编码"定冠词开头")
读"cat": x变成 [0.8, 0.05, ..., 0.3] (编码"动物主语")
读"sat": x变成 [0.6, 0.9, ..., 0.1] (编码"过去时动作")
...
每步:x = A·x_old + B·new_token
A: 遗忘/保持旧信息
B: 融入新信息
为什么维度少也够用?
1. 维度诅咒的反面:
- 64维空间的”容量” = 2
- 连续值表示比离散符号表示信息密度高得多
2. 任务只需”充分统计量”:
- 预测下个词,不需要知道历史每个字
- 只需要: 当前主题 + 语法状态 + 最近几个词
- 这些信息可以用几十维编码
3. 训练目标强制压缩:
Loss = CrossEntropy(predict_next_token(x), true_next_token)
模型学到:x必须包含"预测下个词的最小充分信息"
- 无关信息被丢弃(因为增加x维度也不降低loss)
- 冗余信息被合并
- 最终x是最优压缩
数学类比:主成分分析(PCA)
- PCA用50维可以表示1000维数据,损失很小
- SSM的状态就是”动态PCA”,实时压缩历史
✨ 关键洞察: SSM的状态不是”记忆每个词”,而是”学会提取任务相关的统计特征”。容量来自表示的连续性,不是维度本身。
🎯 其他重要优化(精简版)
5. RoPE - 旋转位置编码
问题: 绝对位置编码(learned position embedding)外推性差 方案: 用旋转矩阵编码相对位置信息
核心思想:
# 传统: Embedding + PositionEmbedding
x = token_emb + pos_emb
# RoPE: 旋转Query和Key
def apply_rope(x, position):
θ = position / 10000^(2i/d) # 频率
rotation_matrix = [[cos(θ), -sin(θ)],
[sin(θ), cos(θ)]]
return rotation_matrix @ x
Q_rope = apply_rope(Q, pos)
K_rope = apply_rope(K, pos)
Attention = Q_rope @ K_rope.T # 自然编码相对位置!优势:
- 外推到更长序列
- 不增加参数
- 旋转不变性
应用: Llama, GPT-Neo, PaLM
6. GQA - 分组查询注意力
问题: Multi-Head Attention中,KV-Cache太大(每个头独立K,V) 方案: 多个Query头共享一组KV
例子:
标准MHA:
32个头 → 32组K,V → KV-Cache = 32 × [seq, d_head]
GQA (4组):
32个Q头,分8组 → 每组4个Q头共享1组KV
KV-Cache = 4 × [seq, d_head] (节省8倍!)
权衡:
- MHA: 每头独立表示 → 表达力强但贵
- MQA: 所有头共享KV → 便宜但弱
- GQA: 中间方案,性能接近MHA,成本接近MQA
应用: Llama-2, Mistral
7. Speculative Decoding - 推测解码
问题: 大模型生成慢(每个token都要跑整个模型) 方案: 小模型快速生成多个候选,大模型批量验证
流程:
1. 小模型(快): "I think the answer is probably" (生成5个token)
2. 大模型(慢): 验证5个token,接受前3个,拒绝后2个
3. 小模型继续从第3个token后生成
4. 重复...
加速原理:
- 大模型验证5个token = 1次forward (并行)
- 如果逐个生成 = 5次forward
- 若接受率60%, 平均加速 = 3x
应用: Medusa, SpecInfer
8-11. 其他技术速览
| 技术 | 核心思想 | 加速/节省 |
|---|---|---|
| PagedAttention | 虚拟内存管理KV-Cache | 减少碎片,提升显存利用 |
| Beam Search | 维护top-k路径 | 生成质量提升 |
| Discretization | SSM的连续→离散转换 | Euler/Tustin方法 |
| Complex SSM | 用复数状态空间 | SSM表达力提升 |
📊 推理优化技术对比
| 技术 | 优化目标 | 加速倍数 | 内存节省 | 精度损失 |
|---|---|---|---|---|
| Flash Attention | 内存+速度 | 2-4x | O(N²)→O(N) | 0% |
| KV-Cache | 生成速度 | 10-100x | +4MB/token | 0% |
| Quantization | 模型大小 | 2-4x | 4x(INT8) | ~1% |
| SSM/Mamba | 序列长度 | 3-10x(长序列) | 常数 | ~2% |
| GQA | KV-Cache内存 | 1x | 4-8x | <1% |
| Speculative Decoding | 生成速度 | 2-3x | 0 | 0% |
🎯 本章总结
推理优化的核心矛盾:速度 vs 质量 vs 内存
不同场景选择:
| 场景 | 推荐技术组合 |
|---|---|
| 云端API(A100) | Flash Attention + KV-Cache + Speculative Decoding |
| 边缘设备(手机) | INT4量化 + Distillation + 小模型 |
| 超长文档(128k) | Mamba/SSM + Sparse Attention |
| 实时对话 | KV-Cache + Beam Search |
| 批量推理 | PagedAttention + INT8量化 |
核心原则:
- 先Profile后优化:找瓶颈(内存?计算?IO?)
- 组合技术:Flash + KV + 量化可叠加
- 权衡取舍:1%精度换10倍速度,值吗?
下一站: 05-生成模型轨道 - 探索VAE, GAN, Diffusion的创造力!