PagedAttention

核心贡献

提出 PagedAttention：将 KV cache 分页存储（类比操作系统虚拟内存），解决大模型推理时的显存碎片化和预分配浪费问题

构建 vLLM 系统：基于 PagedAttention 实现近零显存浪费（< 4%）和灵活的 KV cache 跨请求共享

量化内存问题规模：现有系统由于碎片化和过度预留，浪费 60%–80% 的显存；LLaMA-13B 单序列 KV cache 可高达 1.7 GB

显著提升推理吞吐量：相比 FasterTransformer/Orca 提升 2–4×，相比 HuggingFace Transformers 提升最高 24×

支持多种高级解码策略：Parallel Sampling 和 Beam Search 可共享 KV cache，内存节省最高 55%，吞吐提升最高 2.2×

背景与问题动机

LLM 推理（autoregressive decoding）需要缓存每个 Transformer 层的键值对（Key-Value Cache），以避免重复计算。这带来了严峻的内存管理挑战：

KV cache 体积大：对 LLaMA-13B，单序列可达 1.7 GB

大小动态变化：序列长度在生成过程中不断增长，且各请求间差异极大

现有方案的弊端：为每个序列预分配连续显存块，导致：

• 内部碎片（Internal Fragmentation）：预分配最大长度但实际使用更短
• 外部碎片（External Fragmentation）：不同大小的显存块难以复用
• 现有系统显存利用率约 20%–40%

方法：PagedAttention

核心思想

将 KV cache 按固定大小的物理块（Block）管理，每个请求维护一张逻辑块到物理块的映射表（Block Table），逻辑地址连续但物理地址可以不连续——与操作系统的虚拟内存分页完全同构。

序列的逻辑 KV cache:
  [Block 0] [Block 1] [Block 2] ...（逻辑连续）
      ↓ Block Table 映射
  物理块 #7,  物理块 #3,  物理块 #15 ...（物理不连续）

分块 Attention 计算

设每个 block 包含 $B$ 个 token 的 KV 对。对于查询向量 $q_i$，PagedAttention 分块计算 attention：

$A_{ij} = \frac{\exp(q_i^\top k_j / \sqrt{d})}{\sum_{t=1}^{|s|} \exp(q_i^\top k_t / \sqrt{d})}, \quad o_i = \sum_j A_{ij} v_j$

其中每个物理块存储 $K_j = (k_{(j-1)B+1}, \ldots, k_{jB})$ 和 $V_j = (v_{(j-1)B+1}, \ldots, v_{jB})$。计算时按块逐一读取，最终通过 online softmax 归约得到输出。

Block Table 结构

显存浪费仅发生在最后一个不满的块，在实践中平均浪费 < 4%

KV Cache 共享机制

Parallel Sampling（并行采样）：同一 prompt 生成多个输出序列时，prompt 对应的 KV cache 物理块可被所有输出序列引用计数共享，写时复制（Copy-on-Write）保证正确性。

Beam Search：多个 beam 共享相同前缀的物理块，分叉点后各自独立分配新块。

Prefix Caching：系统提示（System Prompt）等固定前缀在不同请求间复用同一批物理块，无需重复计算。

vLLM 系统架构

┌─────────────────────────────────────┐
│           vLLM Scheduler            │  ← 请求队列管理、抢占策略
├─────────────────────────────────────┤
│         KV Cache Manager           │  ← Block Table 分配/释放/共享
├──────────────┬──────────────────────┤
│   Worker 0   │   Worker 1 ...      │  ← 多 GPU 张量并行执行
│ (GPU 0)      │ (GPU 1)             │
└──────────────┴──────────────────────┘

Scheduler：实现 continuous batching，动态将等待队列中的请求加入运行批次

KV Cache Manager：全局管理物理块的分配、引用计数、写时复制

Worker：执行 PagedAttention kernel，通过 Block Table 访问非连续物理块

抢占策略：当显存不足时，将低优先级序列的 KV cache swap 到 CPU 或直接重计算

实验结果

对比基线

高级解码策略收益

生产部署

LMSYS Chatbot Arena 部署数据：

日均处理 30,000 请求，峰值 60,000

相比原 HuggingFace 后端，GPU 用量减少 50%

与相关工作的区别

核心差异在于：Orca 已解决 batching 调度问题，但未解决内存碎片问题；PagedAttention 在 Orca 思想基础上，进一步解决了显存的高效管理。

与本研究的关联

直接关联：本工作与模型增长/渐进式训练主要针对训练阶段；vLLM 针对推理阶段，两者正交互补。在模型增长后的部署阶段，可直接采用 vLLM 作为推理引擎。

启发点：PagedAttention 中”按需动态分配、共享前缀”的思想，与渐进式训练中”按需扩展网络容量”有相似的系统设计哲学

实用价值：研究增长后模型的推理效率时，KV cache 显存占用是重要约束，PagedAttention 提供了目前最优的显存管理方案

局限性与未来工作

分块管理引入少量计算开销（block 地址解析），但实测可忽略不计

Block size 需要权衡：过小增加 overhead，过大增加内部碎片

当前实现对稀疏注意力模式（如 sliding window attention）的支持有限

跨机器（多节点）的 KV cache 迁移尚未完全优化

引用

@inproceedings{kwon2023efficient,
  title={Efficient Memory Management for Large Language Model Serving with {PagedAttention}},
  author={Kwon, Woosuk and Li, Zhuohan and Zhuang, Siyuan and Sheng, Ying and Zheng, Lianmin and Yu, Cody Hao and Gonzalez, Joseph E. and Zhang, Hao and Stoica, Ion},
  booktitle={Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles (SOSP)},
  year={2023}
}