技术笔记

vLLM 推理原理 从 KV Cache 到分页与连续批处理

用共享书架和排队的例子理解 KV Cache、PagedAttention 与连续批处理,再手算上下文长度和 KV 头数怎样影响显存。

作者与贡献:基于官方文档与论文整理的原理笔记;公式为理论估算,配套 GPU 实测见独立实验记录。

大模型推理通常不只是“加载权重,然后跑一次前向”。对聊天模型来说,生成过程会持续增加 token;服务端还要让长短不一的请求共享 GPU。理解 vLLM,可以先看清两件事:生成中保留哪些数据,下一轮让哪些请求参与计算。

系列入口:vLLM 与 DeepSpeed 学习路线

一次生成有两个阶段

假设输入是“请用一句话解释 KV Cache”。模型先处理输入,再逐步输出回答。

阶段 在做什么 可以怎样理解
Prefill 处理输入 token,计算并保存其 KV 等状态 先把题目读一遍,做好索引
Decode 根据已有内容继续生成,逐步扩展缓存 每写下一部分,就更新笔记

Prefill 常偏计算密集,Decode 在许多常见服务负载下更受访存影响,但具体瓶颈仍取决于模型、批量和硬件。vLLM 的分块 Prefill 会把较长的输入处理拆小,与 Decode 工作一起调度。vLLM 优化与调优

flowchart TB
    A[输入文本转为 token] --> B[Prefill 处理输入并建立 KV]
    B --> C[产生首个输出 token]
    C --> D[Decode 读取历史 KV 并扩展缓存]
    D --> E{是否结束}
    E -->|否| D
    E -->|是| F[结束请求并回收资源]

这里展示的是常规自回归生成流程;推测解码等机制会改变具体执行组织。

KV Cache 缓存的究竟是什么

在注意力计算里,当前 Query 会与历史 Key 计算相关性,再结合 Value 形成输出。常规自回归推理会保存历史 token 的 Key 和 Value,后续不必反复重新计算它们。

它类似阅读时给已经看过的内容做好索引。索引不等于结论:保存 KV 不代表直接保存了一份可以复用给任意问题的答案。PagedAttention 论文关注的正是这类缓存的显存管理。PagedAttention 论文

对于各层形状相同、保留完整历史的普通注意力,单条序列可以作如下容量估算:

$$ M_{KV} \approx 2\times L\times T\times H_{KV}\times d_h\times b $$

符号 含义
$2$ Key 和 Value 两份
$L$ 注意力层数
$T$ 已缓存的 token 数,包括仍保留的输入和生成内容
$H_{KV}$ KV 头数,GQA 中通常小于 Query 头数
$d_h$ 每个头的维度
$b$ 每个缓存元素的字节数,例如 FP16 为 2

取一组假设参数:$L=32$,$T=4096$,$H_{KV}=32$,$d_h=128$,$b=2$,得到 $2,147,483,648$ 字节,即 2 GiB。如果只把 KV 头数改为 8,则为 0.5 GiB。

这不是某个指定模型的实测。多个活跃序列可先按各自缓存长度求和;共享前缀、分页取整和运行时策略会改变实际分配。滑动窗口、MLA、混合架构也不能直接套用这组假设。公式的用途是解释趋势:更长的历史、更多活跃请求,会增加缓存需求。

PagedAttention 怎样整理这些缓存

想象一个不断收到新书的共享书架。如果每位读者必须预留一整段连续空间,有些位置暂时空着,另一些位置又放不下新书。分页思路把缓存分成固定大小的块,用映射表记录逻辑顺序,物理位置不必连续。缓存扩展时可以继续分配块。PagedAttention 论文

逻辑 KV 块通过块表映射到不连续的物理显存块

图中 A 的三个逻辑块对应物理块 2、5、1。数字表示地址映射示例,不是实际 vLLM 的块大小或布局。最后一个块也可能未完全用满,因此分页并不等于完全消除浪费。

分页的价值是让缓存分配更灵活,并为合适场景下的共享提供基础;它不会把已经需要保存的 KV 元素凭空压缩成零。

阅读源码时还要注意版本:官方的 Paged Attention 设计页明确标注为历史文档,不能将其中旧内核描述直接当作当前实现。历史设计页说明

连续批处理解决请求长短不齐

假设 A 要输出 3 步,B 要输出 6 步,C 已在等待。固定一批请求直到整批结束,可能让已经结束的 A 留下空位。按迭代组织的连续批处理则可以让新请求加入后续调度。

生成迭代示意 参与请求 发生了什么
第 1 至 3 轮 A、B 两个请求都未结束
第 4 轮起 B、C A 结束后,C 有机会进入

这张表只解释调度思想,真实调度还受 token 预算、缓存容量和请求状态约束。vLLM V1 采用持续调度组织推理工作。vLLM V1 指南

因此,分页管理回答“缓存放在哪里”,连续批处理回答“这轮处理哪些请求”。两者配合才能理解服务系统;只看某一个内核的速度还不够。

前缀缓存能省掉哪部分工作

如果多次请求都以相同的长文档开头,Automatic Prefix Caching 可以在命中时复用共同前缀的 KV,减少重复 Prefill。它要求可复用的前缀匹配,不能把“语义相近”直接当成命中;新问题的生成仍需要 Decode。vLLM 前缀缓存

对学习笔记问答系统,可以先固定“系统说明+同一篇资料”,只改变末尾问题,观察命中前后的 Prefill 耗时。然后更换资料前部,检查缓存收益是否消失。这样能把缓存机制与回答质量分开分析。

看性能时要同时记录什么

指标 通俗含义 容易漏掉的条件
TTFT 发出请求后,等到首个输出 token 的时间 排队、输入长度、缓存冷热
ITL 连续输出 token 之间的间隔 需要流式输出和一致测量口径
输出吞吐 单位时间生成多少输出 token 并发、输出长度、测量窗口
峰值显存 负载期间最高占用 allocated、reserved 与进程显存口径不同

vLLM 提供 TTFT、排队、Prefill、Decode 等服务指标。客户端还包含网络与接收开销,不能把客户端总延迟直接当成 GPU 计算时间。vLLM 生产指标

下一步可以做一个小验证:同一模型下保持输入和输出长度一致,逐步增加并发,观察吞吐提高时 TTFT 是否也变差。这比只记录“感觉回答快了”更有解释力。

继续阅读:DeepSpeed 与 ZeRO 训练显存到底花在哪里 · 最小实验

END OF NOTE✳
← 返回文章列表