掩码引导的块扩散语言模型 KV 缓存淘汰

arXiv:2610.06996v1 公告类型：新 摘要：块扩散语言模型在整个生成过程中保留大型键值（KV）缓存，并在每个去噪步骤中对其进行注意力计算，这限制了内存容量和生成速度。降低这些开销需要决定哪些过去的 token 用于当前块的去噪（选择），以及哪些 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$。