三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

PagedAttention显存管理算法

PagedAttention显存管理算法

PagedAttention 是一种借鉴操作系统虚拟内存分页(Paging)机制的 LLM 显存管理算法,主要用于解决大语言模型推理阶段 KV Cache(键值缓存)导致的显存碎片化与高并发受限问题。

1. 核心概念

在传统 Transformer 推理中,为了避免重复计算历史 Token 的注意力矩阵,系统会保存所有历史 Token 的 Key 和 Value 向量(即 KV Cache)。传统方式要求这些向量在 GPU 显存中占用连续的空间。

PagedAttention 的核心思想在于将逻辑上的连续显存与物理上的离散显存解耦

  • KV Block(键值块):将 Token 序列的 KV Cache 切分为固定大小的逻辑块(例如每个 Block 包含 16 个 Token)。

  • Block Table(块表):维护类似操作系统页表的映射关系,记下逻辑 KV Block 与物理 GPU 显存块的对应关系。

  • 按需分页计算:在计算 Attention 时,CUDA 内核根据 Block Table 动态调取非连续物理显存中的 KV Block 进行矩阵运算。

2. 核心作用

  • 显存按需分配:在 Prompt 处理和新 Token 生成过程中,不再预先分配最大文本长度的显存,而是产生多少 Token 就动态申请多少物理 Block。

  • 实现显存共享:对于相同的前缀(如 System Prompt 或多轮对话历史),多个请求可以通过映射到同一组物理 Block 实现显存复用。

  • 物理显存解耦:解除对 GPU 连续大块显存的依赖,大幅提升显存利用效率。

3. 能解决什么实际问题

传统 KV Cache 管理痛点PagedAttention 的解决效果
预分配浪费:按最大长度(如 4K/32K)提前划定空间,实际未生成部分全部闲置零预留开销:只为已生成的实际 Token 分配物理内存,显存浪费率从 60%~80% 降至 4% 以下
内存碎片化:请求长度不一且频繁创建/销毁,导致严重的外碎片与内碎片零块外碎片:所有 Block 均为固定尺寸(如 16 Tokens),内部仅在最后一个 Block 产生微小碎片
并发吞吐瓶颈:显存很快被少数请求撑爆,Batch Size 无法拉高吞吐量提升 2-4 倍:节省出的显存可容纳更多并行 Batch,显著降低单 Token 服务成本
复杂采样开销大:Beam Search、Parallel Sampling 需全量复制 KV CacheCopy-on-Write 零拷贝:分支生成时仅复制物理指针,仅当某分支产生新 Token 时才为其分配独立 Block

4. 运用到的主流项目

PagedAttention 已成为当下 LLM 高性能推理引擎的标配底层技术:

  1. vLLM(首创者)

    UC 伯克利团队提出 PagedAttention 的原生框架,是当前部署开源大模型(如 Llama 3、Qwen 2.5)最主流的高并发推理引擎之一。

  2. Hugging Face TGI (Text Generation Inference)

    Hugging Face 的企业级推理服务框架,吸收并集成了基于块管理的分页 KV Cache 优化。

  3. TensorRT-LLM

    NVIDIA 官方推出的推理解析库,在其 Paged KV Cache 机制中采用了与 PagedAttention 相同的思想以适配 H100/A100/L40S 等硬件。

  4. SGLang

    针对结构化文本生成与复杂 Agent 调用的框架,在其底层的 RadixAttention 中深度扩展了 PagedAttention 的前缀共享能力。

  5. LMDeploy

    OpenMMLab 推出的部署工具套件,其 TurboMind 推理引擎支持基于分页管理的 KV Cache 机制。

  6. Ollama / LocalAI

    在后端集成高性能引擎(如 vLLM 或定制化后端)时,均依赖 PagedAttention 提供本地多并发对话支持。

← 返回列表