面向内存受限推理的随机稀疏注意力

自回式解码在长序列情境下变得带宽受限，因为生成每个 token 需要从 KV 缓存中读取所有 $n_k$ 个 key 和 value 向量。我们提出随机加法无乘注意力（SANTA），一种稀疏化 value 缓存访问的方法，通过从 post-softmax 分布中抽取 $S \ll n_k$ 个索引，对仅这些 value 行进行聚合。这一方法为 post-softmax value 聚合提供无偏估计，同时将 value 阶段的乘-累加运算替换为 gather-加法。我们引入分层抽样，以设计方差降低、面向 GPU 的变体，展示在 NVIDIA RTX 6000 Ada 上，SANTA 对 FlashInfer 和 FlashDecoding 实现的 decode 步骤注意力 kernel 加速 1.5 倍，同时在 32k token 上与基线保持一致的准确率。最后，我们提出 Bernoulli $qK^\mathsf{T}$ 抽样作为稀疏化得分阶段的补充技术，减少 key 特征访问通过随机三值查询。两种方法与上游技术如三值量化、低秩投影和 KV 缓存压缩等正交。它们指向稀疏、无乘数且高效的推理。我们将该 kernel 在 https://github.com/OPUSLab/SANTA.git 上开源。