提出EntmaxKV,让长文本生成更快更准。
EntmaxKV: Support-Aware Decoding for Entmax Attention

- 用支持感知的候选选择提前跳过无关缓存页。
- 在100万上下文长度下速度提升5.43倍,误差更低。
- 适合长文本生成、大模型推理等需要高效注意力的场景。
长上下文生成受KV缓存内存流量限制,因每个生成词都需访问线性增长的缓存。现有稀疏解码方法虽可降低开销,但针对softmax注意力设计,其密集尾部导致截断会丢失非零概率质量。而α-entmax能产生精确零值,使稀疏解码从近似变为支持恢复:只要候选集包含entmax支持集,解码即保持精确。尽管最近已有高效的entmax训练内核,但自回归解码瓶颈仍未解决——仍需在知道稀疏性前加载完整缓存。本文提出EntmaxKV,一种原生支持entmax的稀疏解码框架,可在加载前利用查询感知页评分与支持感知候选选择实现提前稀疏化。通过分析被丢弃概率质量δ,证明输出误差由δ控制,当支持集被完全恢复时误差消失。进一步引入高斯感知的entmax选择器,基于轻量页统计估计阈值,自适应调整候选预算。实验表明,EntmaxKV在相同缓存预算下丢弃更少概率质量,保留更多支持词,输出误差低于基于softmax的稀疏解码。在长上下文与语言建模基准上,其性能接近全缓存entmax,仅使用极小部分缓存,于100万上下文长度下相较全注意力基线达3.36倍(softmax)和5.43倍(entmax)加速。
原文摘要 · Abstract (English)
Long-context decoding is increasingly limited by KV-cache memory traffic since each generated token attends over a cache whose size grows linearly with context length. Existing sparse decoding methods reduce this cost by selecting subsets of tokens or pages, but are designed for softmax attention, whose dense tails make any truncation discard nonzero probability mass. In contrast, $α$-entmax produces exact zeros, turning sparse decoding from dense-tail approximation into support recovery: if the selected candidates contain the entmax support, sparse decoding remains exact. While recent entmax kernels enable efficient training, they do not address the autoregressive decoding bottleneck, where dense inference still streams the full KV cache before sparsity is known. In this work, we introduce EntmaxKV, an entmax-native sparse decoding framework that exploits sparsity before KV pages are loaded. EntmaxKV combines query-aware page scoring, support-aware candidate selection, and sparse entmax attention. We analyze truncation error through the dropped probability mass $δ$, showing that output error is controlled by $δ$ and vanishes when the entmax support is recovered. We further introduce a Gaussian-aware entmax selector that estimates the entmax threshold from lightweight page statistics, adapting the selected budget to the score distribution. Empirically, EntmaxKV drops less probability mass, retains more support tokens, and achieves lower output error than softmax-based sparse decoding at matched KV budgets. On long-context and language modeling benchmarks, it closely matches full-cache entmax while using a small fraction of the KV cache, achieving up to $3.36\times$ (softmax) and $5.43\times$ (entmax) speedup over full attention baselines at 1M context length. Code available at: https://github.com/deep-spin/entmaxkv.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。