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

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. 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× across GLM-5.2, DeepSeek-V4-Flash, and Hy-4-Preview at a context length of 1M tokens, without compromising model capabilities.