现在位置: 首页 / AIGC / KVcache优化 / 正文

大模型推理过程中KV Cache 分层/层次化处理方案介绍

  • admin
  • AIGC
  • 1天前
  • 16热度
  • 0评论

一、KV Cache 为什么会随会话增长而膨胀

在自回归解码时，每生成一个 token，模型的每一层都会产生对应的 Key 与 Value 向量并追加到 KV Cache 中。
如果一次会话的上下文长度为 （包括用户输入、模型已生成的 token），模型层数为 ，每层的注意力头数为 ，每个头的隐藏维度为 ，使用的数值精度为  字节（FP16 = 2 B），则 KV Cache 的显存占用可近似为

（“2” 表示 Key 与 Value 两套向量）。
当同一台服务器上需要 并发维护多个会话 时，显存需求会呈 线性叠加，很快超过单卡 GPU 的 HBM 容量，尤其在 100 K‑1 M token 超长上下文场景下更是如此。

二、主流的 KV Cache 分层/层次化处理方案

方案 关键思路 典型实现/论文 适用场景 主要优缺点
层级（Layer‑wise）Offload 按层把 KV Cache 从 GPU 迁移到 CPU；只保留当前计算层的 KV 在显存中。 LayerKV（Ant Group） 中等并发、显存紧张的服务器 能显著降低峰值显存（仅保留少数层），但每层切换会产生 PCIe 传输开销，需要异步流‑拷贝来隐藏延迟。
头部（Head‑wise）Offload 将注意力头划分为 “在 GPU 上保留” 与 “立即 offload 到 CPU”。 HEADINFER（Head‑wise Offloading） 对显存极度受限、但对部分头的计算依赖强的模型 进一步压缩显存（只保留关键头），但实现更复杂，需在每个 token 级别调度。
块/分页（Chunk / Paged）KV 把 KV Cache 划分为固定大小的块，按需分页到 CPU 或磁盘；常配合 PagedAttention、vLLM 的 PagedKVCache。 PagedAttention（Kwon et al., 2023）、vLLM 实现 超长上下文（≥100 K token）且并发量大 通过块级预取实现显存‑CPU 重叠，分页粒度可调；但块大小过小会导致 PCIe 吞吐受限。
压缩（Low‑rank / Quant） 对 KV 向量做 SVD、低秩投影或 INT8/4‑bit 量化，直接减小存储体积。 PALU（低秩投影）、LeanKV、ShadowKV（SVD+压缩） 对延迟敏感、显存极限的单卡部署 压缩率高但会引入近似误差，需要在精度与内存之间权衡。
持久化/长记忆池 将 KV Cache 持久化到 CPU DRAM、页锁定内存，甚至 NVMe SSD，形成跨会话共享的 “长记忆”。 中国信息通信研究院报告（多级 KV‑Cache）、Cheops‑LLM 对 SSD Offload 研究、Compute or Load KV Cache?（CPU + NVMe） 多轮对话、需要跨会话复用上下文的在线服务 CPU → GPU 带宽是瓶颈；NVMe 能提供 TB 级容量但延迟最高，需要分层缓存（GPU → CPU → SSD）进行分层调度。
混合调度框架 同时结合 层/头 offload、分页、压缩 与 持久化，在运行时根据带宽、负载动态决定 KV 的存放位置。 FlexGen、DeepSpeed‑Zero‑Inference、TURNIP（任务图+内存图）、LayerKV 的异步 CUDA 流调度 大规模服务平台、需要兼顾吞吐与延迟 体系结构最完整，但实现复杂，需要专门的调度器与监控。

