💡 一句话总结:LaProx 把”哪些 KV token 该被淘汰”这个看似启发式的问题,重写成了一个有闭式解的矩阵近似问题——这是过去三年 KV Cache 领域最干净的一次范式跃迁。
论文背景
Reformulating KV Cache Eviction Problem for Long-Context LLM Inference(arxiv 2605.07234)于 2026 年 5 月 8 日上线,作者来自清华、香港城市大学和 Cohere AI 的联合团队。论文聚焦的问题非常具体:当上下文长度超过 64K,KV Cache 显存占用超过模型权重本身,必须做淘汰;现有方法是否做对了?
KV Cache 淘汰问题的演进史
要理解 LaProx 为什么是范式转变,先看过去三年这个领域的关键路径:
| 年份 | 方法 | 核心假设 | 局限 |
|---|---|---|---|
| 2023 | H2O | 累计注意力分数高的 token 重要 | 长尾 token 被错杀 |
| 2024 上 | SnapKV | 仅看最近 window 的注意力 | 多跳推理失效 |
| 2024 下 | PyramidKV | 不同层用不同压缩率 | 层间策略割裂 |
| 2025 | ChunkKV | 按语义块淘汰 | 块边界难定义 |
| 2026 | LaProx | 输出端矩阵近似最优 | 每层引入分解开销 |
前四类方法的共同点:都在输入端回答”哪些 token 重要”。LaProx 第一次把视角换到输出端——“淘汰之后,模型的输出还能不能逼近原始输出”。
数学重铸:从打分到优化
原问题表述(H2O / SnapKV 范式)
给定 token 序列 X,注意力权重矩阵 A,目标是选一个 token 子集 S 使得:
S* = argmax_S Σ_{i ∈ S} score(i, A)
其中 score 是启发式函数(累计注意力、最近 window 注意力等)。
这种表述有个根本问题:score 函数和最终输出之间没有直接联系,只是经验上相关。
LaProx 的重铸
LaProx 把问题改写为:找一个 KV 子集,使得淘汰后的注意力输出矩阵 Õ 最接近原始输出矩阵 O:
S* = argmin_S || A_S V_S − A V ||_F
其中 ||·||_F 是 Frobenius 范数,A_S V_S 是只保留 S 中 token 后的注意力输出。
这个表述的关键收益:直接优化我们关心的量(输出准确性),而不是绕一圈走启发式。
闭式近似解
完整问题是 NP-hard(子集选择),LaProx 给出了一个基于**列子集选择问题(CSSP)**的近似算法:
1. 对 A 做 leverage score 采样,初步定位关键列
2. 在采样列上做 QR 分解,得到正交基
3. 用正交基重构 V 的近似,估计每个 token 的边际贡献
4. 贪心选择贡献最大的 token 直至达到预算
整个流程单层 0.3ms 开销(A100 70B 模型),可以摊到 prefill 阶段并行计算,不影响推理延迟。
实验结果亮点
LongBench 综合表现
128K 上下文 Llama-3-70B,KV Cache 压缩率 50%:
| 任务 | 原模型 | SnapKV | PyramidKV | LaProx |
|---|---|---|---|---|
| MultiFieldQA | 51.3 | 47.2 | 48.1 | 50.8 |
| HotpotQA (多跳) | 56.7 | 49.4 | 51.0 | 55.3 |
| 2WikiMQA | 46.1 | 41.8 | 42.7 | 45.2 |
| Passage Count | 9.2 | 7.8 | 8.1 | 9.0 |
| 平均 | 40.8 | 36.6 | 37.5 | 40.1 |
LaProx 在所有任务上都最接近原模型,多跳问答(HotpotQA、2WikiMQA)领先尤其明显——这正是传统淘汰方法的痛点。
极限压缩下的鲁棒性
把压缩率推到 75%(只保留 1/4 KV):
Llama-3-70B, 128K 上下文, 75% 压缩
SnapKV: LongBench 平均 28.4 (↓ 30%)
PyramidKV: LongBench 平均 31.2 (↓ 23%)
LaProx: LongBench 平均 37.6 (↓ 8%)
LaProx 在极限压缩下退化幅度只有 SnapKV 的 1/4,这对显存受限场景(消费级 GPU 跑长上下文)非常关键。
推理速度对比
A100 80GB,128K 上下文,batch=4:
| 方法 | KV 预算 | 首 token 延迟 | 吞吐 (tok/s) |
|---|---|---|---|
| Full KV | 100% | 1.8s | 142 |
| SnapKV | 50% | 1.2s | 218 |
| LaProx | 50% | 1.3s | 211 |
| LaProx | 25% | 0.9s | 287 |
LaProx 50% 时只比 SnapKV 慢 3%,但准确率高得多;推到 25% 时吞吐反而高出 SnapKV 50%,因为 KV 总量更小,attention 计算更快。
与其他范式的关系
LaProx + 滑动窗口
论文 6.2 节做了组合实验:把 LaProx 用在滑动窗口(StreamingLLM)的边界外,效果比纯 LaProx 再涨 1.5 分。背后逻辑是滑动窗口保证局部连贯性,LaProx 负责挑选远距离关键 token。
LaProx + 量化
KV Cache 8bit 量化 + LaProx 50% 淘汰,在 H100 上把 Llama-3-70B 128K 推理的 KV 显存从 40GB 压到 4GB——单卡推理 1M 上下文不再是天方夜谭。
LaProx + 注意力变体
Linear Attention(Mamba2、RWKV-7)本身就不需要 KV Cache,和 LaProx 没关系。但混合架构(DeepSeek V4 的 MHC、Jamba 的混合层)需要在 attention 层做淘汰,LaProx 直接可用。
工程落地路径
如果你想在生产环境用上 LaProx,推荐三步走:
Step 1:评估场景适配
- 上下文长度: 32K-512K → 推荐
- 上下文长度: 小于等于 8K → 不推荐(收益太小)
- 上下文长度: 大于等于 1M → 配合分层存储
- 任务类型: 多跳问答、长文档摘要 → 强烈推荐
- 任务类型: 单轮对话、简单 QA → 可选
Step 2:用参考实现跑通基线
from laprox import LaProxCache
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3-70B-Instruct",
cache_config=LaProxCache(
budget_ratio=0.5,
sample_method="leverage_score",
chunk_size=128
)
)
# 后续推理流程不变
output = model.generate(input_ids, max_new_tokens=512)
Step 3:自定义内核(可选)
参考实现用 Python + PyTorch 跑,单层 leverage score 采样大概 1ms。生产环境建议用 Triton 重写关键内核,论文 7.4 节给出了 Triton 模板,开销能压到 0.3ms。
局限与未指出的问题
LaProx 不是银弹,论文里没明说但需要警惕:
- 采样随机性:leverage score 采样有随机种子敏感性,连续两次同样输入可能淘汰不同 token,下游测试结果可能不稳定
- 批次大小敏感:batch 越大,单卡 leverage score 计算开销越高,目前没看到 batch 大于 16 的实测
- 训练-推理一致性:模型训练时见过完整 KV,推理时只看部分 KV,理论上有 distribution shift,长任务上可能积累误差
社区已经开始针对这些问题打补丁——5 月 12 日上线的 LaProx-Det(确定性版本)就是回应第一个问题。
总结:方法论意义大于工程价值
LaProx 的工程数字很漂亮,但更重要的是它的方法论暗示:KV Cache 优化这条赛道,启发式打分的时代结束了。任何后续工作如果还停留在”设计更好的 score 函数”层面,都得回答 LaProx 的灵魂拷问:你的 score 和最终输出有什么数学关系?
接下来 6 个月,预期会看到一波”输出感知”为关键词的 KV 优化论文。建议关注的方向:
- 输出感知 + 跨层联合优化
- 输出感知 + 训练时引入 distillation
- 输出感知 + 推测解码联合
KV Cache 这条赛道又被重新打开了。