Motivation
训练一个大模型很贵,但服务一个大模型可能更贵。
训练是一次性开销——几千张 GPU 跑几周就结束了。但推理是 7×24 小时不间断的:每一次用户输入、每一个 API 调用,都需要模型实时响应。OpenAI 每天处理数十亿个 token,推理成本占运营总成本的大头。如果推理效率翻倍,意味着成本直接减半——或者同样预算下服务两倍的用户。
推理和训练的优化方向截然不同。训练追求总计算吞吐——尽快把所有数据跑完;推理则需要同时优化两个相互矛盾的指标:
- 吞吐量(Throughput):每秒生成多少个 token?决定了能服务多少并发用户
- 延迟(Latency):每个用户等多久?TTFT(Time To First Token,首 token 延迟)和 TPOT(Time Per Output Token,逐 token 延迟)是关键指标
更重要的是,推理的计算特性和训练完全不同。我们在第一篇中分析过:自回归解码(autoregressive decoding)本质上是矩阵-向量乘法,算术强度只有约 1 FLOPs/Byte,GPU 算力利用率不到 1%。推理不是算力瓶颈,是显存带宽瓶颈——而 KV Cache 是其中最大的那块显存。
本文以 SGLang 为例,深入剖析推理系统的四个核心技术:
- KV Cache——推理显存的大头,为什么它会成为瓶颈
- PagedAttention——借鉴 OS 虚拟内存思想,解决 KV Cache 的显存碎片问题
- RadixAttention——SGLang 的核心创新,用 Radix Tree 实现 KV Cache 复用
- Continuous Batching——迭代级调度,最大化 GPU 利用率
- Chunked Prefill——解决长 prompt 阻塞解码的问题
读完之后,你应该能理解为什么 vLLM 和 SGLang 能比朴素实现快 2-10 倍,以及它们之间的设计差异。
前置知识
- GPU 显存模型与分布式通信基础(第 1 篇)——特别是 HBM 带宽瓶颈和 Roofline 模型
- Transformer 注意力机制的基本理解(Self-Attention、Q/K/V 投影)
- 了解自回归生成的基本流程(逐 token 生成)
KV Cache 基础
为什么需要 KV Cache
自回归语言模型在生成第 $t$ 个 token 时,Attention 的计算公式是:
$$\text{Attention}(Q_t, K_{1:t}, V_{1:t}) = \text{softmax}\left(\frac{Q_t \cdot K_{1:t}^T}{\sqrt{d_k}}\right) V_{1:t}$$
关键点:当前 token 的 Query 需要和所有已生成 token 的 Key、Value 做内积。如果每次都从头计算所有 token 的 K 和 V,复杂度是 $O(t^2)$——生成 1000 个 token 就要做 $1 + 2 + \cdots + 1000 = 500500$ 次 attention 计算。
KV Cache 的核心思想:之前 token 的 K 和 V 不会变,只需要计算一次,缓存起来复用。
flowchart TD
subgraph NO["自回归解码(无 KV Cache)"]
N1["Step 1: 输入 [t1] → 计算 K1,V1 → Attention → 输出 t2"]
N2["Step 2: 输入 [t1,t2] → 重新计算 K1,V1,K2,V2 → Attention"]
N3["Step 3: 输入 [t1,t2,t3] → 重新计算 K1,K2,K3... → Attention"]
NN["Step N: 输入 [t1,...,tN] → 重新计算所有 K,V → O(N²) 总计算"]
N1 --> N2 --> N3 -.-> NN
end
subgraph YES["自回归解码(有 KV Cache)"]
Y1["Step 1: 输入 [t1] → 计算 K1,V1 → 缓存 → Attention → 输出 t2"]
Y2["Step 2: 输入 [t2] → 计算 K2,V2 → 追加缓存 → Attention"]
Y3["Step 3: 输入 [t3] → 计算 K3,V3 → 追加缓存 → Attention"]
YN["Step N: 输入 [tN] → 计算 KN,VN → 追加缓存 → O(N) 总计算"]
Y1 --> Y2 --> Y3 -.-> YN
YNOTE["每步只计算 1 个 token 的 K,V,复用缓存中的所有 K,V"]
end
NO --> YES
style NO fill:#fff3cd,stroke:#856404
style YES fill:#d4edda,stroke:#155724
KV Cache 将总计算量从 $O(N^2)$ 降到了 $O(N)$——这是推理优化的第一性原理,所有现代推理引擎都必须实现它。
KV Cache 有多大
每个 token 的 KV Cache 大小为:
$$\text{KV per token} = 2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times \text{dtype_size}$$
其中因子 2 是因为要存 K 和 V 两份。以几个主流模型为例(FP16,2 bytes per element):
| 模型 | $n_{\text{layers}}$ | $n_{\text{heads}}$ | $d_{\text{head}}$ | KV per token | 2048 tokens |
|---|---|---|---|---|---|
| LLaMA-7B | 32 | 32 | 128 | 0.5 MB | 1.0 GB |
| LLaMA-13B | 40 | 40 | 128 | 0.8 MB | 1.6 GB |
| LLaMA-70B | 80 | 64 | 128 | 2.5 MB | 5.0 GB |
| GPT-3 175B | 96 | 96 | 128 | 4.5 MB | 9.2 GB |
一个 LLaMA-70B 的请求如果生成 2048 个 token,仅 KV Cache 就占 5 GB。如果你想同时服务 16 个这样的请求,KV Cache 就需要 80 GB——一整张 A100 的显存。而且模型权重本身还要 140 GB(FP16)。
这就是推理显存的核心矛盾:模型权重是固定开销,KV Cache 是动态开销,而 KV Cache 随并发请求数和序列长度线性增长,很容易成为瓶颈。
能同时 batch 多少请求,取决于 KV Cache 能放多少——batch size 直接决定吞吐量。所以 KV Cache 的内存管理效率,就是推理系统的性能上限。
朴素分配的浪费
最直观的做法是:为每个请求预分配 max_seq_len 大小的连续显存来存 KV Cache。问题是,你不知道一个请求到底会生成多少 token——可能是 10 个,也可能是 2000 个。
| |
配套代码中的实验表明,当请求长度分布不均匀时(现实中几乎总是如此),朴素预分配的显存浪费率高达 60-80%:
5 requests, max_seq_len=2048
Actual lengths: [37, 152, 8, 420, 91]
Memory allocated: 160.0 MB
Memory actually used: 22.1 MB
Waste: 86.2%
--> This is why we need PagedAttention!
显存浪费 86%——这意味着原本能同时处理 7 个请求的 GPU,实际上只能处理 1 个。浪费的不仅是内存,更是吞吐量和收入。

