VLLM中PagedAttention如何降低显存碎片？

Question

在使用VLLM的PagedAttention机制时，为何能够有效降低显存碎片？传统注意力机制在处理变长序列时需预分配连续显存，导致大量内部碎片。PagedAttention受操作系统分页内存管理启发，将Key-Value缓存切分为固定大小的“页”，按需动态分配物理显存块，逻辑上连续但物理上非连续存储。这种方式避免了长序列因找不到大块连续显存而失败的问题，显著提升显存利用率。请问这种分页式KV缓存如何具体实现显存的离散分配与地址映射，从而减少碎片并支持高效并行访问？

Answer 1

一、显存碎片问题的背景与传统注意力机制的局限性

在大语言模型（LLM）推理过程中，自回归生成阶段需要维护每个已生成 token 的 Key-Value（KV） 缓存，用于加速后续 attention 计算。传统实现中，如 PyTorch 或 Hugging Face Transformers，默认为每个请求预分配一段连续的显存空间来存储其完整的 KV 缓存。

• 当序列长度可变时，系统必须按最大可能长度预留空间，导致大量内部碎片（internal fragmentation）。
• 例如：若最大上下文长度为32k，但多数请求仅使用8k，则剩余24k空间被浪费。
• 更严重的是，随着显存长期运行，大块连续空间逐渐耗尽，即使总空闲显存充足，也无法满足长序列请求，造成外部碎片（external fragmentation）。

这种“一次性预分配+连续存储”的模式严重限制了服务吞吐量和资源利用率。

二、PagedAttention 的核心思想：分页式 KV 缓存设计

PagedAttention 受操作系统虚拟内存管理中的分页机制启发，将逻辑上连续的 KV 缓存切分为固定大小的“页”（page），每页通常包含若干个 token 的 K 和 V 向量。

特性	传统 Attention	PagedAttention
KV 存储方式	连续显存块	离散页表映射
内存分配策略	静态预分配	动态按需分配
碎片类型	内部 & 外部碎片高	显著降低
扩展性	差（受限于最大长度）	好（支持弹性扩容）
并行访问能力	依赖连续布局	通过页表索引并发读取

每个页大小一般设为 16 或 32 个 token，统一管理，便于 GPU 内存池高效复用。

三、显存离散分配与地址映射的具体实现机制

为了支持物理上离散但逻辑上连续的 KV 缓存访问，PagedAttention 引入了类似于页表（Page Table）的数据结构：

1. 每个请求对应一个 逻辑页序列，表示其所需的所有页的顺序。
2. 每个逻辑页通过页表条目映射到一个或多个 物理页块（位于 GPU 显存的不同位置）。
3. 页表本身存储在主机内存或显存中，由调度器维护，并在核函数执行时传入。
4. GPU 核函数在计算 attention 时，根据当前 token 的位置计算所属的逻辑页号，再查页表获得物理地址偏移。
5. 利用 CUDA 的 stride 访问模式，批量加载多个非连续页的数据到共享内存或寄存器中。
6. 所有页具有相同尺寸，确保内存访问模式可预测，利于编译器优化和 warp 并行执行。
7. 内存池预先划分固定大小的页框（frame），采用 slab 分配器风格进行快速分配/回收。
8. 支持跨请求共享公共前缀页（如 prompt 共享），进一步提升缓存命中率。
9. 页表支持多头注意力结构，每个 head 可独立寻址同一组页。
10. 通过 pinned memory 和异步传输优化页间数据迁移开销。

// 简化的页表结构定义（C++ 伪代码）
struct PageTableEntry {
    int physical_block_id;  // 物理块编号
    int offset_in_block;    // 块内偏移（可选）
    bool valid;
};

struct KVCachePage {
    float* k_data;  // 指向实际显存地址
    float* v_data;
    size_t token_count;
};

四、减少碎片与支持高效并行访问的技术路径

通过以下技术组合，PagedAttention 实现了显存利用率与计算效率的双重提升：

graph TD A[请求到达] --> B{是否已有页表?} B -- 是 --> C[追加新页] B -- 否 --> D[创建新页表] C --> E[从内存池分配空闲页] D --> E E --> F[更新页表映射] F --> G[执行Attention核函数] G --> H[通过页表查物理地址] H --> I[CUDA并发读取多个页] I --> J[聚合结果输出]

• 消除内部碎片：仅分配实际使用的页数，避免为短序列预留过多空间。
• 缓解外部碎片：小页单位使碎片粒度细化，提高内存整合能力。
• 支持动态扩展：可在生成过程中动态添加新页，无需初始预估长度。
• 并行访问优化：现代 GPU 支持 scatter/gather 类型访问，结合页表可实现高效的非连续内存读取。
• 批处理友好：不同请求的页可交错分布，统一调度 batched attention kernel。
• 容错与回收机制：完成请求后立即释放页资源，供其他请求复用。

此外，vLLM 框架还引入了 PagedAttention V2，进一步优化页内 layout，支持 block-wise 计算，减少中间激活开销。