大语言模型推理优化:从 KV Cache 到投机解码的系统化实践

随着大语言模型参数量从百亿迈向万亿,推理效率已成为决定 AI 服务可行性的核心瓶颈。本文系统化梳理了 LLM 推理优化的三大核心方向——内存管理、解码加速与系统工程——并深入剖析了 PagedAttention、投机解码、连续批处理等关键技术的原理与实现。

一、推理瓶颈的根源

大语言模型的自回归生成特性决定了其推理过程具有两个根本性挑战:内存密集与计算密集。

以 Llama-2-70B 模型为例,仅模型权重就占用约 140GB(FP16)。更关键的是,自回归解码过程中需要维护一个 KV(Key-Value)缓存,其大小随序列长度线性增长。对于 batchSize=32、seqLen=4096 的场景,KV Cache 可膨胀至数十 GB,成为显存的主要消费者。

从计算角度看,自回归解码的每一步生成一个 token,每步都需要访问全部模型权重。这种"每 token 权重访问数"(arithmetic intensity)极低的特点,使得 LLM 推理通常是内存带宽受限(memory-bandwidth bound)而非计算受限(compute bound)。

由此引出三个核心优化方向:

  1. 内存管理:压缩 KV Cache,提高显存利用率
  2. 解码加速:减少生成每个 token 的计算开销
  3. 系统工程:提高吞吐量和并行效率

二、KV Cache 优化:从静态分配到虚拟内存管理

2.1 传统静态分配的缺陷

在 vLLM 出现之前,大多数推理框架采用静态分配策略:为每个序列预留最大长度的连续 KV Cache 空间。这导致两个严重问题:

  • 内部碎片:序列实际长度远小于最大长度,但预留空间无法被他人使用
  • 外部碎片:不同序列交错释放后,显存中出现大量不连续的空洞

实测表明,静态分配的 KV Cache 利用率通常仅有 20%-40%,这意味着超过一半的显存被浪费。

2.2 PagedAttention:操作系统的智慧

vLLM 的核心创新是将操作系统的虚拟内存管理思想引入 KV Cache 管理。其核心设计包括:

  • 分页机制:将 KV Cache 切分为固定大小的 block(通常 16-32 tokens),类似内存页
  • 页表映射:每个序列维护一个 block table,记录逻辑 block 到物理 block 的映射
  • 按需分配:仅在需要时分配新的物理 block,消除内部碎片
  • 写时复制:多个序列共享同一 prompt 时,通过 CoW 机制避免数据拷贝

配合 Modified PagedAttention 的分块并行策略,PagedAttention 将 KV Cache 利用率提升至 96% 以上,同时使吞吐量提升 2-4 倍。

2.3 KV Cache 量化压缩

进一步降低 KV Cache 内存占用的手段是量化:

  • Per-channel INT8 量化:对 K 矩阵按 channel 维度做 INT8 量化,精度损失极小
  • FP8 KV Cache:H100 原生支持 FP8,在不损失精度前提下将 KV 体积减半
  • GQA-aware 压缩:利用 Grouped Query Attention 中 KV head 数量远少于 Q head 的特点,进一步压缩 KV 维度

三、投机解码(Speculative Decoding)

3.1 核心思想

投机解码是一种"先猜后验"的策略:使用一个小型草稿模型(draft model)快速生成候选 tokens,然后用大模型一次性验证这些 tokens 的正确性。

直觉理解:大模型的单步推理虽然慢,但可以并行处理多个 token;草稿模型虽然快但可能不准确。将两者结合,可以在保证输出分布与原模型完全一致的前提下,减少大模型的解码步数。

3.2 算法流程

标准投机解码的每一步迭代:

  1. 草稿阶段:draft model 自回归生成 γ 个候选 tokens(通常 γ=3-7)
  2. 验证阶段:target model 一次前向传播,并行计算所有位置的 logits
  3. 接受-拒绝采样:逐位比较草稿模型与目标模型的分布,接受一致位置,拒绝分歧位置并从修正分布中重新采样

关键数学保证:通过精心设计的接受-拒绝策略,最终采样的分布与原始大模型完全相同(exact),而非近似。

3.3 变体与优化

  • 树状投机解码:构建多叉候选树而非线性链,增加候选多样性
  • Medusa:在目标模型头部添加多个并行预测头,无需独立草稿模型
  • EAGLE / EAGLE-2:利用目标模型的隐状态直接预测特征序列,消除独立小模型依赖
  • Lookahead Decoding:用 Jacobi 迭代松弛生成草稿,无需额外模型

