ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

KV Cache 显存瓶颈如何破?分级存储与 PagedAttention 实战

KV Cache 显存瓶颈如何破?分级存储与 PagedAttention 实战 1. 为什么 KV Cache 会成为大模型推理的瓶颈1.1 从自回归生成说起每生成一个 token 都要回看历史大模型推理的本质是自回归生成。给定一段提示词模型预测出第一个新 token然后把这个 token 拼回输入序列再预测下一个如此循环。问题就出在这个拼回的动作上Transformer 的注意力机制要求当前 token 和序列中所有历史 token 计算注意力分数而每个历史 token 都要先经过 Key 和 Value 的线性投影。如果每个解码步都重新计算全部历史 token 的 K 和 V计算量会随序列长度平方级增长显存带宽也会被反复读写拖垮。KV Cache 的思路很直接把已经算过的 Key 和 Value 缓存下来后续解码步直接复用不再重算。这样一来每步只需要计算新 token 的 Q、K、V然后拿新 Q 去和缓存里的全部 K 做点积。这个优化带来的收益是巨大的。以 LLaMA-2-7B 为例不做 KV Cache 时生成第 1000 个 token 需要重算前 999 个 token 的 K、V算力浪费极其严重做了 KV Cache 之后每步的注意力计算量基本恒定只有线性投影部分随序列增长。实测下来开启 KV Cache 能让长序列生成的吞吐提升数倍甚至一个数量级。但缓存不是免费的。它把计算压力转化成了显存压力而且这个压力随序列长度、批大小、层数线性增长。这就是为什么 KV Cache 的显存占用会成为长上下文推理的核心瓶颈。1.2 算一笔账KV Cache 到底吃掉多少显存KV Cache 的显存占用公式并不复杂KV Cache 大小 2 × batch_size × seq_len × num_layers × num_kv_heads × head_dim × dtype_bytes其中 2 代表 Key 和 Value 两份num_kv_heads 是 KV 头数GQA/MQA 下小于 Q 头数dtype_bytes 是每个元素的字节数FP16 为 2FP8 为 1。拿一个具体配置算一下。假设模型是 70B 级别80 层采用 GQAnum_kv_heads8head_dim128FP16 精度batch_size1序列长度 32K2 × 1 × 32768 × 80 × 8 × 128 × 2 8.6 GB单条 32K 序列就要 8.6 GB 显存而模型权重本身在 FP16 下约 140 GB。如果 batch_size 提到 16KV Cache 直接飙到 137 GB和权重相当。这意味着在长上下文、高并发场景下KV Cache 的显存占用会超过模型权重成为真正的显存杀手。更麻烦的是KV Cache 是动态增长的。请求进来时你不知道它最终会生成多长只能按最大长度预留或者动态分配。预留会浪费显存动态分配会带来碎片。这个矛盾催生了各种分级存储和分页管理的方案。1.3 分级存储的核心动机把冷数据挪到便宜的地方显存贵、容量小、带宽高内存便宜、容量大、带宽低SSD 更便宜、容量更大、带宽更低。KV Cache 的特点是越靠近当前解码位置的 token被访问的频率越高越早生成的 token被访问频率越低但并非完全不访问——注意力机制每次都要看全部历史。这就形成了一个天然的冷热分层最近生成的 KV 是热数据必须放在显存较早的 KV 是温数据可以放在内存极早的 KV 是冷数据可以放在 SSD 甚至更远的存储。分级存储要解决的核心问题就是在保证注意力计算正确性的前提下把冷数据从显存挪走腾出空间给新请求或更长序列。这个思路和 CPU 的缓存层级L1/L2/L3/内存本质相同只是对象从指令数据变成了 KV 张量。理解了这一点后面所有的工程方案都是围绕如何高效地在层级间搬运数据和如何决定哪些数据该在哪一层展开的。2. 分级存储的架构设计与关键取舍2.1 三层存储的职责划分与容量规划一个典型的分级存储架构会划分三层GPU 显存HBM、主机内存DRAM、本地 SSD。三层的职责和容量规划需要根据硬件配置和业务场景来定。GPU 显存层存放当前活跃请求的 KV Cache以及最近生成的若干 token 的 KV。这一层的容量由 GPU 型号决定A100 80GB、H100 80GB、H200 141GB 是常见配置。显存层要预留一部分给模型权重、激活值和临时缓冲区实际可用于 KV Cache 的可能只有 20-40 GB。主机内存层存放被换出的温数据。一台推理服务器通常配 512GB 到 1TB 内存除去系统和进程开销可用的可能有 400-800 GB。这一层的带宽在 20-50 GB/s 量级比 HBM 的 2-3 TB/s 低两个数量级所以换入换出要尽量批量、异步。SSD 层存放冷数据。NVMe SSD 的顺序读带宽能到 5-7 GB/s随机读在几十万 IOPS。这一层适合存放那些已经生成完毕、但可能被后续请求复用的前缀 KV比如系统提示词、few-shot 示例、长文档的编码结果。容量规划的经验法则是显存层放得下当前并发请求的活跃 KV内存层放得下最近一段时间内可能被复用的 KVSSD 层放得下所有值得缓存的前缀 KV。三层之间的容量比例大致是 1:10:100。2.2 换入换出的触发时机什么时候该搬数据分级存储最难的不是搬运本身而是决定什么时候搬、搬哪些。搬早了浪费带宽搬晚了显存爆掉。常见的触发策略有三种。第一种是水位触发。给显存层设一个高水位比如 90%和一个低水位比如 70%。当占用超过高水位时启动换出把最冷的 KV 搬到内存直到降到低水位。这种策略简单可靠但反应有延迟高并发下可能来不及。第二种是请求驱动。新请求进来需要分配显存时如果空间不够就按 LRU最近最少使用或 LFU最不经常使用淘汰一部分 KV 到内存。这种策略更及时但淘汰决策频繁开销较大。第三种是预测触发。根据历史请求的序列长度分布预测未来一段时间内的显存需求提前换出。这种策略最理想但预测不准时反而添乱。实际系统里通常是组合使用水位触发做兜底请求驱动做精细控制预测触发做优化。我见过的一个生产系统是这样配的显存占用超过 85% 启动后台换出线程每次换出 5% 的冷数据新请求分配失败时立即触发同步换出保证请求不阻塞。2.3 数据搬运动作的代价PCIe 带宽是隐形天花板分级存储的搬运走的是 PCIe 总线。PCIe 4.0 x16 的理论带宽是 32 GB/s实际有效带宽在 25 GB/s 左右PCIe 5.0 x16 翻倍到 64 GB/s实际约 50 GB/s。这个带宽要和模型推理本身的数据传输共享。搬运动作的代价不只是带宽还有延迟和 CPU 开销。一次 cudaMemcpy 的延迟在微秒级但如果是异步拷贝加同步等待延迟会累积。如果搬运逻辑用 CPU 线程做还要考虑线程调度和锁竞争。一个容易被忽略的点是KV Cache 的搬运不是连续大块内存拷贝而是分散在各个层的各个头之间。如果不做内存布局优化搬运会变成大量小拷贝效率极低。所以好的分级存储实现会把 KV Cache 按层、按头组织成连续的大块搬运时整块搬减少拷贝次数。实测数据在 PCIe 4.0 上搬运 1 GB 的 KV Cache 大约需要 40-50 毫秒含开销。如果每生成一个 token 都要搬运那吞吐直接崩掉。所以搬运必须是批量的、异步的、和计算重叠的。3. 从 PagedAttention 到分层卸载的实操实现3.1 PagedAttention把 KV Cache 切成块来管理vLLM 提出的 PagedAttention 是 KV Cache 管理的一个里程碑。它的核心思想借鉴了操作系统的虚拟内存分页把 KV Cache 切成固定大小的 block比如 16 个 token 一块每个 block 在物理显存里可以不连续通过一个 block table 做逻辑到物理的映射。这样做的好处有三个。第一消除了内存碎片因为 block 大小固定分配和回收都是整块操作。第二支持前缀共享多个请求如果有相同的前缀可以共享同一批 block显存占用大幅下降。第三为分级存储提供了天然的搬运单位block 就是最小的换入换出粒度。PagedAttention 的 block table 结构大致是这样每个序列维护一个 block table记录该序列用到的所有 block 的物理地址。注意力计算时kernel 根据 block table 去 gather 对应的 K、V。这个 gather 操作在 GPU 上做开销可控。代码层面vLLM 的 block 管理逻辑在BlockSpaceManager里核心接口是allocate、free、swap_in、swap_out。swap_out 就是把 block 从 GPU 拷到 CPUswap_in 反之。这个设计直接为分级存储留好了接口。3.2 分层卸载的实现GPU-CPU-SSD 三级流水线在 PagedAttention 的基础上做分层卸载需要实现三个组件显存管理器、内存管理器、SSD 管理器以及一个调度器决定 block 的层级归属。显存管理器维护一个空闲 block 池分配时从池里取释放时还回池里。当池子空了触发换出。内存管理器类似但容量大得多。SSD 管理器用文件或内存映射的方式存储 block每个 block 一个文件或一段偏移。调度器的策略是核心。一个可用的策略是给每个 block 打一个热度标签热度由最近被访问的时间和频率决定。新生成的 block 热度最高放在显存随着时间推移热度衰减降到阈值以下就换到内存再降就换到 SSD。访问时如果 block 不在显存触发换入同时提升热度。换入换出要异步做。用一个独立的 CUDA stream 做拷贝和计算 stream 重叠。拷贝完成用 event 通知计算 stream 等待 event。这样搬运的延迟可以被计算掩盖掉一部分。一个实操细节换出时要选那些确定一段时间内不会被访问的 block。对于自回归生成当前序列的早期 block 虽然热度低但每次注意力计算都要访问所以不能换出。真正能换出的是那些已经生成完毕、序列已经结束、但可能被后续请求复用的前缀 block。这个区分很关键搞错了会导致频繁换入换出性能反而下降。3.3 前缀缓存复用分级存储的最大价值场景分级存储最大的价值场景是前缀缓存复用。很多应用场景下大量请求共享相同的前缀系统提示词、角色设定、few-shot 示例、长文档背景。如果每个请求都重新计算这些前缀的 KV浪费巨大。有了分级存储可以把这些公共前缀的 KV 算一次存在 SSD 或内存里后续请求直接复用。复用时只需要把前缀 block 换入显存然后从最后一个前缀 token 开始解码。这个机制的实现要点是前缀匹配。新请求进来先拿它的 token 序列去前缀缓存里查找到最长匹配的前缀复用对应的 block。匹配可以用哈希做把 token 序列分块哈希存一个哈希到 block 的映射表。实测效果在一个 RAG 场景里系统提示词加检索文档约 8K token如果每个请求都重算首 token 延迟在 2 秒左右用了前缀缓存后首 token 延迟降到 300 毫秒以内因为只需要换入 8K token 的 KV不需要重算。注意前缀缓存的命中率高度依赖请求的相似度。如果请求前缀各不相同缓存命中率低反而增加了管理开销。上线前一定要用真实流量测命中率低于 30% 就要重新评估是否值得。4. 性能调优与常见问题排查4.1 换入换出策略的参数调优分级存储的性能对参数很敏感。几个关键参数需要根据硬件和负载调。block 大小太小则管理开销大太大则碎片多、换出粒度粗。经验值是 16 到 64 个 token 一块。16 适合短序列高并发64 适合长序列低并发。水位线高水位设太高容易 OOM设太低频繁换出。建议高水位 85%-90%低水位 60%-70%。两者差距要够大避免在边界反复震荡。换出批量每次换出的 block 数量。太小则换出频繁太大则单次延迟高。建议一次换出显存容量的 5%-10%。热度衰减系数控制热度随时间下降的速度。衰减太快则冷数据被过早换出衰减太慢则显存里堆满不再用的数据。这个参数最好用实际负载的访问模式来拟合。异步拷贝的 stream 数量1 到 2 个足够太多会争抢 PCIe 带宽。调优时建议用固定负载做基准测试每次只改一个参数记录吞吐、延迟、显存占用三条曲线。我一般会画一张参数-性能的敏感度表找出性能平台区然后把参数设在平台区中间留出波动余量。4.2 常见问题速查表问题现象可能原因排查方向解决方法显存频繁 OOM水位线设太高换出不及时看显存占用曲线是否锯齿状降低高水位增大换出批量吞吐突然下降换入换出和计算争抢带宽看 PCIe 利用率是否打满限制换出速率错峰搬运首 token 延迟高前缀缓存未命中或换入慢看缓存命中率和换入耗时优化前缀匹配预热常用前缀生成质量异常KV 换入换出时数据损坏对比换入前后的 KV 值检查拷贝逻辑和同步机制内存占用持续增长block 泄漏释放逻辑有 bug看 block 池的空闲数加引用计数定期审计SSD 写入量过大冷数据被频繁换入换出看 SSD 读写监控提高热度阈值减少无效换出4.3 踩过的坑与实操心得第一个坑是同步换入导致的停顿。早期实现里换入是同步的注意力计算要等换入完成。结果长序列生成时每几步就卡一下。后来改成异步换入加预取提前把下一步可能用到的 block 换进来停顿才消失。预取的逻辑是根据当前解码位置预测未来几步会访问哪些 block提前发起拷贝。第二个坑是 block 碎片化。虽然 PagedAttention 消除了外部碎片但内部碎片还在。如果一个序列最后只用了半个 block那半个 block 就浪费了。高并发下这个浪费会累积。解决办法是用更小的 block或者做 block 合并。我试过把 block 从 32 降到 16显存利用率提升了约 8%但管理开销增加了 5%净收益为正。第三个坑是 SSD 寿命。如果冷数据换入换出太频繁SSD 的写入量会很快耗尽寿命。一个 1TB 的消费级 SSDTBW 可能只有 600如果每天写入 500GB一年多就废了。所以 SSD 层要尽量只放真正冷的数据或者用企业级 SSD。更稳妥的做法是 SSD 层只做只读缓存写入用追加方式减少擦写。第四个坑是前缀缓存的失效。如果前缀缓存用哈希做键模型更新或 tokenizer 变化后旧缓存的哈希对不上会全部失效。所以缓存要带版本号模型更新时清空或迁移。提示分级存储的调试最好从单请求开始逐步加压。先用一条长序列验证换入换出正确性再用小批量验证并发下的调度最后用真实流量压测。每一步都要对比开启和关闭分级存储的结果确认收益为正。5. 分级存储的边界与后续演进方向分级存储不是银弹它有明确的适用边界。当序列长度在 4K 以内、并发在个位数时KV Cache 本身占用不大分级存储的管理开销可能超过收益。当请求前缀高度不重复时前缀缓存命中率低分级存储的价值也有限。真正能吃到红利的场景是长上下文16K 以上、高并发、前缀有复用。从技术演进看几个方向值得关注。一是更细粒度的量化把 KV Cache 压到 FP8 甚至 INT4从源头减少显存占用减轻分级存储的压力。二是注意力机制的改进比如滑动窗口注意力、稀疏注意力减少需要缓存的 KV 数量。三是存储介质的演进CXL 内存池化让内存层可以更大更灵活为分级存储提供更好的中间层。但无论底层怎么变分级存储的核心思想不会过时按访问频率分层把合适的数据放在合适的介质上。这个思想从 CPU 缓存到数据库缓冲池再到 KV Cache一脉相承。理解了这一点具体实现只是工程细节的差异。我在实际部署中的体会是分级存储的收益和复杂度是正相关的。简单的 LRU 换出就能拿到大部分收益再往上每提升一点都要付出成倍的工程代价。所以建议先用最简单能跑的方案上线用监控数据说话确认瓶颈真的在 KV Cache 显存上再逐步加码。盲目上复杂方案很可能调参调到怀疑人生收益还不明显。
返回列表