PaperlayerView exact revision on arXiv

Source-linked AI summary

Sample-Guided Exact Top-K Selection for Long-Context Sparse Attention

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

arXiv:2609.08450v1cs.DC

TL;DR

Exact Top-K selection remains a growing cost for long-context sparse attention, and production radix selectors locate their first actionable boundary only after a complete-row pass. HPC-Ops Top-K uses fixed-stride samples to propose a boundary, then certifies candidates and refines exactly. It outperforms verified external exact implementations across the evaluated operator configurations and improves performance on framework-derived traces.

  • Problem

    Exact Top-K must still examine full materialized score rows as context grows, while production radix selection finds its first actionable boundary after a complete-row pass and needs another row-scale traversal.

    Source: §1 · p. 1

  • Method

    HPC-Ops Top-K uses a fixed-stride current-row view to propose a coarse boundary, then certifies admitted candidates in a complete-row pass and exactly refines the unresolved frontier.

    Source: §8 · p. 26

  • Results

    HPC-Ops Top-K is 1.29–1.75× faster than the best verified external exact implementation across the operator matrix, with a 1.55× geometric-mean speedup and 1.36× and 1.48× speedups on framework-derived traces.

    Source: §8 · p. 26

  • Takeaways & Limitations

    Sampling serves as a proposal mechanism rather than an approximation: complete-row certification determines validity and activates recovery while preserving exact Top-K semantics.

    Source: §8 · p. 26

    Prior-based selectors lack a previous-step prior on first-chunk prefill, single offline calls, or saved-score replay, and then fall back to multi-pass partitioning traffic.

    Source: §2.2 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

Exact Top-K can remain a critical-path cost as indexed score rows grow, because production radix selection discovers an actionable boundary only after a complete-row pass. HPC-Ops Top-K uses a current-row sample to propose a boundary, then certifies and refines it exactly; it is fastest in all 20 evaluated operator configurations and also speeds up two framework-derived traces.

1 Introduction

  • Exact Top-K remains costly as ragged score rows grow, because every valid score must be examined even when fixed K bounds downstream attention work.Production radix selection discovers its first actionable boundary only after a complete-row histogram, requiring another row-scale traversal for candidate formation.Source: §1 · pp. 1–2
  • Fixed-stride views of the current row can locate a useful coarse boundary across ragged lengths without resolving the exact rank boundary.The complete-row pass remains authoritative for exactness.Source: §1 · p. 2
  • HPC-Ops Top-K proposes a coarse boundary, certifies and collects admitted candidates in a complete-row pass, then performs exact FP32 refinement on the unresolved frontier.Underfilled proposals trigger a nested boundary and, if needed, exact recovery before output is committed.Source: §1 · p. 2
  • The GPU portfolio fuses candidate formation with complete-row certification and dispatches among persistent, KV-split, and direct-exact kernels.The implementation uses graph-stable ragged-row dispatch and a shared exact FP32 refinement contract.Source: §1 · p. 2
  • 1.29–1.75× faster than the fastest external baseline across 20 operator configurations, with a 1.55× geometric-mean speedup; framework-derived traces show 1.36× and 1.48× speedups.The evaluation uses indexer scores from Hy4-Preview and verified exact implementations from vLLM, TensorRT-LLM, SGLang, FlashInfer, and PyTorch.Source: §1 · p. 2

2 Background

The background frames exact Top-K as a standalone cost that grows with context length, then compares selector mechanisms by when they obtain a usable boundary and what resources they trade for lower score traffic. Production radix selection’s first boundary arrives only after a full-row pass, structurally requiring another traversal before classification.

2.1 Cost of Standalone Top-K Selection

  • The selector operates on materialized FP32 score rows whose valid lengths vary, while sparse attention consumes only the K selected positions.The contract returns K distinct valid indices when L > K, and uses -1 sentinels for unused slots when L ≤ K.Source: §2.1 · p. 3
  • Deployment requires ragged-row support, padded-pitch alignment, CUDA-graph capture, batch-independent workspace, and compatibility with arbitrary score producers and offline replay.These requirements rule out host-side path selection based on runtime statistics.Source: §2.1 · p. 3
  • For fixed K and data representations, the ratio of selector score traffic to selected-position payload grows linearly with L, and each added full-row traversal increases it proportionally.Req normalizes valid-row score reads by one complete scan; exact refinement after compact reduction is counted separately.Source: §2.1 · pp. 3–4

