Paper

FlexiCache:利用注意力头的时间稳定性管理分层 KV Cache

FlexiCache 根据注意力头的时间稳定性差异,在 GPU 与主机内存间分层放置 KV 页面,降低显存占用、重排开销与在线推理延迟。

FlexiCache: Leveraging Temporal Stability of Attention Heads for Efficient KV Cache Management
Nazmul Takbir, HamidReza Alikhani Koshkak, Nikil Dutt, Sangeetha Abdu Jyothi
Department of Computer Science, University of California, Irvine · Proceedings of Machine Learning and Systems (MLSys 2026)

论文概览

FlexiCache 研究的是长上下文、长生成场景下的 decode KV Cache 管理问题。它的出发点不是简单地判断“哪些 token 不重要”,而是进一步观察:不同 KV head 的重要页面集合在时间维度上的稳定程度不同。有些 head 连续多个 decode step 都关注几乎相同的一组页面;另一些 head 的关注范围则持续漂移。

基于这一差异,FlexiCache 不再对所有注意力头采用统一的缓存策略,而是把 head 分为 stable 与 unstable 两类:unstable head 的完整 KV 保留在 GPU,stable head 只在 GPU 保存当前 top-K 页面,其余完整 KV 下沉到主机内存。系统再以不同频率重新排序两类 head,并仅回传 stable head 中新晋升的页面。

论文发表于 MLSys 2026,实现基于 vLLM。作者报告其在保持约 99% 稠密注意力效果的同时,将长上下文请求的 GPU KV 占用最多降低约 70%,离线吞吐提升 1.38–1.55 倍,在线 TPOT 降低 1.6–2.1 倍。

研究背景与问题

长上下文 LLM 的 decode 阶段同时受到两种压力:

  1. KV Cache 随输入长度与生成长度线性增长,压缩了可容纳的并发请求数量;
  2. 每个 decode step 都要访问历史 KV,注意力计算与显存带宽开销随序列增长而扩大。

现有方案通常沿两条路线处理:

  • 永久丢弃低重要性 KV,例如 StreamingLLM、SnapKV。这类方法能够直接减少显存,但被删除的 token 后续即使重新变得重要也无法恢复,长生成中更容易损失质量。
  • 保留完整 KV,但每步只计算 top-K 页面,例如 Quest。这能够降低注意力计算量,却仍要求完整 KV 常驻 GPU,无法释放主要的容量压力。

LServe 采用混合路线:部分 head 使用流式掩码并永久删除远端 KV,部分 head 动态选择页面。FlexiCache 认为问题的关键不只是“哪些页面重要”,还包括:每个 head 的重要页面集合能稳定多久。这个时间稳定性决定了页面是否适合下沉,以及需要多频繁地重新选择和加载。

核心洞察:注意力头具有不同的时间稳定性

对某一层、某一 KV head,FlexiCache 记录每个 decode step 选出的 top-K 页面集合,并比较相隔若干步的集合重叠度。为消除随机重叠带来的偏差,论文使用 Random-Corrected Overlap(RCO):RCO 为 1 表示两个 top-K 集合完全一致,为 0 表示其重叠程度接近随机抽样。

系统在一个长度为 W 的窗口内聚合 RCO,得到每个 head 的时间稳定性分数。实验显示:

  • 同一层内,不同 KV head 的稳定性差异显著;
  • 某些 head 长期保持高重叠,另一些 head 始终处于低重叠状态;
  • 这种不稳定 head 的身份在不同任务间高度重合,因此更像是模型内部属性,而不是单个请求或任务的偶然现象。

在 Llama-3.1-8B-Instruct 上,不同任务识别出的 unstable head 集合平均重合度约为 0.83;Mistral-7B、Mistral-Small-24B 和 Qwen2.5-32B 上分别约为 0.83、0.76 和 0.81。作者因此采用一次离线 profiling:将稳定性最低的 25% head 归为 unstable,其余 75% 归为 stable。

核心方法

1. 稳定性差异化放置

FlexiCache 对两类 head 采用不同的 GPU/CPU 放置策略:

Head 类型 GPU 中保存的 KV 主机内存中的 KV 重排频率
Unstable head 完整 KV Cache 不依赖主机副本参与选择 每个 decode step
Stable head 当前 top-K 页面 完整 KV Cache 默认每 16 个 step

