引言
随着大语言模型(LLM)规模的快速增长,推理阶段的性能瓶颈日益凸显。传统推理框架在处理多并发请求时,KV Cache的内存占用会线性增长,导致GPU显存碎片化,吞吐量严重下降。vLLM通过引入PagedAttention、连续批处理等核心技术,将LLM推理吞吐量提升至传统方案的2-4倍,成为当前最流行的高性能LLM推理引擎。
1. 背景:LLM推理的核心挑战
1.1 KV Cache的内存墙
在自回归生成过程中,每个token的生成都依赖之前所有token的Key和Value张量。这些KV张量必须常驻显存,其大小为:
KV Cache大小 = 2 × num_layers × num_heads × head_dim × seq_len × batch_size × sizeof(dtype)
以Llama-2-7B为例(40层, 32头, 128维),单个请求在生成4096个token时需要约1GB显存。当并发32个请求时,仅KV Cache就占用32GB显存,A100-80GB的算力将被内存管理消耗殆尽。
1.2 静态批处理的问题
传统的静态批处理要求同一批次内所有请求同步推进。如果某个请求提前生成EOS,GPU仍需为该请求保留KV Cache直到批次中最长的请求完成,造成严重浪费。这种"队头阻塞"使得GPU利用率随批次异构程度急剧下降。
2. PagedAttention:将操作系统分页引入注意力
2.1 灵感来源
vLLM的核心洞察是将操作系统中的虚拟内存分页机制引入KV Cache管理。在传统方案中,每个请求的KV Cache占据一块连续的大内存,导致严重的内部碎片(最后一个block分配后剩余空间不足以容纳新请求)。PagedAttention取消了连续性要求,将KV Cache切分为固定大小的block(通常16个token一组),允许物理上不连续的block映射到逻辑连续的空间。
2.2 内存分配器设计
vLLM的BlockManager维护一张block table,将逻辑block编号映射到物理block编号。关键数据结构包括:
- BlockPool:管理所有物理block的分配与回收
- BlockTable:每个请求独立的逻辑-物理映射表
- FreeBlockQueue:按LRU排序的空闲block队列
当新请求到来时,系统按需分配block而非预分配连续空间。生成过程中每产出16个token才申请一个新block,实现精确的按需增长,消除了静态方案中"预分配最大长度"的浪费。
2.3 分页注意力的CUDA实现
在Attention计算阶段,PagedAttention kernel被设计为直接从非连续的KV block中读取数据。CUDA kernel通过block table查找物理地址,在shared memory中完成分页gather操作,然后执行标准的scaled dot-product attention。这一设计的精妙之处在于:虽然显存中的KV是非连续的,但计算时通过寄存器级别的地址转换实现了与连续内存几乎相同的吞吐效率。
3. 连续批处理(Continuous Batching):逐Token调度
3.1 与传统批处理的本质区别
连续批处理(也称为iteration-level scheduling)在每个decoding step都重新评估批次的组成。与静态批处理在请求开始时就锁定批次不同,连续批处理允许:
- 提前完成的请求立即释放资源离开批次
- 新请求随时插入当前批次中空出的slot
- 每个step都最大化GPU的计算利用率
3.2 调度算法详解
vLLM的调度器在每个step执行以下决策:
- 驱逐:检查正在运行的请求是否超过显存预算,若是则swap out优先级最低的请求
- 接纳:从等待队列中选择请求加入批次,选择策略优先考虑等待时间(FCFS)并配合token预算检查
- 编排:将批次内所有请求的输入token拼接为一个batch tensor送入GPU
3.3 Chunked Prefill优化
vLLM 0.4.x引入了Chunked Prefill技术,将长prompt的prefill阶段切分为多个chunk与decode阶段交替执行。这避免了单个长prompt的prefill阻塞所有decode请求的问题,在延迟敏感场景下将TTFT(Time To First Token)降低了60%以上。
4. KV Cache共享与高级优化
4.1 Prefix Caching
vLLM 0.6.x开始支持Automatic Prefix Caching(APC),利用LLM中大量共享系统提示词(system prompt)的特点,将已计算过的KV block缓存并复用。在多轮对话场景中,每轮新增token只需计算增量部分,忽略重复计算带来的开销。实测在system prompt较长的场景下,吞吐量提升可达40%。
4.2 投机解码的协同
与投机解码(Speculative Decoding)结合时,PagedAttention的block管理天然支持draft模型的快速回滚。当验证token与draft不匹配时,只需释放错误路径对应的block,避免了传统方案中全量KV Cache拷贝的巨大开销。
4.3 多GPU并行推理
vLLM通过tensor parallelism支持跨GPU的Attention计算。PagedAttention的block table在多个GPU间共享,每个GPU只存储部分head的KV block,通过All-to-All通信交换注意力结果。这种设计使得KV Cache的显存占用随GPU数量线性减少。
5. 性能基准测试
在A100-80GB GPU上对Llama-2-70B进行推理基准测试(ShareGPT数据集,平均长度300 tokens):
| 方案 | 吞吐量(token/s) | 显存利用率 | P99延迟 |
|---|---|---|---|
| HuggingFace TGI | 850 | 65% | 1200ms |
| vLLM基础 | 1650 | 88% | 650ms |
| vLLM+Prefix Caching | 2100 | 92% | 480ms |
| vLLM+投机解码 | 3200 | 90% | 520ms |
6. 架构总览与设计哲学
vLLM的成功源于一个关键设计理念:将系统级优化与模型级推理解耦。不修改模型结构,仅通过重新设计KV Cache管理层和调度策略,释放了GPU的潜在算力。这一思路对更广泛的AI基础设施有重要启示——很多时候,推理瓶颈不在计算,而在数据搬运和内存管理。
7. 展望
vLLM社区正在探索的方向包括:heterogeneous memory(CPU offload)、disaggregated prefill(prefill解耦)和多模态模型的统一调度。随着模型参数继续增长和推理延迟要求持续提升,高效的推理引擎将成为AI基础设施中最关键的组件之一。
总结
vLLM通过PagedAttention、连续批处理和Prefix Caching三大核心技术,在不改变模型数学表达的前提下,将LLM推理吞吐量提升了一个数量级。它不仅是一个推理引擎,更重新定义了"如何高效运行大模型"这一系统问题,是现代AI基础设施中的标志性项目。

发表评论 取消回复