Reformulating KV Cache Eviction Problem for Long-Context LLM Inference¶

会议: NeurIPS 2026
arXiv: 2605.07234
论文: https://arxiv.org/abs/2605.07234
领域: LLM 效率
关键词: 大语言模型推理、KV Cache 驱逐、长文本推理、矩阵乘法近似、全局预算分配

一句话总结¶

针对长文本大语言模型推理中的显存瓶颈与传统注意力启发式驱逐策略的局部性缺陷，本文将 KV Cache 驱逐重构为层级矩阵乘法近似优化问题，提出了显式建模注意力分布与投影后值状态乘积交互的 LaProx 准则，并通过跨层归一化实现了无需离线标定的全模型统一全局驱逐。

研究背景与动机¶

随着大语言模型（LLM）上下文窗口拓展至数十万甚至数百万 token，长文本在多轮对话、长文档理解、代码生成以及跨文档推理等任务中展现出卓越能力。然而自回归生成过程中，历史 token 的键值缓存（KV Cache）随着上下文长度线性膨胀，不仅吞噬了宝贵的 GPU 显存，造成极大的内存带宽压力，更导致解码延迟剧烈攀升。为了在有限显存预算下部署长文本模型，学术界与工业界广泛采用动态 KV Cache 驱逐策略，通过评估并丢弃不重要的 token 缓存以保持恒定的缓存规模。

然而，现存主流驱逐机制（如 H2O、SnapKV、CAKE、AdaKV 等）主要依赖注意力权重的累加或局部窗口池化。这种注意力权重启发式方法存在两个根本性理论缺陷。一方面，注意力权重仅仅反映了 query 与 key 的相关性分布，而注意力层的最终输出本质上是注意力概率与投影后值状态（$V W_O$）的矩阵乘积。由于 Softmax 归一化限制，注意力权重的数值范围极其狭窄，而 $V W_O$ 具有更宽泛的幅值分布且随网络深度增加而逐渐放大，仅凭注意力峰值无法忠实代表 token 对层输出的实际贡献。另一方面，现有方法普遍采用独立的单头驱逐或依赖启发式先验划分预算，忽视了多头注意力（MHA）通过输出投影矩阵在线性相加层面融合成统一表示的事实，导致大量在局部头内排名高但全局投影贡献微弱的“局部赢家”浪费了宝贵的缓存容量。

针对这一本质矛盾，本文跳出传统的“单头局部注意力均值统计”思路，从多头注意力的数学结构出发，将保留关键 KV 缓存抽象为在固定显存预算下逼近各层注意力输出的矩阵乘法近似优化问题。核心 idea：将 KV Cache 驱逐重构为以 $A \times (V W_O)$ 矩阵乘积重构误差最小化为目标的近似问题，提出以注意力列范数与投影值行范数乘积为核心的 LaProx 指标，并结合跨层分数归一化实现跨头、跨层的无标定模型级全局动态驱逐。

方法详解¶

整体框架¶

LaProx（Layer Approximated Cache）旨在不调整预训练权重、不引入微调与离线校准的前提下，在 Prefill 阶段结束时完成全局 KV Cache 的高保真压缩。其核心流程分为两个阶段：驱逐重要性指标计算（Eviction Indicator）与模型级全局驱逐决策（Eviction Action）。

首先，在单个注意力层内，算法利用矩阵乘法的结合律预先计算值状态与输出投影矩阵的乘积 $H = V W_O$，随后结合观测窗口内的注意力权重矩阵 $A$，基于蒙特卡洛矩阵乘法近似理论，将各 token 在各头上的重要性显式定义为注意力列向量与投影值行向量的 $L_2$ 范数乘积。其次，在跨头与跨层聚合阶段，算法将所有注意力头的得分拉平，通过层级相对重要性归一化消除深浅层之间的绝对数值尺度偏差，最后在全模型全局预算约束下统一执行 Top-K 选择，自适应地将更多缓存容量倾斜给高信息密度的注意力头与关键网络层。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["长文本输入 Prefill 阶段<br/>生成 Q, K, V 缓存与观测窗口 A"] --> B["矩阵乘法近似指标计算<br/>融合注意力列范数与投影值行范数: ||A[:, i]|| x ||V W_O[i, :]||"]
    B --> C["层内多头得分展平<br/>映射至统一输出表示空间: p_(l, k)"]
    C --> D["跨层分布归一化<br/>消除层间幅值绝对尺度差: s_(l, j) = p_(l, j) / sum(p)"]
    D --> E["模型级全局 Top-K 驱逐<br/>按全模型预算自适应分配各层各头保留容量"]
    E --> F["压缩 KV Cache 解码<br/>保持恒定显存与高效自回归生成"]

