Guess-Verify-Refine: Data-Aware Top-K for Sparse-Attention Decoding on Blackwell via Temporal Correlation

Abstract: Sparse-attention decoders rely on exact Top-K selection to choose the most important key-value entries for each query token. In long-context LLM serving, this Top-K stage runs once per decode query and becomes a meaningful latency bottleneck even when the indexer and attention kernels are already highly optimized. We present \textbf{Guess-Verify-Refine (GVR)}, a data-aware exact Top-K algorithm for sparse-attention decoding on NVIDIA Blackwell. GVR exploits temporal correlation across consecutive decode steps: it uses the previous step's Top-K as a prediction signal, computes pre-indexed statistics, narrows to a valid threshold by secant-style counting in 1-2 global passes, verifies candidates with a ballot-free collector, and finishes exact selection in shared memory. We connect this behavior to the Toeplitz / RoPE structure of DeepSeek Sparse Attention (DSA) indexer scores and validate the design on real DeepSeek-V3.2 workloads integrated into TensorRT-LLM. GVR achieves an average \textbf{1.88x} single-operator speedup over the production radix-select kernel, with up to \textbf{2.42x} per layer per step, while preserving bit-exact Top-K outputs. In controlled TEP8 min-latency deployment, it improves end-to-end TPOT by up to \textbf{7.52%} at 100K context, with larger gains at longer contexts and smaller but still positive gains under speculative decoding. While implemented and validated in the current TensorRT-LLM DSA stack on Blackwell, the same principle may extend to sparse-attention decoders whose decode-phase Top-K exhibits temporal stability.

• The paper introduces the GVR algorithm, leveraging temporal correlation in autoregressive decoding to warm-start threshold estimation for exact Top-K selection.
• It employs secant-method and histogram-based refinement techniques to achieve up to 2.42× speedup and significantly reduce end-to-end inference latency.
• The approach is integrated into TensorRT-LLM on Blackwell, demonstrating robustness across varied layer correlations and sequence lengths.

Guess-Verify-Refine: Data-Aware Top-K for Sparse-Attention Decoding on Blackwell via Temporal Correlation

Motivation and Background

Sparse-attention mechanisms in LLMs necessitate exact Top-K selection to identify key-value (KV) entries of maximal importance per query token. In extended-context inference (32K–128K+ tokens), this Top-K operation arises as a dominant latency bottleneck, persisting even after aggressive optimization of indexer and attention kernels. Prior GPU Top-K research, notably the radix-select kernels in TensorRT-LLM, treats the selection problem as distribution-agnostic, unconcerned with the underlying temporal structure present in autoregressive decoding. However, empirical analysis demonstrates strong temporal correlation between consecutive Top-K index sets—an untapped signal with the potential to accelerate exact selection without resorting to approximation.

In the DeepSeek Sparse Attention (DSA) paradigm, each decoding step requires ranking up to hundreds of thousands of indexer scores, selecting the $K = 2048$ highest, and using these entries in subsequent attention computation.

Figure 1: Indexer Top-K selection, with $K{=}5$ in a toy example. For DSA the typical regime is $K{=}2048$, $N$ up to $128$K+.

The "Guess-Verify-Refine" (GVR) algorithm capitalizes on the temporal correlation induced by Toeplitz/RoPE positional structures, translating prior-step Top-K output into a warm-start threshold estimate for efficient exact selection.

Temporal Correlation: Empirical and Theoretical Foundations

Empirical measurements on DeepSeek-V3.2 workloads show pronounced Top-K overlap between decoding steps (35–50% for layers 20–60), considerably elevated in deep layers relative to shallow ones. This result is visualized below and forms the cornerstone for the GVR predictor.

Figure 2: Raw Top-K overlap (hit ratio) between consecutive decoding steps for various layers, with deeper layers exhibiting 35–50% average overlap.

Theoretical justification stems from the Toeplitz structure created by RoPE positional encoding: the attention score matrix is constant along diagonals ($g(\Delta)$), and the query-adjacent score landscape shifts minimally with each decode step. YaRN extensions further enhance this temporal persistence by preserving high-magnitude peaks at large relative positions—enabling Top-K sets to stably span both neighboring and distant tokens.

Figure 3: Temporal correlation under an offset+1 shift; Toeplitz score matrix structure and RoPE frequency analysis underpin Top-K stability.

By leveraging either precomputed structural priors or real prior-step Top-K indices, GVR produces threshold estimates tightly coupled to the true Top-K boundary.

GVR Algorithm: Guess, Verify, Refine

The GVR procedure proceeds in four phases within a single CTA:

