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

2026/07/26 by Ziqi Yin, Jianyang Gao, Peiqi Yin +2
#cs.LG

paper · pdf

Abstract

Indexer-TopK, the operation to compute the scores and select the top-k candidates, is widely used by sparse attention algorithms 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 study, inspired by the curse of dimensionality phenomenon, we first observe that sparse attention scores exhibit a score concentration phenomenon, where scores tend to fall within a narrow range. Based on this observation, we propose LITETOPK, an efficient fused Indexer-TopK kernel. LITETOPK first samples a small subset of data to estimate query-data score ranges, then partitions candidates into bins accordingly. This organization allows the LITETOPK kernel to maintain a tight approximate threshold online, write back only promising candidates, reduce unnecessary I/O and memory overhead while preserving exact Top-k correctness. Building on LITETOPK, we further propose LITEDSA, which exploits the similarity of top-k candidate sets among neighboring tokens. LITEDSA packs neighboring tokens' candidates for joint computation and masks out extra scores for each query, thereby reducing memory traffic while preserving correctness. Experimental results in a real-world deployment environ ment with eight B200 GPUs show that LITETOPK+LITEDSA accelerates the prefill stage of GLM 5.2 by 1.35x, with no performance loss and lower memory overhead.

Citations

• DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models
• The Sparse Frontier: Sparse Attention Trade-offs in Transformer LLMs
• Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
• Exploiting Sparsity for Long Context Inference: Million Token Contexts on Commodity GPUs
• FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving
• TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention
• Embedding And Clustering Your Data Can Improve Contrastive Pretraining
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• PQCache: Product Quantization-based KVCache for Long Context LLM Inference
• Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference
• SnapKV: LLM Knows What You are Looking for Before Generation
• The Faiss library
• SGLang: Efficient Execution of Structured Language Model Programs
• SparQ Attention: Bandwidth-Efficient LLM Inference
• Efficient Memory Management for Large Language Model Serving with PagedAttention
• Rethinking the Role of Token Retrieval in Multi-Vector Retrieval
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Memory-efficient Transformers via Top-k Attention
• Billion-scale similarity search with GPUs
• Pointer Sentinel Mixture Models

Related

• StreamIndex: Memory-Bounded Compressed Sparse Attention via Streaming Top-k
• Efficient Low Rank Attention for Long-Context Inference in Large Language Models
• Learn from the Past: Fast Sparse Indexing for Large Language Model Decoding
• FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel
• Guess-Verify-Refine: Data-Aware Top-K for Sparse-Attention Decoding on Blackwell via Temporal Correlation