LiteTopK: Exploiting the Curse of Dimensionality for a Fused Indexer-TopK Kernel in Long-Context Sparse Attention
Ziqi Yin, Jianyang Gao, Peiqi Yin, Jiangneng Li, Gao Cong
Indexer-TopK, which computes scores and selects the top-$k$ candidates, is a key primitive in long-context sparse attention for large language models and vector retrieval for recommendation systems and vector databases. However, existing GPU-based Indexer-TopK kernels like DeepSeek Sparse Attention (DSA) typically separate scoring from selection and materialize the full score matrix in high-bandwidth memory. This strategy is inefficient for large score matrices, such as those encountered in long-context sparse attention, due to excessive global-memory traffic, costly synchronization, and prohibitive memory overhead. In this study, motivated by the curse of dimensionality, we observe that sparse attention scores exhibit a concentration phenomenon, with scores falling within a narrow range. Based on this observation, we propose LiteTopK, an efficient fused Indexer-TopK kernel. LiteTopK samples a small subset of data to estimate query-data score ranges, then partitions candidates into bins accordingly. This design allows LiteTopK to maintain a slightly relaxed threshold online and write back only promising candidates, reducing unnecessary I/O and memory overhead while preserving exact top-$k$ correctness. Building on LiteTopK, we further propose LiteDSA, which packs the overlapping candidates of neighboring query tokens into a shared key/value list for joint attention and uses per-query masks to preserve each token's original top-$k$ membership while reducing memory traffic. We integrate LiteTopK into three widely used LLM training and inference systems: vLLM, SGLang, and Megatron-LM. Experiments in a production environment with eight NVIDIA B200 GPUs demonstrate prefill speedups of up to $1.49\times$ across GLM-5.2, DeepSeek-V4-Flash, and Hy-4-Preview at a context length of 1M tokens, without compromising model capabilities.
尚无 AI 解读，展示原文摘要