关键设计¶

1. 输出感知矩阵乘法近似驱逐指标：以注意力与投影值联合范数刻画真实贡献
传统方法只考虑注意力权重 $A$，完全忽略了值向量 $V$ 和输出投影矩阵 $W_O$ 的加权调制作用。从多头注意力的计算结构来看，令 $W_O = [W_O^1, W_O^2, \dots, W_O^h]$ 为按头切分的输出投影矩阵，$H^i = A^i V^i$ 为第 $i$ 个注意力头的输出，则整个 MHA 层的最终输出可精确分解为各头线性贡献之和：$\text{Output} = \sum_{i=1}^h H^i W_O^i = \sum_{i=1}^h A^i (V^i W_O^i)$。这意味着 token 的实际物理影响必须经由 $V W_O$ 的线性变换投射到隐空间中。若简单地用标量幅值重新缩放均值注意力，会破坏矩阵乘法原有的行列对齐几何关系。本文将缓存驱逐严谨形式化为采样最优子集以逼近矩阵乘积 $A \times (V W_O)$ 的优化问题。根据蒙特卡洛随机矩阵乘法理论，当采样概率正比于对应列向量与行向量的欧氏范数乘积时，乘积估计的重构方差达到最小。因此，第 $i$ 个 token 在该头中的重要性得分 $p_i$ 严格定义为： $$p_i \propto \|\boldsymbol{A}[:, i]\|_2 \cdot \|\boldsymbol{V W_O}[i, :]\|_2$$ 其中 $\|\boldsymbol{A}[:, i]\|_2$ 刻画了当前 token 在查询观测窗口内被关注的强度聚集度，而 $\|\boldsymbol{V W_O}[i, :]\|_2$ 刻画了该 token 在经过输出投影后注入残差流的向量特征范数大小。该指标能够准确惩罚那些虽然注意力权重较高但特征模长极小、以及模长巨大但未被查询关注的边缘 token，确保保留真正主导层输出的关键信息。

2. 层级加法可比性与跨层归一化：消除层间绝对尺度偏差
在传统的单头驱逐框架中，各头独立按照局部排名截断 Top-K，这导致部分头即使整体贡献微弱，仍被迫保留预设配额的“局部赢家”；而承载复杂全局上下文的注意力头却因配额受限丢失关键 token。由于公式中引入了 $W_O^i$，各头原本独立的子空间被投影映射到了同一维度的隐藏层空间中，使得头与头之间的得分具备了直接相加的物理意义。然而，当进一步将选择拓展至全模型跨层范围时，Transformer 结构中随着层数加深残差累加导致的特征幅值逐渐放大，深层 $V W_O$ 的模长往往显著大于浅层。若直接使用原始得分进行全局 Top-K 比较，会导致浅层缓存被无节制地过度驱逐，引发早阶语义特征崩溃。为此，LaProx 设计了简洁有效的层级归一化机制，将每层内展平的 token 驱逐得分转化为层内相对重要性概率分布： $$s_{l, j} = \frac{p_{l, j}}{\sum_k p_{l, k}}$$ 该操作严格保全了每一层内部不同头、不同 token 之间的相对排序与贡献权重，同时抹平了层间的绝对数值量纲差距，使得模型浅层与深层之间的重要性得分具备完全可比的全局物理含义。

3. 全局自适应动态预算分配：突破单头局部决策的模型级 Top-K 联合选择
在获得跨层、跨头归一化评分 $s_{l,j}$ 后，LaProx 在全模型范围内执行统一的 Top-K 驱逐。与需要使用特定校准集在离线阶段进行网格搜索或微调学习头预算权重的基线方案（如 HeadKV、PyramidKV、CAKE）不同，LaProx 完全免除超参数调节，也不需要任何离线校准数据。在推理过程中，模型面对具体输入提示词时，信息自适应地根据 $s_{l,j}$ 的全局分布自发沉淀到最具信息量的层和头上。对于检索密集型任务，负责长程记忆检索的“检索头”和核心交互层会自动保留多达数百个 token；而对于上下文冗余的钝化头或纯局部归纳头，缓存则被大幅压缩。这种纯数据驱动的全局动态选择，使得模型能够在极其严苛的极低显存预算下（如平均每头仅保留 128 或 256 token），仍能逼近 FullKV 的全量注意力输出精度。

损失函数 / 训练策略¶

