06 推理优化——KV Cache、量化与投机解码

摘要:

训练一个大语言模型需要数百万美元,但真正的成本大头在推理——一个被数百万用户使用的模型,每天处理数十亿 token,推理成本远超一次性的训练成本。本文深入剖析 LLM 推理阶段的核心优化技术:KV Cache 消除自回归生成中的重复计算,权重量化(GPTQ/AWQ/GGUF)将模型压缩到 4-bit 甚至更低,投机解码(Speculative Decoding)用小模型加速大模型的生成,以及 Continuous BatchingPagedAttention 等系统级优化。这些技术共同将 LLM 的推理吞吐量提升数倍到数十倍,使大规模部署成为可能。


第 1 章 LLM 推理的两个阶段

1.1 Prefill 与 Decode

LLM 的推理过程分为两个截然不同的阶段:

Prefill(预填充)阶段:模型一次性处理整个输入 prompt。所有 token 同时参与 Self-Attention 计算,充分利用 GPU 的并行能力。这个阶段是计算密集型的——GPU 的算力是瓶颈。

Decode(解码)阶段:模型逐个 token 生成输出。每一步只生成一个 token,需要与之前所有 token 的 Key/Value 做注意力计算。这个阶段是访存密集型的——每一步的计算量很小(只处理一个 token),但需要从显存中读取大量的 KV Cache 数据,GPU 的显存带宽是瓶颈。

阶段处理 token 数瓶颈GPU 利用率占总推理时间
Prefill所有输入 token(并行)计算(FLOPs)~10-30%
Decode每步 1 个 token(串行)显存带宽(Memory Bandwidth)~70-90%

Decode 阶段占据了绝大部分推理时间,且 GPU 利用率极低——这就是 LLM 推理优化的核心目标:提升 Decode 阶段的效率

1.2 推理的核心指标

  • TTFT(Time To First Token):从输入到生成第一个 token 的延迟,主要由 Prefill 阶段决定
  • TPOT(Time Per Output Token):每个输出 token 的生成时间,由 Decode 阶段决定
  • 吞吐量(Throughput):每秒处理的总 token 数(输入 + 输出),通常以 tokens/s 衡量
  • 总延迟(Latency)

不同应用场景对指标的优先级不同:聊天应用优先 TTFT(用户期望快速看到首字);批量处理优先吞吐量(总处理速度)。


第 2 章 KV Cache——消除重复计算

2.1 KV Cache 的动机

02 GPT 架构——Decoder-Only 的自回归语言模型 中我们知道,Decode 阶段每一步生成一个 token,该 token 需要与之前所有 token 做 Self-Attention 计算。

如果没有 KV Cache,每生成一个新 token,模型都需要重新计算整个序列(包括所有已生成的 token)的 Key 和 Value——这意味着第 步需要处理 个 token,总计算量为 为总序列长度)。这显然是极度浪费的——前 个 token 的 Key 和 Value 在第 步已经计算过了,完全相同。

KV Cache 的核心思想:将每一步计算出的 Key 和 Value 缓存在 GPU 显存中,后续步骤直接复用,无需重复计算。

2.2 KV Cache 的工作机制

在 Prefill 阶段,模型处理完整个输入序列后,将所有层的 Key 和 Value 矩阵缓存起来。

在 Decode 阶段的第 步:

  1. 只对新生成的 token 计算 (通过一次线性投影)
  2. 追加到缓存中:
  3. 计算 与整个 的注意力:

这样每步只需计算一个 token 的 Q/K/V 投影(而非整个序列),将 Decode 的计算量从 降低到

2.3 KV Cache 的显存消耗

KV Cache 的代价是显存——需要为每个请求的每个 token 存储所有层的 Key 和 Value。

对于一个 LLaMA-7B 模型(32 层,32 头,,BF16):

其中 代表 K 和 V 两个矩阵,最后的 代表 BF16 的 2 字节。