三、分层存储的典型实现细节

  1. GPU → CPU（页锁定内存）

    • 速度最快的离显存方案，PCIe ×16（约 30 GB/s）可在 256 token 块上实现 0.5 ms 左右的拷贝（见）。
    • 需要 异步 CUDA 流 与 批量拷贝，避免每个 token 单独传输导致的 kernel 启动开销。

  2. CPU → NVMe SSD

    • 适用于 长期 KV 保存（如 prefix cache、跨会话记忆池），容量可达 TB 级。
    • 访问延迟在 0.1‑0.5 ms（NVMe PCIe 4.0），吞吐约 2‑3 GB/s，必须配合 块预取 与 读写合并（Cheops‑LLM 研究表明，块大小 4‑8 KB 能在 SSD 上达到 80% 带宽利用率）。

  3. 压缩与选择性加载

    • 低秩投影（PALU）把每个 KV 向量从 128 维压到 32 维，显存降低 75% 且误差 < 0.2%（实验在 Llama‑3‑8B 上完成）。
    • 稀疏/重要性筛选（ShadowKV）只保留 “关键 token” 的 KV，其他 token 采用 SVD 近似或直接丢弃，进一步削减带宽需求。

四、实际部署时的容量规划示例

假设部署 Llama‑3‑8B（层数 L = 32，头数 H = 32，head‑dim D = 128），使用 FP16，单卡显存 24 GB。

上下文长度 S 仅 GPU KV（不压缩）显存需求 加入 层‑wise Offload（仅保留 4 层） 再加 INT8 KV 量化 结果
2 048 token ≈ 3.2 GB ≈ 0.4 GB（仅 4 层） ≈ 0.2 GB 完全适配
8 192 token ≈ 12.8 GB ≈ 1.6 GB ≈ 0.8 GB 仍在显存余量内
64 000 token ≈ 100 GB ≈ 12.5 GB（仍超） ≈ 6.3 GB（仍超） 必须 分页/SSD Offload

从表中可以看出：当上下文长度进入万级以上，仅靠 GPU + 量化 已不足，需要 分页 + CPU/SSD 分层（如 LayerKV + PagedAttention）才能保证不 OOM。

五、推荐的系统级实现路线

  1. 基础层：在 GPU 上保留 当前层的 KV（Layer‑wise）或 关键头的 KV（Head‑wise），使用 异步 CUDA 流 实现计算‑拷贝重叠。
  2. 中间层：把 已完成层的 KV 按 256‑512 token 块 异步 offload 到 页锁定 CPU 内存，利用 PCIe 带宽最大化吞吐。
  3. 持久层：对 跨会话、需要长期保存的 prefix cache，采用 NVMe SSD 作为后备存储，配合 块预取 与 LRU/热点调度，实现 TB 级容量。
  4. 压缩层（可选）：在 KV 进入 CPU/SSD 前做 低秩投影 或 INT8 量化，进一步降低带宽与存储开销。
  5. 调度器：使用 LayerKV、FlexGen、DeepSpeed‑Zero‑Inference 或 TURNIP 的 任务图 + 内存图 机制，动态决定每块 KV 的存放位置，以满足 吞吐‑延迟 双目标。

六、结论

  • KV Cache 的显存需求随会话长度线性增长，单卡 GPU 在超长上下文（≥10 K token）或高并发会话下必然出现显存瓶颈。
  • 业界已经形成 多层次、分层 offload + 压缩 + 持久化 的完整生态：
    • 层/头级 offload（LayerKV、HEADINFER）降低峰值显存；
    • 分页/块级缓存（PagedAttention、vLLM）实现显存‑CPU 动态调度；
    • 压缩/低秩投影（PALU、LeanKV、ShadowKV）在不显著牺牲精度的前提下进一步削减内存；
    • CPU + NVMe 持久化（多级 KV‑Cache、Cheops‑LLM）提供 TB 级容量，支撑跨会话长记忆。
  • 在实际部署时，建议先在 GPU + CPU（页锁定） 之间实现层/头级 offload 与块分页；当业务需要 跨会话长期记忆 时，再引入 NVMe SSD 作为后备层，并配合 低秩压缩 以控制带宽开销。

通过上述分层方案，既能保持 高吞吐、低延迟，又能突破显存上限，实现数十万甚至上百万 token 的长上下文会话。