← 研踪 PaperTrailarXiv:2607.11976 · v4
LiteTopK: Exploiting the Curse of Dimensionality for a Fused Indexer-TopK Kernel in Long-Context Sparse Attention
首次提交 2026-07-13 11:53
版本更新 2026-10-27 10:38
雷达观测 2026-10-29 00:35
摘要
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.
作者备注：15 pages
观测到的版本
相邻记录可能跳过中间版本；首次发现 v2 及以上也会标为版本更新。
v4 · 2026-10-27 10:38 · 最新观测版本
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. 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.
查看此版本解读 →查看此版本原文 ↗