Back to blog

Efficient Memory Management for Large Language Model Serving with PagedAttention

通过 PagedAttention 算法实现高效 KV cache 内存管理,构建 vLLM 高吞吐量 LLM 服务系统

Efficient Memory Management for Large Language Model Serving with PagedAttention

论文信息: arXiv:2309.06180 [cs.LG] 12 Sep 2023

作者: Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica

机构: UC Berkeley, Stanford University, UC San Diego

会议: SOSP 2023

代码: https://github.com/vllm-project/vllm

许可: CC BY 4.0


一、论文概述

1.1 研究背景

高吞吐量的 LLM 服务需要同时处理大量请求。然而,现有系统面临 KV cache 内存管理的挑战:

挑战说明
KV cache 巨大13B 模型单个 token 需要 800KB,单请求最大 1.6GB
内存碎片现有系统仅 20.4%-38.2% 的 KV cache 内存有效使用
无法共享现有系统无法利用 KV cache 共享机会
动态增长KV cache 随时间动态增长和收缩,生命周期不可预测

现有系统的内存浪费:

浪费类型说明
预留浪费为最大序列长度预留空间,实际使用可能远小于此
内部碎片预分配的连续内存块中未使用的部分
外部碎片分配器产生的碎片,无法被其他请求利用

1.2 核心贡献

贡献说明
PagedAttention受 OS 虚拟内存启发的注意力算法,支持非连续 KV cache 存储
vLLM 系统基于 PagedAttention 的高吞吐量分布式 LLM 服务引擎
近零内存浪费通过分页管理消除内存碎片和预留浪费
KV cache 共享支持跨请求的内存共享,进一步减少内存使用

二、核心思想

2.1 问题定义

如何高效管理 LLM 服务中的 KV cache 内存,以提高服务吞吐量?

现有方案的不足:

  1. 连续内存存储导致严重的内存碎片
  2. 预分配策略浪费大量内存空间
  3. 无法利用不同解码算法中的内存共享机会

2.2 解决方案概述

PagedAttention:受 OS 虚拟内存和分页技术启发的注意力算法

核心思想:

  • 将 KV cache 分割为固定大小的块(类似 OS 的页)
  • 块可以存储在非连续的物理内存中
  • 通过块表映射逻辑块到物理块
  • 支持按需分配和动态增长

三、技术架构

3.1 现有系统问题

内存布局 图 1:左图展示 13B 模型在 A100 上的内存分布,右图展示 vLLM 相比现有系统的吞吐量提升。

内存浪费 图 2:现有 LLM 服务系统中 KV cache 内存浪费比例。

现有内存管理 图 3:现有系统内存管理的三种浪费:预留、内部碎片、外部碎片。

关键发现:

  • 现有系统仅 20.4%-38.2% 的 KV cache 内存有效使用
  • 13B 模型单个请求的 KV cache 可达 1.6GB
  • 连续内存存储导致严重的内存碎片

3.2 vLLM 系统架构

vLLM 架构 图 4:vLLM 系统架构,包含中央调度器、KV cache 管理器、GPU/CPU 块分配器。

系统组件:

组件说明
中央调度器协调分布式 GPU worker 的执行
KV cache 管理器管理分页的 KV cache 内存
块引擎在 GPU/CPU 上分配和管理物理块
PagedAttention 内核执行分页注意力计算

3.3 PagedAttention 算法

PagedAttention 图 5:PagedAttention 算法示意图,KV 块存储在非连续物理内存中。

核心公式:

标准注意力: aij=exp⁡(qi⊤kj/d)∑t=1iexp⁡(qi⊤kt/d),oi=∑j=1iaijvja_{ij} = \frac{\exp(q_i^\top k_j / \sqrt{d})}{\sum_{t=1}^{i} \exp(q_i^\top k_t / \sqrt{d})}, o_i = \sum_{j=1}^{i} a_{ij} v_j

PagedAttention(分块计算): Aij=exp⁡(qi⊤Kj/d)∑t=1⌈i/B⌉exp⁡(qi⊤Kt1/d),oi=∑j=1⌈i/B⌉VjAij⊤A_{ij} = \frac{\exp(q_i^\top K_j / \sqrt{d})}{\sum_{t=1}^{\lceil i/B \rceil} \exp(q_i^\top K_t \mathbf{1} / \sqrt{d})}, o_i = \sum_{j=1}^{\lceil i/B \rceil} V_j A_{ij}^\top