LaProx 是一种纯推理期（Inference-time）无需训练（Training-free & Tuning-free）的算法。它不改变任何预训练权重，不引入辅助分类器或门控网络，仅在 Prefill 阶段针对历史上下文执行一次性的矩阵范数计算、层级归一化与 Top-K 索引提取。在实现上，为了与 Grouped Query Attention（GQA）架构原生兼容，各 Query 组内部采用组内 Query 注意力权重的均值进行对齐；所有矩阵操作均基于 FlashAttention-2 高效算子加速，引入的 Prefill 延迟开销微乎其微。

实验关键数据¶

主实验¶

论文在 LongBench 长文本评测基准（涵盖 16 个数据集、6 大核心任务领域：单文档 QA、多文档 QA、长文本摘要、少样本学习、合成检索、代码补全）以及 Needle-In-A-Haystack（NIAH）大海捞针基准上进行了全方位评测，测试基座涵盖 Meta-Llama-3.1-8B-Instruct、Mistral-7B-Instruct-v0.3 以及 Qwen3-8B。在极其严苛的固定缓存容量 $B_{\text{total}} = 128L$（平均每头仅保留 128 token）设置下，LongBench 各模型各任务的综合性能对比如下表所示。

模型与配置	单文档QA (Avg)	多文档QA (Avg)	摘要 (Avg)	少样本 (Avg)	合成 (Avg)	代码 (Avg)	LongBench 全局均分
Llama-3.1-8B-Instruct
FullKV (无压缩上限)	43.56	47.93	29.19	69.34	54.19	61.92	49.51
StreamingLLM	25.01	36.02	19.37	54.77	53.83	54.30	38.86
SnapKV	32.53	41.04	20.54	58.69	53.21	56.14	42.31
AdaKV	33.31	41.43	21.26	59.77	53.69	57.66	43.12
CAKE	33.25	41.44	21.27	59.20	53.69	55.77	42.77
CriticalKV	33.44	41.42	21.15	59.35	53.69	56.32	42.88
LaProx (本文)	36.71	43.19	22.26	63.34	54.20	59.04	45.19
Mistral-7B-Instruct-v0.3
FullKV (无压缩上限)	41.18	42.32	29.43	70.66	51.75	61.97	48.01
StreamingLLM	23.51	29.40	18.12	57.20	42.75	53.60	36.08
SnapKV	31.93	34.38	20.45	58.72	49.75	55.28	40.41
AdaKV	32.23	35.38	21.07	59.23	49.50	56.22	40.95
CAKE	33.62	35.78	21.55	59.43	51.00	55.73	41.54
CriticalKV	32.60	35.43	20.91	59.95	50.00	55.25	41.07
LaProx (本文)	37.43	36.98	22.22	65.55	50.50	58.26	44.00
Qwen3-8B
FullKV (无压缩上限)	43.86	47.92	27.42	58.94	50.50	51.62	46.17
StreamingLLM	20.67	27.43	16.76	46.29	38.00	45.30	31.25
SnapKV	30.02	37.75	17.40	54.38	50.00	47.36	38.33
AdaKV	32.73	38.22	17.27	53.05	49.50	46.35	38.47
CAKE	32.48	42.01	18.47	50.63	52.00	45.75	39.14
CriticalKV	33.36	36.84	18.07	55.44	49.50	48.04	39.15
LaProx (本文)	35.21	42.93	19.07	57.78	50.00	50.04	41.57

消融实验¶

为了严格验证指标构建（$V$ 与 $W_O$ 的逐级必要性）以及预算分配策略（展平与归一化的交互影响），论文在 Mistral-7B-Instruct-v0.3 上开展了细粒度消融评测（预算固定为 $B_{\text{total}} = 128L$）。

