大模型推理优化：从自回归生成到 KV Cache 与量化

从 GPT 自回归生成开始，逐步讲清楚 Prefill、Decode、KV Cache、GQA、分页式缓存管理、INT8 量化、校准集和离群值。

预计阅读 7 分钟 最后更新
  • AI
  • 大模型
回到卷帙

大模型推理优化：从自回归生成到 KV Cache 与量化

前面的文章已经讲过 GPT 的基本路径：

文本 -> token -> embedding -> attention -> logits -> 下一个 token

如果还没有读，可以先看：

这篇文章继续往工程方向走：模型结构理解了以后，它真正推理时为什么慢，为什么占显存，为什么要做 KV Cache、Prefill、Decode、量化和缓存管理。

先从一句话生成开始

假设 prompt 是：

我喜欢

GPT 会预测下一个 token。假设预测出：

学习

下一步输入就变成：

我喜欢学习

模型再预测下一个 token。

这个过程叫自回归生成：

已有 token -> 预测下一个 token -> 拼回序列 -> 继续预测

它有一个天然限制：第 10 个 token 必须等第 9 个 token 生成后才能预测。

所以生成阶段不像训练那样可以完全并行。

朴素推理会重复计算

在 Attention 里，每个 token 会产生 Query、Key、Value。

公式是：

Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

生成新 token 时，历史 token 已经固定了。比如已经生成：

我 喜欢 学习

下一步要预测第 4 个 token。历史三个 token 的 Key 和 Value 不会改变。

如果每一步都从头算一遍历史 token 的 K/V，就会浪费大量计算。

这时自然引出 KV Cache。

KV Cache 的直觉

KV Cache 的思想很简单：

历史 token 的 Key 和 Value 算过一次后缓存起来，后面直接复用。

生成第 tt 个 token 时，只需要：

  1. 为当前新 token 计算新的 Q,K,VQ,K,V。
  2. 把新的 K,VK,V 追加到缓存。
  3. 用当前 QQ 去读取缓存里的所有历史 K,VK,V。

流程可以画成：

flowchart LR
  P[Prompt tokens] --> A[Prefill: 一次性计算 prompt 的 K/V]
  A --> C[(KV Cache)]
  C --> D[Decode step]
  D --> E[计算新 token 的 Q/K/V]
  E --> F[追加新 K/V 到缓存]
  F --> C
  E --> G[当前 Q 读取历史 K/V]
  G --> H[输出下一个 token]

KV Cache 没有改变模型数学结构，它只是避免重复计算历史 K/V。

Prefill：先处理完整 prompt

推理通常分成两个阶段：Prefill 和 Decode。

Prefill 阶段负责处理用户输入的完整 prompt。

特点：

  • 输入 token 多。
  • 可以并行处理 prompt 中所有 token。
  • 会生成每一层的初始 KV Cache。
  • 主要影响首 token 延迟，也就是 Time To First Token。

例如用户输入 2000 个 token，模型要先读完这 2000 个 token，建立上下文表示和 KV Cache，然后才能开始输出第一个 token。

Decode：逐 token 生成

Decode 阶段负责一个 token 一个 token 往外生成。

特点：

  • 每一步通常只处理一个新 token。
  • 必须依赖上一步生成结果。
  • 每一步都会读取历史 KV Cache。
  • 很容易受显存带宽影响。

如果用户要求生成 500 个 token，decode 就要循环很多次。

所以优化推理时要分清：

指标主要受什么影响
TTFT，首 token 延迟Prefill
TPOT，每输出 token 延迟Decode
吞吐batch 调度、显存带宽、缓存管理

KV Cache 为什么占显存

KV Cache 是用显存换计算。

它的大小大致和下面这些量成正比：

batch_size × num_layers × seq_len × num_kv_heads × head_dim × 2 × bytes_per_element

其中乘以 2 是因为同时存 Key 和 Value。

如果使用 FP16，每个元素 2 字节：

KV bytes=B×L×S×Hkv×D×2×2\text{KV bytes} =B\times L\times S\times H_{kv}\times D\times 2\times 2

含义是：

  • BB：batch size。
  • LL：层数。
  • SS：序列长度。
  • HkvH_{kv}：KV head 数。
  • DD：每个 head 的维度。