需要强调的是,unstable head 虽然完整驻留 GPU,但实际 decode attention 仍只计算 top-K 页面。完整驻留的目的是让系统能够每步重新选择,而不产生频繁的主机到 GPU 传输。

2. 基于 MinMax 元数据的页面评分

FlexiCache 沿用 Quest/LServe 的页面评分思路。每个 KV 页面维护逐维 key 最小值与最大值;给定当前 query,系统用这些边界估计该页面可能产生的最大点积,从而选出 top-K 页面。

关键设计在于:MinMax 元数据与完整 KV 页面解耦。即使 stable head 的 KV 页面已经被卸载到主机内存,对应的 MinMax 元数据仍保留在 GPU。因此系统可以在不读取完整页面的前提下,对所有历史页面重新评分,并判断哪些离线页面应该晋升回 GPU。

3. 周期性重排与增量晋升

Unstable head 每步重新计算 top-K;stable head 默认每 16 步重排一次,并在两个重排点之间复用上一轮 top-K 集合。重排后,系统不会重新加载整个 top-K,而只传输“当前 top-K 减去上一轮 top-K”得到的新晋升页面。

在 Llama-3.1-8B、2048-token budget 的设置中,stable head 的 GPU top-K KV 总量约为 192 MB,而每次重排实际晋升的数据约为 44–67 MB,即 top-K 总量的 23%–35%。这说明稳定性确实减少了集合变化,但变化量仍不可忽略,周期性重排仍然必要。

4. 传输与计算重叠

对 stable head,prefill 后的大批量 GPU→CPU 卸载在低优先级 CUDA stream 上异步执行;decode 继续进行,卸载完成后才释放对应 GPU block。后续页面填满时,只进行一次增量卸载。

CPU→GPU 的页面晋升更复杂,因为同一请求在加载期间继续 decode 可能读取到不一致的 block。FlexiCache 暂停正在 reload 的请求,但继续执行 batch 中其他请求。由于 stable head 每 16 步才重排,任一时刻通常只有约 1/16 的请求受到影响。

系统还实现了直接访问 pinned host KV 的 CUDA UVA 搬运 kernel,避免 CPU 端 gather-copy-scatter;同时通过低优先级 stream、限制 grid size 和切分短 kernel 来降低传输 kernel 对主计算 SM 的占用。

系统设计

FlexiCache 在 vLLM 上增加了四类关键组件:

  • Stability-aware reranker:根据 head 类别决定 top-K 评分频率;
  • Sparse decode kernel:每个 head 只对选中页面执行 decode attention;
  • Hierarchical block allocator:分别管理 GPU 与主机 KV block;
  • FlexiCache scheduler / KV transfer module:调度请求暂停、异步卸载和页面晋升。

Head 级 block table

标准 PagedAttention 通常让全部层和 head 共享一个请求级 block table。但 FlexiCache 中不同 layer-head pair 的页面会独立晋升和驱逐,因此它将 block table 扩展为 [batch, layer, head, logical_block]

这种细粒度映射会带来明显的控制面开销,论文实现了三项优化:

  1. Dirty tracking:只同步发生变化的 block-table 区域,而不是每步传输完整四维表;
  2. Physical block reuse:重排时将被驱逐页面的物理 block 直接重新分配给晋升页面,避免 CPU allocator 的频繁分配与回收;
  3. Null block:被驱逐的逻辑位置仍保留表项,但指向不可解引用的 null block,使 block table 保持规整、便于向量化处理。

实验设置

项目 配置
GPU 单张 NVIDIA H100 94 GB HBM
CPU 2× AMD EPYC 9554,单颗 64 核
主机内存 1.1 TB DDR5,默认划分 180 GB 给 KV Cache
互连 PCIe 5.0,峰值双向带宽 64 GB/s
软件 vLLM、PyTorch 2.7.0、CUDA 12.8、cuDNN 9.5
模型 Llama-3.1-8B-Instruct、Mistral-7B-Instruct-v0.2;稳定性与准确率扩展到 Mistral-Small-24B、Qwen2.5-32B
数据集 LongBench 16 个任务、L-Eval 多个长上下文/长生成任务
系统基线 vLLM Triton Flash-Decoding
精度对比 Dense attention;LServe
默认策略 25% unstable heads,stable head 每 16 步重排,top-K budget 为 1024 或 2048 token

