PagedAttention 内存管理深度解析
PagedAttention 是 vLLM 的核心创新,借鉴操作系统的虚拟内存分页机制来管理 KV Cache,彻底解决了传统推理系统中显存碎片化和浪费的问题。
传统方案的痛点
在传统推理系统中,KV Cache 采用预分配连续显存的方式:
- 内部碎片:预分配最大序列长度,短序列浪费大量显存
- 外部碎片:不同长度的请求释放后产生不连续的空闲区域
- 内存预留:必须为最坏情况预留显存,实际利用率低
以 Llama-2-7B 为例,最大序列长度 4096、batch size 32 时,KV Cache 需要约 16GB 显存,但平均利用率仅 60%。
分页机制设计
PagedAttention 将 KV Cache 划分为固定大小的块(Block),每个块存储固定数量 Token 的 Key 和 Value 向量:
python
# PagedAttention 块管理示意
class BlockManager:
def __init__(self, block_size=16, num_blocks=1024):
self.block_size = block_size # 每块存储的 Token 数
self.num_blocks = num_blocks
self.free_blocks = list(range(num_blocks)) # 空闲块列表
self.block_tables = {} # 请求 -> 块映射表
def allocate(self, request_id, num_tokens):
"""为请求分配 KV Cache 块"""
num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
if num_blocks_needed > len(self.free_blocks):
return False # 显存不足
allocated = []
for _ in range(num_blocks_needed):
block_id = self.free_blocks.pop()
allocated.append(block_id)
self.block_tables[request_id] = allocated
return True
def free(self, request_id):
"""释放请求占用的所有块"""
if request_id in self.block_tables:
for block_id in self.block_tables[request_id]:
self.free_blocks.append(block_id)
del self.block_tables[request_id]块大小选择
块大小(block_size)是关键参数,需要在灵活性和管理开销间平衡:
| 块大小 | 内部碎片率 | 管理开销 | 推荐场景 |
|---|---|---|---|
| 8 | ~3.9% | 高 | 短序列为主 |
| 16 | ~4.7% | 中 | 通用场景(默认) |
| 32 | ~5.5% | 低 | 长序列为主 |
最佳实践
vLLM 默认 block_size=16,对于平均序列长度在 500-2000 的场景表现最优。如果输入以超长上下文为主,可适当增大到 32。
Copy-on-Write 机制
PagedAttention 支持 Copy-on-Write(COW),在 Beam Search 等场景中多个候选序列可以共享前缀的 KV Cache:
- 多个序列共享同一组物理块
- 仅当序列产生分歧时才复制需要修改的块
- 大幅减少 Beam Search 的显存占用
Swap 机制
当 GPU 显存不足时,vLLM 可以将部分 KV Cache 块交换到 CPU 内存:
python
# 启用 Swap 的配置
from vllm import LLM
llm = LLM(
model="meta-llama/Llama-2-7b-hf",
swap_space=4, # CPU swap 空间大小(GB)
gpu_memory_utilization=0.85,
)Swap 的代价
Swap 机制虽然可以避免 OOM,但 CPU-GPU 数据传输会引入显著延迟。建议仅在显存确实紧张时启用,且优先考虑减小 batch size。
显存利用率分析
PagedAttention 相比传统方案的显存利用率提升:
- 传统连续分配:有效利用率 50-70%
- PagedAttention:有效利用率 85-95%
- 配合 Continuous Batching:整体吞吐量提升 2-4x