跳到主要内容
推理优化

第2章：推理引擎核心技术

掌握 PagedAttention、Continuous Batching、Prefix Cache 和 Chunked Prefill 四大推理引擎核心技术

PagedAttention Continuous Batching Prefix Cache Chunked Prefill

现代 LLM 推理引擎(vLLM、SGLang、TensorRT-LLM)能把 GPU 利用率从 20% 拉到 80%+,靠的是四个组合拳:PagedAttention 解决 KV 碎片、Continuous Batching 让 GPU 永不空转、Prefix Cache 复用前缀、Chunked Prefill 解 Prefill/Decode 互扰。本章逐个深入,理解为什么这套技术栈能把吞吐做到 20× 提升。

📑 目录

1. PagedAttention:KV 的虚拟内存

1.1 传统 KV 分配的问题

每个请求一来就预留 max_seq_len 大小的连续 KV Cache:

Request A (实际 100 token):  ████____________(预留 4096 个 slot)
Request B (实际 4000 token): ██████████████__(几乎用满)
Request C (300 token):        ████__________(浪费严重)

问题:

  • 内部碎片:实际用不到 max_len,大量空间浪费
  • 外部碎片:请求结束释放后,空闲块大小不一,新请求难以放入
  • 预留过大:典型场景 GPU 内存利用率只有 20-40%

1.2 PagedAttention 思想

借鉴操作系统虚拟内存:把 KV Cache 切成固定大小的 page(典型 16 token / page),用 Page Table 映射逻辑序列到物理页。

Request A 的逻辑序列:
  Logical block 0 (token 0-15)   → Physical page 5
  Logical block 1 (token 16-31)  → Physical page 12
  Logical block 2 (token 32-47)  → Physical page 3

Physical pool:
  Page 0  Page 1  Page 2  Page 3  ... Page N
   [free] [used]  [used]  [A.2]      [A.0]

1.3 收益

  • 几乎零内部碎片:每页 16 token,粒度很细
  • 无外部碎片:页大小固定,任何空闲页都能用
  • 内存利用率 80-90%(vs 传统 ~40%)

vLLM 论文实测:比 HuggingFace 快 20×,比 FasterTransformer 快 2-4×。

1.4 GPU 实现挑战

间接寻址(通过 page table)破坏了合并访存模式,需要专门的 paged attention CUDA kernel。

// 简化伪代码
__global__ void paged_attn(...) {
    for (int page_idx = 0; page_idx < num_pages; page_idx++) {
        int physical_page = page_table[req_id][page_idx];
        // 从 kv_pool[physical_page] 读 K, V
        // 用 online softmax 累加
    }
}

实际工业级实现见 vLLM 的 csrc/attention/paged_attention.cu 或 FlashInfer 的 paged decode kernel。

1.5 Copy-on-Write:Beam Search 优化

Beam Search 需要从一个父序列分裂多个候选——传统做法复制整个 KV,浪费严重。PagedAttention 用 CoW:多个 beam 共享前缀页,只在分歧处复制。

2. Continuous Batching:动态拼车

2.1 Static Batching 的问题

传统做法:等齐一个 batch 才开始,所有请求同时结束才能换新 batch。

Time →
Batch 1: ████████████████████ (8 个请求,长度 200-2000 不等)
                              ↑ 必须等最长的完成
Batch 2: ████████____

问题:

  • 短请求等长请求结束,latency 高
  • batch 内部分位置已生成完毕,GPU 空算
  • 新请求要等到下个 batch,排队等待长

2.2 Continuous Batching(Iteration-level Scheduling)

每个 iteration 重新组 batch:已完成的请求立刻退出,新请求随时插入。

Time →
slot 0: ████ ___ ████████ ████ ███           (短-空-长-短-短)
slot 1: ██████████ ____ █████ ██             (长-空-短-短)
slot 2: ██ █████████████ ███████             (短-长-中)
slot 3: ████████ █████ ████ ████ ███         (中-中-短-短-短)
       ↑ 每一列都是一次 forward

GPU 始终满载,实测 throughput 提升 5-23×(vLLM 数据)。

2.3 Iteration-level Scheduler 伪代码

class Scheduler:
    waiting = []      # 排队的新请求
    running = []      # 正在 decode 的请求

    def step(self):
        # 1. 检查 running 中已完成的,移除
        finished = [r for r in self.running if r.is_done()]
        self.running = [r for r in self.running if not r.is_done()]
        for r in finished: self.return_to_user(r)

        # 2. 从 waiting 中取新请求(KV 空间够就加入)
        while self.waiting and self.has_kv_space():
            r = self.waiting.pop(0)
            self.run_prefill(r)
            self.running.append(r)

        # 3. 当前 running 集合一起 decode 一步
        self.run_decode_step(self.running)

2.4 与 PagedAttention 的协同

Continuous Batching 需要频繁分配/释放 KV——而 PagedAttention 的页粒度恰好解决了这个动态分配问题。两者是天作之合,vLLM 的核心设计就是把它们组合起来。

3. Prefix Cache:复用前缀 KV

3.1 场景

很多应用中,多个请求共享同一段 system prompt 或 few-shot 例子:

请求 A: [SYSTEM PROMPT][用户问题 A][...]
请求 B: [SYSTEM PROMPT][用户问题 B][...]
请求 C: [SYSTEM PROMPT][用户问题 C][...]

如果每次都重新 Prefill SYSTEM PROMPT,浪费了 90% 的算力。

3.2 Prefix Cache 机制

