Skip to content

PagedAttention 内存管理深度解析

PagedAttention 是 vLLM 的核心创新,借鉴操作系统的虚拟内存分页机制来管理 KV Cache,彻底解决了传统推理系统中显存碎片化和浪费的问题。

PagedAttention内存管理示意

传统方案的痛点

在传统推理系统中,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

相关资源

最近更新