2.2 Approaches to Exact Selection

  • Ordered-state methods reach Req = 1, but shared-memory and register footprints cap supported ranks at 256 for bitonic networks and 2048 for queues.A queue is faster only below K ≈256 in the cited identical-hardware measurement.Source: §2.2 · p. 4
  • Alternatives trade traffic for different resources: threshold search needs on-chip rows, delegate reduction depends on a tuned subrange, and sampled methods vary in exactness guarantees.SampleSelect and Qrita retain deterministic exact completion; Prof-K uses user-specified high-probability recovery.Source: §2.2 · p. 5
  • History-based GVR can approach one threshold-search traversal on correlated decode work, but lacks prior-step state on prefill starts, offline calls, or saved-matrix replay.With nsearch threshold-search traversals, its row-scaled traffic satisfies Req ≥ nsearch + 1.Source: §2.2 · p. 5
  • The unresolved target is long-row support with regular GPU execution, current-row state, and traffic near one valid-row scan; radix’s first boundary arrives too late to achieve it.The proposed direction is to test whether a current-row view can locate a compact tail before mandatory row inspection.Source: §2.2 · p. 6, §2.3 · p. 7

2.3 The Production Radix Bottleneck

  • The first radix histogram and boundary-dependent classification require 2L reads before candidate shrinkage; round fusion reduces a four-digit schedule’s worst-case reads from 8L to 5L.Atomics and synchronization affect each pass’s throughput but do not remove this structural reread.Source: §2.2 · p. 5, §2.3 · pp. 6–7

3 Motivation

The measurements show that coarse boundaries can retain a compact upper tail and that fixed-stride views track global upper-tail ranks across ragged rows. The theoretical basis motivates separating early candidate filtering from authoritative complete-row certification and exact refinement.

3.1 Empirical Observations

  • 100× the exact adjacent score gap retains a median 1.03K values, while 95% of rows retain no more than 1.15K.At 500×, the median is 1.18K and the 95th percentile is 2.02K; tied Kth-boundary rows are excluded because the normalization requires a positive gap.Source: §3.1 · p. 8
  • Even 1/256 views remain within 5% of requested rank across the sweep, while views through 1/64 remain within 1.3%.The sweep covers five requested ranks from K through 2K and every phase of the selected rows.Source: §3.1 · p. 8
  • At requested rank 1.5K, 1/64 and 1/128 views span 1.23K–1.77K and 1.09K–1.92K, respectively, across the fifth-to-95th-percentile ranges.Their medians remain within 0.03K of the request; 95% of absolute errors are within 0.32K and 0.50K, respectively.Source: §3.1 · p. 8
  • Endpoint deciles contain 52–74% of exact top-K indices versus 20% under uniform positions, and endpoint-structured rows have 14–28% lower RMS rank error.The lower error is relative to the exchangeable reference across all five view sizes.Source: §3.1 · p. 8, Figure 3 · p. 9

3.2 Theoretical Basis

  • Under positional exchangeability, phase order statistics center near the requested global rank, while reducing view size increases uncertainty only at a square-root rate.Local rank geometry gives C(ϵ) − K ≈ ϵ; at ϵ = 500, this predicts a tail near 1.24K, versus a measured 1.18K median.Source: §3.2 · p. 10

3.3 Design Implications

  • The partial view supplies an early boundary, but the complete-row stage verifies tail sufficiency before exact ordering, so sampling error changes retained work rather than the final output.This coarse-to-exact division makes the early estimator’s remaining work, not reproduction of the final cutoff, the relevant design target.Source: §3.3 · p. 10

4 Method

HPC-Ops Top-K separates coarse boundary proposal from exact ranking: a fixed-stride view proposes a compact tail, and a complete-row pass certifies and forms it before FP32 refinement. Bounded recovery handles underfilled proposals, while persistent and KV-split mappings adapt the work to row-level parallelism.

4.1 Algorithm Overview

  • A fixed-stride partial view proposes a row-local coarse boundary, while the mandatory complete-row traversal validates it and forms the admitted tail.The partial view determines remaining work; complete-row certification and FP32 refinement determine the returned indices.Source: §4.1 · p. 11

