针对扩散语言模型设计了基于掩码令牌的缓存淘汰机制，显著提升长文本推理效率。
Mask Tokens as Prophet: Fine-Grained Cache Eviction for Efficient dLLM Inference
- 利用注意力权重识别并淘汰低重要性提示词，实现细粒度缓存管理。
- 仅用256对缓存（不足5%令牌数）保持94%完整缓存性能，32k长度下提速31倍。
- 无需训练，适配资源受限场景下的长上下文推理任务。
扩散大语言模型（dLLMs）通过并行解码提供了优于主流自回归模型（ARMs）的潜力，但需付出巨大计算与内存开销。特别是dLLMs中双向注意力的缓存机制需要大量内存，限制了其在资源受限环境下处理长上下文的能力。现有缓存淘汰策略专为ARMs设计，忽略了dLLMs的独特特性，导致性能不佳。为此，我们提出MaskKV——一种无需训练、专为dLLMs设计的缓存淘汰框架，聚焦于掩码令牌的作用。MaskKV包含两项关键创新：(1) 基于掩码查询的评分机制，利用注意力权重识别并淘汰每个头中较不重要的提示词；(2) 自适应缓存预算分配策略，减少中间层分配，集中资源于偏好提示的头。在LLaDA上启用MaskKV后，将KV缓存压缩至仅256对（低于总令牌数5%），在LongBench上仍保持94%的全缓存性能，并在32k提示长度下实现最高31倍加速。代码已公开于https://github.com/jianuo-huang/MaskKV。
原文摘要 · Abstract (English)
Diffusion large language models (dLLMs) present a promising alternative to dominant autoregressive models (ARMs) by the ability of parallel decoding at the expense of substantial computation and memory costs. Specifically, the cache mechanism for bidirectional attention in dLLMs demands large memory footprint, restricting their ability to handle long contexts under resource-limited settings. Existing cache eviction strategies are designed for ARMs and ignore the unique characteristics of dLLMs, thus leading to unsatisfactory performance. To address these challenges, we introduce MaskKV, a training-free cache eviction framework tailored to dLLMs, focusing on the effect of mask tokens in dLLMs. MaskKV is built on two key innovations: (1) a mask-query guided scoring mechanism that leverages attention weights to identify and evict less critical prompt tokens for each head; (2) an adaptive cache budgeting strategy that improves efficiency by reducing allocation in intermediate layers and concentrating resources on prompt-preferring heads. On LLaDA with MaskKV, compressing the KV cache to only 256 pairs (less than 5% of tokens) retains 94% of the full-cache performance on LongBench and achieves up to 31x acceleration at 32k prompt length. The code is publicly available at: https://github.com/jianuo-huang/MaskKV
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。