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

2026-07-12

AI Summary

This work addresses the inefficiencies of existing Indexer-TopK kernels on GPUs, which suffer from frequent global memory accesses, high synchronization overhead, and substantial memory consumption. Leveraging the concentration of pairwise distances in high-dimensional spaces—a manifestation of the “curse of dimensionality”—the authors propose, for the first time, a fused Indexer-TopK kernel that dynamically estimates score ranges via sampling, performs online bucket-based candidate filtering, and writes back only potential Top-k results. This approach significantly reduces I/O and memory overhead while preserving result accuracy, enabling efficient sparse attention indexing. When deployed in GLM 5.2, the method achieves a 1.2× speedup during the prefill phase and markedly lowers memory usage.

Abstract

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.

Problem

• Indexer-TopK
• sparse attention
• memory overhead
• global memory traffic
• GPU kernel

Innovation

• curse of dimensionality
• fused kernel
• sparse attention
• Top-k selection
• memory efficiency