4.2 Phase 1: Coarse Boundary Localization

  • The sample rank targets eK = ρK, and a deterministic row-phased view of approximately Lr/s values locates a coarse boundary across ragged lengths.A nested deeper boundary from the same histogram supports bounded recovery without another sample traversal.Source: §4.2 · p. 12

4.3 Phase 2: Fused Validation and Candidate Formation

  • When the admitted tail contains at least K values, complete-row certification establishes a Top-K solution; underfill triggers a secondary band or exact coarse recovery.The fused pass also persists candidates and builds the leading exact-digit histogram, yielding the common refinement interface.Source: §4.3 · pp. 13–14

4.4 Phase 3: Exact Radix Refinement

  • Theorem 4.2 guarantees K distinct valid indices satisfying the exact Top-K contract for every admissible row with L > K.FP32 radix refinement commits groups ahead of the rank boundary and continues only through the unique unresolved boundary group.Source: §4.1 · p. 11, §4.4 · p. 14

4.5 Execution Mappings

  • Persistent execution assigns one CTA per row, while KV-split partitions validation across segments and uses a finisher to initialize the same refinement state.KV-split targets low row counts by shortening validation span, trading for additional aggregate sample work.Source: §4.5 · p. 15, §4.6 · p. 16

4.6 Complexity Analysis

  • The certified sampled path has row-scaled traffic Req = 1 + 1/s before frontier refinement, and primary-boundary-certified rows avoid secondary scans and exact-recovery passes.Observed primary underfill is rare, and the secondary boundary makes exact recovery rarer still.Source: §4.6 · pp. 15–16

5 Implementation

The implementation fuses boundary validation, candidate formation, and leading-histogram construction, then refines only unresolved frontiers. A shape-stable portfolio selects persistent, KV-split, or direct-exact kernels while keeping row state and dispatch compatible with graph capture.

5.1 Implementation Overview

  • The kernel family shares a common FP32 refinement interface, while a shape-stable planner selects execution mapping from captured dimensions.Sampled long-row kernels use persistent or KV-split execution; row-local and eight-CTA kernels cover complementary regimes.Source: §5.1 · p. 16, §5.5 · p. 19

5.2 Phase 1: On-Chip Boundary Localization

  • Sampled CTAs construct coarse boundaries with a regular row-phased view, shared histogram, and in-place prefix; unusable views take exact coarse localization.The histogram allocation is cleared and reused for the leading exact-digit histogram in Phase 2.Source: §5.2 · pp. 16–17

5.3 Phase 2: Fused Full-Row Validation

  • The vectorized complete-row loop fuses endpoint classification, candidate storage, admitted-count measurement, and leading-digit histogram construction.The principal path uses one complete-row traversal; additional row traffic is confined to recovery.Source: §5.3 · pp. 17–18

5.4 Phase 3: Frontier-Only Radix Refinement

  • Each radix round gathers scores only for unresolved indices, routing them to output, discard, or the next boundary frontier while building its histogram.Shared arrays ping-pong between frontiers, with global buffers extending capacity when needed.Source: §5.4 · p. 18

5.5 Execution Mappings and Dispatch

  • Persistent execution suits abundant rows, KV-split raises parallelism for a few long rows, and direct-exact mappings serve complementary shapes.The host planner fixes kernel family, view density, split count, and workspace before graph capture, while live row counts and valid lengths remain device inputs.Source: §5.5 · pp. 18–19

6 Evaluation

The evaluation tests exact Top-K on Hy4-Preview FP32 indexer scores, from standalone operator shapes to framework traces and sampled-path ablations. HPC-Ops leads the operator matrix and framework traces, while the ablations characterize where sampling helps and how underfilled proposals are recovered.

6.1 Experimental Setup

  • The evaluation uses NVIDIA H20 GPUs and Hy4-Preview FP32 scores, validates returned rows under a set-valued exactness contract, and times operator execution with CUDA-Graph measurements.The standalone operator uses K = 2048 and preserves native ragged row lengths; input preparation, correctness checks, and compilation are excluded from timing.Source: §6.1 · pp. 20–21, Table 3 · p. 20

