LLM 推理优化深度实战：从 PagedAttention 到 Speculative Decoding 的工程架构
大语言模型（LLM）推理是 AI 落地中最昂贵的环节。一个 70B 参数量的模型，单次推理可能消耗数十 GB 显存，而在线服务要求的 P99 延迟往往需要控制在 2 秒以内。如何在有限的 GPU 资源上同时实现高吞吐和低延迟，是过去两年系统优化领域最激烈的战场之一。本文将从显存管理、批处理策略和解码算法三个维度，深入拆解现代 LLM 推理引擎的核心技术，并给出可落地的工程实践方案。
一、显存瓶颈：KV Cache 管理的工程困局
Transformer 自回归推理的核心瓶颈不在计算，而在显存。每一步 decode 都需要读取之前所有 token 的 Key 和 Value 张量，这就是 KV Cache。对于一个层数 L、隐藏维度 d、序列长度 n、批次大小 b 的模型，KV Cache 的显存占用约为：
KV Cache ≈ 2 × L × d × n × b × sizeof(dtype)以 Llama-2-70B 为例（L=80, d=8192），处理一个 2048 token 的序列，FP16 下单卡 KV 就需要约 512MB。当并发请求达到 100 时，仅 KV Cache 就需要 50GB——这是大多数单卡 GPU 无法承受的。
更棘手的是，传统推理引擎会为每个请求预留其最大长度的连续显存空间。实际场景中，序列长度分布高度不均：有的请求只需要几十个 token，有的需要几千个 token。提前为每个请求分配最大长度的显存空间，意味着大量显存被浪费在预留但未使用的空间上。实测数据显示，传统方案中 KV Cache 的有效利用率往往不足 25%。
二、PagedAttention：从连续分配到虚拟内存管理
2023 年，UC Berkeley 的 Kwon 等人在论文 "Efficient Memory Management for Large Language Model Serving with PagedAttention" 中提出了革命性的解决方案，其思想直接借鉴了操作系统中的虚拟内存和分页机制。
2.1 核心思想
PagedAttention 的核心是：将 KV Cache 视为"内存页"，按需分配，物理上不要求连续。
具体实现上，系统预先将 GPU 显存划分为固定大小的 block（通常 16 或 32 个 token 的 KV 数据）。每个请求维护一个 block table，逻辑上连续的序列在物理上可以散布在任意 block 中。当序列长度增长时，只需追加新的 block；当请求完成时，回收这些 block 供其他请求使用。
# PagedAttention 伪代码：block table 管理
class BlockTable:
"""每个请求的逻辑到物理 block 映射"""
def __init__(self, block_size: int = 16):
self.block_size = block_size # 每个 block 存储多少个 token 的 KV
self.logical_blocks: list[int] = [] # 物理 block 索引列表
self.num_tokens: int = 0
def append_block(self, physical_block: int):
"""序列增长时追加一个物理 block"""
self.logical_blocks.append(physical_block)
def get_physical_block(self, logical_idx: int) -> int:
"""根据逻辑位置获取物理 block"""
return self.logical_blocks[logical_idx]
@property
def num_blocks(self) -> int:
"""需要的 block 数量 = ceil(num_tokens / block_size)"""
return (self.num_tokens + self.block_size - 1) // self.block_size2.2 显存利用率提升
这种分页管理带来的最直接的收益是显存利用率接近 100%。因为：
- 消除内部碎片：只有最后一个 block 可能存在未填满的情况，浪费最多一个 block（通常仅占总分配的 3-6%）。
- 消除外部碎片：物理 block 池化使用，不存在传统方案中因请求长度差异大导致的显存"空洞"。
- 精确按需分配：只分配实际需要的 block，不使用预分配最大长度的策略。
实测对比显示，PagedAttention 可以将系统吞吐提升 2-4 倍，同时保持与基线方案相同的延迟水平。
2.3 Copy-on-Write 与共享前缀
PagedAttention 的另一大亮点是支持 Copy-on-Write（COW）机制。在 system prompt 共享的场景下（如多轮对话、few-shot prompting），多个请求的前缀 KV Cache 可以共享同一个物理 block。只有在某个请求需要修改前缀时，才触发实际的拷贝操作。
# 共享 system prompt 的 COW 示例
system_prompt_blocks = [10, 11, 12, 13] # system prompt 占用的物理 block
# 两个用户请求共享相同的 system prompt
request_a = BlockTable(shared_blocks=system_prompt_blocks.copy())
request_b = BlockTable(shared_blocks=system_prompt_blocks.copy())
# 当 request_a 的用户输入部分开始写入时，只需为新内容分配 block
request_a.append_block(alloc_new_block()) # 写入用户 token 的 KV
request_b.append_block(alloc_new_block()) # 写入用户 token 的 KV
# system prompt 部分仍然共享，无需拷贝三、Continuous Batching：突破静态批处理吞吐瓶颈
在 PagedAttention 之前，推理引擎主要使用 Static Batching：将一批请求一起送入 GPU，等所有请求都完成解码后再返回结果。假设一批中有 10 个请求，其中 9 个在 100 token 内完成，第 10 个需要 2000 token——那么前 9 个请求的 GPU 计算在等待中白白浪费，整体吞吐被最慢的请求拖累。
Continuous Batching（也称为 Iteration-Level Scheduling）在每个 decode 步骤的间隔中检查队列，动态地将已完成请求移出、将新请求插入。其效果可以通过以下对比理解：
Static Batching:
Batch-1: [Req1: 100 tokens] [Req2: 2000 tokens] [Req3: 50 tokens]
GPU 利用率: ████████░░░░░░░░░░░░░░░ (前50步100%, 中间1500步只有Req2在跑)
Continuous Batching:
Step 0-50: [Req1] [Req2] [Req3] ← 3个请求并行
Step 50: [Req3完成, Req4加入] ← Req3 退出, Req4 进入
Step 50-100: [Req1] [Req2] [Req4] ← GPU 持续满载
Step 100: [Req1完成, Req5加入]
Step 100+: [Req2] [Req5] ← 持续高利用率vLLM 的实测数据显示，在现实流量模式下（请求到达服从泊松分布，长度服从重尾分布），Continuous Batching 比 Static Batching 提升吞吐 23 倍。
3.1 调度策略的权衡
Continuous Batching 并非没有调度开销。当 batch 中混合了 Prefill（首 token 生成，计算密集）和 Decode（逐 token 生成，访存密集）阶段时，需要仔细设计调度策略：
# 简单但有效的调度策略: Swapping-aware Scheduling
class ContinuousBatchScheduler:
def schedule(self, waiting_queue: list[Request], running_batch: list[Request]):
"""
调度逻辑:
1. 优先保证 running batch 中 decode 请求的 GPU 时间
2. 空余算力用于预填充 (prefill) 等待队列中的请求
3. 如果显存不足，swap out 低优先级 decode 请求到 CPU
"""
gpu_budget = self.max_batch_size
kv_budget = self.available_kv_blocks
# 1. 保留 running batch 中的 decode 请求（不能丢，它们已经在生成中）
scheduled = [req for req in running_batch if req.phase == 'decode']
remaining_slots = gpu_budget - len(scheduled)
remaining_blocks = kv_budget - sum(req.kv_blocks for req in scheduled)
# 2. 从等待队列中选择请求进行 prefill
for req in waiting_queue:
if remaining_slots <= 0 or remaining_blocks < req.estimated_kv_blocks:
continue
scheduled.append(req)
remaining_slots -= 1
remaining_blocks -= req.estimated_kv_blocks
return scheduled四、Speculative Decoding：用小模型跑赢大模型的解码速度
即便解决了显存和批处理问题，大模型推理还有一个根本性的瓶颈：decode 阶段的 GPU 利用率极低。每步 decode 只生成一个 token，70B 模型的数十亿参数只为了计算一个 token 的 softmax，计算访存比（arithmetic intensity）接近 1——这是 GPU 上最糟糕的工作模式。
Speculative Decoding（Leviathan et al., 2023）的核心洞察是：让一个小得多的 Draft 模型先生成候选 token，然后让大模型一次性验证这些候选。
4.1 直觉理解
这就像写文章时先让助手起草草稿，然后你一次性审阅修改。虽然你（大模型）每一步能做的判断比助手（小模型）准，但你审阅 N 个草稿 token 的时间，可能比你自己写 N 个 token 的时间还少。
数学表达：设大模型一步耗时 $T_L$，小模型一步耗时 $T_S$（$T_S \ll T_L$），小模型生成的候选长度为 $\gamma$，最终接受率为 $\alpha$。则每次 speculation 的期望时间：
$$E[t] = T_S \times \gamma + T_L$$
而等效的传统 decode 时间为：
$$E[t_{baseline}] = T_L \times (\alpha \times \gamma + 1)$$
当 $\alpha \times \gamma \times T_L > \gamma \times T_S + T_L$ 时，speculative decoding 就带来了加速。由于 $T_S \ll T_L$ 且 $\alpha$ 通常在 0.7-0.9 之间，这种加速在大多数场景下非常可观。
4.2 验证算法
Speculative Decoding 的关键在于验证阶段如何高效地使用大模型。大模型不需要重新计算整个序列——它只需要一次前向传播，同时验证所有候选 token：
import torch
def speculative_verify(draft_tokens, model_large, prefix_kv,
prefix_logits=None, temperature=0.0):
"""
draft_tokens: [batch, gamma] - draft 模型生成的候选 token
model_large: 目标大模型
prefix_kv: 前缀的 KV Cache（已计算好的）
返回: 每个位置接受的 token 数量和具体的接受/拒绝决策
"""
gamma = draft_tokens.shape[1]
# 1. 大模型一次性计算所有候选 token 的 logits
# 使用 prefix_kv 作为前缀，draft_tokens 作为当前输入
# 这只需要一次 forward pass
logits_large = model_large.forward_with_kv(
draft_tokens, past_key_values=prefix_kv
) # [batch, gamma, vocab_size]
# 2. 获取 draft 模型对应位置的 logits（需要保留推理时的概率分布）
logits_draft = model_small.get_logits() # [batch, gamma, vocab_size]
# 3. 逐位置验证
accepted = 0
for i in range(gamma):
prob_large = softmax(logits_large[:, i, :] / temperature)
prob_draft = softmax(logits_draft[:, i, :] / temperature)
# 接受概率（两种等价的采样方法）
# 方法 A: Rejection sampling
coin_flip = torch.rand(prob_large.shape[0])
accept_prob = torch.clamp(
(prob_large.gather(1, draft_tokens[:, i:i+1]) -
prob_draft.gather(1, draft_tokens[:, i:i+1])) /
(prob_draft.gather(1, draft_tokens[:, i:i+1]) + 1e-8),
min=0.0
)
if (coin_flip < accept_prob).all():
accepted += 1
else:
# 在位置 i 拒绝：从调整后的分布中采样一个 token
adjusted_prob = torch.clamp(prob_large - prob_draft, min=0)
adjusted_prob = adjusted_prob / adjusted_prob.sum(dim=-1, keepdim=True)
# 重采样...
break
return accepted, final_tokens4.3 Medusa：多头解码的工程突破
Speculative Decoding 的一个限制是 Draft 模型的选择。如果小模型和大模型差异太大，接受率 $\alpha$ 会很低，导致额外的计算负担反而拖慢速度。
Medusa（Cai et al., 2024）提出了一种优雅的解法：直接在大模型之上训练多个并行的预测头，不需要独立的 Draft 模型。
Normal Decode:
[Base Model] → token₁ → token₂ → token₃ → token₄
Medusa:
[Base Model + Head₁] → token₁
[Base Model + Head₂] → token₂ (预测位置+2)
[Base Model + Head₃] → token₃ (预测位置+3)
[Base Model + Head₄] → token₄ (预测位置+4)
↓
[Verification] 一次性验证每个 Medusa Head 是一个轻量的 MLP 模块，附加在大模型的最后一层之后。它根据当前位置的 hidden state 直接预测未来第 k 个位置的 token。最终验证时，大模型的前向传播可以共享绝大部分计算——验证 N 个候选 token 与验证 1 个 token 的计算量几乎相同。
Medusa 在保持大模型输出分布完全一致的前提下，实现 2-3 倍的解码加速。
五、工程实践：vLLM 的架构解析
vLLM 是目前最流行的开源 LLM 推理引擎，它将 PagedAttention 和 Continuous Batching 深度集成。理解其架构有助于我们在实际部署时做出正确的配置选择。
5.1 核心组件
┌─────────────────────────────────────────────────────────┐
│ vLLM Architecture │
├─────────────────────────────────────────────────────────┤
│ │
│ ┌──────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │ FastAPI │───▶│ Scheduler │───▶│ GPU Model │ │
│ │ Server │ │ │ │ Runner │ │
│ └──────────┘ │ ┌──────────┐ │ │ │ │
│ │ │ Prefill │ │ │ ┌──────────┐ │ │
│ │ │ Queue │ │ │ │ Block │ │ │
│ │ └──────────┘ │ │ │ Manager │ │ │
│ │ ┌──────────┐ │ │ └──────────┘ │ │
│ │ │ Decode │ │ │ ┌──────────┐ │ │
│ │ │ Batch │ │ │ │ Paged │ │ │
│ │ └──────────┘ │ │ │ Attention│ │ │
│ └──────────────┘ │ └──────────┘ │ │
│ └──────────────┘ │
└─────────────────────────────────────────────────────────┘5.2 关键配置参数
# vLLM 启动配置示例：Llama-2-70B 部署
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Llama-2-70b-chat-hf \
--tensor-parallel-size 4 \ # 4-way TP 分布在 4 张 A100-80G
--gpu-memory-utilization 0.90 \ # 预留 10% 显存给 system/PagedAttention 开销
--max-num-batched-tokens 8192 \ # 每 batch 最大 token 数（控制 prefill 粒度）
--max-num-seqs 256 \ # 最大并发序列数
--enable-speculative-decoding \ # 启用 speculative decoding
--speculative-model \ # draft 模型路径
"meta-llama/Llama-2-13b-chat-hf" \
--num-speculative-tokens 5 \ # 每个 speculation 的候选 token 数
--block-size 16 # PagedAttention 的 page 大小5.3 性能调优实践
在实际部署中，有几个容易被忽视但对性能影响巨大的参数：
gpu-memory-utilization：不宜设置过高（建议 0.85-0.92）。PagedBlock 管理器需要少量显存做 bookkeeping，同时 CUDA 上下文本身也需要几百 MB。设置过高会导致运行时 OOM。
max-num-batched-tokens：这是单次 forward 中最多的 token 总数，直接影响 prefill 阶段的延迟。对于 latency-sensitive 场景，建议设置为 2048-4096；对于 throughput-only 场景，可以适当增大。
chunked-prefill：vLLS 2024 引入的新特性，允许将一个长 prefill 请求拆分为多个 chunk，与 decode 步骤交错执行。这解决了长 prefill "阻塞" decode 请求的问题，对 P99 延迟有显著改善。
六、前沿趋势：推理引擎的下一个战场
6.1 分布式推理与异构调度
当模型规模从 70B 扩展到 405B 甚至更大时，单节点 8 卡已经不够。未来的推理系统需要在多节点间分布 KV Cache，实现跨节点的 PagedAttention。SGLang、 LMDeploy 等引擎已经在这个方向上有了初步探索。
6.2 量化与压缩的协同优化
GPTQ、AWQ、SmoothQuant 等量化技术可以将模型从 FP16 压缩到 INT4/INT8，直接降低 KV Cache 的显存占用。尤其是 KV Cache 量化——在 Attention 计算中使用低精度的 KV，精度损失极小（<0.1 个 perplexity 点），但显存和带宽节省可达 4 倍。
6.3 投机解码的变体与进化
Tree-based Speculative Decoding 将线性的 candidate 序列扩展为候选树结构，在一次验证中从多个分支中选择最优路径。而 Prompt Lookup Decoding 利用用户 prompt 中已有的短语直接作为 draft，在特定场景下实现零额外开销的投机解码。
七、总结
LLM 推理优化的核心是一个系统工程问题——它不只是算法层面的创新，更是操作系统原理（分页、调度）与硬件特性（GPU 访存模式、Tensor Core）的深度结合。从 PagedAttention 到 Continuous Batching 再到 Speculative Decoding，每一步优化都解决了一个明确的工程瓶颈，而它们组合在一起的收益远超各部分之和。
对于实践者来说，理解这些底层机制的价值在于：当你在 vLLM 中看到一个 OOM 错误时，你会知道该调高 gpu-memory-utilization 还是减小 max-num-seqs；当你发现 P99 延迟抖动时，你会怀疑是不是 chunked-prefill 没有开启；当你的 Draft 模型接受率很低时，你会检查两个模型的分布差异是否过大。
这就是系统工程的魅力：原理告诉你原因，数据指导你优化。
发表评论 取消回复