离线吞吐实验从 L-Eval 随机采样 500 个请求,输入长度为 10k–30k token,输出长度为 50–1500 token。在线实验同样使用 500 个请求,输入为 10k–30k,输出为 20–2000,并按 Poisson 过程到达。

主要结果

准确率保持

在 LongBench 16 个任务上,FlexiCache 相对 dense attention 的平均得分比例为:

  • Llama-3.1-8B:约 1.00;
  • Mistral-7B:约 0.99。

在 L-Eval 上,2048-token budget 对 Llama-3.1-8B、Mistral-7B、Mistral-Small-24B、Qwen2.5-32B 的平均保持率分别约为 0.99、0.99、1.001 和 0.998。

消融结果更能说明机制作用:不进行周期性 reranking 时,L-Eval 保持率会下降到约 0.88–0.944;不区分 unstable head、将所有 head 都视为 stable 时,则约为 0.95–0.978。也就是说,周期性重新发现重要页面为漂移 head 保留完整 GPU KV都不可缺少。

显存与 kernel 性能

  • 序列长度超过 20k、top-K budget 为 1024 token 时,GPU KV 占用降低超过 70%;理论上随序列增长逐渐接近 75%,对应仅 25% unstable head 保留完整 KV。
  • batch size 为 40、输入为 10k token 时,decode kernel 最多加速约 4 倍。
  • 稳定性感知 reranking 相比所有 head 每步重排,在 batch size 40 时快约 2.44 倍。

离线与在线服务

  • 论文摘要报告离线 token throughput 提升 1.38–1.55 倍;随着输出长度增加,decode 占比提高,收益更明显。
  • 固定输出长度 500 时,Llama-3.1-8B 和 Mistral-7B 的请求吞吐提升约 1.35–1.46 倍。
  • 对 Mistral-Small-24B 与 Qwen2.5-32B,端到端吞吐分别提升约 1.37 倍与 1.33 倍。
  • 在线负载达到 0.4 req/s 时,FlexiCache-1024 的平均 TPOT 为 34.6 ms,vLLM 为 71.5 ms,改善约 2.1 倍;FlexiCache-2048 为 44.9 ms,改善约 1.6 倍。

FlexiCache 对 TTFT 的直接优化有限,但它降低了每个 decode 请求的 GPU KV 占用,因而在高负载下延缓显存饱和与排队,使 TTFT 不会像基线那样突然恶化。

与相关工作的区别

与 SnapKV、StreamingLLM 等永久压缩方法相比,FlexiCache 不永久删除 stable head 的历史 KV,而是把完整副本保存在主机内存,因此后续可以重新晋升曾经不重要的页面。

与 Quest 相比,两者都利用 query-aware top-K 页面选择,但 Quest 仍把完整 KV 放在 GPU;FlexiCache 用 head 稳定性决定哪些完整 KV 可以下沉,从而真正释放显存。

与 LServe 相比,LServe 将部分 head 转换为 streaming head 并永久移除其远端 KV;FlexiCache 采用可恢复的分层存储,更适合长生成中注意力焦点发生漂移的场景。

与 FlexGen、MoE-Lightning 等分层卸载系统相比,FlexiCache 的放置粒度是 layer-head-page,并显式依据页面重要性和 head 的时间稳定性,而不是把同一层或同一批请求的 KV 视为均匀对象。

与 LMCache 相比,LMCache 主要服务于跨请求 prefix reuse 和 prefill/decode 解耦,以降低重复 prefill 与 TTFT;FlexiCache 管理的是已经进入 decode 的活跃请求 KV。

局限性

1. 仍然依赖大容量主机 DRAM

FlexiCache 释放的是 GPU 容量,但 stable head 的完整 KV 仍线性保存在主机内存,实验默认预留 180 GB pinned DRAM。它并未解决总 KV 容量随上下文和并发增长的问题,只是将主要容量压力转移到了更便宜的层级。

