用概率化方法提升长文本推理效率，显著加速大模型生成
SOCKET: SOft Collision Kernel EsTimator for Sparse Attention
- 用软碰撞核替代传统哈希的硬匹配，实现更精准的注意力筛选
- 在多个长文本任务中表现超越或持平现有方法，内存占用更低
- 适合追求高效长序列生成的开发者与研究者使用
在长上下文推理中利用稀疏性是扩展大语言模型的关键，因为注意力机制主导了自回归解码的计算开销。稀疏注意力通过限制计算范围来降低开销，但其效果依赖于推理时高效的评分与选择。本文重新审视局部敏感哈希（LSH），提出SOCKET——一种软碰撞核估计器，将硬桶匹配替换为概率化、相似度感知的聚合方式。传统LSH仅提供二值碰撞信号，影响排序质量且需大量内存；而软LSH在多个哈希表中累积渐进式碰撞证据，能以更少内存保持top-k顺序。这一设计将LSH从候选生成器转变为可解释的稀疏注意力评分核。基于此，SOCKET无需额外投票机制即可高效选词，在多个长文本基准测试中表现优异。结合定制CUDA评分核与Flash Decode Triton后端，SOCKET吞吐量最高达FlashAttention的1.5倍。
原文摘要 · Abstract (English)
Exploiting sparsity during long-context inference is key to scaling large language models, as attention dominates the cost of autoregressive decoding. Sparse attention reduces this cost by restricting computation to a subset of tokens, but its effectiveness depends on efficient scoring and selection at inference time. 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 yields binary collision signals that limit ranking quality and require substantial memory to perform well. In contrast, soft LSH accumulates graded collision evidence across hash tables, preserving top-k ordering with significantly less memory. This reframes LSH from a candidate generator into a principled scoring kernel for sparse attention. Leveraging this property, SOCKET enables efficient token selection without ad hoc voting and matches or surpasses prior sparse attention methods across multiple long-context benchmarks. With a custom CUDA scoring kernel and a Flash Decode Triton backend, SOCKET achieves up to 1.5$\times$ higher throughput than FlashAttention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。