LiteTopK:利用维度诅咒加速长上下文稀疏注意力

LiteTopK: Exploiting the Curse of Dimensionality for a Fused Indexer-TopK Kernel in Long-Context Sparse Attention

精选理由

这个新方法挺聪明,用高维空间特性优化注意力计算,在GLM 5.2上实测提速1.2倍还省内存。

AI 摘要

LiteTopK是一种针对长上下文稀疏注意力中Indexer-TopK操作的高效融合内核。它利用高维空间中向量距离集中在狭窄范围的特性,通过采样小部分数据估计查询-数据分数范围,并在线将候选结果分桶。该设计减少了全局内存流量和同步开销,同时保持精确Top-k正确性。实验在GLM 5.2的预填充阶段实现1.2倍加速,内存开销更低。

原文 · arXiv: DeepSeek

LiteTopK: Exploiting the Curse of Dimensionality for a Fused Indexer-TopK Kernel in Long-Context Sparse Attention

Indexer-TopK, the operation to compute the scores and select the top-k candidates, is widely used by sparse attention kernels in large language models and vector retrieval in recommendation systems and vector databases. However, existing GPU-based Indexer-TopK kernels like DeepSeek Sparse Attention (DSA) remain inefficient due to excessive global memory traffic, costly synchronization, and prohibitive memory overhead. In this work, we exploit the curse of dimensionality in high-dimensional spaces, where distances between high-dimensional vectors tend to concentrate within a narrow range, to design LITETOPK, a novel and efficient fused Indexer-TopK kernel. LITETOPK first samples a small subset of data to estimate query-data score ranges, then uses these estimates to partition candidate results into bins online. This organization allows the LITETOPK kernel to maintain a tight approximate threshold, write back only promising candidates, reduce unnecessary I/O, substantially lower memory overhead, and still preserve exact Top-k correctness. Experimental results show that LITETOPK accelerates the prefill stage of GLM 5.2 by 1.2x in real-world deployment scenarios while incurring lower memory overhead.