1. Guess (Phase 1–2): Uses predicted indices (prev-step Top-K) to compute statistics and interpolate a threshold (via secant-style search) close to $x_{(K)}$.
2. Verify (Phase 3): Collects candidates above threshold, caching counts to bypass redundant data passes.
3. Refine (Phase 4): Executes histogram-based exact selection and tie-breaking snap iterations to produce exactly $K$ indices.

   Figure 4: GVR Top-K algorithm: Guess (threshold via interpolation), Verify (candidate collection), Refine (histogram-based selection), all within a single CTA.

The secant-method threshold search is visualized as iteratively narrowing brackets to center around the candidate set with $K$–$C$ elements, with empirical convergence in 1–2 passes for high-correlation layers.

Figure 5: Phase 2 secant-method threshold search: quick bracket narrowing via interpolation using prior Top-K signal.

Histogram-based selection enables rapid tie-breaking and partitioning of candidate elements; the snap refinement loop exhibits fast convergence under most data distributions.

Figure 6: Phase 4 details: min/max scan, histogram build, warp-parallel bin search, snap tie-breaking iterations yielding exact Top-K.

Implementation on Blackwell and Dispatch Logic

GVR is implemented as a single-CTA kernel, maximizing L2 cache reuse and intra-phase communication efficiency. Shared memory layout is expanded ($K{=}5$060KB per CTA), allowing persistent candidate buffers.

Integration into TensorRT-LLM is architecture- and workload-aware: GVR is preferred when prior Top-K signal and scratch buffers are available, falling back to radix-select otherwise.

Figure 7: Full decode-stage Top-K dispatch: Python selects CuTE DSL Top-K or C++ kernel; GVR takes precedence with valid preIdx, else radix-select fallback triggers.

Performance Evaluation and Numerical Results

Single-Operator Speedup

Synthetic benchmarks (random Q/K + static RoPE prior) demonstrate break-even performance at sequence length $K{=}5$1K, scaling to 1.75$K{=}5$2 speedup at $K{=}5$3K. On real DeepSeek-V3.2 decoding data, GVR achieves average 1.88$K{=}5$4 single-operator speedup over radix-select, peaking at 2.42$K{=}5$5 layer-wise; worst-case layers (lognorm, weak temporal correlation) still yield 1.59$K{=}5$6.

Figure 8: Kernel latency vs. sequence length: both kernels scale as $K{=}5$7, with GVR's fitted slope about 42% of baseline.

Figure 9: Per-layer kernel latency (average over 17 decode steps): GVR achieves 1.59–2.04$K{=}5$8 speedup, consistent across DeepSeek-V3.2 layers.

End-to-End Latency Impacts

In batch-1 min-latency TEP8 serving (DeepSeek-V3.2-Exp on B200$K{=}5$98), GVR reduces end-to-end TPOT by 5.47% (64K) to 7.52% (100K) context length; gains are monotonic with sequence length and positive across speculative settings.

Figure 10: End-to-end TPOT reduction: GVR provides greater gains as context length grows and maintains benefit under speculative decode.

Robustness and Ablation

The prediction signal’s quality directly drives speedup: even random selection ($K{=}2048$0) enables 1.44$K{=}2048$1 improvement via architectural optimizations. High-correlation layers average 1.94$K{=}2048$2, low-correlation layers 1.65$K{=}2048$3. Ballot-free collection and count-cache optimizations ensure efficiency regardless of preIdx accuracy.

Practical and Theoretical Implications

GVR demonstrates that data-aware heuristics exploiting temporal correlation can yield substantial gains in exact Top-K selection without compromising correctness. Integrating GVR as a standard primitive in sparse-attention serving pipelines (TensorRT-LLM) enables significant end-to-end latency reduction, especially at long context—crucial for emerging agentic and code-centric workloads.

The approach is shown robust across distributional regimes, and because RoPE/YaRN underlies many sparse-attention LLMs, GVR’s principle is likely applicable across architectures where decode-phase Top-K exhibits temporal stability. For layers or models with weaker correlation, GVR gracefully degrades to parity with data-agnostic radix-select methods.

Future Directions

Prospective enhancements include:

• Adaptive dispatch for short sequence or batch-parallel regimes
• Extension of GVR to multi-CTA and ultra-long contexts
• Explore alternative prediction signals (prefill/MTP-aware)
• Cross-architecture validation (Ampere/Hopper/Blackwell+)
• Direct application to hierarchical indexer variants and other sparse-attention routing mechanisms

Practical deployment would benefit from further SMEM optimization and broader integration in batch-serving frameworks.

Conclusion

Guess-Verify-Refine establishes an exact, data-aware Top-K selection primitive for sparse attention decoding, leveraging temporal correlation to substantially lower both kernel and end-to-end inference latency without loss of accuracy. Its architectural and workload-informed optimizations deliver stable and robust speedup in realistic long-context LLM deployments, marking a step toward more efficient inference primitives tightly coupled to LLM decoder dynamics. Further exploration of more general signals and wider applicability is suggested by the results and the theoretical framework presented.