通过随机采样减少长文本推理的内存访问，提升速度且不损失精度。
Stochastic Sparse Attention for Memory-Bound Inference
- 从软最大分布中采样少量值向量，仅读取部分缓存数据
- 32k上下文下解码速度提升1.5倍，端到端延迟快1.25倍
- 适用于长文本生成，适合追求高效推理的开发者
自回归解码在长上下文场景下受带宽限制，因每生成一个词元需读取所有 $n_k$ 个键值向量。我们提出随机加性无乘法注意力（SANTA），通过从后软最大分布中采样 $S \\< n_k$ 个索引，仅聚合对应值行，实现对后软最大值聚合的无偏估计，并将值阶段的乘加操作替换为取值与求和。引入分层与系统采样设计出方差更低、适配GPU的变体。在 Llama-3.1-8B-Instruct 32k-token 上评估，S$^2$ANTA 在保持基线准确率的同时，相较 FlashInfer 与 FlashDecoding 实现最高 1.5 倍的解码步注意力核加速；批量长上下文生成中，端到端解码延迟最高降低 1.25 倍。进一步提出伯努利 $qK^\mathsf{T}$ 采样作为互补技术，通过随机二值查询稀疏化得分阶段，减少键特征访问。两项方法均与上游量化、低秩投影、KV缓存压缩及选择方法兼容，共同指向稀疏、无乘法器、节能的推理方向。代码已开源：https://github.com/OPUSLab/SANTA.git
原文摘要 · Abstract (English)
Autoregressive decoding becomes bandwidth-limited at long contexts, as generating each token requires reading all $n_k$ key and value vectors from KV cache. We present Stochastic Additive No-mulT Attention (SANTA), a method that sparsifies value-cache access by sampling $S \ll n_k$ indices from the post-softmax distribution and aggregates only those value rows. This yields an unbiased estimator of the post-softmax value aggregation while replacing value-stage multiply-accumulates with gather-and-add. We introduce stratified and systematic sampling to design variance-reduced, GPU-friendly variants. Evaluated on Llama-3.1-8B-Instruct at 32k-token contexts, S$^2$ANTA matches baseline accuracy while achieving up to $1.5\times$ decode-step attention-kernel speedup over FlashInfer and FlashDecoding on an NVIDIA RTX 6000 Ada. In batched long-context generation, these kernel gains translate to up to $1.25\times$ end-to-end decode-latency speedup. Finally, we propose Bernoulli $qK^\mathsf{T}$ sampling as a complementary technique to sparsify the score stage, reducing key-feature access through stochastic ternary queries. Both methods are complementary to upstream quantization, low-rank projection, KV-cache compression, and KV-cache selection methods. Together, they point toward sparse, multiplier-free, and energy-efficient inference. We open-source our kernels at: https://github.com/OPUSLab/SANTA.git
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。