一个 4096 token 长度的序列的 KV Cache 约为 。如果同时服务 32 个并发请求,KV Cache 就需要 64 GB——超过了单块 A100 的显存。

模型KV Cache / token4K 上下文32K 上下文128K 上下文
LLaMA-7B512 KB2 GB16 GB64 GB
LLaMA-70B (GQA, 8 KV头)640 KB2.5 GB20 GB80 GB
Mistral-7B (GQA, 8 KV头)128 KB0.5 GB4 GB16 GB

这正是 GQA(Grouped-Query Attention)如此重要的原因——它通过减少 K/V 头的数量,将 KV Cache 大小降低数倍。

核心概念

KV Cache 是 LLM 推理中”以空间换时间”的经典策略。它将 Decode 阶段从”计算密集型”转变为”访存密集型”——瓶颈从 GPU 算力转移到了显存带宽和显存容量。后续的大量优化技术(PagedAttention、量化 KV Cache、GQA)都是围绕 KV Cache 的显存压力展开的。


第 3 章 权重量化——模型压缩

3.1 量化的基本原理

量化(Quantization)是将模型权重从高精度浮点数(如 FP16/BF16,16 bit)转换为低精度整数(如 INT8 或 INT4),从而减小模型体积和加速推理。

最简单的量化方式是线性量化:将 FP16 值映射到 INT8 区间

其中 是量化位数。反量化时:

量化带来的好处是双重的:

  • 模型体积减半到四分之一:FP16 → INT8 减半,FP16 → INT4 减到四分之一
  • 推理加速:更小的模型意味着更少的显存读取,在访存密集的 Decode 阶段直接转化为速度提升

3.2 训练后量化(Post-Training Quantization)

训练后量化不需要重新训练模型,只需要一小批校准数据来确定量化参数(scale 和 zero_point)。

GPTQ(Generative Pre-trained Transformer Quantization)

GPTQ(Frantar et al., 2023)是最早在 LLM 上成功实现高质量 4-bit 量化的方法。核心思想是逐层量化:一次处理一个线性层,利用最优脑损伤(Optimal Brain Damage)的思想,按照每个参数对输出误差的影响大小,选择最优的量化顺序和补偿方式。

GPTQ 使用少量校准数据(通常 128 条文本)计算每一层的 Hessian 矩阵(衡量参数重要性),然后按照重要性从低到高量化参数,每量化一个参数就调整其他未量化参数来补偿误差。

AWQ(Activation-aware Weight Quantization)

AWQ(Lin et al., 2024)观察到一个关键现象:不是所有权重同等重要——少数”关键权重”对应着激活值中的少数”关键通道”。这些关键通道的激活值幅度特别大,对应权重的量化误差会被放大。

AWQ 的策略是:对关键权重(对应高激活值通道)乘以一个大于 1 的缩放因子,使其值范围更大、量化精度更高;对不关键的权重缩小,牺牲其精度换取整体更好的量化效果。

GGUF(GPT-Generated Unified Format)

GGUF 是 llama.cpp 生态使用的量化格式,支持 CPU 推理。它使用混合精度量化——不同层可以使用不同的量化位数(如注意力层用 8-bit,FFN 用 4-bit),在精度和体积之间灵活权衡。

GGUF 的命名约定如 Q4_K_M 表示:4-bit 量化 + K-quants 分组方法 + Medium 精度配置。

量化方法量化位数校准数据需求推理框架速度质量
GPTQ4-bit / 8-bit少量(128 条)ExLlama, vLLM快(GPU)
AWQ4-bit少量vLLM, TensorRT-LLM快(GPU)略优于 GPTQ
GGUF2-8 bit 混合llama.cpp中(CPU/GPU)灵活
BitsAndBytes4-bit (NF4) / 8-bitTransformers中(GPU)好(QLoRA 使用)

3.3 量化对质量的影响

