vAttention

概览

论文通过 CUDA 提供的虚存管理 API，解决了在用户态管理显存分页的问题。对比 FlashAttention 和 FlashInfer 中的 PagedAttention kernel，分别获得了 22%、29% 的性能提升。

论文地址：vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention

动机

vLLM 通过引入 PagedAttenion，通过分页管理显存，来节省 KV Cache。但是这带来了几个明显的问题：

1. 需要重写 GPU kernel，难以复用现有代码
2. CPU 侧推理框架代码复杂，难以维护
3. 性能开销
   1. GPU 侧：在 FlashAttention、FlashInfer、TensroRT-LLM 等框架中，PagedAttention 的性能比不分页的 Attention kernel 差 10%～28% 不等
   2. CPU 侧：管理页表也带来了明显的开销

方法

基本观察

对推理框架有这样两个观察：

1. KV Cache 需求可预测：在 decode 单个 token 的粒度，每轮迭代所需要的显存是可预测的
2. KV Cache 不需要高内存分配带宽：单个 token 的 KV Cache 大小并不大，因此 KV Cache 的分配需求（以 MB/s 计）受限于解码速度，处于一个较低的水平 基于这两点观察，高效的动态 KV Cache 分配机制成为可能。

系统设计

下文用到的 CUDA 虚存 API：https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html

1. 系统启动阶段，预分配所有显存
   1. 虚存：按模型所需要的最大 context length 分配，分配的虚存可以远远大于实际的显存大小
   2. 物理显存：按 GPU 能分给 KV Cache 的最大量分配
2. 推理阶段，动态映射虚存和物理显存
   1. 在每次 decode a token 之前将虚存映射到物理显存
   2. 解码结束后取消虚存和物理显存的映射
3. 系统关闭阶段，释放所有的虚存和物理显存

优化

1. CUDA 默认的 page 是 2MB 大，可以通过扩展 GPU 驱动的方式，构造一些能分配 64KB、256KB 大小的页面
   1. 我个人认为这个方式是不可维护的
2. 显存分配和计算重叠。第 轮 decode 计算时，可以并行分配第 轮的显存
3. 显存回收推迟。某个 decode 请求完全结束后，不立刻取消虚存和物理显存的映射，而只是清零显存，分配给下一个新的请求。这样只有在下一个新的请求 decode 长度大于本次请求长度时，才需要额外的虚存-物理显存映射
   1. 反过来，如果下一次请求的 decode 长度更短，会造成显存浪费。因此设定系统的可用显存降低到某个阈值以下后，再回收物理显存