arXiv cs.AI· Gleb Molodtsov, Ekaterina Alimaskina, Evgeny Uskov, Artur Zagitov, Aleksandr Beznosikov·· 10 天前AI 评分24
MaskAhead：块扩散语言模型中的掩码引导 KV Cache 淘汰方法
Mask-Guided KV Cache Eviction in Block Diffusion Language Models
AI 导读
针对块扩散语言模型生成时 KV cache 占用大、速度慢的问题，研究者提出免训练方法 MaskAhead，用单一掩码查询排序机制同时完成当前块的 KV 选择与未来块的淘汰。
正文 · AI 翻译
摘要：块扩散语言模型在生成过程中会保留一个大型键值（KV）缓存，并在每个去噪步骤中对其进行注意力计算，从而限制了内存容量和生成速度。要降低这些成本，就需要决定哪些过去的 token 用于当前块的去噪（选择），以及哪些保留在内存中供未来块使用（逐出）。我们提出 MaskAhead，一种无需训练的方法，通过单一的基于掩码查询的排序机制同时解决这两个任务。当前块的掩码指导选择，而对即将到来的掩码块的探测指导逐出。两者都根据 KV 条目对注意力输出的估计贡献进行排序。我们的量化变体 Q-MaskAhead 直接从低位 KV 计算选择和注意力，在很大程度上保留了被选中的条目。在 Fast-dLLM-v2、DreamReasoner 和 LLaDA2.0-mini 上的实验涵盖了长生成推理、长提示问答和大海捞针检索。在长提示问答上，与稠密推理相比，MaskAhead 平均将 KV 内存减少 $9.5\times$，平均 F1 损失为 1.2 个点。Q-MaskAhead 将减少量提高到 $20.1\times$，平均 F1 损失为 2.3 个点。在 batch-32 系统配置中，MaskAhead 相比稠密推理实现了 $1.23\times$ 的端到端加速和 $1.68\times$ 的解码阶段加速。
| 学科： | 机器学习（cs.LG）；人工智能（cs.AI）；计算与语言（cs.CL） |
| 引用为： | arXiv:2610.06996 [cs.LG] |
| （或使用 arXiv:2610.06996v1 [cs.LG] 查看此版本） | |
| https://doi.org/10.48550/arXiv.2610.06996 由 DataCite 颁发的 arXiv DOI（待注册） |
提交历史
来自：Gleb Molodtsov 先生 [查看电子邮件]
[v1]
2026年10月4日星期日 07:52:31 UTC（715 KB）
来源：arXiv cs.AI · arxiv.org