1. 从一次显存爆炸说起KV Cache 分级存储到底在解决什么第一次在单卡 24G 显存上跑 32K 上下文的推理服务时我盯着nvidia-smi里那条几乎垂直向上的显存曲线心里只有一个念头这玩意儿怎么涨得比房价还快。模型权重加载完占了 14G剩下 10G 本来以为够用结果请求一进来batch 稍微开大一点显存直接 OOM。后来把 batch 压到 1上下文拉到 16K勉强能跑但吞吐低得可怜延迟也飘得厉害。问题就出在KV Cache上。KV Cache 是大模型自回归生成时的核心缓存结构。每生成一个 token模型都要拿当前 token 的 Query 去和前面所有 token 的 Key、Value 做注意力计算。如果每步都重新算一遍历史 token 的 K 和 V计算量会随序列长度平方级增长推理根本没法用。所以工程上会把历史 token 的 Key 和 Value 缓存下来后续每步只算新 token 的 K、V然后拼到缓存里。这个缓存就是 KV Cache。它带来的好处是计算量从 O(n²) 降到接近 O(n)但代价是显存占用随上下文长度线性增长。以常见的 7B 模型为例假设 32 层、隐藏维度 4096、32 个注意力头、FP16 精度每个 token 的 KV Cache 大小大约是2 × 层数 × 头数 × 头维度 × 2 字节。头维度是 4096/32128算下来每 token 约2 × 32 × 32 × 128 × 2 512KB。16K 上下文就是 8GB32K 就是 16GB。这还只是一个请求batch 一开显存直接见底。KV Cache 分级存储要解决的就是这个矛盾上下文越来越长、并发越来越高但显存容量和带宽是有限的。思路很直接——把 KV Cache 按访问频率和时效性分层热数据放显存温数据放主机内存冷数据放本地高速存储甚至可以把不再需要的部分直接丢弃或压缩。这样既保住了推理速度又突破了单卡显存的物理上限。这套方案适合谁如果你在做本地大模型部署、私有化推理服务、长上下文应用开发或者单纯想搞懂 vLLM、SGLang 这些推理框架内部到底怎么管显存那 KV Cache 分级存储就是绕不开的一环。它不是什么高深论文里的概念而是工程落地时实打实要面对的显存管理问题。2. 拆解 KV Cache为什么它成了显存杀手2.1 从注意力机制看 KV Cache 的必然性Transformer 的自注意力计算可以简单理解为每个 token 生成一个 Query、一个 Key、一个 Value。当前 token 的 Query 要和所有历史 token 的 Key 做点积得到注意力权重再对 Value 加权求和。没有缓存时第 t 步生成需要重新计算前 t-1 个 token 的 K 和 V总计算量是 12...n也就是 O(n²)。KV Cache 的做法是第 t 步只计算当前 token 的 Q、K、V然后把新的 K、V 追加到缓存里。注意力计算时直接读取缓存中的历史 K、V。这样每步计算量基本恒定总计算量降到 O(n)。代价就是缓存需要一直占着显存而且随着序列增长不断膨胀。这里有个容易被忽略的点KV Cache 的显存占用和batch size是线性关系。单请求 16K 上下文占 8Gbatch 开到 8 就是 64G单卡根本放不下。很多人在本地部署时发现模型权重明明只占十几 G但一跑长上下文就爆显存根源就在这里。2.2 显存占用的精确计算与瓶颈定位要谈分级存储先得把账算清楚。KV Cache 的单 token 占用公式KV Cache per token 2 × num_layers × num_heads × head_dim × dtype_bytes其中 2 是因为要存 Key 和 Value 两份。以 LLaMA-2 7B 为例32 层、32 头、头维度 128、FP16 占 2 字节单 token 就是2 × 32 × 32 × 128 × 2 524288 字节 ≈ 512KB。16K 上下文约 8GB32K 约 16GB。如果是 70B 模型80 层、64 头、头维度 128单 token 约2 × 80 × 64 × 128 × 2 2.6MB。8K 上下文就是 20GB 以上单卡 A100 80G 也就勉强跑一个请求。这就是为什么长上下文推理必须做显存优化。瓶颈不只是容量还有带宽。注意力计算每步都要读取整个 KV Cache显存带宽决定了读取速度。HBM 带宽虽然高但面对几十 GB 的缓存读取延迟依然可观。分级存储如果设计不好把热数据放到低速介质上推理速度会断崖式下跌。2.3 分级存储的核心思想让数据待在它该待的地方分级存储的本质是按访问频率和时效性分配存储介质。刚生成的 token 的 KV 会被后续很多步访问属于热数据必须放显存。较早生成的 token随着注意力窗口推移访问频率下降可以挪到主机内存。如果上下文特别长很久以前的部分可能再也不被访问就可以放到本地 SSD 甚至直接丢弃。这和 CPU 的缓存层级L1/L2/L3逻辑一致越靠近计算单元容量越小、速度越快、成本越高。KV Cache 分级存储就是把显存当 L1主机内存当 L2本地高速盘当 L3形成一个容量和速度的梯度。关键问题是什么时候搬、搬多少、搬到哪里。搬得太频繁带宽开销吃掉收益搬得太少显存还是不够。这就需要一套策略通常结合注意力窗口、请求优先级、序列位置来决策。3. 分级存储的工程实现从策略到落地3.1 分层策略设计热、温、冷三层的划分依据实际落地时我一般把 KV Cache 分成三层热层显存存放最近 N 个 token 的 KV以及当前活跃请求的全部 KV。N 的取值和注意力窗口、模型层数有关通常覆盖最近 2K 到 4K token。温层主机内存存放超出热层但仍在上下文窗口内的 KV。主机内存容量大DDR4/DDR5 带宽虽然远低于 HBM但通过 PCIe 4.0/5.0 传输批量搬运时开销可控。冷层本地高速存储存放极长上下文中很少访问的部分或者已经完成生成但可能被后续请求复用的前缀 KV。NVMe SSD 顺序读写能到 5GB/s 以上适合大块搬运。划分依据主要是注意力局部性。大量实验和工程实践表明解码阶段注意力权重往往集中在最近一段 token 和少数关键 token 上早期 token 的注意力权重衰减明显。所以把最近 token 放显存老 token 外迁对生成质量影响很小。注意分层不是越细越好。层数太多搬运逻辑复杂调试困难收益可能被开销抵消。我建议先从两层显存内存做起跑通后再考虑加冷层。3.2 搬运时机与触发条件别让搬运拖慢推理搬运时机的选择直接决定性能。常见触发条件有几种显存水位触发当显存占用超过阈值比如 85%开始把最老的 KV 块搬到内存。这是最直接的策略但可能造成突发搬运影响延迟。序列位置触发当序列长度超过热层容量自动把超出部分外迁。这种策略比较平滑适合长上下文场景。请求优先级触发高优先级请求的 KV 尽量留显存低优先级请求的 KV 优先外迁。适合多租户服务。预取与回写结合在注意力计算前预判哪些 KV 会被访问提前从内存取回显存计算完再把不再热的块写回内存。我实测下来序列位置触发 显存水位兜底的组合比较稳。纯水位触发容易在长上下文时频繁搬运纯位置触发又可能在短请求多的时候浪费显存。搬运粒度也很关键。按 token 搬太碎PCIe 事务开销大按块搬比如 128 或 256 token 一块效率更高。vLLM 的 PagedAttention 就是把 KV Cache 按块管理块大小通常 16 或 32分级存储可以复用这套块管理机制。3.3 数据结构与索引怎么快速找到每一块 KV分级存储后KV Cache 不再是一块连续显存而是分散在不同介质上。需要一个索引结构记录每个 token 或每个块的 KV 当前在哪一层、具体偏移是多少。我一般用一张块表Block Table来管理字段说明block_id块编号全局唯一layer_id所属层token_range覆盖的 token 范围location当前位置GPU/CPU/DISKoffset在对应介质中的偏移last_access最近访问时间戳ref_count引用计数用于回收注意力计算时先查块表如果块在显存就直接读在内存就触发搬运回显存在磁盘就异步加载。块表本身放显存或锁页内存保证查询低延迟。实操心得块表用哈希表还是数组取决于块数量。块数量少几千用数组加偏移计算更快块数量大几十万用哈希表避免稀疏浪费。我一般用两级结构按 layer 分数组层内用哈希。3.4 与推理框架的集成方式KV Cache 分级存储不是独立模块要嵌到推理框架的注意力计算流程里。以 vLLM 为例它的 PagedAttention 已经把 KV Cache 切成块分级存储可以作为一个块调度器插在块分配和注意力计算之间。具体集成点块分配时优先分配显存块显存不足时分配内存块并记录位置。注意力计算前检查所需块是否在显存不在则触发搬运。计算完成后更新块的访问时间标记冷块择机外迁。请求结束时根据策略决定回收还是保留前缀缓存。如果是自己写推理代码可以在 KV Cache 的append和read两个接口上做拦截。append时判断是否需要外迁read时判断是否需要回迁。这样对上层模型代码侵入最小。4. 实操过程从零搭一个简易分级 KV Cache4.1 环境准备与基础推理代码先准备一个能跑的最小推理环境。我用 PyTorch Transformers 加载一个小模型比如 Qwen2.5-1.5B 或 LLaMA-3.2-1B方便在单卡上验证。显存 8G 以上即可。import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name Qwen/Qwen2.5-1.5B tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained( model_name, torch_dtypetorch.float16, device_mapcuda ) model.eval()先跑一个普通生成观察 KV Cache 显存增长input_ids tokenizer(你好请介绍一下大模型推理。, return_tensorspt).input_ids.cuda() with torch.no_grad(): out model.generate(input_ids, max_new_tokens512, do_sampleFalse) print(tokenizer.decode(out[0], skip_special_tokensTrue))用torch.cuda.memory_allocated()打印生成前后的显存差就能看到 KV Cache 的实际占用。4.2 手动管理 KV Cache 的分层存储Transformers 的past_key_values就是 KV Cache。默认它在显存里。我们可以手动把它的一部分搬到 CPU。def split_kv_cache(past_key_values, keep_recent512): 把 KV Cache 按层拆分最近 keep_recent 个 token 留 GPU其余放 CPU。 past_key_values: tuple of (key, value)每个 shape [batch, heads, seq_len, head_dim] new_past [] cpu_past [] for layer_k, layer_v in past_key_values: seq_len layer_k.shape[2] if seq_len keep_recent: new_past.append((layer_k, layer_v)) cpu_past.append(None) else: gpu_k layer_k[:, :, -keep_recent:, :].contiguous() gpu_v layer_v[:, :, -keep_recent:, :].contiguous() cpu_k layer_k[:, :, :-keep_recent, :].cpu() cpu_v layer_v[:, :, :-keep_recent, :].cpu() new_past.append((gpu_k, gpu_v)) cpu_past.append((cpu_k, cpu_v)) return tuple(new_past), cpu_past这个函数把每层 KV 切成 GPU 部分和 CPU 部分。生成下一步时需要把 CPU 部分临时取回 GPU 参与注意力计算算完再放回去。实际实现要改注意力计算逻辑这里只演示数据搬运。4.3 搬运开销实测与参数调优搬运开销主要来自 PCIe 传输。PCIe 4.0 x16 理论带宽 32GB/s实际有效带宽约 20-25GB/s。搬运 1GB KV 大约 40-50ms。如果每步都搬延迟直接爆炸。所以必须批量搬运 异步传输。用torch.cuda.Stream做异步拷贝stream torch.cuda.Stream() def async_move_to_gpu(cpu_tensor): gpu_tensor torch.empty_like(cpu_tensor, devicecuda) with torch.cuda.stream(stream): gpu_tensor.copy_(cpu_tensor, non_blockingTrue) return gpu_tensorkeep_recent这个参数需要调。太小搬运频繁太大显存省不下来。我实测 1.5B 模型、16K 上下文keep_recent1024时显存占用从 6G 降到 2.5G生成速度下降约 15%。keep_recent2048时显存降到 3.5G速度下降约 8%。具体取值要看模型大小和延迟容忍度。keep_recent显存占用生成速度下降适用场景5122.0G25%显存极度紧张10242.5G15%平衡20483.5G8%延迟敏感40965.0G3%显存充足4.4 冷层落盘与前缀缓存复用如果上下文特别长或者多个请求共享相同前缀比如系统提示词可以把冷层 KV 落盘下次直接加载。落盘用torch.save或safetensors按块存块表记录文件偏移。import os def save_kv_block(kv_tensor, path, block_id): os.makedirs(path, exist_okTrue) file_path os.path.join(path, fblock_{block_id}.pt) torch.save(kv_tensor.cpu(), file_path) return file_path def load_kv_block(file_path): return torch.load(file_path, map_locationcpu)前缀缓存复用的关键是块哈希。把前缀 token 序列哈希成 key块表里查是否已有对应 KV。有就直接挂到新请求上省去重新计算。这在多轮对话和 RAG 场景里收益很大。注意落盘 KV 要考虑精度和压缩。FP16 落盘占空间大可以转 FP8 或 INT8 再存加载时反量化。但量化会引入误差长上下文下可能累积需要评估。5. 常见问题与排查技巧实录5.1 生成结果变差或重复搬运导致的注意力错位分级存储后如果搬运逻辑有 bug注意力计算可能读到错误的 KV导致生成乱码或重复。排查时先关掉分级确认原始模型输出正常再逐步开启分层观察哪一层引入问题。常见原因块表偏移算错读到了相邻块的 KV。搬运时没有同步注意力计算读到了未完成的拷贝。序列位置编码RoPE在搬运后没有正确对齐。实操心得在注意力计算前加一个校验对比搬运前后 KV 的 checksum。不一致就报错能快速定位。5.2 速度不升反降搬运开销超过收益如果搬运太频繁PCIe 带宽成为瓶颈速度可能比不分层还慢。排查方法用torch.cuda.Event测每步搬运耗时和计算耗时。如果搬运耗时占比超过 30%说明策略太激进。调大keep_recent或增大搬运块大小减少搬运次数。另一个坑是锁页内存。CPU 内存如果不是锁页的non_blocking拷贝实际是同步的速度很慢。用pin_memory()把 KV 所在内存锁页cpu_k cpu_k.pin_memory()5.3 显存碎片与 OOM块管理没做好分级存储后显存里是离散的块容易碎片化。跑一段时间后明明总空闲显存够但分配不出连续块导致 OOM。解决办法用固定大小的块避免变长分配。实现块回收和整理把空闲块合并。预留一部分显存作为缓冲不要占满。vLLM 的 PagedAttention 就是为解决这个问题设计的块大小固定按需分配碎片很少。自己实现时可以参考。5.4 多请求并发下的公平性与隔离多请求同时跑如果都往内存搬内存带宽会被抢。需要做配额和优先级每个请求分配显存配额超了就必须外迁。高优先级请求保留更多热层。内存搬运队列按优先级调度避免低优先级请求饿死。我一般用简单的令牌桶限流每个请求每步最多搬 N 个块保证公平。问题现象可能原因排查方法解决方向生成乱码搬运错位校验 KV checksum检查块表和偏移速度下降搬运过频测搬运/计算耗时比调大 keep_recentOOM显存碎片打印块分配情况固定块大小回收并发卡顿带宽抢占监控 PCIe 利用率配额优先级结果重复位置编码错对比搬运前后位置对齐 RoPE6. 进阶方向压缩、量化与调度协同分级存储只是显存优化的一环。实际生产里我会把它和KV Cache 量化、稀疏注意力、请求调度结合起来。量化把 FP16 压到 INT8显存直接减半搬运量也减半。稀疏注意力只保留重要 token 的 KV进一步降低缓存量。调度层根据请求的上下文长度和优先级动态决定分层策略。比如一个 RAG 请求前缀是长文档这部分 KV 可以量化后放内存甚至磁盘生成时只把最近问题和回答部分的 KV 留显存。这样单卡能同时服务更多请求。另一个方向是跨请求共享。多个请求如果有相同系统提示词这部分 KV 只存一份大家共享。块表里用引用计数管理请求结束减引用零引用才回收。这在企业私有化部署里很实用系统提示词往往很长共享后显存省一大截。我在实际项目里踩过最大的坑是搬运和计算没有 overlap。一开始同步搬延迟高得没法用。后来改成双缓冲一块在算一块在搬用 CUDA Stream 并行速度才回来。这个思路和数据库的预取、CPU 的流水线是一个道理核心就是让搬运和计算同时进行别互相等。最后分享一个小技巧调试分级存储时先把所有搬运日志打出来包括块 ID、源位置、目标位置、耗时。跑一个短请求看日志是否符合预期。确认无误后再关日志跑性能。这样能省很多瞎猜的时间。