这解释了为什么长上下文推理很吃显存：序列越长，KV Cache 线性增长。

GQA 为什么能降低缓存压力

在 ROPE与GQA 里已经讲过 GQA。这里从推理角度再看一次。

传统 MHA 中，每个 attention head 都有自己的 K/V。这样效果强，但 KV Cache 大。

MQA 让所有 query head 共享一组 K/V，缓存最小，但表达能力可能下降。

GQA 介于两者之间，让一组 query head 共享一组 K/V。

结构KV head 数KV Cache 压力取舍
MHA多高表达能力强
MQA1低压缩最激进
GQA中间值中常用折中

所以 GQA 不只是模型结构技巧，也是推理显存优化手段。

分页式 KV 管理

真实服务里，请求长度不一样，生成长度也不一样。

如果每个请求都提前分配一整块最大长度缓存，会浪费显存；如果频繁申请和释放不同大小的缓存，又容易产生碎片。

分页式或 block 式 KV 管理的思路是：

  • 把 KV Cache 切成固定大小的块。
  • 请求需要更多上下文时，再追加块。
  • 请求结束后，把块还给缓存池。
  • 调度 batch 时按块管理不同请求的历史上下文。

它解决的是服务层面的内存管理问题，不改变 Attention 公式。

量化：让模型更小、更省带宽

除了 KV Cache，模型权重本身也占大量显存。量化就是把高精度数值转换成低精度表示。

例如从 FP16 变成 INT8。

一个简单线性量化可以写成：

xint8=round(xs)x_\text{int8}=\text{round}\left(\frac{x}{s}\right)

反量化时：

x^=s⋅xint8\hat{x}=s\cdot x_\text{int8}

这里 ss 是 scale。

量化的收益：

  • 模型占用更小。
  • 权重读取带宽更低。
  • 某些硬件上 INT8 计算更快。
  • 可以容纳更大 batch 或更长上下文。

为什么量化会损失精度

量化是用有限整数刻度近似连续浮点数，所以会有误差。

影响精度的关键因素有：

  1. 激活范围是否稳定。
  2. 校准集是否接近真实输入。
  3. 是否存在离群值。
  4. 哪些层保留浮点计算。
  5. 只量化权重，还是权重和激活都量化。

离群值很麻烦。假设大多数激活在 [−1,1][-1,1]，但少数值到了 20。为了覆盖 20，scale 会变大，导致 [−1,1][-1,1] 内的大量普通值可用刻度变少。

这就是为什么大模型量化经常要特殊处理 outlier，或者保留部分敏感层为 FP16 / BF16。

校准集为什么重要

后训练量化通常需要校准集。

校准集的作用是估计真实推理时的激活范围。

如果校准集和真实输入差异很大，scale 就会估错：

  • 激活可能被截断。
  • 刻度可能浪费。
  • 下游任务精度下降。

所以校准集必须尽量贴近真实业务数据。

这个思想和 机器学习建模 里讲的数据分布问题是一致的：训练、校准、测试和真实环境的分布差异，都会影响模型表现。

推理优化是在做取舍

推理优化不是越激进越好。它总是在延迟、吞吐、显存和精度之间取舍。

方法收益代价
KV Cache减少重复计算增加显存
GQA / MQA减少 KV 显存可能影响表达能力
INT8 量化降低模型大小和带宽可能损失精度
分页式 KV减少碎片，提升调度系统实现更复杂
动态 batch提升吞吐单请求延迟可能波动

交互式聊天更关心首 token 延迟和每 token 延迟；批量离线生成更关心吞吐；长文档问答更关心上下文长度和 KV Cache 显存。

总结

大模型推理优化可以按这条线理解：

自回归生成很慢
  -> 历史 K/V 会重复计算
  -> KV Cache 用显存换计算
  -> Prefill 和 Decode 拆分不同阶段
  -> GQA/MQA 降低 KV 压力
  -> 分页式管理减少显存碎片
  -> INT8 量化降低模型和带宽成本

这篇文章不是单独讲部署技巧，而是接在 GPT 和 Attention 之后：只有先理解 token 如何通过 Attention 读取上下文，才能真正理解为什么 KV Cache 是推理系统里的核心结构。

尺素

0 封尺素

正在取尺素...