大语言模型推理优化：从 KV Cache 到投机解码

引言：推理阶段的性能挑战

大语言模型（LLM）的推理阶段往往比训练阶段更具挑战性。训练可以离线进行，而推理需要实时响应用户请求。随着模型规模增长到数百亿甚至数千亿参数，单次推理的计算量和内存占用急剧增加，如何在有限的 GPU 资源下提升吞吐量和降低延迟，成为工程落地的关键问题。

推理优化的核心目标有两个：降低延迟（Latency）和提升吞吐量（Throughput）。延迟影响用户体验，吞吐量决定服务成本。本文将深入探讨当前主流的 LLM 推理优化技术。

一、KV Cache：推理加速的基础

1.1 为什么需要 KV Cache？

Transformer 模型在生成每个新 token 时，需要计算当前 token 与所有历史 token 的注意力。如果不使用缓存，每生成一个 token 都要重新计算所有历史 token 的 Key 和 Value 向量，计算复杂度为 O(n²)，其中 n 是序列长度。

KV Cache 的核心思想是：将已经计算过的 Key 和 Value 向量缓存起来，在生成新 token 时直接复用，避免重复计算。这样每步生成只需计算当前 token 的 K、V，然后与历史缓存拼接，计算复杂度降为 O(n)。

1.2 KV Cache 的内存开销

KV Cache 虽然提升了计算效率，但带来了显著的内存开销。缓存大小可以通过以下公式估算：

KV Cache 大小 = 2 × 层数 × 头数 × 头维度 × 序列长度 × 批大小 × 数据类型字节数

以 LLaMA-2-70B 为例：80 层、64 个头、128 维头维度、bf16 精度（2 字节），单个请求的 KV Cache 约为 2 × 80 × 64 × 128 × 2 = 2.6 MB/token。如果序列长度为 4096，单个请求需要约 10.7 GB 显存，这甚至超过了模型权重本身的占用。

1.3 KV Cache 优化策略

• MQA（Multi-Query Attention）：所有注意力头共享同一组 K、V，大幅减少缓存大小（如 Falcon、PaLM 模型）
• GQA（Grouped-Query Attention）：将注意力头分组，每组共享 K、V，在性能与缓存大小间取得平衡（如 LLaMA-2-70B、Mistral）
• KV Cache 量化：将 KV Cache 从 FP16 量化到 INT8 甚至 INT4，减少内存占用和带宽压力
• KV Cache 驱逐：当缓存超过阈值时，丢弃不重要的历史 KV（如 H2O、StreamingLLM）

二、投机解码：用小模型加速大模型

2.1 投机解码的核心思想

自回归生成是逐 token 串行的，每步都需要完整的模型前向传播。投机解码（Speculative Decoding）提出了一种巧妙的加速方法：使用一个小的"草稿模型"快速生成多个候选 token，然后用大模型并行验证这些 token 的正确性。

如果草稿模型的预测准确，大模型可以一次前向传播验证多个 token，从而加速生成；如果预测错误，则回退到常规解码。这种方法在保持输出质量（与大模型完全一致）的前提下，理论上可以带来 2-3 倍的加速。

2.2 投机解码的数学原理

投机解码基于拒绝采样（Rejection Sampling）理论。设草稿模型为 q，目标模型为 p，生成 k 个候选 token 的过程如下：

1. 草稿模型自回归生成 k 个候选 token：x₁, x₂, ..., xₖ
2. 目标模型并行计算这 k 个位置的概率分布 p(x|context)
3. 对于每个位置 i，以概率 min(1, p(xᵢ)/q(xᵢ)) 接受该 token
4. 如果拒绝，从修正后的分布中重新采样
5. 一旦某个位置被拒绝，后续候选全部丢弃，从当前位置继续

这种机制保证了最终输出与目标模型的分布完全一致，不会引入质量损失。

2.3 投机解码的实践挑战

• 草稿模型选择：需要与目标模型"足够相似"才能有高接受率。通常使用同系列的小模型（如 LLaMA-2-7B 作为 LLaMA-2-70B 的草稿模型）
• 接受率与加速比：接受率越高，加速效果越好。但草稿模型太小会降低接受率，太大则减少加速收益
• 动态调整：可以根据上下文难度动态调整草稿长度 k，简单场景用更长的草稿，复杂场景用更短的草稿

三、连续批处理：提升吞吐量

