SOCKET: 用于稀疏注意力的软碰撞核估计器
SOCKET: SOft Collision Kernel EsTimator for Sparse Attention
AI总结：
SOCKET通过引入软碰撞核估计器改进稀疏注意力，利用概率性相似性聚合替代传统二元碰撞信号，提升长上下文推理效率和内存利用率。
摘要：
在长上下文推理中利用稀疏性是扩展大语言模型的关键，因为注意力主导了自回归解码的成本。稀疏注意力通过限制计算到部分token来降低此成本，但其效果依赖于推理时高效的评分和选择。我们重新审视局部敏感哈希（LSH）并引入SOCKET，一种软碰撞核估计器，用概率性、相似性感知的聚合替代硬桶匹配。传统LSH产生二元碰撞信号，限制排名质量并需要大量内存才能表现良好。相比之下，软LSH在哈希表中积累分级碰撞证据，保留top-k顺序的同时显著减少内存使用。这将LSH从候选生成器转变为稀疏注意力的原理性评分核。利用这一特性，SOCKET实现了高效的token选择，无需随意投票，并在多个长上下文基准上匹配或超越了先前的稀疏注意力方法。通过定制的CUDA评分内核和Flash Decode Triton后端，SOCKET的吞吐量比FlashAttention高至1.5倍。
英文摘要：
Exploiting sparsity is key to efficient long-context inference, as attention dominates the cost of autoregressive decoding. Sparse attention reduces this cost by restricting computation to a subset of tokens, but its effectiveness hinges on fast and accurate token scoring and selection at inference time. Data-agnostic approaches offer an attractive way to perform this selection, but often incur substantial memory overhead to maintain high recall. We revisit Locality-Sensitive Hashing (LSH) and introduce SOCKET, a SOft Collision Kernel EsTimator that replaces hard bucket matches with probabilistic, similarity-aware aggregation. Traditional LSH relies on binary collision signals, providing limited information for ranking tokens and necessitating many hash tables for accurate retrieval. In contrast, soft LSH accumulates graded collision evidence across hash tables, closely preserving the true top-$k$ ordering with significantly less memory. This reframes LSH from a candidate-generation mechanism into a principled scoring kernel for sparse attention. Building on this insight, SOCKET enables efficient token selection without ad hoc voting and matches or outperforms existing sparse attention methods across multiple long-context benchmarks and diverse language models. With a custom set of CUDA/Triton kernels for scoring, selection, and attention, SOCKET achieves up to approximately $1.5\times$ higher throughput than FlashAttention. Code is open-sourced at https://github.com/amarka8/SOCKET.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