量化不可避免地引入精度损失,但不同位数的影响差异巨大:

  • 8-bit 量化:几乎无损——在大多数基准上与 FP16 的差异小于 1%
  • 4-bit 量化:轻微损失——在知识密集型任务(如 MMLU)上可能下降 1-3%,在推理任务上影响更小
  • 3-bit 量化:明显损失——质量开始显著下降,仅适用于对精度要求不高的场景
  • 2-bit 量化:严重损失——当前技术下质量不可接受

生产避坑

量化对不同能力的影响不均匀。经验表明,量化对”事实性知识”的损伤大于对”推理能力”的损伤——因为事实知识依赖于权重中存储的精确信息,而推理能力更多依赖于模型的结构性模式。因此,如果你的应用场景对事实准确性要求极高(如医疗、法律),建议使用 8-bit 量化而非 4-bit。

3.4 KV Cache 量化

除了权重量化,KV Cache 也可以被量化以节省显存。KV Cache 量化通常使用 8-bit(FP8 或 INT8),对质量的影响很小,但可以将 KV Cache 的显存占用减半。

vLLM 和 TensorRT-LLM 都支持 KV Cache 的 FP8 量化。


第 4 章 投机解码——用小模型加速大模型

4.1 自回归生成的根本瓶颈

LLM 推理的根本瓶颈在于自回归生成的串行性——每一步只生成一个 token,但每步都需要从显存读取完整的模型权重。对于大模型,每步的计算量很小(处理一个 token),但显存读取量很大(读取整个模型),GPU 的算力严重闲置。

从硬件指标看,A100 GPU 的峰值算力为 312 TFLOPS(BF16),但显存带宽只有 2 TB/s。一个 7B 模型(14 GB,BF16)每步需要读取全部参数,耗时约 ms。但这一步的实际计算量极小——大量时间花在”等数据从显存传到计算单元”上。

这个瓶颈被称为内存墙(Memory Wall):GPU 的计算能力远超数据供给能力。

4.2 投机解码的核心思想

投机解码(Speculative Decoding,Leviathan et al., 2023; Chen et al., 2023)的核心思想来自一个观察:大模型生成的很多 token 其实是”简单”的——小模型也能正确预测

比如”The capital of France is Paris”这个输出,前几个词(“The capital of”)几乎任何语言模型都能正确生成。只有关键信息词(“France”、“Paris”)可能需要大模型的能力。

投机解码的策略是:

  1. 用一个小型草稿模型(Draft Model,如 68M 参数)快速生成 个候选 token( 通常为 4-8)
  2. 将这 个候选 token 一次性送入大型目标模型(Target Model)进行验证——大模型在一次前向传播中同时处理 个位置
  3. 从左到右检查:如果草稿 token 与大模型的预测一致(或满足某个接受准则),就接受;遇到第一个不一致的位置就拒绝,从大模型的分布中重新采样

sequenceDiagram
    participant Draft as "草稿模型 (68M)"
    participant Target as "目标模型 (70B)"
    participant Output as "最终输出"

    Draft->>Draft: "快速生成 K=4 个 token"
    Draft->>Target: "候选: [The, capital, of, France]"
    Target->>Target: "一次前向传播验证 4 个位置"
    Target->>Output: "接受 3 个: [The, capital, of]"
    Target->>Output: "拒绝第 4 个, 从目标分布采样"

关键性质:投机解码的输出分布与直接用大模型生成完全相同——它不是近似,而是精确的。接受-拒绝机制保证了这一点。

4.3 加速比分析

投机解码的加速比取决于两个因素:

  • 接受率 :草稿 token 被大模型接受的概率。越高越好—— 意味着平均每 5 个草稿 token 有 4 个被接受
  • 草稿模型速度与大模型速度的比值:草稿模型越快越好

理论加速比约为:

