MAGE：块扩散LLM中首个All-[MASK]块已知如何查找位置

块扩散语言模型（Block diffusion LLMs）作为语言生成的新兴范式，其使用键值（KV）缓存导致在长上下文场景下内存访问成为主要瓶颈。尽管动态稀疏注意力已被广泛探索，但现有针对自回归LLM的方法依赖于近似重要性估计，在适应块扩散模型时表现不佳。本文识别出块扩散模型的一个关键机会：在首次All-[MASK]去噪步骤中的注意力能够可靠地预测重要的KV条目和预算需求，使MAGE能够在每个块上执行一次精确的注意力传递，并将其用于无训练的稀疏去噪。在LongBench和Needle-in-a-Haystack等长上下文基准测试中，MAGE在仅使用少量KV预算的情况下实现了接近无损的精度，并带来了高达3-4倍的端到端加速，持续优于基于自回归的稀疏注意力基线。一种轻量级的微调策略进一步强化了[MASK]引导的模式，成本极低，仅需在单个NVIDIA H100 GPU上进行数小时的训练即可适用于1.5B和7B模型。