EntmaxKV：Entmax Attention 的支持感知解码

长上下文解码越来越受到 KV 缓存内存流量的限制，因为每个生成的词元都参与大小随上下文长度线性增长的缓存。现有的稀疏解码方法通过选择标记或页面的子集来降低这种成本，但它们是为 softmax 注意力而设计的，其密集的尾部使得任何截断都会丢弃非零概率质量。相反，$α$-entmax 产生精确的零，将稀疏解码从密集尾近似转变为支持恢复：如果所选候选者包含 entmax 支持，则稀疏解码保持精确。虽然最近的 entmax 内核能够实现高效的训练，但它们并没有解决自回归解码瓶颈，即在稀疏性已知之前，密集推理仍然会流式传输完整的 KV 缓存。在这项工作中，我们介绍了 EntmaxKV，一种 entmax 原生稀疏解码框架，它在加载 KV 页面之前利用稀疏性。 EntmaxKV 结合了查询感知页面评分、支持感知候选选择和稀疏 entmax 注意力。我们通过下降的概率质量 $δ$ 来分析截断误差，表明输出误差受 $δ$ 控制，并在 entmax 支持恢复时消失。我们进一步引入了一个高斯感知的 entmax 选择器，它根据轻量级页面统计数据估计 entmax 阈值，使所选预算适应分数分布。根据经验，在匹配的 KV 预算下，EntmaxKV 比基于 softmax 的稀疏解码下降更少的概率质量，保留更多的支持标记，并实现更低的输出错误。在长上下文和语言建模基准测试中，它与全缓存 entmax 非常接近，同时使用了一小部分 KV 缓存，在全注意力基线上实现了高达 $2.88\times$ (softmax) 和 $5.5\times$ (entmax) 的加速，并且在 1M 上下文长度下比 softmax 和 entmax top-k 实现了超过 $2\times$ 的加速。