四、连续批处理与调度优化

4.1 传统静态批处理的局限

最简单的批处理策略是等待一批请求到达后统一处理。问题在于:一批请求中最长的序列决定了所有请求的等待时间,导致大量 Padding 浪费。

4.2 Continuous Batching(Orca 方案)

Orca 提出的连续批处理(iteration-level scheduling)在每次迭代结束时动态调整批次组成:

  • 已完成的请求立即移出批次
  • 新到达的请求可以立即加入当前批次
  • 配合 PagedAttention 的内存管理,实现近乎 100% 的 GPU 利用率

与动态批处理相比,连续批处理消除了 Padding 开销。实测在请求到达率较高时,吞吐量提升可达 10 倍以上。

4.3 优先级调度与公平性

在实际部署中,还需考虑:

  • Preemption:抢占式调度,低优先级请求在显存不足时腾出空间
  • SLA-aware scheduling:根据延迟 SLO 动态调整 batch size
  • Chunked prefilling:将长 prompt 的 prefill 阶段与 decode 阶段交错执行,避免长 prompt 阻塞短请求

五、量化加速:从权重压缩到混合精度

5.1 推理量化的三种粒度

  • 仅权重量化(W8A16/W4A16):权重 INT8/INT4,激活 FP16。适用于矩阵乘为主的计算
  • 权重量化 + 激活量化(W8A8):两端均 INT8,需要校准步骤确定量化参数
  • FP8 推理:H100/B200 硬件原生支持,无需量化校准

5.2 GPTQ / AWQ / SmoothQuant 对比

当前主流的 4-bit 权重量化方案:

  • GPTQ:基于 Optimal Brain Quantization,逐列校准,压缩比高但推理时需反量化
  • AWQ:感知激活的显著通道,保护重要权重通道,精度更优
  • SmoothQuant:将激活的量化难度转移到权重端,兼顾两端效率
  • GGUF / llama.cpp:针对 CPU 优化的量化格式,支持混合位宽

5.3 稀疏化与 MoE

  • Structured Sparsity:NVIDIA 2:4 结构化稀疏,硬件加速下 2x 加速
  • Mixture of Experts:每个 token 仅激活部分专家(如 Mixtral 8x7B 仅激活 2 个专家),等效降低 70% 的推理计算量

六、推理引擎工程实践

6.1 SGLang 的 RadixAttention

SGLang 提出的 RadixAttention 利用 Radix Tree 自动管理前缀缓存。在多轮对话和少样本场景下,共享系统提示的 KV Cache 可节省 60%-80% 的显存和计算。

6.2 TensorRT-LLM

NVIDIA 官方推理引擎,提供:

  • GPU 定制 CUDA kernel(fused attention、fused MoE)
  • In-flight batching 与 Paged KV Cache
  • FP8/INT4 量化支持
  • 强大的编译器优化(layer fusion、constant folding、形状特化)

6.3 分布式推理

对于超大规模模型(如 100B+ 参数),单 GPU 无法容纳,需要:

  • Tensor Parallelism:将权重切分到多 GPU,需要高速 NVLink 互联
  • Pipeline Parallelism:不同层分配到不同 GPU,存在"流水线气泡"
  • Distributed Serving(DistServe):将 Prefill 与 Decode 阶段分离到不同实例,避免互相干扰

七、2026 年展望

推理优化技术仍在快速演进中:

  • 统一内存架构:CPU-GPU 统一内存(如 Apple Silicon、Grace Hopper)打破显存物理限制
  • NPU 推理:端侧 NPU 支持 INT4/FP4 推理,手机端可运行 7B 级别模型
  • 推测性自验证:LLM 自身具备验证能力,新一代"验证式解码"可完全消除草稿模型
  • 编译器-硬件协同设计:如 FlashAttention 系列展示的那样,算法创新与硬件特性深度协同将持续产生突破

推理优化不仅是一个工程问题,更是一个需要算法、系统、硬件三个层面协同创新的交叉领域。随着模型规模持续增长和部署场景日益丰富,高效推理技术将成为 AI 普惠化的关键基石。

点赞(0) 打赏

评论列表 共有 0 条评论

暂无评论
立即
投稿

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部