其中:

  • BB:KV 块大小(block size)
  • KjK_j:第 jj 个 key 块
  • VjV_j:第 jj 个 value 块

3.4 Block Table 映射

Block Table 图 6:Block Table 翻译机制,逻辑块到物理块的映射。

映射过程:

  1. 请求的 KV cache 表示为一系列逻辑 KV 块
  2. 每个逻辑块映射到一个物理块
  3. 物理块可以存储在非连续的内存位置
  4. 块表记录每个请求的逻辑-物理块映射

3.5 内存管理机制

双请求管理 图 7:vLLM 同时管理两个请求的内存。

关键机制:

  • 按需分配:仅在需要时分配物理块
  • 动态增长:KV cache 可以动态增长,无需预分配
  • 块级共享:支持跨请求的块级内存共享
  • 引用计数:跟踪每个物理块的引用次数

四、核心创新

4.1 创新点总结

创新点说明理论/实验依据
PagedAttention分块注意力算法,支持非连续存储内存效率实验
分页内存管理受 OS 虚拟内存启发的内存管理内存浪费减少实验
Copy-on-Write块级写时复制机制并行采样和 Beam Search 实验
统一内存共享支持多种解码算法的内存共享吞吐量提升实验

4.2 解码算法支持

并行采样

并行采样 图 8:并行采样中的 copy-on-write 机制。

机制:

  • 多个输出序列共享提示的 KV cache
  • 当需要修改共享块时,使用 copy-on-write
  • 仅复制需要修改的块,减少内存开销

Beam Search 图 9:Beam Search 中的块共享模式。

机制:

  • 不仅共享提示块,还共享其他公共块
  • 动态调整共享模式
  • 释放不再需要的块,分配新块

共享前缀

共享前缀 图 10:共享前缀的机器翻译示例。

机制:

  • 预定义共享前缀的物理块
  • 用户请求映射到共享的物理块
  • 仅计算用户任务输入部分

五、实验结果

5.1 实验设置

模型配置:

模型GPU总内存参数大小KV cache 内存最大 KV 块数
13BA10040GB26GB12GB15.7K
66B4×A100160GB132GB21GB9.7K
175B8×A100640GB346GB264GB60.1K

工作负载:

  • ShareGPT:用户与 ChatGPT 的对话,输入平均 8.4× 更长
  • Alpaca:GPT-3.5 生成的指令数据集

基准系统:

  • FasterTransformer:优化延迟的推理引擎
  • Orca (Oracle):假设已知输出长度的上界
  • Orca (Pow2):输出预留最多 2× 空间
  • Orca (Max):预留最大序列长度空间

5.2 基本采样性能

OPT-13B ShareGPT 图 12a:OPT-13B 在 ShareGPT 上的性能。

OPT-66B ShareGPT 图 12b:OPT-66B 在 ShareGPT 上的性能。

OPT-175B ShareGPT 图 12c:OPT-175B 在 ShareGPT 上的性能。

关键发现:

  • vLLM 在 ShareGPT 上可维持 1.7×-2.7× 更高的请求率(相比 Orca Oracle)
  • 相比 Orca (Max) 提升 2.7×-8×
  • 相比 FasterTransformer 提升高达 22×

5.3 批处理请求数对比

批处理 ShareGPT 图 13a:ShareGPT 上的平均批处理请求数。

批处理 Alpaca 图 13b:Alpaca 上的平均批处理请求数。

关键发现:

  • OPT-13B 在 ShareGPT 上,vLLM 同时处理 2.2× 更多请求(相比 Orca Oracle)
  • 相比 Orca (Max) 处理 4.3× 更多请求

并行采样内存节省 图 14a:并行采样的内存节省。

Beam Search 内存节省 图 14b:Beam Search 的内存节省。

内存节省:

  • 并行采样:6.1%-30.5% 内存节省
  • Beam Search:37.6%-66.3% 内存节省

性能提升:

  • 基本采样:vLLM 比 Orca (Oracle) 提升 1.3×
  • Beam Search (width=6):提升 2.3×

5.5 共享前缀性能

共享前缀 图 16:共享前缀的翻译工作负载性能。

