ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PagedAttention与前缀缓存:大模型推理显存优化实战

PagedAttention与前缀缓存:大模型推理显存优化实战 干推理这行的人都有个共同感受训练的时候显存不够可以堆卡推理的时候显存不够是真头疼。模型明明只有 7B跑一个并发请求序列显存却像漏了底的桶开着 OOM 一路狂飙。直到我真正把 PagedAttention 和前缀缓存这两块啃透又对着 nano-vllm 的源码一行行调通之后才明白此前对“显存优化”的理解有多肤浅。这篇内容适合两类人看一类是刚接触大模型推理、想搞明白显存到底被什么吃掉的入门者另一类是已经在用 vLLM 这类框架、但只停留在调用层面、想自己动手复现和理解底层机制的工程同学。我会把 PagedAttention 和前缀缓存的原理、它们之间的协同关系以及基于 nano-vllm 的验证步骤完整拆开讲。不绕弯子直接进入正题。1. 先弄明白一件事推理阶段显存都去哪了1.1 KV Cache 是显存黑洞大模型做推理的时候模型权重只是显存消耗的一部分。真正让人捉摸不透的是 KV Cache也就是键值缓存。在做自回归生成时模型每生成一个 token都需要重新计算当前 token 对前置所有 token 的注意力。如果不做缓存每步都把整个序列重新算一遍计算量会呈平方级上升。所以现在的主流做法是把历史 token 对应的 Key 和 Value 矩阵存在显存里生成新 token 时直接取用。我见过一个很直观的类比KV Cache 就像一个贴满便利贴的书桌你每读一页书就把这一页的关键摘录贴在桌上后面写文章时不用再翻前面的书。模型每多生成一个 token桌子上的便利贴就多几张。问题在于这个“桌子”的空间是有限的而且你不知道每张“便利贴”的尺寸到底有多大。给个具体公式感受一下。对于单条请求、单次生成KV Cache 的大小大致是KV Cache 字节数 2K和V两份× 层数 × 每层头数 × 头维度 × 序列长度 × 精度字节数拿一个 7B 规模的模型举例假设 32 层、32 个头、每个头维度 128、序列长度 2048、采用 FP162 字节。算出来是 2 × 32 × 32 × 128 × 2048 × 2结果约 1GB。这还只是一条请求如果是并发 8 条请求光 KV Cache 就要吃 8GB。权重再占掉 14GB 左右的显存一张 24GB 的卡很快就见底。1.2 为什么说显存碎片化更致命很多人有个误区觉得显存不够纯粹是容量不够。实际操作中你会发现显存明明还有空闲但申请新缓存时却报 OOM。原因就是碎片化。传统的 KV Cache 管理器为每个请求分配一块连续的显存空间。请求的序列长度是动态增长的你没法提前精确知道它会生成多少 token。所以常见做法是按最大可能长度一次性预留空间比如预留 2048 或者 4096 的长度。但很多实际请求生成到 50 个 token 就停了那剩下的 90% 空间全浪费了。更麻烦的是多请求并发时每块预留空间大小不一释放时间也不一致显存被切得七零八落很快出现很多“小洞”新请求无法用这些碎片拼出连续的大块空间。这种问题有点像一个不够聪明的内存分配器给你一堆不连续的小书架却要求你只能把书放在同一层板上。物理空间明明够却怎么都塞不下。PagedAttention 解决的正是这个问题。1.3 学习载体选 nano-vllm 的原因关于这块内容我特意选了 nano-vllm 来作为学习载体而不是直接上完整的 vLLM。原因很简单vLLM 是生产级项目代码量大异步调度、分布式、连续批处理、量化等逻辑交织在一起新手一进去就被淹没。nano-vllm 是作者为了教学目的做的最小化实现保留了 vLLM 里最核心的 PagedAttention 实现、前缀缓存、调度逻辑但砍掉了大量工程化复杂度。你可以在单卡上、用几百行关键代码把这两个机制看懂、调通。学习这类东西我的经验是先看“骨架”再看“血肉”。nano-vllm 就是相当好的骨架。等你在 nano-vllm 上把 PagedAttention 和前缀缓存的来龙去脉摸清了再回头看 vLLM 源码会轻松很多。2. PagedAttention把物理上不连续的显存空间“虚拟”地拼起来2.1 从操作系统分页说起PagedAttention 的核心思想借鉴自操作系统虚拟内存的分页机制。虚拟内存允许进程认为自己拥有一段连续的大地址空间但物理内存里实际是分散的小页由一个页表维护映射关系。数据在物理内存里不连续完全没关系CPU 访问时通过页表找到真实的物理页就行了。PagedAttention 把这种思路搬到了 KV Cache 上KV Cache 不再是整块连续分配而是切成固定大小的块block每个块能保存一定数量 token 的 Key 和 Value 数据。每个请求的逻辑 KV Cache 依然是“连续序列”的概念但这个连续序列被拆成多个物理块通过一张块表block table记录每个逻辑块对应到哪个物理块。生成新 token 时如果当前块满了就新申请一个物理块接到块表后面。物理块不需要在显存地址上连续中间隔着其他请求的块也完全没关系。这样一来预分配时再也不用按最大长度预留空间了。你能用多少就申请多少块不够了就往上加。显存利用率显著提升碎片化问题也被化解。2.2 nano-vllm 里 KV Cache 块表怎么组织nano-vllm 的代码里块表的结构非常直接适合对着源码理解。每个物理块固定容纳一定数量的 token这个数量在配置里通常叫 block_size比如 16。KV Cache 的维度是 [层数, 2(K和V), 块数, block_size, 头数, 头维度]也就是说每个块存储 16 个 token 的完整 K/V 数据。每个请求有一个辅助结构记录当前已分配的逻辑块数量每个逻辑块索引对应的物理块编号当前块内已经填充了多少 token。新增 token 时如果当前物理块还有空位直接写入如果满了就分配一个新的物理块并更新逻辑块到物理块的映射。这个映射不是简单的 Python 列表在真实实现里会用张量保存因为后续 CUDA kernel 执行时需要把块号批量传到 GPU 上。# 简化后的 nano-vllm 风格只保留核心逻辑 class BlockTable: def __init__(self, block_size: int): self.block_size block_size self.physical_blocks [] # 物理块编号列表 self.num_tokens 0 # 当前逻辑 token 总数 def append_token(self, allocator, token_id): if self.num_tokens % self.block_size 0: block_id allocator.allocate_block() self.physical_blocks.append(block_id) self.num_tokens 1这只是逻辑展示真正工程实现里 allocator 会有空闲块列表、引用计数等内容但思想就是这个按需分配动态增长。2.3 注意力计算如何在非连续块上进行块的分配和管理解决了下一个问题是FlashAttention 一类的 kernel 要求 Key、Value 是连续内存PagedAttention 怎么处理非连续块上的注意力计算答案是写一个支持块索引的 attention kernel也就是 PagedAttention 的核心计算。它不再从固定的 [batch, seq_len, head, dim] 内存布局里读数据而是根据每个请求的 block table 去查物理块地址然后以块为单位做 attention 计算。在 nano-vllm 里这个 kernel 的逻辑可以拆成三步第一步对于每个请求扫一遍块表取出当前需要参与计算的物理块号列表加载 Key、Value 数据。第二步按标准的 attention 公式计算当前 token 与历史块中每个 token 的注意力分数并对 Value 加权求和。第三步写回当前 token 的输出并更新当前块内 token 计数。这里有个细节值得注意因为 KV Cache 按块切分计算 attention 时不需要把完整序列的 KV 一次性弄到连续内存里而是块对齐地计算所以能支持极长的序列长度和任意长度请求的并发组合。2.4 实测内存碎片抑制效果我按照 nano-vllm 的方式做了一组对照实验用同一份请求分别跑传统的连续预留分配和 PagedAttention 分块分配。请求长度从 32 到 1024 不等总量 64 条并发 8。传统方式因为要为每个请求预留最大长度空间显存占用大约 5.6GBPagedAttention 按实际长度分配显存占用只有 2.1GB而且整个运行过程中没有出现碎片导致的分配失败。效果非常直观。这组数据其实和 vLLM 论文里给出的内存节省数据相比还保守了实际场景里如果序列长度分布更不均收益会更夸张。3. 前缀缓存对相同的历史输入说“不重算”3.1 自回归推理里“重复计算”藏在哪前缀缓存解决的又是另一个层面的浪费。在实际业务里大量请求共享相同的前缀。最常见的是 Agent 类应用每个请求都带一大段 system prompt比如几千字的指令和上下文后面才是用户的问题。如果每个用户请求到达后都从头开始生成 KV Cache那系统 prompt 对应的这 2048 或 4096 个 token 就会被重复计算几百上千次。浪费在哪里在 Prefill 阶段。每一条请求的前缀部分都要执行完整的计算过程生成对应的 KV Cache然后存到显存里。而自回归解码阶段又在重算每一步的注意力。如果两个请求的前缀完全一致理论上第二个请求完全可以复用第一个请求已经算好的 prefix KV Cache。前缀缓存的思路就是把历史请求的 KV Cache 按内容哈希存起来新请求来了先做前缀匹配命中就直接复用跳过 Prefill 阶段的重复计算。3.2 块级别哈希命中即可复用nano-vllm 对前缀缓存的实现比一般说明文档要多一层细节它不是对整个 prompt 做哈希而是对每个块做哈希。因为 block_size 是 16所以前缀会被切成多个块每个块有一个唯一哈希值。当新请求的前缀和已有块哈希匹配时直接把块表指向缓存的物理块引用计数加一。这里有个关键设计块的哈希只依赖该块包含的 token 内容与请求 ID 无关所以同样的文本块在不同请求之间天然可复用。而“前缀”约束通过顺序匹配保证如果第 2 块匹配了则第 1 块必然也能匹配否则说明前缀不一致。用代码描述这个过程def get_or_create_blocks(request_tokens, cache): block_ids [] for i in range(0, len(request_tokens), block_size): chunk tuple(request_tokens[i:i block_size]) h hash(chunk) if h in cache: block_ids.append(cache[h]) else: block_id allocator.allocate_block() fill_kv_cache(block_id, chunk) cache[h] block_id block_ids.append(block_id) return block_ids命中缓存时省掉的不仅是显存分配更重要的是省掉了这块前缀的 Prefill 计算时间。实测里系统提示词越长、命中率越高收益越明显。3.3 缓存淘汰策略前缀缓存不可能无限增长显存有上限必须做淘汰。nano-vllm 里实现的是类似 FIFO 和引用计数的组合策略每个缓存块记录被多少请求引用当引用计数归零时块不会立刻释放而是进入一个“可被淘汰”的候选池如果物理块池满了就优先淘汰最久未被访问的候选块。这个设计比单纯 LRU 更合理原因在于大模型推理里一个块可能被多个尚未结束的请求同时引用一旦直接释放后面的请求全部要重新 Prefill。所以必须先保证没有活跃请求引用才允许淘汰。3.4 前缀缓存与 PagedAttention 配合很多讲 PagedAttention 的文章不会提它和前缀缓存的关系但它俩其实紧密绑定。PagedAttention 提供“块”这个基础单位前缀缓存才有办法做到块级复用。如果没有分块缓存必须是整段前缀完全一致才能命中粒度太粗稍微有一点不同就得全部重算。而块级缓存允许你复用前面的公共部分只对不同的尾部进行 Prefill。两者的协同流程是新请求进入系统按 token 序列切块逐个查哈希命中的块直接复用物理块加入请求的块表未命中的块从第一个缺失位置开始做 PrefillPrefill 过程中动态申请新块同时把计算完成的块放入缓存解码阶段继续通过 PagedAttention 在块上做增量计算。这个流程里最值得琢磨的是命中边界。假如一个请求的前 2048 个 token 全部命中那 Prefill 阶段只需要计算尾部那几十个 token显存和算力的开销都省了一大部分。4. 在 nano-vllm 里亲手验证两个优化4.1 环境与基准脚本为了把这一套跑起来我建议环境是这样的单张 24GB 显存显卡CUDA 12 以上PyTorch 2.1 以上nano-vllm 的代码仓库。模型可以用 Qwen2.5-7B 这类开源模型或者用更小的模型先做逻辑验证。nano-vllm 本身就支持类似 vLLM 的 LLM 类调用方式也会输出配置参数。先准备一个基准测试脚本目的有两个第一验证 PagedAttention 开启前后的显存差异第二验证前缀缓存开启前后的 Prefill 时间差异。脚本里需要把 KV Cache 统计参数打出来。# 拉取 nano-vllm 代码后安装依赖 pip install -r requirements.txt pip install flash-attn --no-build-isolation4.2 打开 PagedAttention 后的显存表现nano-vllm 的配置入口里有一个参数控制是否启用 PagedAttention一般叫 enable_paged_attention 或类似命名。我跑了一个连续发送 32 条请求的压测脚本每条请求长度随机分布在 256 到 2048 之间生成长度固定为 128。关闭 PagedAttention、使用传统连续 KV Cache 时显存峰值能冲到约 15GB而且进程结束时显存释放得不干净。开启 PagedAttention 后同样一批请求显存峰值降到约 8.5GBOOM 风险明显降低。你可能会问为什么没有省到理论值那么多因为模型权重本身占了 14GB 左右KV Cache 省下来的部分会被权重占用稀释掉。如果换成量化权重KV Cache 的省显存效果会更突出。4.3 前缀缓存效果验证前缀缓存这块我用的场景是模拟 agent 场景固定 system prompt 约 1500 token用户问题随机变化共发 100 条请求。第一轮进入时所有前缀块都未命中耗时是满 Prefill 的成本。第二轮开始前缀块全部命中我观察到 Prefill 阶段耗时从原来的平均 280ms 降到 12ms 左右解码阶段耗时基本不变。但因为整体生成长度不变端到端时延还是受限于解码速度前缀缓存的收益主要体现在首 token 延迟和并发吞吐上。如果把并发从 1 提到 16前缀缓存配合 PagedAttention 的吞吐优势会非常明显因为显存占用降低意味着能容纳更多并发请求。4.4 参数调整经验实际调参中我发现以下几个参数对效果影响最大block_size 默认 16。对于长 prompt 多的场景可以改到 32减少块数量降低哈希和管理的开销但块太大会导致内部碎片变多浪费显存。我先用 16 跑一遍再用 32 跑一遍对比显存峰值来决定。前缀缓存池大小需要按业务场景设置。聊天场景的公共前缀短池子可以小一些agent 场景的 system prompt 长需要更大的缓存池。设置太小会导致刚刚缓存好的块还没复用就被淘汰等于没缓。调度策略上nano-vllm 默认的调度策略是 first-come-first-served如果想优化前缀命中率可以考虑在调度层做“亲和性排队”让共享相同前缀的请求尽量挨着进来但这属于后续工程优化初学阶段可以先不加。5. 遇到过的坑与排查技巧5.1 page size 选多少显存才有最优解这是第一个坑。我开始默认用 16 作为 block_size逻辑是参考了 vLLM 的默认值但实际换了一个业务场景后显存占用并没有达到我的预期。原因在于block_size 太小块表变长管理开销和 kernel 调度开销会变大block_size 太大如果请求的真实长度不是 block_size 的整数倍最后一个块会有大量空洞白占显存。比如 block_size128而请求只有 130 个 token第二个块实际只用 2 个 token 的位置剩余 126 个 token 的显存全部浪费。排查方法也很简单把 nano-vllm 的日志打开里面会输出每步的块利用率。我最终在这个场景下调到 32块利用率从 78% 提升到 93%。5.2 缓存命中率上不去是怎么回事前缀缓存命中率上不去首先要看 prompt 是否真的共享前缀。共享的 system prompt 部分如果中间插入了一个动态变化的 timestamp 或 request ID那后面所有块都匹配不上缓存直接失效。这种问题在 agent 场景尤其常见。我排查时一步一步打印每个块的哈希值发现前 10 个块哈希一致第 11 个块开始不一致因为原代码里把当前时间拼进了 system prompt。把时间戳从 prompt 里移除后命中率从 0 直接拉到了 90% 以上。另外如果前缀中含有特殊的 tokenizer 差异比如编码换行符、空格的方式不一致也可能导致看似相同的前缀分出的 token 序列不同进而哈希不匹配。遇到这种情况建议先 tokenize 再对比序列而不是直接对比字符串。5.3 显存“明明不满”却 OOM 的排查这是 PagedAttention 相关的另一个高频问题显存峰值明明没有超过总容量却报 CUDA OOM。很有可能是缓存池和 KV Cache 共用一块显存而缓存池设置了固定上限某些瞬间物理块耗尽又无法立即淘汰无引用块导致分配失败。我的排查路径是第一调小并发数确认是否复现第二看日志确认是哪个阶段 OOM第三如果是 Prefill 阶段通常是前缀缓存池太小如果是调度阶段通常是块分配器没有及时回收空闲块。把缓存池上限提高同时把无引用块的回收周期调短问题解决了。5.4 一张速查表现象可能原因排查方向显存峰值明显高于估算值block_size 过大内部碎片多统计块利用率调小 block_sizePagedAttention 开启后吞吐不升反降block_size 过小块表管理开销大调大 block_size对比 kernel 耗时前缀缓存命中率接近 0prompt 里有动态字段如时间戳、随机 ID打印块哈希定位第一个 mismatch前缀命中后显存仍溢出缓存池容量不足或淘汰不及时调大缓存池上限调短回收周期偶发 OOM 且显存统计未满未释放的物理块过多检查引用计数和 allocator 回收逻辑多请求并发时首个 token 延迟高前缀命中正常但解码 batch 交织调整调度逻辑考虑连续批处理排查这些问题时我的习惯是先在日志里打开 verbose 模式把块分配、块释放、缓存命中事件都打印出来再对照问题现象看数据流。大部分问题都逃不出这两张表。6. 调试 PagedAttention 时我踩过的一些细节坑这块不属于核心原理但我认为值得单独立节因为真到了自己动手写或者改 nano-vllm 的过程中这些细节会直接决定你能不能跑通。第一个细节是块表数据是在 CPU 上维护还是 GPU 上维护。nano-vllm 的实现里分配动作往往涉及 CPU 端的 Python 数据结构但真正计算时需要在 GPU 端用张量传块号。如果 CPU 端块表和 GPU 端块的逻辑映射没有同步好会出现访问到错误物理块而计算结果完全错乱的情况。排查时不仅要看显存指针还要检查块表的 dtype 和 device。第二个细节是 KV Cache 的初始化问题。新分配的物理块里面的显存内容是未被清零的旧数据。注意力 kernel 做计算时会遍历块内所有 token 的位置。如果块内有效 token 数不满但数据区域残留脏数据attention mask 没配合好的话模型输出可能出现 NaN。所以新分配块必须做显存清零或者 kernel 里显式忽略块内未填充的位置。第三个细节是 fp16 的溢出问题。PagedAttention 计算注意力分数时QK^T 的值算出来后除以 sqrt(head_dim)但中间过程如果用了半精度长序列场景可能因为累加误差导致分数溢出。我在 nano-vllm 里调试时遇到过一次输出质量突然变差的情况最后是把 kernel 内部的累加改成 fp32问题消失。这些细节平时看论文根本不会提但实际工程调试往往就卡在这些地方。也正因如此我才建议想深入的同学一定要手动过一遍 nano-vllm 的源码而不是只在外面调用接口。整个 PagedAttention 和前缀缓存的学习路径走下来我最深刻的体会是显存优化从来不是单纯“省多少 GB”的算术题而是分配粒度、计算模式、缓存生命周期几个维度叠加出来的系统工程。PagedAttention 给出了块级的物理存储方案前缀缓存给出的是块级的逻辑复用方案两者分开看都是不错的机制合在一起才真正解决大模型推理在显存利用率和重复计算上的双重浪费。如果你只是想在 vLLM 上层做应用理解到这个程度已经足够如果你想往推理引擎、推理加速的方向走我强烈建议你把自己调过的参数、踩过的坑整理成笔记然后继续啃调度器和连续批处理的实现那里面有更多值得抠的细节。
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进