驱逐指标组成	值向量 $V$	输出投影 $W_O$	LongBench 平均得分	相比基线变化
传统注意力均值 (SnapKV baseline)	✗	✗	40.53	基准
仅加入值状态范数 (Attention $\times \\|V\\|$)	✓	✗	41.69	+1.16
完整 LaProx 指标 (Attention $\times \\|V W_O\\|$）	✓	✓	41.84	+1.31

预算分配机制	跨头展平 (Head-Flatten)	跨层展平 (Layer-Flatten)	跨层归一化 (Normalize)	LongBench 平均得分	机制说明
均匀逐头分配 (无自适应)	✗	✗	✗	41.84	每头固定配额，仅采用 $V W_O$ 指标
仅跨头自适应分配	✓	✗	✗	42.96	各层内部自适应，层间均分
朴素跨层跨头直接展平	✓	✓	✗	41.14	深浅层数值量纲不一，性能反向暴跌
完整全局分配 (LaProx)	✓	✓	✓	44.00	归一化后跨层跨头全局 Top-K，性能最优

关键发现¶

输出投影 $W_O$ 是跨头可比性的关键数学桥梁：消融表明，仅考虑 $V$ 虽有提升，但无法真正刻画对残差流的最终注入量；必须结合 $W_O$ 才能使跨头特征映射到同一度量空间。在 Table 6 的对比中，若直接使用未结合 $V W_O$ 的纯注意力分数做全局分配（$A_G$ 得分 43.05），其效果与均分基线（$A_L$ 得分 42.31）相差无几，甚至显著落后于使用 LaProx 指标但均分配额的变体（$L_L$ 得分 43.26），直接证实了 $V W_O$ 是实现有效全局预算调配的前提条件。
跨层归一化解决了深层数值压制浅层的问题：如果不加归一化直接对所有层的分数进行全局排序，分数会由于深层模长天然偏大而严重倾斜至深层，浅层缓存被过度裁汰，导致平均分跌至 41.14（甚至低于未做自适应分配的 41.84）。通过简单的层内求和归一化，即刻使得分攀升至 44.00。
极低预算与复杂多目标任务下优势显著：在 Needle-In-A-Haystack 的多针检索（4 针 4 目标 4N-4T）及极低显存（256 token）场景下，LaProx 相比 CriticalKV、AdaKV 和 CAKE 实现了 1.5 倍至 3 倍的精度恢复率；在 NVIDIA H100 上的 128K 长文本测试中，LaProx 将解码吞吐加速了 2.3 倍，显存从 63.3GB 直降至 47.5GB。

亮点与洞察¶

将驱逐问题提升至矩阵乘法逼近理论高度：跳脱出长期以来把 KV 驱逐当成“滑动窗口+启发式注意力累加”的经验主义陷阱，借助蒙特卡洛矩阵乘法近似理论，为“什么样的 token 具有高输出贡献”提供了坚实严格的数学根基。
实现了全模型视角的无参全局自适应分配：过去无论是分层金字塔设计（PyramidKV）还是注意力方差调节（CAKE），都需要人工引入超参数或依靠标定集预先拟合权重。LaProx 利用 $V W_O$ 赋予的加法可比性与层级归一化，首次实现了完全无需调参、无需标定集、完全由当前提示词动态激活模式决定的全局分配机制。
跨架构泛化能力极强：由于注意力的数学输出分解特性在标准稠密模型与混合专家模型（MoE）中完全一致，LaProx 无需任何架构修改即可即插即用部署于 Qwen3-30B-A3B 等 MoE 架构上，性能退化控制在 3% 以内。

局限与展望¶

理论近似形式依然较为简明：当前为了兼顾工程实现开销，指标计算主要依托一阶列向量与行向量的 $L_2$ 范数相乘，未来可探索更高阶的随机矩阵近似方法或考虑注意力 Softmax 结构下的非线性扰动。
归一化方案可进一步细化：目前采用简单的层内求和归一化处理层间尺度不均衡，未来可针对不同深度层的残差缩放因子或 RMSNorm 权重进行理论对齐，以进一步挖掘跨层分配潜力。
动态更新与增量驱逐开销：当前评测主要在 Prefill 结束后执行一次性集中压缩，自回归生成阶段随着步数延长仍需维护周期性微量淘汰策略，如何在多轮持续交互中维持近零开销的矩阵范数增量追踪值得深入探索。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从矩阵乘法近似与跨头线性相加视角系统重构了 KV Cache 驱逐范式，打破了长久以来的局部启发式定式。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 大主流开源模型及 30B MoE 架构，测试跨越 LongBench 16 个数据集、NIAH、RULER 及 InfiniteBench，并配有完整的矩阵输出余弦相似度与硬件效率实测。
写作质量: ⭐⭐⭐⭐⭐ 数学证明严密自洽，从动机分析、公式推导到算法落地层层递进，消融逻辑高度清晰。
价值: ⭐⭐⭐⭐⭐ 即插即用、无需训练与标定，为超长上下文 LLM 推理系统的显存压减与加速提供了兼具理论深度与实用价值的全新基础支撑。

预算分配机制	跨头展平 (Head-Flatten)	跨层展平 (Layer-Flatten)	跨层归一化 (Normalize)	LongBench 平均得分	机制说明
均匀逐头分配 (无自适应)	✗	✗	✗	41.84	每头固定配额，仅采用 \(V W_O\) 指标
仅跨头自适应分配	✓	✗	✗	42.96	各层内部自适应，层间均分
朴素跨层跨头直接展平	✓	✓	✗	41.14	深浅层数值量纲不一，性能反向暴跌
完整全局分配 (LaProx)	✓	✓	✓	44.00	归一化后跨层跨头全局 Top-K，性能最优