用于内存限制推理的随机稀疏注意力

自回归解码在长上下文中会受到带宽限制，因为生成每个词元需要从 KV 缓存中读取所有 $n_k$ 键和值向量。我们提出了随机加性无多重注意力（SANTA），这是一种通过从 post-softmax 分布中采样 $S \ll n_k$ 索引来稀疏值缓存访问并仅聚合那些值行的方法。这产生了 post-softmax 值聚合的无偏估计器，同时用收集和相加替换值阶段乘法累加。我们引入分层和系统抽样来设计方差减少、GPU 友好的变体。在 32k-token 上下文中的 Llama-3.1-8B-Instruct 上进行评估，S$^2$ANTA 匹配基线精度，同时在 NVIDIA RTX 6000 Ada 上比 FlashInfer 和 FlashDecoding 实现高达 $1.5\times$ 解码步骤注意内核加速。在批量长上下文生成中，这些内核增益可转化为高达 $1.25\times$ 的端到端解码延迟加速。最后，我们提出伯努利 $qK^\mathsf{T}$ 采样作为稀疏评分阶段的补充技术，通过随机三元查询减少关键特征访问。这两种方法都是上游量化、低秩投影、KV 缓存压缩和 KV 缓存选择方法的补充。它们共同指向稀疏、无乘数且节能的推理。我们在以下位置开源内核：https://github.com/OPUSLab/SANTA.git