PageAttention

PagedAttention (分页注意力机制)

💡 核心痛点：为什么需要 PagedAttention？

在大语言模型（LLM）的推理过程中，模型需要把过去计算好的 Key 和 Value 状态保存下来（即 KV Cache），以免在生成下一个 Token 时重复计算。这通常会占用极大的 GPU 显存空间。

传统的应对方式是预分配完整的连续内存：

  • 最大化预分配：例如大模型设置的最大上下文长度为 8192，系统一旦接收到请求，就会直接为其划出一块长达 8192 个 Token 的连续显存位置。
  • 内存碎片化（严重的内存浪费）：如果用户实际的对话连同生成的回答最终只占用了 700 个 Token，那么剩下的 7000 多个空位的显存就被白白占着却没被使用。据统计，传统连续 KV Cache 有多达 80%~90% 是被这样浪费掉的“内部碎片”，其他排队的并发任务也无法使用这部分空间。

🚀 解决思路：借鉴操作系统的虚拟内存分页

PagedAttention 的设计灵感非常直接：像操作系统管理虚拟内存分页一样来动态管理 KV Cache。

  • 非连续存储（Blocks）：KV Cache 不再强求存储在一整块连续的显存中。系统将多余的显存切分成一个个固定尺寸的**“物理块”**（Physical Blocks），每个物理块只能容纳固定数量（例如 16 个）的 Token。
  • 按需动态分配：随着新的 Token 被生成出来，系统先尝试把它们塞进当前的块中。只有当当前这一个块被写满时，系统才会去内存池里找一个新的空物理块分配给该请求，做到真正的按需分配。
  • 块表映射（Block Table）：面对这些天女散花的内存块，系统专门通过一张所谓的“块表”（Block Table）来统一管理。只要查询这张表，系统就能得知一个请求完整的上下文具体是由哪几个不连续的物理块拼凑出来的（这跟操作系统里的“页表”是完全一样的道理）。

📈 带来的巨大收益

  1. 极低的碎片率：内存浪费仅仅发生在最后一个未填满的物理块剩下的那几个空位上。整体的显存碎片率可以直接锐减到 1% 以下。
  2. 成倍提升系统吞吐量 (Batch Size)：由于以往被白白浪费的显存现在被彻底压榨出来迎接新的用户请求，同样的单卡环境下并发处理容量甚至能扩大 7 倍以上。
  3. 高效的内存共享机制：
    • 因为请求被“打包”成了抽象的离散块，如果有不同的请求碰巧包含了完全相同的前缀或者系统提示词（System Prompt），它们的块表可以直接指向同一些物理块！这被称为 Prompt Cache。
    • 在进行 Beam Search（集束搜索）时，生成的多条分支可以安全地共享同一个前缀的内存块，而无需对整个上下文状态进行大段拷贝。

⚙️ 底层实现细节窥探（vLLM）

在以 vLLM 为代表的高性能推理框架中，由于数据在物理显存上不再是连续放置的，原本经典的 Attention 计算方式无法直接跑起来，因此底层需要精妙的定制版 CUDA Kernel 代码：

  • 底层 Kernel 的指针（如 k_ptr）在从全局显存读取数据时，需要配合 Block Table 边查表边跳转。
  • 对并发线程（Thread Groups 和 Warps）的切分做了深度配合：每个线程负责只读其分配的某个局部矩阵块，以达到最佳的显存读取合并（Memory Coalescing）效应抵消跳转表带来的开销。

📝 参考资料: