ARTICLE DETAIL

资讯详情

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

大模型原理之 KV Cache 分级存储

大模型原理之 KV Cache 分级存储 摘要KV Cache 分级存储的本质是以带宽换容量利用注意力访问的时间局部性将冷数据下沉至低速大容量存储热数据保留于显存。关键词KV Cache分级存储驱逐策略换入换出预取长上下文一、KV Cache 的定义自回归生成过程中每产生一个新 token均需其与全部历史 token 执行注意力运算。为避免对历史 Key/Value 的重复计算系统将历史 token 的 Key 与 Value 予以缓存此即 KV Cache。二、为什么需要分级存储显存墙是长上下文推理的核心瓶颈。KV Cache 的大小随上下文长度线性增长KV 大小 ≈ 2 × 层数 × 序列长度 × 隐藏维度 × 精度字节数 \text{KV 大小} \approx 2 \times \text{层数} \times \text{序列长度} \times \text{隐藏维度} \times \text{精度字节数}KV大小≈2×层数×序列长度×隐藏维度×精度字节数以 LLaMA-2 70B80 层、隐藏维度 8192、64 头、FP16 精度为例上下文长度KV Cache 占用单卡 HBM80GB4k~10 GB可容纳32k~80 GB勉强容纳128k~320 GB无法容纳序列长度从 4k 增至 128kKV Cache 由 10 GB 膨胀至 320 GB而单卡 HBM 容量仅 80 GB。高并发与多轮对话使问题进一步加剧。KV Cache 随请求数线性叠加10 个并发 128k 会话即需 3.2 TB远超单卡乃至单机容量上限多轮对话中历史 token 随轮次持续累积同样推高 KV Cache 峰值。除容量约束外经济性亦构成约束。HBM 单位成本约为 DRAM 的 10 倍、SSD 的百倍量级全量 KV 驻留显存代价过高。综上长上下文推理、高并发服务、多轮对话等场景对 KV Cache 容量提出极高要求容量与经济性双重约束下无法全量驻留显存分级存储因而被视为突破显存墙的主流方案。三、存储层级3.1 三级存储架构分级存储以带宽代价换取有效显存容量数倍至数十倍的扩展其载体为三级存储架构命中未命中命中未命中换入换出驱逐驱逐生成请求L1 GPU HBM热层 ~3TB/s 几十GBAttention 计算L2 CPU DRAM温层 ~100GB/s 几百GBL3 SSD/网络冷层 ~几GB/s TB级图 1 三级存储架构与数据流向层级介质带宽容量存储内容L1 热层GPU HBM~3 TB/s数十 GB最近或最常被注意的 tokenL2 温层CPU DRAM~100 GB/s数百 GB较久未访问、可能被回看的 tokenL3 冷层SSD/网络~数 GB/sTB 级长期未被访问的 token3.2 与相邻技术的关系前缀缓存跨请求复用相同前缀的 KV降低 prefill 阶段开销但不解决单请求内 KV 的容量问题。PagedAttention通过分页管理消除显存碎片提升显存利用率但不扩展总容量。分级存储从时间维度对 KV 进行分层管理直接解决容量瓶颈。四、核心机制4.1 驱逐策略与下沉对象注意力分数驱动H2O累计注意力得分较高的 token 为 heavy-hitter予以保留其余下沉 [5]。时间局部性最近 N 个 token 保留于显存更早的 token 下沉。语义重要性SnapKV依据聚类或语义信息选取代表性 token 保留 [6]。关键权衡驱逐过激导致精度下降驱逐不足则无法有效释放显存。4.2 换入换出与数据迁移换出HBM 至 DRAM/SSD 的数据迁移在生成过程中异步批量执行不阻塞计算。换入DRAM/SSD 至 HBM 的数据迁移在 token 被注意力访问时同步执行。批量搬运累积至一定批量后统一迁移摊薄带宽开销。异步执行换出操作与计算重叠隐藏迁移延迟。4.3 预取机制与提前迁移预测未来将被访问的 KV提前从 DRAM 迁移至 HBM。预测依据注意力模式具有可预测性时间局部性与 heavy-hitter 稳定性。命中时免除等待预测失败则产生无效迁移浪费带宽。4.4 压缩机制与体积缩减量化FP16 至 INT8/INT4体积减半或缩减至四分之一。稀疏化仅存储重要 token 的 KV。与分级正交可先压缩后分级或先分级后压缩。4.5 机制协同示例以下伪代码展示驱逐、换入换出、预取三项机制在生成循环中的协同方式defgenerate_with_tiered_cache(request):forstepinrange(max_steps):# 预取预测下一步需要的 KV提前从 DRAM 搬回 HBMprefetch_kv(predicted_tokens)# 前向计算HBM 命中则直接算未命中则换入outputattention(query,hbm_cache)# 驱逐HBM 满了按策略选冷 token 下沉到 DRAM/SSDifhbm_cache.is_full():victimsselect_victims(eviction_policy)swap_out(victims,target_tierDRAM)# 换入计算中发现需要的 token 在 DRAM/SSD搬回 HBMifneeded_kvnotinhbm_cache:swap_in(needed_kv,source_tierDRAM)四项机制由推理框架自动执行开发者仅需配置层级容量配比与驱逐阈值等参数。五、代表项目下表从核心机制、存储层级、驱逐策略、预取能力与适用场景五个维度对代表性项目进行对比。项目核心机制层级驱逐策略预取适用场景vLLM [1]PagedAttention 按页管理消除显存碎片2 级页级换出无中等长度序列FlexGen [2]线性规划求解最优换入换出调度3 级LP 调度无单请求长上下文InfiniGen [3]学习注意力模式预测未来访问2 级注意力驱动有多轮对话Mooncake [4]KV 中心分离架构跨请求/跨卡共享多级全局调度有高并发与长上下文六、工程实践6.1 层级容量配比HBM 与 DRAM 容量比约为 1:4 至 1:8热层小、温层大。SSD 作为兜底容量不设上限。依据并发数与平均序列长度估算。6.2 驱逐阈值调优阈值过激精度下降困惑度上升。阈值过松无法有效释放显存。采用离线标注集与线上 A/B 测试确定。6.3 常见问题换出带宽瓶颈SSD 带宽仅数 GB/s需控制换出频率。延迟抖动驱逐与换入时机不稳定导致 TPOT 毛刺。预取预测失效无效迁移浪费带宽预测模型需持续校准。七、结论分级存储的本质并非更大的显存而是以带宽换容量。注意力访问的时间局部性决定了大部分 KV 在大部分时间内并不被需要将其下沉至廉价存储、在被访问时迁移回显存是长上下文推理在有限显存条件下运行的关键。未来研究方向包括更精确的预取预测、更细粒度的迁移控制、以及与压缩量化技术的联合优化。参考文献[1] W. Kwon et al.“Efficient Memory Management for Large Language Model Serving with PagedAttention.”SOSP, 2023.[2] Y. Sheng et al.“FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU.”ICML, 2023.[3] W. Lee et al.“InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management.”OSDI, 2024.[4] R. Qin et al.“Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving.”ACM Trans. Storage, 2025.[5] Z. Zhang et al.“H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models.”NeurIPS, 2023.[6] Y. Li et al.“SnapKV: LLM Knows What You are Looking for Before Generation.”arXiv, 2024.
返回列表