LiteTopK:利用维度灾难优化稀疏注意力融合Indexer-TopK内核

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

精选理由

这论文用维度诅咒的视角,搞了个LiteTopK算子,比DeepSeek的DSA更快更省内存,GLM 5.2预填充直接快1.2倍。

AI 摘要

本文针对长上下文稀疏注意力中的Indexer-TopK操作,指出现有GPU内核如DeepSeek Sparse Attention(DSA)存在全局内存开销大、同步成本高问题。作者利用高维空间中向量距离分布集中的特性,提出LiteTopK内核:通过采样小部分数据估计查询-数据分数范围,在线分桶,维持紧近似阈值并回写有潜力候选。实验在GLM 5.2的预填充阶段实现1.2倍加速,同时降低内存开销,且保持精确Top-k结果。

原文 · 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.