Abstract: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.
| Comments: | 7 figures, 17 tables |
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2602.06283 [cs.LG] |
| (or arXiv:2602.06283v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2602.06283 arXiv-issued DOI via DataCite |
From: Sahil Joshi [view email]
[v1]
Fri, 6 Feb 2026 00:41:44 UTC (924 KB)
[v2]
Fri, 8 May 2026 00:20:43 UTC (878 KB)
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。