2. 对高带宽单机环境依赖较强

实验使用 H100、PCIe 5.0 和双路 EPYC。低带宽 PCIe、NUMA 跨 socket、多 GPU 共享链路或主机内存争用都可能放大 reload 开销。论文尚未给出这些条件下的敏感性分析。

3. 静态二分类较粗

系统通过一次离线 profiling 固定 25% unstable heads,并统一使用 16-step 重排周期。论文已经观察到稳定性本身是连续谱,因此二分类是工程简化,而不是最优策略。模型微调、量化或注意力实现变化后,也可能需要重新 profiling;Llama-3.1-8B 的一次 profiling 约需 2 小时。

4. 页面粒度与模型范围有限

论文以 16-token page、1024/2048-token top-K budget 为主,重点评估传统 dense/GQA 模型。对于 DSA 等原生 token-level sparse attention,模型内部已经有更细粒度的索引器,MinMax 页面评分未必仍是合适路径。

5. 生产规模验证不足

主要结果来自单 GPU、10k–30k 输入和最多 2000-token 输出。系统尚未覆盖百万 token、长时间多轮 agent 会话、PD 分离、多节点迁移以及集群级负载均衡。系统基线也主要是 vLLM,缺少与更多 decode KV offloading 系统的端到端对比。

对大模型推理加速和存储优化的启发

稳定性可以直接转化为缓存生命周期

FlexiCache 最有价值的地方,不是再次证明注意力具有稀疏性,而是把时间稳定性转化为放置层级、重排周期和预取频率。重要性回答“现在需要什么”,稳定性回答“多久后需要重新判断”。后者更接近缓存系统中的 reuse distance 和生命周期信号。

对 HBM–DRAM–SSD / NPU–ASU 架构的启发

FlexiCache 当前是 GPU–DRAM 两级设计,但其策略可以自然扩展:

  • unstable head 或短期频繁变化的 KV 留在 HBM/近端 DRAM;
  • stable head 的当前 top-K 留在 HBM;
  • stable head 的完整冷 KV 下沉到 ASU DRAM 或 SSD;
  • 轻量索引、重要性摘要与映射元数据保留在 HBM,避免为了判断重要性先读取完整 KV。

不过,FlexiCache 每 16 步仍可能晋升 44–67 MB 数据。对 SSD 层而言,这种页面级批量回迁仍然过粗,必须进一步依赖 token-level 选择、异步预取、多请求 I/O 合并和更长时间尺度的预测,才能把尾延迟隐藏在 decode 计算之后。

原生稀疏模型可以省掉一部分评分开销

对于带有原生 top-K indexer 的模型,系统不必再用 MinMax 估计页面重要性。模型索引器已经输出候选 token,存储系统可以直接消费这些结果,把研究重点转向:top-K 命中率、跨层预测、不同时间尺度的保留策略以及 HBM/DRAM/SSD 间的数据迁移。

稳定性不应只做二分类

更合理的下一步是把 head/token 的稳定性映射为多个时间层级:每步重排、短周期重排、长周期重排以及纯冷数据。这样可以形成与多级存储介质延迟更匹配的策略,而不是 stable/unstable 的单一阈值。

今夜白的观察

FlexiCache 是一篇逻辑完整、系统实现扎实的工作。它把三件以往常被分开讨论的事情连在了一起:稀疏注意力降低计算、时间稳定性降低选择频率、分层缓存降低 GPU 容量需求。特别是 MinMax 元数据常驻 GPU、完整 KV 下沉的设计,清楚地区分了“用于决策的轻量索引”和“真正参与 attention 的重数据”。

但它仍然是一种以主机 DRAM 为安全后备层的设计。完整 KV 并没有消失,系统收益很大程度上建立在 H100、PCIe 5.0 和 180 GB pinned host memory 上。对下一代 AI-centric storage 架构而言,真正困难的问题会从“能否卸载”转向“如何让更慢、更便宜的 SSD 层也能及时返回下一步真正需要的稀疏 KV”。

因此,这篇论文对后续工作的主要启发不是直接复制 stable/unstable 二分类,而是采用它的分析范式:先量化访问集合随时间变化的速度,再让缓存层级和预取机制匹配这一变化速度。

参考资料