对 prompt 的前缀做 hash,如果命中已有的 KV Cache,直接复用,跳过 Prefill 这部分。

新请求来:
  1. 把 prompt 按 page 大小切分
  2. 对每个 page 的 token 做 hash
  3. 在 hash table 中查询是否有命中
  4. 命中的页直接共享物理页(增加引用计数)
  5. 未命中的部分才走 Prefill

3.3 收益

System prompt 1000 token,问题 100 token,共享率 90%:

  • 朴素 Prefill:1100 token,~50 ms
  • Prefix Cache 命中:只 Prefill 100 token,~5 ms——TTFT 降低 10×

3.4 vLLM 用法

from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-2-7b", enable_prefix_caching=True)

# 第一次请求,正常 Prefill
out1 = llm.generate(["[SYSTEM]...[USER]问题 A"], sampling)
# 第二次,前缀命中,只算用户问题部分
out2 = llm.generate(["[SYSTEM]...[USER]问题 B"], sampling)

4. RadixAttention:更聪明的前缀树

vLLM 的 Prefix Cache 是按 page hash 匹配,只能复用”完全相同的”前缀。RadixAttention(SGLang)用 Radix Tree(基数树),可以复用任意公共前缀。

4.1 直觉

Tree:
  root → [SYSTEM 你是助手]
           ├── [USER 计算...]
           │     ├── [对 1+1]
           │     └── [对 2+2]
           └── [USER 翻译...]
                 └── [hello]

每个节点是一段 token,节点间共享路径。新请求来,沿树匹配最长公共前缀,后续才 Prefill。

4.2 收益

复杂多轮对话场景,RadixAttention 的命中率比 Prefix Cache 高 30-50%。

4.3 SGLang 用法

import sglang as sgl

sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000"))

@sgl.function
def multi_turn(s):
    s += "[SYSTEM] 你是助手\n"
    s += "[USER] 1+1=?\n"
    s += "[ASSISTANT]" + sgl.gen("ans1", max_tokens=10)
    s += "[USER] 2+2=?\n"
    s += "[ASSISTANT]" + sgl.gen("ans2", max_tokens=10)

第二个 USER 后的 prefill 会复用第一轮的 KV——多轮场景加速尤其明显。

5. Chunked Prefill:解 PD 互扰

5.1 问题

Continuous Batching 让 Prefill 和 Decode 请求混在一个 batch 里。但 Prefill 是 compute bound 的大算子(几千 token 一次),会严重拖慢同 batch 的 Decode——Decode 请求被迫等 Prefill 算完才能拿到下一 token。

5.2 Chunked Prefill

把长 Prefill 切成多个 chunk(典型 512 token),每个 iteration 只算一个 chunk:

Iteration 1: [Decode A][Decode B][Prefill C 的 chunk 1]
Iteration 2: [Decode A][Decode B][Prefill C 的 chunk 2]
Iteration 3: [Decode A][Decode B][Prefill C 的 chunk 3] ← C 完成,转 Decode

TPOT 显著改善——Decode 请求每次只被一个 chunk 阻塞,而非整个 Prefill。

5.3 vLLM / SGLang 配置

llm = LLM(
    model=...,
    max_num_batched_tokens=4096,    # 一个 iteration 总 token 上限
    enable_chunked_prefill=True,
)

5.4 与 PD 解耦的关系

Chunked Prefill 是”温和”方案:仍在同 GPU 池处理 P/D。 PD 解耦(第 6 章)是”激进”方案:分到不同 GPU 池。

两者互补:小集群用 Chunked Prefill,超大集群上 PD 解耦。

6. 四大技术的协同

vLLM 的整体架构:

请求到达
   ↓
Scheduler(Iteration-level)
   ├── 查 Prefix Cache:命中部分跳过 Prefill
   ├── 长 Prefill 切 Chunk
   └── 组合 Prefill chunks + Decode 请求成一个 batch
   ↓
PagedAttention 分配 KV 物理页
   ↓
执行一个 iteration
   ↓
完成的请求退出,新请求加入,循环

四个技术任一缺失,整体性能就会降一个档:

缺失后果
无 PagedAttentionKV 碎片严重,并发上不去
无 Continuous BatchingGPU 利用率仅 20-30%
无 Prefix Cache重复 Prefill,TTFT 高
无 Chunked Prefill长 Prefill 拖垮 Decode P95

🌟 vLLM、SGLang、TensorRT-LLM 都实现了这四件套——它们的差异在 RadixAttention(SGLang)、量化集成(TensorRT-LLM)、生态成熟度(vLLM)等方面。

✅ 自我检验清单

  • PagedAttention 类比:能用操作系统虚拟内存类比解释 PagedAttention 设计
  • 碎片问题:能解释传统 KV 分配的内部碎片和外部碎片来源
  • Continuous Batching:能用图示对比 Static Batching 和 Continuous Batching 的时间利用率
  • Iteration-level Scheduler:能默写 scheduler 一个 step 的伪代码(回收完成 / 加入新请求 / 一起 decode)
  • Prefix Cache 命中:能算出 system prompt 1000 token、问题 100 token 时的 TTFT 提升倍数
  • Prefix Cache vs RadixAttention:能解释 RadixAttention 在多轮对话中为什么命中率更高
  • Chunked Prefill 收益:能解释为什么不切 chunk 会让 Decode P95 翻倍
  • 四件套协同:能说出每个技术缺失会导致什么具体性能问题
  • vLLM 配置:能正确启用 prefix caching + chunked prefill + max_num_batched_tokens

📚 参考资料