SOCKET: SOft Collision Kernel EsTimator for Sparse Attention

TL;DR

SOCKET introduces a probabilistic hashing-based kernel for sparse attention, improving token selection efficiency and throughput in long-context language model inference.

Contribution

It replaces hard collision signals with soft, similarity-aware aggregation, enabling more efficient and accurate sparse attention scoring.

Findings

SOCKET matches or surpasses prior sparse attention methods on benchmarks.

Achieves up to 1.5× higher throughput than FlashAttention.

Uses a custom CUDA kernel and Triton backend for efficiency.

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.