在实践中,投机解码通常能获得 2-3 倍的加速。接受率高度依赖于任务类型——模板化的输出(如代码、格式化文本)接受率高,创意性输出(如故事续写)接受率低。

4.4 Medusa——无需草稿模型的投机解码

Medusa(Cai et al., 2024)避免了使用独立的草稿模型——它在大模型本身上添加多个轻量级的”Medusa Head”(额外的 FFN 层),每个 Head 负责预测不同未来位置的 token。

Medusa 的优势是不需要独立的草稿模型(避免了草稿模型与目标模型之间的分布差异问题),但需要额外的训练来学习 Medusa Head。


第 5 章 系统级优化

5.1 Continuous Batching

传统的批处理(Static Batching)将多个请求打包成一个 batch 同时处理。问题是:不同请求的输出长度不同,短请求生成完毕后必须等长请求完成,GPU 在等待期间空闲。

Continuous Batching(也叫 Iteration-level Scheduling)改变了这种模式:在每个 Decode 步之后,如果某个请求已完成(生成了 <eos>),立即将其替换为队列中的新请求。这样 GPU 始终在处理满载的 batch,不存在等待浪费。

Continuous Batching 可以将吞吐量提升 2-5 倍,是 vLLM、TensorRT-LLM 等推理框架的标配功能。

5.2 PagedAttention

PagedAttention(Kwon et al., 2023,vLLM 的核心技术)借鉴了操作系统的虚拟内存分页机制来管理 KV Cache。

传统方式为每个请求预分配一整块连续的 KV Cache 显存——按最大可能长度分配。如果请求的实际输出长度远小于最大长度,就会造成严重的显存碎片化和浪费。实测表明,传统方式的 KV Cache 显存浪费率高达 60-80%。

PagedAttention 的做法:

  1. 将 KV Cache 分割为固定大小的(Block,如每页 16 个 token 的 K/V)
  2. 每个请求维护一个页表(Page Table),记录逻辑页到物理页的映射
  3. 物理页在显存中不需要连续——就像操作系统的虚拟内存一样
  4. 新 token 的 K/V 写入当前页,页满后分配新页

这消除了显存碎片和预分配浪费,实际上可以将同时服务的请求数提升 2-4 倍(因为同样的显存可以容纳更多的 KV Cache)。

设计哲学

PagedAttention 的设计精妙地展示了操作系统原理在 AI 系统中的应用。虚拟内存分页是操作系统解决物理内存碎片化的经典方案,PagedAttention 将同样的思想应用于 GPU 显存中的 KV Cache 管理。这提醒我们:AI 系统工程不仅需要机器学习知识,还需要扎实的系统设计功底。

5.3 Prefix Caching

在很多应用场景中,不同请求共享相同的前缀——比如所有请求都有相同的 System Prompt,或者 RAG 场景中多个请求检索到了相同的文档片段。

Prefix Caching(前缀缓存)将共享前缀的 KV Cache 缓存并在请求之间复用,避免重复计算。vLLM 的 Automatic Prefix Caching(APC)功能可以自动检测和复用共享前缀。

对于有长 System Prompt 的聊天应用,Prefix Caching 可以将 TTFT 降低数倍。


第 6 章 推理框架生态

6.1 主流推理框架对比

框架核心特性硬件支持适用场景
vLLMPagedAttention, Continuous Batching, Prefix CachingNVIDIA/AMD GPU通用在线服务
TensorRT-LLMNVIDIA 官方优化, 高度定制的 CUDA kernelNVIDIA GPU追求极致性能
llama.cppCPU + Metal/CUDA, GGUF 量化CPU/Apple Silicon/GPU本地部署, 边缘设备
SGLangRadixAttention (高级前缀缓存), 结构化生成NVIDIA GPU复杂 prompt 工程
ExLlamaV2GPTQ 量化优化, 高效的 4-bit kernelNVIDIA GPU量化模型推理
Ollamallama.cpp 封装, 一键部署CPU/GPU本地快速体验