3.1 传统批处理的局限

传统 LLM 推理服务采用静态批处理（Static Batching）：等待一批请求到达后，一起送入模型处理，整批全部完成后才返回结果。这种方法存在两个问题：

• 队头阻塞：如果某个请求生成很长，其他短请求必须等待
• 资源浪费：批次中所有请求必须等到最长的完成，GPU 利用率低

3.2 连续批处理（Continuous Batching）

连续批处理（也称为 In-flight Batching 或 Dynamic Batching）的核心思想是：在生成过程中动态调整批次。当一个请求完成生成后，立即从等待队列中取一个新请求加入当前批次，而不需要等待整个批次完成。

这种方法的优势：

• 消除队头阻塞：短请求完成后立即释放资源，长请求不影响其他请求
• 提高 GPU 利用率：批次始终保持满载，减少空闲时间
• 降低平均延迟：新请求可以更快被调度，无需等待批次凑齐

3.3 PagedAttention：内存管理的革新

vLLM 提出的 PagedAttention 技术借鉴了操作系统的虚拟内存分页机制。传统 KV Cache 管理采用连续内存分配，导致严重的内存碎片和浪费。PagedAttention 将 KV Cache 分割成固定大小的"页"（Page），每个页存储固定数量的 token 的 KV 向量：

• 按需分配：只分配实际需要的页，避免预分配造成的浪费
• 非连续存储：页可以分散在内存中，通过页表索引，消除碎片
• 共享与复制：多个请求可以共享相同的页（如共享 prompt），写入时才复制（Copy-on-Write）

PagedAttention 通常可以将 GPU 利用率提升 2-4 倍，是 vLLM 高性能的核心秘诀。

四、量化与压缩：减少计算与内存

4.1 权重量化

将模型权重从 FP16/BF16 量化到 INT8、INT4 甚至更低精度，可以显著减少内存占用和计算量：

• INT8 量化：内存减半，精度损失较小，适合大多数场景
• INT4 量化：内存降至 1/4，精度损失较明显，需要量化感知训练或后训练量化（PTQ）技术
• GPTQ、AWQ、GGUF：主流的量化方案，各有优劣

4.2 稀疏化与剪枝

研究表明，LLM 中存在大量冗余参数。通过剪枝移除不重要的权重或注意力头，可以在几乎不损失性能的情况下减少模型大小：

• 结构化剪枝：移除整个注意力头或前馈层，保持硬件友好性
• 非结构化剪枝：移除单个权重，需要稀疏计算库支持
• 动态稀疏：根据输入动态激活部分参数（如 MoE 混合专家模型）

五、工程实践：主流推理框架

5.1 vLLM

vLLM 是目前最流行的 LLM 推理框架，核心特性包括：

• PagedAttention 内存管理
• 连续批处理
• 支持多种量化方案
• 兼容 Hugging Face 模型生态
• 支持张量并行和多 GPU 推理

5.2 TensorRT-LLM

NVIDIA 官方推出的高性能推理框架，针对 NVIDIA GPU 深度优化：

• 高度优化的 CUDA 内核
• 支持 FP8、INT4 等低精度推理
• 与 Triton Inference Server 集成
• 支持投机解码

5.3 Hugging Face TGI 与 Transformers

Hugging Face 提供的推理工具链，适合快速部署和实验：

• Transformers 库：模型加载与基础推理
• Text Generation Inference (TGI)：生产级推理服务
• Optimum：针对特定硬件的优化版本

六、未来趋势

• 多模态推理优化：随着视觉-语言模型的普及，需要针对图像、视频等模态设计专门的优化策略
• 长上下文推理：支持 100K+ token 的上下文窗口，需要更高效的注意力机制和内存管理
• 边缘端推理：在移动设备、嵌入式设备上运行 LLM，对模型压缩和量化提出更高要求
• 推理与训练融合：在线学习、持续学习的场景下，推理和训练需要更紧密的协同

结语

LLM 推理优化是一个快速发展的领域，从 KV Cache 到投机解码，从连续批处理到 PagedAttention，每一项技术都在推动大模型从实验室走向生产环境。理解这些优化技术的原理，不仅能帮助开发者选择合适的工具和配置，更能启发我们思考如何设计更高效的模型架构。

随着硬件技术的进步和算法创新，我们有理由相信，未来 LLM 推理将变得更加高效、便捷，真正成为人人可用的基础设施。