PagedAttention
从操作系统借鉴的智慧
PagedAttention 是 vLLM(OSDI'23)的核心贡献,它的灵感来自一个成熟了 50 年的技术——操作系统的虚拟内存。
在操作系统中,进程看到的是虚拟地址空间——连续的、独立的。但物理内存是以固定大小的页(page) 分配的,通过页表(page table) 将虚拟地址映射到物理地址。进程不需要拿到连续的物理内存,操作系统在幕后把碎片化的物理页拼成看似连续的虚拟空间。
PagedAttention 把这个思想完美移植到 KV Cache 管理上:
flowchart LR
subgraph OS["OS 虚拟内存"]
A1["进程"]
A2["虚拟页"]
A3["物理页帧"]
A4["页表"]
A5["按需分配(缺页中断)"]
A6["页大小固定(4KB)"]
A7["CoW(fork)"]
A8["页面置换(swap)"]
end
subgraph PA["PagedAttention"]
B1["请求(sequence)"]
B2["逻辑 KV block"]
B3["物理 KV block"]
B4["block table"]
B5["按需分配(生成新 token 时)"]
B6["block 大小固定(16 tokens)"]
B7["CoW(beam search)"]
B8["preemption(换出请求)"]
end
A1 -.-> B1
A2 -.-> B2
A3 -.-> B3
A4 -.-> B4
A5 -.-> B5
A6 -.-> B6
A7 -.-> B7
A8 -.-> B8
style OS fill:#cce5ff,stroke:#004085
style PA fill:#d4edda,stroke:#155724
Block 分配机制
PagedAttention 将 GPU 上的 KV Cache 显存划分为一个固定大小的 block 池。每个 block 能存储固定数量的 token 的 KV(通常 16 或 32 个 token)。每个请求有一个 block table,记录其逻辑 block 到物理 block 的映射。
flowchart TD
subgraph POOL["物理 Block 池(GPU 显存)"]
B0["B0<br/>Seq0"]
B1["B1<br/>Seq1"]
B2["B2<br/>Free"]
B3["B3<br/>Seq0"]
B4["B4<br/>Seq1"]
B5["B5<br/>Free"]
B6["B6<br/>Seq0"]
B7["B7<br/>Free"]
end
subgraph T0["Seq 0 Block Table"]
L0_0["Logic 0 → B0"]
L0_1["Logic 1 → B3"]
L0_2["Logic 2 → B6"]
end
subgraph T1["Seq 1 Block Table"]
L1_0["Logic 0 → B1"]
L1_1["Logic 1 → B4"]
end
L0_0 --> B0
L0_1 --> B3
L0_2 --> B6
L1_0 --> B1
L1_1 --> B4
NOTE["物理 block 不需要连续,逻辑上连续即可"]
style POOL fill:#cce5ff,stroke:#004085
style B2 fill:#f8f9fa,stroke:#6c757d
style B5 fill:#f8f9fa,stroke:#6c757d
style B7 fill:#f8f9fa,stroke:#6c757d
工作流程:
- 新请求到达:分配一个空的 block table,不预分配任何 block
- 生成第 1 个 token:从 free pool 拿一个 block,存入 KV,更新 block table
- 当前 block 填满(写满 16 个 token 的 KV):再从 free pool 拿一个新 block
- 请求结束:归还所有 block 到 free pool,立即可被其他请求复用
| |
从 86% 浪费到 4% 浪费
PagedAttention 的内存浪费只来自一个地方:每个请求最后一个 block 的内部碎片(internal fragmentation)。如果 block_size=16,最后一个 block 平均只用了一半,浪费约 $\frac{1}{2} \times \frac{1}{N_{\text{blocks}}}$ 的显存。
配套代码中的对比实验结果:
Config: 32 heads, head_dim=128, max_seq=2048, block_size=16
10 requests with lengths: [47, 183, 12, 891, 256, 5, 1024, 73, 330, 15]
Metric Naive Paged Actual
-----------------------------------------------------------------------
KV cache memory (MB) 320.0 45.5 44.3
Memory utilization 13.8% 97.4% 100.0%
Waste (MB) 275.7 1.2 0.0
Waste (%) 86.2% 2.6% 0.0%
朴素方式浪费 86%,PagedAttention 只浪费 2.6%。这意味着同样的 GPU 显存可以 batch 多 2-4 倍的请求,直接把吞吐量提高 2-4 倍。

Copy-on-Write:高效的序列分叉
在 beam search 和 parallel sampling 场景下,一个请求会分叉出多个候选序列。这些序列共享前缀的 KV Cache——如果每个候选都复制一份前缀的 KV,显存开销会翻倍。
PagedAttention 的解决方案是 Copy-on-Write(CoW),同样借鉴自操作系统的 fork() 系统调用:
flowchart TD
subgraph BEFORE["fork 前"]
S0a["Seq 0: Block A → Block B → Block C (ref=1)"]
end
subgraph AFTER["fork 后(Seq 0 + Seq 1)"]
S0b["Seq 0: Block A → Block B → Block C"]
S1b["Seq 1: Block A → Block B → Block C"]
REF["Block C ref=2(共享,无拷贝!)"]
end
subgraph WRITE["Seq 1 写入新 token 到 Block C"]
S0c["Seq 0: Block A → Block B → Block C (ref=1)"]
S1c["Seq 1: Block A → Block B → Block C' (ref=1, 此时拷贝)"]
NOTE2["只有被修改的 block 才拷贝,前面的 block 仍共享"]
end
BEFORE --> AFTER --> WRITE
style BEFORE fill:#cce5ff,stroke:#004085
style AFTER fill:#fff3cd,stroke:#856404
style WRITE fill:#d4edda,stroke:#155724
核心逻辑在 append_token 中:写入前检查 block 的引用计数,如果 >1(被共享),就先拷贝一份新 block,再写入:
| |
CoW 的收益在 beam search 中尤其显著:beam_size=4 时,4 个候选序列共享大量前缀 KV,只有最后一个 block(正在生成的部分)需要独立存储。显存开销接近 1 个序列而非 4 个。
RadixAttention(SGLang 核心创新)
从 PagedAttention 到 Prefix Caching
PagedAttention 解决了单个请求内部的显存碎片问题。但它忽略了一个更大的优化机会:不同请求之间的 KV Cache 复用。
在实际的 LLM 应用中,大量请求共享相同的前缀:
flowchart TD
subgraph S1["场景 1: System Prompt"]
SP["共享前缀: You are a helpful assistant..."]
UA["+ 用户 A: What is machine learning?"]
UB["+ 用户 B: Write a Python function for..."]
UC["+ 用户 C: Translate this text..."]
SP --> UA
SP --> UB
SP --> UC
end
subgraph S2["场景 2: Few-shot Learning"]
FS["共享前缀: Here are some examples: × 5 样例"]
Q1["+ Input: new query 1"]
Q2["+ Input: new query 2"]
FS --> Q1
FS --> Q2
end
subgraph S3["场景 3: Multi-turn Chat"]
MT["共享前缀: System + Turn 1 + Turn 2"]
T3["+ User Turn 3(新消息)"]
MT --> T3
end
style S1 fill:#cce5ff,stroke:#004085
style S2 fill:#d4edda,stroke:#155724
style S3 fill:#fff3cd,stroke:#856404
如果 system prompt 有 500 个 token,每个请求都重新计算和存储这 500 个 token 的 KV Cache,那么 100 个并发请求就重复存了 100 份——完全相同的数据。
Prefix Caching 的核心思想:如果两个请求的 token 序列有相同的前缀,它们的 KV Cache 在这些位置上完全一样,可以直接共享,无需重新计算。
vLLM 后来也加入了 prefix caching 功能,但它的实现基于预定义的前缀匹配——你需要显式声明哪些请求共享前缀。SGLang 的 RadixAttention 提供了一种更灵活、更优雅的方案。
Radix Tree:索引任意前缀的数据结构
SGLang(LMSYS,2024)的核心创新是用 Radix Tree(基数树) 来索引和管理所有已缓存的 KV Cache。
Radix Tree 是一种压缩前缀树(compressed trie)。普通 trie 每个节点代表一个字符,radix tree 将只有一个子节点的路径压缩成单个节点,减少树的深度。在 SGLang 中,每个节点代表一段 token 序列,节点上存储着对应的 KV Cache block 指针。
graph TD
ROOT["Root"]
SYS["System Prompt<br/>KV blocks: B0, B1, B2"]
ML["User: What is ML?<br/>KV blocks: B3, B4"]
PY["User: Write Python<br/>KV blocks: B5, B6"]
AST["Asst: ML is...<br/>KV blocks: B7"]
ROOT --> SYS
SYS --> ML
SYS --> PY
ML --> AST
REQD["新请求 D: System Prompt + What is ML? + new query<br/>→ 匹配 B0-B4 直接复用,只需为 new query 计算新 KV"]
style SYS fill:#cce5ff,stroke:#004085
style ML fill:#d4edda,stroke:#155724
style PY fill:#d4edda,stroke:#155724
style AST fill:#fff3cd,stroke:#856404
style REQD fill:#f8f9fa,stroke:#6c757d

工作流程
当一个新请求到达时,SGLang 的调度器执行以下步骤:
flowchart TD
subgraph FLOW["RadixAttention 处理新请求"]
A["1. PREFIX MATCHING<br/>在 Radix Tree 中查找最长匹配前缀<br/>例: sys_prompt 的 KV blocks 可直接复用"]
B["2. REUSE CACHED KV<br/>匹配到的前缀部分跳过 prefill<br/>只需计算不匹配的后缀部分的 KV"]
C["3. INSERT INTO TREE<br/>请求完成后将新 KV 插入 Radix Tree<br/>后续相同前缀的请求可复用"]
A --> B --> C
end
subgraph TIMELINE["时间对比"]
T1["无 cache: ====== prefill 500 tokens ====== decode"]
T2["有 cache: == prefill 100 tokens == decode<br/>400 tokens 的 KV 从 cache 复用,跳过计算"]
end
FLOW --> TIMELINE
style FLOW fill:#cce5ff,stroke:#004085
style T1 fill:#fff3cd,stroke:#856404
style T2 fill:#d4edda,stroke:#155724
在 system prompt 场景下(几乎所有请求共享同一个 system prompt),prefix caching 可以跳过大部分 prefill 计算,将 TTFT(Time To First Token)降低 50-90%。
LRU 逐出策略
GPU 显存有限,不可能缓存所有见过的前缀。当 block 池用满时,需要逐出一些缓存来给新请求腾空间。
SGLang 使用 LRU(Least Recently Used) 策略:最久没被访问的 KV Cache 节点优先被逐出。
flowchart TD
subgraph NODES["Radix Tree 节点(按最近访问时间排序)"]
N1["sys_prompt + user_A — 10 sec ago ✅ 保留"]
N2["sys_prompt + user_B — 30 sec ago ✅ 保留"]
N3["sys_prompt + user_C — 120 sec ago ⚠️ 候选逐出"]
N4["old_prompt + old_query — 300 sec ago ❌ 优先逐出"]
end
subgraph EVICT["逐出流程"]
E1["1. 从 LRU 链表尾部取最久未访问的节点"]
E2["2. 释放该节点的 KV blocks"]
E3["3. 从 Radix Tree 中删除该节点"]
E4["4. 若父节点变成无引用叶子,递归清理"]
E1 --> E2 --> E3 --> E4
end
style N1 fill:#d4edda,stroke:#155724
style N2 fill:#d4edda,stroke:#155724
style N3 fill:#fff3cd,stroke:#856404
style N4 fill:#f8d7da,stroke:#721c24
LRU 的合理性在于时间局部性:最近被访问的前缀往往会再次被请求到(比如同一个用户的多轮对话,或热门的 system prompt)。
Cache-Aware Scheduling
有了 Radix Tree 之后,调度器可以做一个非常聪明的优化:优先调度那些在 cache 中匹配前缀更长的请求。
flowchart TD
subgraph QUEUE["等待队列"]
RA["Req A: 匹配 500 tokens 缓存 → prefill 仅需 100 tokens"]
RB["Req B: 匹配 0 tokens → prefill 需要 600 tokens"]
RC["Req C: 匹配 300 tokens 缓存 → prefill 仅需 200 tokens"]
end
subgraph SCHED["调度策略对比"]
FCFS["朴素 FCFS: A → B → C"]
CACHE["Cache-aware: A → C → B(优先 cache 命中率高的)"]
end
subgraph BENEFIT["好处"]
B1["1. 高 cache 命中的请求 prefill 更快 → TTFT 更低"]
B2["2. KV 留在 tree 中 → 后续相似请求也命中 → 正反馈循环"]
B3["3. 减少总 prefill 计算量 → 更多 GPU 时间给 decode"]
end
QUEUE --> SCHED --> BENEFIT
style RA fill:#d4edda,stroke:#155724
style RB fill:#f8d7da,stroke:#721c24
style RC fill:#fff3cd,stroke:#856404
style CACHE fill:#d4edda,stroke:#155724
这种调度策略使 SGLang 在 few-shot learning、multi-turn chat 等共享前缀场景下性能远超不做 prefix caching 的系统。
与 vLLM 的对比
vLLM(v0.3+)后来也加入了 prefix caching 功能,但设计思路不同:
| 特性 | SGLang (RadixAttention) | vLLM (Prefix Caching) |
|---|---|---|
| 数据结构 | Radix Tree | Hash Map |
| 前缀匹配 | 自动匹配任意共享前缀 | 基于 token block hash |
| 匹配粒度 | token 级别 | block 级别(block_size 对齐) |
| 逐出策略 | LRU on tree nodes | LRU on blocks |
| Cache-aware 调度 | 原生支持 | 后续版本加入 |
| 多轮对话 | 天然支持(树结构) | 需要 hash 匹配 |
SGLang 的 Radix Tree 方案在灵活性上更胜一筹:
- 任意前缀共享:不限于预定义的 prompt template,任何两个请求只要有共同前缀就能自动匹配
- 树结构天然适配多轮对话:对话的历史消息自然形成一棵树,不同分支代表不同的后续对话
- 渐进式匹配:可以匹配到 token 级别的精确前缀,而非 block 级别的近似匹配
vLLM 的 hash-based 方案在实现简单性和大规模部署稳定性上有优势,两者在实际性能上差距取决于具体的工作负载特征。
Continuous Batching
静态 Batching 的问题
在理解了 KV Cache 的内存管理之后,我们来看另一个维度的优化:调度策略。
最朴素的推理方式是 静态 batching:把一组请求凑成一个 batch,一起做 prefill,然后一起做 decode,等 batch 中所有请求都生成完毕后,才处理下一个 batch。
问题在于:不同请求的输出长度差异巨大。
flowchart TD
subgraph STATIC["静态 Batching — Batch = Req 0-3, 输出长度: 3, 8, 2, 5"]
direction LR
R0["Req 0: ## ## ## .. .. .. .. ..<br/>第 3 步结束"]
R1["Req 1: ## ## ## ## ## ## ## ##<br/>最长,其他人都在等"]
R2["Req 2: ## ## .. .. .. .. .. ..<br/>第 2 步结束"]
R3["Req 3: ## ## ## ## ## .. .. ..<br/>第 5 步结束"]
end
WASTE["## = 有效计算 · .. = GPU 空闲<br/>GPU 利用率: 18/32 = 56.25%<br/>浪费了 43.75% 的 GPU 算力<br/>Req 4, 5 必须等 Req 1 结束才能开始"]
style STATIC fill:#fff3cd,stroke:#856404
style WASTE fill:#f8d7da,stroke:#721c24
当输出长度方差很大时(这在实际中几乎总是如此——有人问"是几点了",有人要写一篇长文),静态 batching 的 GPU 利用率可以低至 20-50%。
迭代级调度
Continuous Batching(Orca, OSDI'22)的核心改进是将调度粒度从 batch 级别 降到 iteration 级别:
每一个 decode step 结束后,立即检查:
- 哪些请求已经结束(生成了 EOS 或达到 max_length)?→ 驱逐,释放 KV Cache
- 等待队列中有没有新请求可以加入?→ 纳入,开始 prefill
flowchart TD
subgraph CB["Continuous Batching — 迭代级调度"]
direction LR
R0["Req 0: ## ## ## — step 3 结束"]
R1["Req 1: ## ## ## ## ## ## ## ## — step 8 结束"]
R2["Req 2: ## ## — step 2 结束"]
R3["Req 3: ## ## ## ## ## — step 5 结束"]
R4["Req 4: step 3 加入 → ## ## ## ##"]
R5["Req 5: step 4 加入 → ## ## ## ## ## ##"]
end
RESULT["每步活跃数: 4 4 4 4 4 3 3 3 1<br/>GPU 几乎一直满载!结束的请求立刻被新请求替换"]
style CB fill:#d4edda,stroke:#155724
style RESULT fill:#cce5ff,stroke:#004085
配套代码中的对比实验显示,continuous batching 在吞吐量上通常有 2-3 倍的提升:
| |

Preemption:优雅应对显存压力
Continuous batching 还引入了一个关键机制:preemption(抢占)。当 KV Cache 显存不够为所有 running 请求分配新 block 时,调度器可以:
- 暂停一个或多个 running 请求(通常选最新加入的——LIFO 策略,因为它们做的工作最少)
- 释放被暂停请求的 KV Cache blocks
- 用腾出的空间继续服务其他请求
- 稍后恢复被暂停的请求(重新做 prefill 来重建 KV Cache)
flowchart TD
A["正常运行<br/>Running: Req 0, 1, 2, 3<br/>KV blocks: 200/256 (78%)"]
B["Req 1 生成长回复, KV 持续增长...<br/>KV blocks: 253/256 (99%) — 即将 OOM!"]
C["触发 Preemption<br/>暂停 Req 3 (LIFO: 最新, 做的工作最少)<br/>释放 Req 3 的 20 个 blocks"]
D["KV blocks: 233/256 (91%)<br/>有空间继续了"]
E["稍后 Req 0/1 结束, 空间释放<br/>→ 恢复 Req 3, 重新 prefill"]
F["替代方案: 不做 preemption → OOM crash → 所有请求失败!"]
A --> B --> C --> D --> E
B -.-> F
style A fill:#d4edda,stroke:#155724
style B fill:#f8d7da,stroke:#721c24
style C fill:#fff3cd,stroke:#856404
style D fill:#d4edda,stroke:#155724
style F fill:#f8d7da,stroke:#721c24
Preemption 的代价是被暂停的请求需要重新做 prefill(重建 KV Cache),增加了该请求的延迟。但比起 OOM 导致所有请求崩溃,这是一个合理的 trade-off。
配套代码中的 preemption demo 显示,即使在非常受限的显存下(32 blocks),preemption 也能保证所有请求最终完成:
Config: 20 requests, 32 KV blocks, max_batch=4
Preemptions triggered: 12
Preemption allows the system to handle load spikes without OOM crashes.
All 20 requests completed: True
Chunked Prefill
长 Prompt 阻塞解码
Prefill 阶段(处理用户输入的 prompt)和 decode 阶段(逐 token 生成)有截然不同的计算特性:
- Prefill:一次性处理整个 prompt(可能数千个 token),是 Compute-bound 的大 GEMM 操作
- Decode:每步只生成 1 个 token,是 Memory-bound 的 GEMV 操作
问题在于:如果一个请求的 prompt 有 8000 个 token,做 prefill 可能需要几百毫秒。在这段时间里,所有正在 decode 的请求都被阻塞了——它们在等 GPU 做完这个长 prefill。
flowchart TD
subgraph NOCHUNK["无 Chunked Prefill"]
GPU1["GPU: ====== Prefill 8K tokens (200ms) ====== → D D D D"]
RA1["Req A: 等待 200ms... 然后才能 decode"]
RB1["Req B: 等待 200ms... 然后才能 decode"]
end
PROBLEM["问题: Req A, B 的 TPOT 从 10ms 暴增到 200ms+<br/>用户体验: 正在打字的对话突然卡住 200ms"]
style NOCHUNK fill:#fff3cd,stroke:#856404
style PROBLEM fill:#f8d7da,stroke:#721c24
对于正在和模型对话的用户来说,TPOT 从 10ms 跳到 200ms 会造成明显的"卡顿感"——字在流式输出时突然停了一下。
分块 Prefill
Chunked Prefill 的解决方案直截了当:把长 prompt 的 prefill 分成多个小块(chunk),每个 chunk 处理固定数量的 token(比如 512 个),每个 chunk 之间穿插一轮 decode step。
flowchart TD
subgraph CHUNK["Chunked Prefill — 8192 tokens, chunk_size=512, 共 16 chunk"]
GPU2["GPU: P1 → D → P2 → D → P3 → D → P4 → D → P5 ..."]
LEGEND["P = prefill chunk (512 tokens, ~12ms) · D = decode step (~3ms)"]
end
subgraph DECODE["Decode 请求视角"]
RA2["Req A: D — D — D — D — D (每 ~15ms 一次)"]
RB2["Req B: D — D — D — D — D (每 ~15ms 一次)"]
end
RESULT2["TPOT: ~15ms(稳定!)vs 无 chunking 时的 200ms 尖峰"]
CHUNK --> DECODE
DECODE --> RESULT2
style CHUNK fill:#d4edda,stroke:#155724
style RESULT2 fill:#cce5ff,stroke:#004085
Trade-off:
- Prefill 总时间变长了:因为每个 chunk 之间插入了 decode step,8K token 的 prefill 从 200ms 变成了约 240ms(TTFT 略增)
- Decode 延迟变稳定了:TPOT 从 200ms 的尖峰降到稳定的 15ms
这是一个 TTFT vs TPOT 的 trade-off:
| 指标 | 无 Chunked Prefill | 有 Chunked Prefill |
|---|---|---|
| TTFT (首 token) | 200ms | ~240ms (+20%) |
| TPOT (逐 token,最坏情况) | 200ms (被阻塞) | ~15ms (稳定) |
| TPOT (P99) | 极差 | 可控 |
| 用户体验 | 偶尔严重卡顿 | 流畅稳定 |
在大多数在线服务场景下,稳定的 TPOT 比略低的 TTFT 更重要——用户更在意"打字流畅"而非"快 40ms 看到第一个字"。
Chunk 大小的选择
Chunk 大小的选择需要平衡:
- 太小(如 64 tokens):prefill 的 GPU 利用率低(小矩阵乘法效率差),TTFT 显著增加
- 太大(如 4096 tokens):接近不分块的情况,decode 仍然会被长时间阻塞
- 典型值:512-1024 tokens,兼顾 prefill 效率和 decode 稳定性
SGLang 和 vLLM 都支持 chunked prefill,并可根据当前 running batch 的状态动态调整 chunk 大小——如果当前没有 decode 请求,就不需要分块,直接做完整的 prefill 以最小化 TTFT。
推理系统全景
将以上所有组件整合起来,一个现代 LLM 推理引擎(如 SGLang)的整体架构如下:
flowchart TD
HTTP["HTTP Server<br/>接收用户请求"]
subgraph SCHED["Scheduler(调度器 — 大脑)"]
WQ["Waiting Queue"]
RT["Radix Tree<br/>(Prefix Cache 索引)"]
STEPS["每个 iteration:<br/>1. 驱逐已完成请求<br/>2. 匹配新请求前缀<br/>3. Cache-aware 排序 + 纳入<br/>4. 必要时 preempt<br/>5. Chunked prefill 调度"]
end
subgraph BM["Block Manager(块管理器 — 内存)"]
POOL2["Physical Block Pool: B0 B1 B2 B3 B4 B5 ..."]
BMOPS["按需分配/释放 · CoW · LRU 逐出"]
end
subgraph ME["Model Executor(模型执行器 — 计算)"]
ATTN["Attention kernel (FlashAttention / FlashInfer)"]
KV["从 Block Table 读取 KV(非连续物理地址)"]
PAR["TP / PP 并行推理"]
end
HTTP --> SCHED --> BM --> ME
style SCHED fill:#cce5ff,stroke:#004085
style BM fill:#fff3cd,stroke:#856404
style ME fill:#d4edda,stroke:#155724
调度器是大脑——它决定每一步执行哪些请求,是 prefill 还是 decode,是否需要 preempt。Block Manager 是内存管理器——它负责分配、释放、共享 KV Cache blocks。Model Executor 是执行引擎——它拿到调度器的决策和 block table 的映射,执行实际的 attention 计算。
三者协同工作,构成了一个高效的推理流水线。
Key Takeaways
KV Cache 是推理的显存瓶颈。每个 token 需要存储 $2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times \text{dtype_size}$ 的 KV 数据。LLaMA-70B 下每个 token 约 2.5 MB,100 个并发请求、每个 2048 tokens 就需要 500 GB 的 KV Cache——远超单卡显存。能 batch 多少请求、能处理多长序列,完全取决于 KV Cache 的管理效率。
PagedAttention 用 OS 虚拟内存思想解决显存碎片。固定大小的 block 按需分配,block table 做逻辑-物理映射,CoW 支持高效的序列分叉。将显存浪费从 60-80% 降到 <5%,同样的 GPU 可以 batch 2-4 倍的请求。
RadixAttention 实现跨请求的 KV Cache 复用。Radix Tree 索引所有已缓存的 token 前缀,新请求自动匹配最长前缀并复用其 KV Cache。结合 LRU 逐出和 cache-aware scheduling,在共享前缀场景下(system prompt、few-shot、multi-turn chat)大幅降低 TTFT。
Continuous Batching 将调度粒度从 batch 级降到 iteration 级。每步驱逐已完成请求、纳入等待请求,GPU 利用率从 20-50% 提升到 80-95%。配合 preemption 机制,优雅应对显存压力,避免 OOM 崩溃。
Chunked Prefill 解决长 prompt 阻塞 decode 的问题。将 prefill 分成小块,穿插 decode step,用略高的 TTFT 换取稳定的 TPOT。chunk 大小的选择需要根据负载特征权衡。
这些技术不是互斥的,而是层层叠加的。现代推理引擎(SGLang、vLLM)同时使用 PagedAttention + Prefix Caching + Continuous Batching + Chunked Prefill,每一层解决不同维度的效率问题,叠加起来实现了比朴素实现 2-10 倍的吞吐量提升。
配套代码
本文配套代码位于 code/03-inference-sglang/:
kv_cache_from_scratch.py— 从零实现 KV Cache 和 PagedAttention Block Manager。包含 4 个部分:(1) 朴素 KV Cache 的浪费演示;(2) PagedAttention block 分配、释放、CoW 的完整实现;(3) 用 block-managed KV cache 跑一个 mini decode loop;(4) 朴素 vs PagedAttention 的内存效率对比。运行方式:python kv_cache_from_scratch.pycontinuous_batching.py— 静态 batching vs continuous batching 的对比模拟。包含 4 个部分:(1) 静态/连续 batching 的吞吐量对比;(2) ASCII 调度时间线可视化;(3) preemption 在显存压力下的表现;(4) 不同 batch size 下的 scaling 分析。运行方式:python continuous_batching.py
所有代码支持 CPU 模式(无需 GPU),方便在任何环境下学习核心逻辑。
参考资料
- Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., SOSP'23 (vLLM) — PagedAttention 的原始论文,将 OS 虚拟内存思想引入 KV Cache 管理
- SGLang: Efficient Execution of Structured Language Model Programs — Zheng et al., 2024 (LMSYS) — RadixAttention 和 cache-aware scheduling 的原始论文
- Orca: A Distributed Serving System for Transformer-Based Generative Models — Yu et al., OSDI'22 — Continuous Batching(iteration-level scheduling)的开创性工作
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao et al., NeurIPS'22 — 推理引擎中 attention kernel 的核心优化
- vLLM Documentation — docs.vllm.ai — vLLM 官方文档
- SGLang Documentation — sgl-project.github.io — SGLang 官方文档
- Sarathi-Serve: Chunked-Prefills for Efficient LLM Serving — Agrawal et al., 2024 — Chunked Prefill 的系统化分析