LLM推理服务的推测解码与PagedAttention内存管理

引言

大语言模型（LLM）的推理服务正面临两大核心瓶颈：自回归解码的顺序性导致GPU计算单元大量空闲，以及KV缓存的内存占用随序列长度线性增长。2024年以来，推测解码（Speculative Decoding）和PagedAttention（分页注意力）以及两者协同优化的推理系统成为学术界和工业界关注的焦点，成功将LLM推理吞吐量提升了2-4倍。

推测解码（Speculative Decoding）原理

传统自回归解码以逐token方式生成输出，每一步都执行完整的模型前向计算，但每个token实际只需要模型计算的一个子集。推测解码的核心洞察是：大多数token的预测准确率足够高，可以用轻量级草稿模型快速生成候选序列，再由目标模型并行验证。

具体流程如下：

1. 草稿阶段：使用小型草稿模型（如目标模型的蒸馏版本或n-gram缓存）自回归生成k个候选token
2. 验证阶段：将整个候选序列和历史上下文一次性输入目标模型执行并行前向传播，通过单次GPU矩阵运算计算所有位置的概率分布
3. 接受/拒绝：逐位置比较草稿模型与目标模型输出的概率分布，以概率min(1, p_target/p_draft)接受该token

理论证明（Leviathan et al., 2023）：假设草稿模型的接受率为α，则推测解码的期望加速比为(1-α^(k+1))/(1-α)，当α=0.8且k=5时加速比可达约3倍。

推测解码的工程实践变体

Medusa（Cai et al., 2024）抛弃了小模型方案，改为在目标模型顶部附加多个独立的预测头（Medusa Head），每个头负责预测未来不同长度的token。由于这些轻量级头与目标模型主干并行执行，草稿生成开销几乎为零，可稳定实现2-3倍加速。

EAGLE/EAGLE-2（Li et al., 2024）利用目标模型的特征层输出（而非logits）作为轻量级自回归头的特征输入，通过复用主干网络的中间表示大幅提升了草稿准确率。EAGLE-2还引入上下文感知的动态草稿长度调整。

LayerSkip则让模型在草稿生成阶段只使用部分Transformer层（如跳过偶数层），在验证阶段再使用完整层，实现了模型内部的推测执行。

PagedAttention：KV缓存的内存管理革命

vLLM提出的PagedAttention借鉴了操作系统中虚拟内存和分页机制来解决KV缓存的内存碎片问题。在传统推理引擎中，每个序列需要预先分配连续的大块内存，但由于序列实际长度差异巨大，大量内存未被使用但仍被预留（浪费比例可达60-80%）。

PagedAttention的设计核心：

• 逻辑块与物理块分离：每个序列维护一个Block Table，逻辑上连续的KV缓存映射到物理上不连续的内存块（通常每块存储16-64个token的KV）
• 按需分配：只有当当前块写满时才分配新物理块，消除预分配浪费
• 内存共享：在beam search和并行采样场景下，父子序列可以通过写时复制（Copy-on-Write）共享相同的物理块

实验表明PagedAttention将GPU内存利用率从约20-40%提升至接近100%，显存浪费控制在4%以内，使单卡上可同时服务数十倍的并发请求。

Continuous Batching与调度策略

传统Static Batching要求等待最慢的请求完成才能返回结果，导致GPU利用率波动剧烈。Continuous Batching（又称Iteration-Level Scheduling）在每个解码步骤都可以动态插入新请求或驱逐已完成请求，配合PagedAttention实现了真正的在线服务。

关键调度策略包括：

• Shortest-Remain-First：优先处理剩余生成token数短的请求，降低尾延迟
• Watermark-based Admission：设置已用内存水印阈值，超过阈值时拒绝新请求或将其排队
• Chunked Prefill：将长输入的Prefill阶段拆分为多个Chunk与Decode阶段交错执行，避免单个超长Prefill阻塞其他正在解码的请求

推测解码与PagedAttention的协同效应

推测解码的并行验证模式让GPU在大batch下仍能保持高利用率，而PagedAttention的内存效率使推理引擎能够同时容纳更多并发请求。两者的结合创造了正反馈：PagedAttention释放的显存空间可以容纳更多草稿序列的并发验证，推测解码带来的计算密度提升又让高并发的KV缓存管理变得有价值。

vLLM v0.5+和TensorRT-LLM已将推测解码集成到PagedAttention框架中，实测在Llama-3-70B模型上使用4块A100可提供高达5000+ tokens/sec的端到端吞吐，相比基线方案提升3.5倍以上。

未来方向

当前研究前沿包括：动态草稿长度的自适应控制、推测解码与Mixture-of-Experts（MoE）架构的协同优化、KV缓存的量化压缩与跨节点分布式管理、以及针对超长上下文（100K+ tokens）的分层推测策略。随着多模态模型的普及，视觉token与文本token的异构推测解码也将成为新的挑战。