6.2 vLLM 的架构

vLLM 是当前最主流的 LLM 推理框架,其核心架构包括:

  • Scheduler:请求调度器,实现 Continuous Batching 和优先级管理
  • BlockManager:基于 PagedAttention 的 KV Cache 显存管理
  • Worker:GPU 上的模型执行引擎,支持张量并行
  • Tokenizer:高效的分词和解码
  • OpenAI-compatible API:提供与 OpenAI API 兼容的 HTTP 接口

vLLM 的部署只需一行命令:

vllm serve meta-llama/Llama-2-7b-chat-hf --tensor-parallel-size 2

第 7 章 实际部署的性能考量

7.1 显存预算分配

部署 LLM 时,GPU 显存需要在以下几部分之间分配:

组成部分7B (BF16)7B (INT4)70B (BF16)70B (INT4)
模型权重14 GB3.5 GB140 GB35 GB
KV Cache(每请求 4K)2 GB2 GB2.5 GB2.5 GB
其他开销~2 GB~2 GB~5 GB~5 GB
单请求总计~18 GB~7.5 GB~147.5 GB~42.5 GB

模型权重是固定开销(所有请求共享),KV Cache 是按请求分配的。因此,最大化并发请求数的关键是减少模型权重占用(量化)和高效管理 KV Cache(PagedAttention + GQA)。

7.2 延迟与吞吐量的权衡

增大 batch size 可以提升吞吐量(GPU 利用率更高),但会增加单个请求的延迟(等待 batch 凑齐的时间,以及 batch 内的请求相互影响)。

实际部署中需要根据 SLA(Service Level Agreement)来平衡:

  • 交互式聊天(低延迟优先):batch size 较小,TTFT < 500ms
  • 批量处理(高吞吐优先):batch size 尽量大,最大化 tokens/s

第 8 章 总结

8.1 优化技术分层

层次优化技术效果
算法层KV Cache, 投机解码, GQA消除冗余计算
模型层量化 (GPTQ/AWQ/GGUF)减小模型体积, 加速访存
系统层Continuous Batching, PagedAttention, Prefix Caching提升 GPU 利用率和并发
硬件层Flash Attention, Tensor Core, NVLink底层计算加速

8.2 下一步


参考文献

  1. Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention”, SOSP 2023
  2. Leviathan et al., “Fast Inference from Transformers via Speculative Decoding”, ICML 2023
  3. Frantar et al., “GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers”, ICLR 2023
  4. Lin et al., “AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration”, MLSys 2024
  5. Cai et al., “Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads”, ICML 2024
  6. Dao et al., “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning”, ICLR 2024
  7. Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, EMNLP 2023
  8. Yu et al., “ORCA: A Distributed Serving System for Transformer-Based Generative Models”, OSDI 2022

思考题

  1. KV Cache 将已生成 token 的 Key 和 Value 缓存起来,避免重复计算。对于一个 70B 参数的模型,生成 4096 个 token 时 KV Cache 占用多少显存?PagedAttention(vLLM 使用)通过’分页’管理 KV Cache——它解决了什么问题?与传统的连续内存 KV Cache 相比,PagedAttention 如何提升显存利用率?
  2. 模型量化(INT8/INT4/GPTQ/AWQ)将 FP16 权重压缩为低精度——减少显存和加速推理。但量化不可避免地带来精度损失。在什么任务类型上量化的精度损失最明显(如数学推理 vs 文本摘要)?GPTQ(逐层量化 + 校准集)和 AWQ(激活感知量化)的量化策略有什么本质区别?
  3. 投机解码(Speculative Decoding)使用一个小模型快速生成候选 token 序列,然后用大模型一次性验证。如果小模型的预测准确率为 80%,投机解码的加速比大约是多少?小模型和大模型的’对齐度’如何影响加速效果?在什么场景下投机解码的收益最大?