6.2 Operator-Level Top-K Performance

  • HPC-Ops is fastest in all 20 operator configurations, beating the best exact external result by 1.29–1.75× with a 1.55× geometric-mean speedup.The advantage widens with allocated row width, from 1.29–1.42× at 128K columns to 1.68–1.75× at 1M.Source: §6.2 · p. 21, Figure 7 · p. 22

6.3 Framework-Level Top-K Performance

  • On the 245K and 430K framework traces, HPC-Ops is exact and fastest, achieving 1.36× and 1.48× speedups over SGLang.The trace replays preserve each invocation’s original row count, allocated width, valid lengths, and order.Source: §6.3 · p. 22, Figure 8 · p. 23

6.4 Sampled-Path Ablations

  • With eK = 1.5K and s = 128, speedup reaches 1.51–1.54× while recovery stays below 0.9%; s = 256 raises recovery and reverses part of the gain.A deeper retention target suppresses remaining recoveries but retains more candidates on every row, producing a broad latency plateau.Source: §6.4 · pp. 23–24, Figure 10 · p. 24
  • Conditional rescue resolves 93–96% of primary misses, and coarse rescue uses 61–62% less recovery-path time than full exact recovery.The secondary boundary reduces residual recovery rates from 0.40% and 0.86% to 0.015% and 0.060% in the two terminal steps.Source: Figure 11 · p. 24, §6.4 · pp. 24–25
  • Matched controls attribute 92% and 97% of the latency difference to sparse boundary localization, with proposal time falling from approximately 18 to 2 µs and 33 to 3.5 µs.The full-row validation and exact tail processing remain comparable across controls.Source: §6.4 · p. 25, Figure 12 · p. 25
  • At 16K and 24K widths exact proposal wins every tested cell; at 64K both sampled variants win all 14 cells, with median speedups of 1.13× for s = 64 and 1.22× for s = 128.Results are mixed from 32K through 56K, supporting a transition range rather than a universal cutoff.Source: §6.4 · p. 25

7 Related Work

The related-work discussion distinguishes standalone selection over materialized score rows from approaches that change the surrounding sparse-attention pipeline. Those approaches fuse or stream selection, reuse indexing work, or move data and execution to another substrate.

Fused and streaming indexer selection

  • Fused and streaming indexer methods avoid materializing the complete score matrix, but optimize the fused indexer pipeline rather than the standalone materialized-row operator.LiteTopK filters scores as produced, while StreamIndex evaluates chunks and merges bounded candidate sets.Source: §7 · pp. 25–26

Reusing indexing work

  • Reuse methods reduce how often indexing runs by exploiting redundancy across nearby queries or layers, while owner layers and per-query refinement can still require selection.The passage characterizes reducing per-selection cost as complementary to reducing the number of selections.Source: §7 · pp. 25–26

Alternative memory and execution substrates

  • Alternative substrates address capacity or data placement rather than materialized-row selection cost: HiSparse uses host memory with a bounded GPU cache, while KARAT proposes near-memory scoring and selection.HiSparse composes with a GPU selector; KARAT replaces the GPU-resident boundary and is evaluated in simulation.Source: §7 · pp. 25–26

8 Conclusion

HPC-Ops Top-K uses a sampled current-row view to propose a compact upper tail, then certifies candidates over the complete row and exactly refines the remaining frontier. The portfolio is faster than verified external exact implementations on operator configurations and framework traces.

8 Conclusion

  • A regular partial view proposes a row-local boundary; a complete-row pass certifies and forms candidates, and FP32 refinement resolves the remaining frontier.Underfilled proposals are recovered before output is committed, preserving exact Top-K semantics on every path.Source: §8 · p. 26
  • The portfolio is 1.29–1.75× faster than the best verified external exact implementation across the operator matrix, with a 1.55× geometric-mean speedup, and reaches 1.36× and 1.48× on framework traces.The reported gains come primarily from replacing complete-row boundary localization with a sparse current-row proposal.Source: §8 · p. 26
  • Sampling is a proposal mechanism rather than an approximation to Top-K: it sets common-path work while certification determines validity and triggers recovery when needed.This trades common-path traffic against rare recovery work while preserving exact Top-K semantics.Source: §8 · p. 26

Related papers

Loading 2609.08450v1…