arXiv：cs.LG· Gleb Molodtsov, Ekaterina Alimaskina, Evgeny Uskov, Artur Zagitov, Aleksandr Beznosikov·· 4 小时前AI 评分34
MaskAhead：块扩散语言模型中的掩码引导 KV Cache 驱逐
Mask-Guided KV Cache Eviction in Block Diffusion Language Models
AI 导读
MaskAhead 是一种免训练方法，用单一掩码查询排序机制同时完成块扩散语言模型的 KV 选择与驱逐。在长提示词问答任务上，它平均减少 9.5× KV 内存，F1 仅下降 1.2 分；量化版 Q-MaskAhead 将内存压缩提升至 20.1×，F1 下降 2.3 分。batch-32 系统测试中，端到端与解码阶段分别提速 1.23× 和 1.68×。
正文
Abstract:Block diffusion language models keep a large key-value (KV) cache throughout generation and attend to it at every denoising step, limiting both memory capacity and generation speed. Reducing these costs requires deciding which past tokens to use for denoising the current block (selection) and which to keep in memory for future blocks (eviction). We propose MaskAhead, a training-free method that solves both tasks with a single mask-query-based ranking mechanism. Current-block masks guide selection, while probes of upcoming masked blocks guide eviction. Both rank KV entries by their estimated contribution to the attention output. Our quantized variant, Q-MaskAhead, computes selection and attention directly from low-bit KV, largely preserving the selected entries. Experiments on Fast-dLLM-v2, DreamReasoner, and LLaDA2.0-mini cover long-generation reasoning, long-prompt question answering, and needle-in-a-haystack retrieval. On long-prompt QA, MaskAhead reduces KV memory by $9.5\times$ on average with a 1.2-point mean F1 loss relative to dense inference. Q-MaskAhead increases the reduction to $20.1\times$ with a 2.3-point mean F1 loss. In a batch-32 systems profile, MaskAhead achieves $1.23\times$ end-to-end and $1.68\times$ decode-stage speedups over dense inference.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL) |
| Cite as: | arXiv:2610.06996 [cs.LG] |
| (or arXiv:2610.06996v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.06996 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Gleb Molodtsov Mr [view email]
[v1]
Sun, 4 Oct 2026 07:52:31 UTC (715 KB)
来源：arXiv：cs.LG · arxiv.org