New sample-based method speeds exact Top-K sparse attention on GPU
Sample-Guided Exact Top-K Selection for Long-Context Sparse Attention
Distributed, Parallel, and Cluster Computing
Summary
Large language models use complex attention mechanisms that need to select the most important tokens from a long context quickly. The authors found that looking at a small, spaced sample of data first can guide the exact selection process more efficiently without losing accuracy. They implemented this method on GPUs and showed it is about 1.3 to 1.7 times faster than previous exact methods across many test cases. This improvement helps make large models work faster during inference.
What this means in practice
- •For software engineers optimizing llms: Speed up exact Top-K token selection in sparse attention to reduce inference time on GPUs for long-context language models.
- •For high-performance computing developers: Incorporate sample-guided Top-K selection techniques to improve GPU operator libraries handling irregular sparse data structures.
Authors
Siran Liu, Yang Xue, Theo Tang, Changxu Shao, Qian Cheng, Haimeng Ren, Donghua Jiang, Haipeng Ming, Lehua Ding, Zhonghan Lin, Shengying Wei, Wei Liu, Kai Liu, Jianchen Zhu
Abstract
Sparse attention bounds downstream attention work by retaining a fixed-size subset of indexed tokens, but its standalone exact Top-$K$ stage must still process materialized score rows whose length grows with context. Production radix selectors discover their first actionable boundary only after a complete-row pass, forcing another row-scale traversal before exact refinement. We observe that locating a compact upper tail requires substantially less resolution than identifying the exact rank boundary, and that fixed-stride partial views of the current row remain calibrated to the corresponding complete-row rank across ragged lengths. We present HPC-Ops Top-K, a sample-guided exact selector for ragged sparse-attention score rows. A fixed-stride view proposes a row-local coarse boundary; the mandatory complete-row pass certifies its sufficiency, forms the admitted candidate set, and initializes exact FP32 refinement over the unresolved frontier. A nested secondary boundary and exact recovery handle underfilled proposals before any output is committed, so sampling controls common-path work but never correctness. The GPU implementation fuses complete-row certification and candidate formation, and combines persistent, KV-split, and direct-exact execution behind graph-capturable ragged-row dispatch. We evaluate HPC-Ops Top-K on indexer scores from Hy4-Preview. It outperforms the fastest verified external exact baseline by $1.29$--$1.75\times$ across 20 operator configurations, with a $1.55\times$ geometric-mean speedup. It further achieves $1.36\times$ and $1.48\times$ speedups on two framework-derived sparse-attention traces. The implementation is available in HPC-Ops, Tencent's open-source high-performance operator library for LLM inference, at https://github.com/Tencent/hpc-ops.