EntmaxKV：面向 Entmax 注意力的支持感知解码

长上下文解码正受到 KV 缓存内存流量的限制，因为每个生成的 token 都要注意其 attention 过的 cache，其大小随上下文长度线性增长。现有的稀疏解码方法通过选择 token 或页面的子集来降低此成本，但这些方法是为 softmax 注意力设计的，其稠密尾部使得任何截断都会丢弃非零概率质量。相反，$\alpha$-entmax 可产生精确的零，使稀疏解码从稠密尾部近似转为支持恢复：如果所选候选者包含 entmax 支持，则稀疏解码保持精确。虽然最近的 entmax 内核支持高效训练，但它们不解决自回归解码瓶颈，在此之前稀疏性尚未知晓，稠密推理仍会流式传输完整的 KV 缓存。在本工作中，我们引入 EntmaxKV，一个面向 entmax 的稀疏解码框架，利用稀疏性在 KV 页面加载前进行 exploitation。EntmaxKV 结合查询感知的页面评分、支持感知的候选者选择以及稀疏 entmax 注意力。我们通过 dropping 的概率质量 $\delta$ 分析截断误差，表明输出误差由 $\delta$ 控制，当 entmax 支持被恢复时即为零。我们进一步引入一个受高斯启发的 entmax 选择器，从轻量级页面统计量估计 entmax 阈值，使所选预算适应得分分布。实验上，EntmaxKV 丢弃的概率质量更少，保留的支持 token 更多，在匹配 KV 预算的情况下实现更低的输出误差。在长上下文和语言建模基准上，它在使用小小比例 KV 缓存的情况下几乎匹配全缓存 entmax，实现了最高可达 $3.36\times$(softmax) 和 $5.43\times$(entmax) 的加速。代码可在 https://github.com/deep-spin/entmaxkv 获取。