关键发现:

  • 1-shot 前缀共享:vLLM 比 Orca (Oracle) 提升 1.67×
  • 5-shot 前缀共享:提升 3.58×
  • 共享更多示例时,提升更显著

5.6 聊天机器人性能

聊天机器人 图 17:聊天机器人工作负载性能。

关键发现:

  • vLLM 可维持 2× 更高的请求率
  • 在长对话场景下优势更明显
  • PagedAttention 有效处理长提示

5.7 消融实验

内核延迟 图 18a:注意力内核延迟对比。

块大小影响 图 18b:不同块大小的端到端延迟。

关键发现:

  • PagedAttention 内核开销可忽略不计
  • 块大小影响性能:较大块减少内核开销但增加碎片
  • 块大小 16 是较好的平衡点

六、代码实现分析

6.1 项目结构

  • 前端: FastAPI,扩展 OpenAI API 接口
  • 引擎: 8.5K 行 Python + 2K 行 C++/CUDA
  • 组件: 调度器、块管理器(Python),PagedAttention 内核(CUDA)
  • 模型: 支持 GPT、OPT、LLaMA 等

6.2 内核优化

优化说明
融合 reshape 和块写入最小化内核启动开销
融合块读取和注意力支持非连续块读取
融合块复制批量处理 copy-on-write 操作

6.3 解码算法实现

方法说明
fork从现有序列创建新序列
append向序列追加新 token
free删除序列

七、相关工作

7.1 LLM 服务系统

系统特点局限性
FasterTransformer优化延迟无调度器,内存管理低效
Orca迭代级调度连续内存存储,无法共享
Triton动态批处理类似 Orca 的内存管理
vLLM分页内存管理近零内存浪费,支持共享

7.2 内存管理技术

技术来源应用
虚拟内存OSKV cache 分页管理
Copy-on-WriteOS并行采样和 Beam Search
块分配OS按需内存分配

八、总结

8.1 核心贡献

  1. PagedAttention:受 OS 虚拟内存启发的注意力算法,支持非连续 KV cache 存储
  2. vLLM 系统:基于 PagedAttention 的高吞吐量分布式 LLM 服务引擎
  3. 内存效率:实现接近零的 KV cache 内存浪费
  4. 内存共享:支持跨请求的 KV cache 共享,进一步减少内存使用

8.2 技术影响

  • 服务成本:通过提高吞吐量降低 LLM 服务成本
  • 系统设计:为 LLM 服务系统提供新的内存管理范式
  • 开源生态:vLLM 成为最流行的 LLM 服务框架之一
  • 后续研究:启发了大量关于 KV cache 优化的研究

8.3 局限性

  1. 块大小权衡:较大块减少内核开销但增加内存碎片
  2. 交换开销:CPU-GPU 内存交换可能成为瓶颈
  3. 调度策略:FCFS 调度可能不是最优的
  4. 模型限制:主要评估了 OPT 和 LLaMA 模型

九、参考资源

9.1 论文链接

9.2 关键图表

图表说明路径
图 1内存布局与吞吐量figure-1-memory-layout.jpg
图 2内存浪费统计figure-2-memory-waste.jpg
图 3现有内存管理figure-3-existing-memory-management.jpg
图 4vLLM 架构figure-4-vllm-overview.jpg
图 5PagedAttentionfigure-5-pagedattention.jpg
图 6Block Tablefigure-6-block-table.jpg
图 7双请求管理figure-7-two-requests.jpg
图 8并行采样figure-8-parallel-sampling.jpg
图 9Beam Searchfigure-9-beam-search.jpg
图 10共享前缀figure-10-shared-prefix.jpg

9.3 相关论文

论文作者年份关系
OrcaYu et al.2022基准系统
FasterTransformerNVIDIA2022基准系统
Megatron-LMShoeybi et al.2020模型并行
Virtual MemoryOS-核心思想来源

9.4 关键技术术语

术语英文说明
PagedAttentionPagedAttention分页注意力算法
KV CacheKey-Value Cache注意力机制的键值缓存
块表Block Table逻辑块到物理块的映射
Copy-on-WriteCopy-on-Write写时复制机制
迭代级调度Iteration-level Scheduling每个迭代更新批次
模型并行Model Parallelism跨 GPU 分布模型参数

分析完成时间:2026年6月24日 分析工具:Claude Code + paper-analyzer skill