ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

KV Cache 显存优化:LMCache 与 vLLM 分层卸载实践

KV Cache 显存优化:LMCache 与 vLLM 分层卸载实践 先说一个我自己的场景。之前在内网服务里把模型的上下文从 32K 拉长到 128K当天下午就有两个 worker 因为 CUDA OOM 自动重启了。第一反应当然是加显存但仔细一看实际上绝大部分显存都被 KV cache 吃掉了而且不同请求之间有一大段公共前缀每次都在重复计算。后来我花了一个周末把 LMCache 和 vLLM 接起来把 KV cache 按访问热度分层卸载到 CPU 内存和 SSD 上128K 上下文才真正跑得动。这篇文章就把这次探索的全过程写清楚KV cache 为什么这么占显存、把缓存从 GPU 搬到 CPU/SSD 在数学上为什么划算、LMCache 在 vLLM 里到底做了什么、具体怎么配置以及我踩过的几个坑。如果你也在被长上下文、多轮对话、RAG 重复前缀这类问题折磨这篇应该能给你一个直接可用的方案。1. 显存刺客KV cache 到底吃了多少显存1.1 一个 token 的 KV cache 有多大KV cache 是预填充阶段算出来的 Key 和 Value 张量目的是让后续每个 token 生成时不用重新计算历史部分的注意力。它的尺寸不取决于模型参数量而取决于层数、注意力头数和 head_dim。以 Llama-3.1-8B 为例32 层、8 个 KV 头GQA、head_dim 为 128FP16 精度下单个 token 的 KV cache 大小是2K 和 V 两张表x 32 层 x 8 个 KV 头 x 128 维度 x 2 字节 128KB注意这是每生成一个 token 就要累积这么多显存。上下文长度 32K 时单个序列的 KV cache 就是 4GB到了 128K 上下文直接变成 16GB。如果并发 8 个请求光是 KV cache 就要 128GB。这还没算模型权重和激活值。70B 级别模型的数字更夸张。Llama-3.1-70B 是 64 层、8 个 KV 头单个 token 是 256KB128K 上下文就是 32GB。想让几十路请求同时跑长上下文显存再大也不够。1.2 显存的分配悖论vLLM 本身对显存使用做了两个关键优化PagedAttention 的块式管理、以及通过gpu_memory_utilization参数把大部分显存预留给 KV cache。但这就产生了一个问题gpu_memory_utilization开得越高留给活跃 KV cache 的空间越多可是模型权重和运行时需要的其他显存就被压得越紧开得低KV cache 空间又不够长上下文立刻 OOM。我当时的服务就是卡在这里。模型权重 16GB8 卡 A100 每张 80GB看起来绰绰有余但 128K 上下文只要同时跑 6 个请求单卡 KV cache 空间就爆了。LMCache 的核心价值就是把这部分 KV cache 按热度拆成几层 GPU 只保留最热的那一小块冷数据放到 CPU 内存和 SSD 上显存的压力一下就解了。2. 卸载到 CPU/SSD 的数学账为什么带宽差这么多还值得做2.1 三层存储的带宽差距先看一组量级数据。GPU HBM 带宽大约是 2-3.5TB/sCPU 内存带宽约为 50-100GB/sNVMe SSD 的顺序读带宽大约 3-7GB/s。从数字上看SSD 比 HBM 慢了差不多三个数量级。如果 decode 阶段每生成一个 token 都要从 SSD 把全部历史 KV cache 读回 GPU那确实行不通生成速度会灾难性地下降。这也是很多人第一次听到“KV cache 卸载到 SSD”时的直觉反应SSD 那么慢怎么可能扛得住2.2 关键在跳过重复的 prefill要理解这里面的账得先分清两个阶段。Prefill预填充阶段是处理 Prompt 的计算量随序列长度增长非常快本质上是大量矩阵乘法极度吃算力。Decode 阶段是逐个生成 token 的虽然每个 token 计算量不大但它需要反复读取整段历史的 KV cache。如果 KV cache 全部存在 GPU 显存里decode 阶段读起来很快所以 vLLM 默认策略没问题。问题是显存不够的时候连 prefill 都做不了。LMCache 的思路反过来了把 KV cache 存到 CPU 内存和 SSD 上下次遇到相同前缀的请求直接从这些慢速存储里把 KV cache 读回来跳过整个 prefill 阶段。虽然读回也要时间但比起重复计算成千上万个 token 的 prefill耗时往往少一个量级。我做了一个大致估算。32K token 的 KV cache在 8B 模型上是 4GB从一块 7GB/s 的 NVMe SSD 顺序读回只需要 0.6 秒左右在 H100 上重新 prefill 32K token按 20K tokens/s 的速度算也要 1.6 秒以上实际因为并行度和显存带宽竞争通常更慢。如果前缀更长比如 128K差异就更明显。2.3 命中次数决定收益上限还有一个容易被忽略的点缓存方案的收益是复利式的。同一个长前缀第一次写入后第二次读回就省了一次 prefill之后每次命中都是在赚钱。RAG 场景里文档库可能被几千个问题反复引用多轮对话里历史对话是天然的重复前缀。命中率越高这套方案越划算。所以评价一个 KV cache 卸载方案核心指标不是单次读回快不快而是命中率 x 单次省下的 prefill 时间是不是远大于存储读写和序列化的开销。LMCache 和 vLLM 结合后很多实测收益就落在缓存命中率和 TTFT首 token 延迟上而不是 decode 速度。3. LMCache 在 vLLM 体系里到底做了什么3.1 vLLM 原生 prefix caching 的边界vLLM 自己是有自动前缀缓存机制的原理是把 KV cache 按 block 粒度组织相同前缀的 block 可以被多个请求复用。但它的缓存默认放在 GPU 显存里显存紧张时这些缓存 block 会被淘汰掉。vLLM 也支持把 KV cache 的一部分放到 CPU 上通过cpu_offload_gb参数。这个参数的作用是给 CPU 上预留一部分空间存 swap-out 的 KV block但它的定位更像页交换主要为了短时缓解显存峰值不是做长远的缓存分层。SSD 一级基本不在考虑范围内。3.2 分层缓存GPU 是 L1CPU 是 L2SSD 是 L3LMCache 做的事情本质上是给 vLLM 加了一个多级缓存层。GPU 显存里保留最活跃的 KV cache冷一点的数据落到 CPU 内存再冷的落到 SSD 磁盘类似 CPU 的 L1/L2/L3 缓存架构。具体到实现层面LMCache 通过 vLLM 的 KV connector 接口接入把 vLLM 中按 block 管理的 KV cache 序列化后存入本地。每个缓存项对应一段 token 序列的 KV 值匹配时按 token 前缀来找。命中后直接把 KV cache 读回 GPUvLLM 就不需要再对这些 token 做 prefill。这里有几个关键设计块粒度存储LMCache 以 KV block 为最小存储单位和 vLLM 的 PagedAttention block 对齐不需要重新切分数据。LRU 淘汰新 KV cache 写入时如果 CPU/SSD 空间不够按照最近最少使用原则淘汰冷数据。跨实例复用同一台机器上可以跑多个 vLLM 实例缓存写入共享目录A 实例算过的前缀B 实例直接复用。这比每台机器各算一遍要省太多。3.3 和 chunked prefill 的配合长上下文场景下LMCache 和 vLLM 的 chunked prefill 配合起来很有意思。vLLM 把 prefill 阶段切成若干 chunk避免单个长请求把 GPU 占死。LMCache 在 chunk 粒度上做缓存命中的 chunk 直接跳过只对未命中的 chunk 做实际计算。这意味着即使同一个请求中只有一部分前缀命中也能获得部分收益。比如一个 128K 的请求前 100K 是文档内容之前已经算过后 28K 是新的问题那只需要算 28K token剩下 100K 全部从 SSD 读回。这个特性对 RAG 特别友好。知识库文档被分割成多个 chunk不同用户的问题引用同一批文档时chunk 级缓存几乎是天然命中的。4. 实操接入把 vLLM 的 KV cache 导到 CPU 和 SSD4.1 版本和安装LMCache 以 Python 包的形式发布安装很简单pip install lmcache需要注意版本匹配。vLLM 的 KV connector 接口一直在变最好用 vLLM 和 LMCache 社区互相适配的版本组合。我当时用的是 vLLM 0.6.x 对应的 LMCache 版本接口相对稳定。如果你用的是更新的 vLLM V1 引擎要确认 LMCache 是否已经适配否则 connector 可能加载不进去。启动 vLLM 时用--kv-transfer-config参数把 KV 传输层切到 LMCachevllm serve meta-llama/Llama-3.1-8B-Instruct \ --gpu-memory-utilization 0.7 \ --max-model-len 131072 \ --kv-transfer-config {kv_connector:LMCacheConnector,kv_role:kv_both}这个配置的含义是KV connector 使用 LMCache当前实例既能写入 KV cache也能读取 KV cache。4.2 缓存路径和环境变量LMCache 默认会在本地目录下创建缓存文件。建议把缓存目录放到独立的数据盘上不要放在系统盘。export LMCACHE_CPU_OFFLOAD_PATH/data/lmcache/cpu export LMCACHE_SSD_OFFLOAD_PATH/data/lmcache/ssd在 vLLM 启动日志里能看到 LMCache 的初始化信息包括 cache 目录和已加载的缓存大小。看到类似LMCache initialized的日志说明 connector 已经接上了。如果目录不存在记得先创建并保证权限正确。我第一次跑的时候忘了建目录LMCache 抛了一个权限异常服务直接起不来。4.3 验证缓存是否真的命中启动之后先跑一个简单测试同样的 Prompt 连续问两次观察日志中是否有 KV cache hit 的记录。LMCache 在命中时会在日志里打印 hit 信息或者你可以在代码里通过 metrics 接口拿到命中率。没有日志输出的话说明缓存没接上大概率是版本不匹配。再测一个长文档问答场景。第一次问某份文档耗时包括了完整 prefill第二次换一个问题继续问同一份文档如果命中TTFT 会明显下降。如果只是跑通 Demo到这里就算完成了。但要真正把显存压力降下来还需要调整--gpu-memory-utilization。我实际用的配置是 0.65-0.75给 GPU KVCache 留一部分同时把大量冷 KV cache 推到 CPU/SSD。开太高的话GPU 缓存空间压缩了 LMCache 的工作可以更充分但同时活跃序列的 KV cache 也容易 OOM。5. 收益实测与适用边界什么时候开什么时候别开5.1 我看到的收益区间在我们内部的长文档问答服务上LMCache 接入后效果比较明显。改造前128K 上下文、并发 4 路时单卡显存很快打满服务持续 OOM 重启改造后A100 单卡可以稳定承载 8-12 路并发TTFT 从 1.8 秒左右降到 0.3-0.5 秒前提是前缀命中率足够高。多轮对话场景收益更大。每轮请求携带全部历史对话假设历史有 64K token每次新问题只需要算新加的一小段。没有缓存时每轮都要为 64K 历史做重复 prefill有缓存后除了第一轮后面的请求基本就是纯 decode 速度。我还测过前缀完全不同的场景比如每次请求的 Prompt 都是随机生成的无关联文本命中率趋近于 0。这种情况下KV cache 写入是纯开销TTFT 反而比不开 LMCache 时慢了 10%-15%主要是序列化和写磁盘的时间。5.2 适合开的场景RAG / 知识库问答文档 chunk 被反复引用天然适合缓存。多轮对话历史对话是稳定的重复前缀。多实例推理多个 vLLM 副本共享缓存目录避免每台机器重复算相同前缀。长上下文 高并发显存成为瓶颈时把冷 KV cache 换到 CPU/SSD能显著提升并发能力。5.3 不适合开的场景一次性短 Prompt上下文只有几百 token缓存收益太小读写开销占比过高。GPU 显存非常充裕所有 KV cache 都能放进显存时卸载反而增加了无谓的序列化和读盘延迟。数据盘性能太差如果只有一块 SATA SSD 或者 IO 被其他业务打满读回速度可能比 prefill 还慢。CPU 内存也非常紧张CPU 内存一旦开始 swap整个节点都会卡顿此时卸载到 CPU 内存收益为负。性能收益和命中的关系大致可以概括为下表场景命中率对 TTFT 的影响对显存的影响RAG 长文档问答高大幅降低大幅缓解多轮对话高大幅降低明显缓解随机短 Prompt低轻微变慢帮助不大GPU 显存充裕依赖命中率可能变慢无意义6. 踩坑实录配置、性能和运维层面的教训6.1 坑一Prompt 前缀里的动态内容让缓存全部失效我们内部有个服务在系统提示词前面拼了一段带时间戳的文本比如“当前时间2025-01-12 14:23:11”。这个时间戳每次请求都在变导致 LMCache 认为前缀完全不一致缓存命中率直接降到 10% 以下。排查的时候看日志发现大量 miss后来把一个请求的完整前缀打出来才意识到问题。解决办法是把动态内容从系统提示词移到用户消息的末尾让公共系统提示词保持完全一致。改动之后命中率提到了 80% 以上。这是一个极其容易被忽略的问题。LMCache 的匹配逻辑是严格按 token 前缀对齐的任何动态字段放在前缀部分都会让整段缓存失效。规范的做法是系统提示词固定不变动态信息一律放到请求尾部。6.2 坑二把缓存目录放在系统盘第一次部署的时候我没指定LMCACHE_SSD_OFFLOAD_PATH默认缓存写到了系统盘。测试到一半整台机器的 IO 全部被打满其他服务的日志、监控全部异常。原因很简单KV cache 的文件尺寸非常大长对话写下来上百 GB 很正常系统盘根本扛不住这种写入压力。后来把缓存路径指到独立 NVMe 数据盘并且在数据盘上做了容量监控才解决。建议一定要用独立的数据盘最好是 TLC 或企业级 SSD不要用 QLC。QLC 盘的写入寿命和一致性表现在 KV cache 高频写入场景下会非常吃力。6.3 坑三gpu-memory-utilization 不是越大越好我第一次优化参数时把gpu-memory-utilization调到 0.95想着尽可能多地给 GPU KVCache 留空间。结果显存反而频繁 OOM因为活跃序列的 KV cache 和 LMCache 读回的缓存会同时存在 GPU 上。后面把参数降到 0.7 左右情况才稳定。我的理解是LMCache 的设计目标是分层缓存如果 GPU 显存里没有一点多余空间接收从 SSD 读回的 KV cache那卸载就没有意义。合理的做法是给 GPU KVCache 一个安全余量同时让它保留一部分高频命中的热点。6.4 坑四多实例共享缓存目录时的写冲突多副本部署时多个 vLLM 实例写同一个缓存目录一个实例在写、另一个实例刚好淘汰同一个文件会出现读取失败或缓存损坏的异常。我们目前的方案是每个实例使用独立子目录再通过上层负载均衡让相同请求尽量落同一个实例。如果要严格共享最好用支持并发读写的文件系统并且注意 LMCache 版本对多进程并发的支持情况。这个领域还在快速迭代不同版本的稳定性差异很大。6.5 坑五CPU 内存装太满触发了 swapCPU 内存也不是无限的。有一次我把LMCACHE_CPU_OFFLOAD_PATH指向 /dev/shm共享内存顺手把大量 KV cache 都放上去结果 CPU 内存用尽操作系统的 swap 开始介入服务整体性能断崖式下跌。KV cache 从 SSD 读回时会先经过 CPU 内存再搬回 GPU所以 CPU 内存是不可避免的中间层。建议给 LMCache 的 CPU 内存使用设置上限不要让它把整台机器的内存吃光。6.6 坑六日志和监控没有提前接入接入 LMCache 之后建议第一时间把命中率和缓存读写量接到监控里。没有指标就没有优化依据我改造初期完全靠日志里的零散信息判断走了不少弯路。运维层面要关注四个指标缓存命中率、SSD 读带宽、CPU 内存占用、磁盘剩余空间。这四个指标基本决定了一套 KV cache 卸载方案的健康度。命中率低是前缀设计问题SSD 带宽打满是存储性能问题CPU 内存高是分层策略问题磁盘满则是淘汰策略问题。我在实际部署中的体会是LMCache 和 vLLM 的组合并不是把 SSD 当作 GPU 显存的替代品而是把 KV cache 当成一份可以被多级存储管理的缓存数据。显存紧张时这套方案确实能让你在同样硬件上跑更长的上下文和更高的并发。不过它是有前提的缓存命中率、存储性能和容量的规划决定了整个方案是锦上添花还是画蛇添足。如果让我重新做一次我会先把请求前缀的稳定性治理好再谈卸载策略这样后续的一切优化才有意义。
返回列表