Pith. sign in

REVIEW 3 major objections 6 minor 1 cited by

Approximate Top-$k$ for Increased Parallelism

T0 review · 3 major / 6 minor · reviewed 2026-08-11 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read This paper argues that replacing exact top-k with a two-stage bucketed approximation unlocks substantial parallelism on ML accelerators, yielding 2-4x speed-ups (over 4x in sparse attention) with little to no downstream task degradation.

desk verdict A transparent, useful evaluation of bucketed approximate top-k that earns its design rules, with a real but narrow robustness gap around periodic input structure. read the letter →

arxiv 2412.04358 v1 pith:3MCNUBXJ submitted 2024-12-05 cs.LG

Oscar Key , Luka Ribar , Alberto Cattaneo , Luke Hudlass-Galley , Douglas Orr This is my paper · ORCID
classification cs.LG
keywords approximatetop-kbucketedselectionGPUparallelismsparseattentionlanguagemodelinferenceknowledgegraphlinkpredictionrecallspeedup
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Exact top-$k$ selection—finding the $k$ largest entries of a vector—does not parallelise well on GPUs because the $k$ largest values must be aggregated across the whole vector, forcing cooperation between threads. This paper argues that relaxing exactness to a bucketed approximation unlocks the parallelism that sparsity workloads need. The algorithm splits the input into $b$ interleaved buckets, independently selects the top $k_b$ values from each bucket, and optionally runs a final exact top-$k$. Using a uniform-distribution recall model and cost models, plus experiments on sparse attention for language models, vocabulary sampling, and knowledge-graph link prediction, the authors report top-$k$ speed-ups of 2–4×, and over 4× in the sparse-attention setting, with little to no loss in downstream accuracy. The design guidance that emerges: when $k$ is a large fraction of $n$, raise $k_b$ while keeping $b \cdot k_b = k$; when $k$ is small, add buckets instead.

What carries the argument

The object that carries the argument is the two-stage bucketed top-$k$: split the $n$ inputs into $b$ interleaved buckets, perform an exact top-$k_b$ within each bucket in parallel, and if the concatenated $b \cdot k_b$ candidates exceed $k$, run a final exact top-$k$ to select the $k$ largest. The quality model is the binomial recall bound of Equation (1), $\mathbb{E}[R(k,b,k_b)] = \frac{1}{k}\left(k_b + \sum_{i=k_b}^{k-1} F(k_b-1; i, 1/b)\right)$, which assumes top values are uniformly spread across buckets; Equation (2) gives the worst case when all top values concentrate in $\lceil bk/n\rceil$ buckets. The implementation keeps per-thread priority queues of size $k_b \le 4$ in registers, so buckets can be processed with no or minimal inter-thread communication.

What would settle it

Take a vector of length $n$ where the top $k$ entries are placed exclusively in positions congruent to one residue class modulo $b$ (so all top values fall in a single interleaved bucket), run the bucketed top-$k$ with given $b$ and $k_b$, and compare recall to Equation (1)'s prediction; the observed recall should drop to the worst-case level of Equation (2) if the uniform assumption is doing the work.

Watch

Extended reading notes

Core claim

The central claim is that bucketed approximate top-$k$ is a practical drop-in replacement for exact top-$k$ in the sparsity methods that machine learning accelerators actually run. The paper establishes this by characterising the algorithm's two design parameters—$b$, the number of buckets, and $k_b$, the number of candidates kept per bucket—through a binomial recall model that upper-bounds expected recall error, and through serial and parallel cost models that show large speed-ups are theoretically available. Empirically, on sparse attention in a large language model with $k = n/16$ to $n/8$, using $k_b = 2$ and $b \cdot k_b = k$ reduces top-$k$ cost by more than 4× with almost no degradation in task performance. On small-$k$ tasks, vocabulary sampling with $n = 128{,}256$, $k = 256$ and knowledge-graph link prediction with $n \approx 2.65$M, $k = 100$, speed-ups between 2× and 4× are achieved while recall error stays low. The paper's main caveat is that the bucket assignment must be interleaved, because real data are correlated along the sequence and contiguous assignment degrades recall.

Load-bearing premise

The whole quality story rests on the assumption that, once buckets are interleaved, each bucket receives about the same number of the true top-$k$ values, so the uniform-distribution recall bound of Equation (1) applies; if the largest values line up in a periodic or positional pattern that matches the bucket stride, recall degrades toward the worst-case bound and the reported 'little degradation' would break.

Editorial extensions

If this is right

  • For sparsity methods with $k$ proportional to $n$, a small per-bucket top-$k_b$ with $b \cdot k_b = k$ can cut top-$k$ cost by more than 4× while preserving downstream accuracy.
  • For small-$k$ settings such as vocabulary sampling and knowledge-graph link prediction, increasing the number of buckets $b$ (with $k_b = 1$) delivers 2–4× speed-ups with low recall error.
  • Using the bucketed approximation in place of exact top-$k$ inside sparse attention adds a further ~10% end-to-end generation speed-up at 40,000-token prompts, moving from 1.9× to 2.1× over dense attention.
  • The recall and cost models give practitioners a way to pick $b$ and $k_b$ for a target accuracy–speed trade-off, without rerunning the downstream task.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If interleaved bucket assignment breaks positional correlation as thoroughly as the correlated-Gaussian experiments suggest, the same bucketing trick could apply to other reduction primitives—argmax, top-$p$ sampling quantiles, or threshold selection—wherever accelerators need more parallelism.
  • The uniform-recall model predicts a concrete failure mode that the paper does not stress-test: inputs whose top values concentrate in one residue class modulo $b$ would fall toward the worst-case bound of Equation (2), so adversarial or periodic data is the regime to probe before deploying bucketed top-$k$.
  • Replacing exact top-$k$ with an approximation during training (as opposed to inference) may compound error across optimisation steps; the paper evaluates inference-time sparsity, so training-time use would need its own stability check.
  • In distributed settings, where exact top-$k$ requires cross-device communication, bucketing could have a larger advantage than the single-accelerator numbers here show; the paper itself flags this as future work.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 6 minor

Summary. The paper studies bucketed approximate top-k algorithms, in which the input is divided into interleaved buckets, a small top-kb is computed per bucket in parallel, and an optional exact top-k merges the candidates. It analyzes the recall/cost trade-off with a binomial model of recall (Appendix C) and abstract cost models (Appendix D), then validates the approach on SparQ attention for LLM inference, LLM vocabulary sampling, and knowledge-graph link prediction, reporting 2-4x speedups of the top-k operation with little downstream degradation. The authors release a CUDA/PyTorch implementation and provide end-to-end Llama 2 generation results.

Significance. If the central claim holds, the paper offers a practical and well-motivated alternative to exact top-k for sparsity workloads on parallel accelerators, with a clear design dichotomy (increase b for k << n, increase kb for k proportional to n). The release of a working PyTorch implementation is a concrete contribution, and the theoretical recall derivation is self-contained and machine-checkable in structure. The empirical coverage is respectable: three downstream tasks plus an end-to-end LLM experiment. The main weakness is that the 'little to no degradation' half of the claim rests on an implicit distributional assumption (top-k values spread across buckets) that is only tested on AR(1) correlated Gaussian data and on three real-world tasks, not on adversarial or periodic concentration patterns that the paper's own worst-case bound describes.

major comments (3)
  1. [Appendix A.1 and Appendix C] The claim that interleaved bucket assignment 'works very well in retaining recall, no matter the degree of correlation' is supported only by an AR(1) multivariate-normal simulation (Figure 6). Interleaving does not decorrelate inputs with periodic structure aligned with the bucket stride: if high scores occur at positions i ≡ r (mod b), the top-k values concentrate in few buckets and recall degrades toward the worst-case bound of Equation (2). Since the central claim of 'little to no degradation' depends on the uniform-spread assumption behind Equation (1), the paper should add stress tests with periodic or adversarially concentrated inputs (e.g., scores peaking at a fixed residue class) or explicitly restrict the claim to inputs without such alignment.
  2. [Appendix D and Figure 17] The cost-model comparison is not parameter-free: Figure 15 states the serial and parallel cost models 'have been aligned to match PyTorch at n = 2^10', and Figure 16 indicates the priority-queue models were 'shifted vertically for sake of visual tracking'. Consequently, the theoretical trade-off curves in Figure 3 and Appendix E are fits to a single hardware/software point, not first-principles predictions. The design conclusions are also tested empirically, so this is not fatal, but the paper should describe the theoretical contribution as model-based guidance and disclose which constants are fitted.
  3. [Section 4.1 and Appendix A] The runtime benchmark in Figure 1 selects, for each configuration of n, k, b, kb, the better of the two implementation modes ('we always select the mode... giving the best performance'). This is an oracle selection that ignores the cost and possible misprediction of the mode-selection heuristic described in Appendix A. Reported speedups therefore represent an upper bound on what a deployed implementation would achieve. The end-to-end result in Figure 5 mitigates this concern, but the paper should state clearly that Figure 1 is an oracle-mode benchmark.
minor comments (6)
  1. [Appendix B.1] The text says 'using PyTorch 3.12 and CUDA 12.1'; PyTorch version numbers are 2.x, so this is presumably Python 3.12 with a PyTorch 2.x release. Please clarify.
  2. [Figure 8] The y-axis label 'Contiguous assignment speedup' is ambiguous; it should state the baseline (presumably speedup over interleaved assignment) in the caption.
  3. [Appendix C, Equation (2)] The expression involves n/b, which is not necessarily an integer; the floor and modulo operations should be defined with explicit rounding or the analysis restricted to cases where b divides n.
  4. [Appendix D.2] The assumption 'If statements are free, but all branches are taken' is confusing, since a free if statement cannot simultaneously have a cost for taking a branch; please rephrase to clarify the intended accounting.
  5. [Abstract and Section 1] The statement that exact top-k requires aggregating 'the k largest values must be aggregated along the vector' is too absolute, since radix-select and other parallel exact algorithms exist (as the paper itself discusses). Suggest softening the motivation.
  6. [Appendix C] The sentence 'differently from the previous analysis where all buckets were always considered equiprobable... Equation (1) can still be seen as a upper bound' conflates the uniform-input assumption with the small-kb/n assumption; separating these two conditions would improve readability.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity: the recall and cost analyses are derived from stated assumptions, and the design conclusions are validated by independent empirical benchmarks rather than by construction.

full rationale

The derivation chain is self-contained. The recall model in Appendix C starts from an explicit uniform-distribution assumption ('As top values are equally likely to be contained in any of the buckets'), derives Equation (1) as a probability calculation, and separately gives the worst-case bound Equation (2) when that assumption is dropped. The statement that real-world recall 'is always sufficiently closely aligned with Equation (1)' is an empirical observation, not a fitted parameter used to produce the prediction. The cost models in Appendix D use operation counts with stated assumptions; Figure 15 notes the serial and parallel models 'have been aligned to match PyTorch at n = 2^10', which is a constant-scale calibration, not a fit of the design conclusion itself. The design guidance (prefer kb>1 for k proportional to n, prefer more buckets for k much smaller than n) is then checked against GPU benchmarks and downstream tasks, so it is not forced by construction. The authors' use of SparQ Attention (Ribar et al. 2024) is a same-group self-citation, but SparQ is a downstream application under test rather than the source of the bucketed top-k derivation; the central speed and quality claims are evaluated against exact torch.topk, PyTorch/RAFT baselines, SQuAD, a repetition task, LLM vocabulary sampling, and PharMeBINet. No equation or fitted quantity is renamed as a prediction, and no uniqueness theorem or prior-work ansatz is invoked to make the design choice forced. Therefore no significant circularity is present.

Assumptions & free parameters 3 free parameters · 4 assumptions · 0 invented entities

The central claim depends on the uniform-distribution recall model and on interleaving to make real data behave uniformly; the cost models use hand-chosen constants and one-point hardware alignment. These are modeling assumptions, not circular fits to downstream accuracy. No new entities are introduced.

free parameters (3)
  • Cost-model alignment constants = matched to PyTorch at n = 2^10; vertically shifted in Figure 16
    Used in Appendix D.3 to compare serial and parallel cost models with measured runtimes; the alignment makes theoretical cost axes hardware-dependent, though it does not affect downstream quality claims.
  • Operation cost weights = 1 for simple ops, 2 for read-modify-write
    Hand-chosen in Appendix D.2 for the abstract execution model; these weights influence which algorithm looks optimal in the serial and parallel cost models.
  • Mode-selection heuristic thresholds = use mode 1 if threads >= lanes or buckets < 64
    Hand-chosen heuristic in Appendix A; benchmarks additionally select the better mode per configuration, so reported speedups are not for a single fixed path.
assumptions (4)
  • domain assumption Input values are uniformly distributed across buckets so each bucket contains approximately the same number of top-k values.
    Used in Appendix C to derive expected recall Equation (1); real data is correlated, so the paper relies on interleaved assignment to restore this assumption.
  • domain assumption Interleaved bucket assignment breaks positional correlations in real inputs.
    Validated on correlated Gaussian data and on the paper's tasks in Appendix A.1, but not stress-tested against adversarial periodic patterns.
  • domain assumption Stage 1 and Stage 2 costs are additive in the cost model.
    Appendix D assumes additivity; Figure 17 notes the model underestimates Stage 2 cost, so this is an approximation.
  • domain assumption The abstract execution model's free operations do not change algorithmic rankings.
    Appendix D.2 assumes linear iteration, fixed-offset addressing, bounds checks, and if statements are free and that all branches are taken; these simplifications make hardware-independent analysis tractable but are not exact for real GPUs.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Approximate Top-$k$ for Increased Parallelism." pith.science (2026). https://pith.science/paper/3MCNUBXJ

@misc{pith2026241204358,
  author       = {Pith},
  title        = {Pith review of: Approximate Top-$k$ for Increased Parallelism},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/3MCNUBXJ}},
  note         = {Machine review of arXiv:2412.04358}
}
abstract

We present an evaluation of bucketed approximate top-$k$ algorithms. Computing top-$k$ exactly suffers from limited parallelism, because the $k$ largest values must be aggregated along the vector, thus is not well suited to computation on highly-parallel machine learning accelerators. By relaxing the requirement that the top-$k$ is exact, bucketed algorithms can dramatically increase the parallelism available by independently computing many smaller top-$k$ operations. We explore the design choices of this class of algorithms using both theoretical analysis and empirical evaluation on downstream tasks. Our motivating examples are sparsity algorithms for language models, which often use top-$k$ to select the most important parameters or activations. We also release a fast bucketed top-$k$ implementation for PyTorch.

Figures

Figures reproduced from arXiv: 2412.04358 by the authors.

Figure 1
Figure 1. Our approximate top-k implementation ( ×+), compared with exact top-k imple￾mentations from PyTorch and RAFT, and a bucketed top-k using torch.argmax, tested in float32 on an H100 PCIe GPU with batch size m = 128. Total bandwidth is the minimum number of bytes transferred by top-k, divided by runtime. Left: Small fixed k = 64; it is faster to retrieve kb = 1 element per bucket, varying the total number b · kb of ele… view at source ↗
Figure 2
Figure 2. Left: An example of a bucketed top-k, with n = 11, k = 4, b = 3 and kb = 2. In Stage 1, n elements are reduced to b · kb elements via b independent top-kb. An optional Stage 2 takes final top-k. Right: The trade-off between top-k runtime duration and downstream task accuracy for SparQ Attention in SQuAD and a sequence repetition task (see Appendix B.2), when using different bucketed top-k settings with batch size m … view at source ↗
Figure 3
Figure 3. Theoretical trade-off curves, using the serial cost model, which computes the count of all [PITH_FULL_IMAGE:figures/full_fig_p003_3.png] view at source ↗
Figures from the paper (17 more)
Figure 4
Figure 4. Figure 4: Bucketed top-k trade-off for LLM vocabulary sampling (left, n = 128,256, k = 256, m = 64) and for knowledge graph link prediction (right, n = 2,653,751, k = 100, m = 128). In both regimes, kb = 1 gives peak performance, but kb = 2 sacrifices some speed for sake of a lo…
Figure 5
Figure 5. Figure 5: End-to-end speed-ups achieved when generating text from Llama 2 7B using SparQ [PITH_FULL_IMAGE:figures/full_fig_p005_5.png]
Figure 6
Figure 6. Figure 6: Comparison of achieved recall vs. exact top- [PITH_FULL_IMAGE:figures/full_fig_p009_6.png]
Figure 7
Figure 7. Figure 7: Comparison of LLM downstream task performance using interleaved and contiguous [PITH_FULL_IMAGE:figures/full_fig_p009_7.png]
Figure 8
Figure 8. Figure 8: Runtime comparison of contiguous vs interleaved assignment on SQuAD and Repetition [PITH_FULL_IMAGE:figures/full_fig_p009_8.png]
Figure 9
Figure 9. Figure 9: Bandwidth comparison with CUDA graphs enabled, in the same configurations used for [PITH_FULL_IMAGE:figures/full_fig_p011_9.png]
Figure 10
Figure 10. Figure 10: Bandwidth comparison for bfloat16 values, in the same configurations used for [PITH_FULL_IMAGE:figures/full_fig_p011_10.png]
Figure 11
Figure 11. Figure 11: PRIORITYQUEUE (serial), with insertion sort, cost m · n · (3k − 1) + O(m). The first sort can be merged into the loop. def topk_radix_select(data, k): # Find kth value kth_value, mask, count_gt = 0, 0, 0 for r in range(31, -1, -1): # * log(n) r_mask = 1 << r kth_value…
Figure 12
Figure 12. Figure 12: RADIXSELECT (serial), cost m · n · (4 log2 n + 4) + O(m log n). Note that we assume a key length of log2 n, to uniquely identify n elements; in practical scenarios, the key length is separate from n. If the k th element may be tied, a second “collect” step may be nece…
Figure 13
Figure 13. Figure 13: SCANMAX (parallel), cost k · (2 log2 n + 3) + O(1). def scan_cumsum(data): s = data.copy() # +1 for i in range(log2(len(s))): # * log(n) s = [s[j] + (s[j-2**i] if j-2**i >= 0 else 0) # | +2 for j in range(len(s))] return s def topk_radix_select_parallel(data, k): # Fi…
Figure 14
Figure 14. Figure 14: RADIXSELECT (parallel), cost log2 n · (2 log2 n + 16) + O(1). 15 [PITH_FULL_IMAGE:figures/full_fig_p015_14.png]
Figure 15
Figure 15. Figure 15: PyTorch top-k runtime as a function of n, k. Left: Batch size m = 1. Right: m = 256. The serial and parallel cost models have been aligned to match PyTorch at n = 210. Practical runtime does not depend strongly on k, and follows the parallel cost model when batch size…
Figure 16
Figure 16. Figure 16: Our priority queue top-k runtime as a function of n, k. Left: Batch size m = 1. Right: m = 256. The serial and parallel cost models have been shifted vertically for sake of visual tracking, since they do not directly predict wall-clock time. Runtime scaling with k is …
Figure 17
Figure 17. Figure 17: A comparison of the relative runtime of our approximate top- [PITH_FULL_IMAGE:figures/full_fig_p018_17.png]
Figure 18
Figure 18. Figure 18: Theoretical trade-off curves under the basic cost model, as [PITH_FULL_IMAGE:figures/full_fig_p019_18.png]
Figure 19
Figure 19. Figure 19: Theoretical trade-off curves under the serial cost model, as [PITH_FULL_IMAGE:figures/full_fig_p020_19.png]
Figure 20
Figure 20. Figure 20: Theoretical trade-off curves under the parallel cost model, as [PITH_FULL_IMAGE:figures/full_fig_p021_20.png]

Discussion (0). Continue with ORCID to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. Accelerating Linear Recurrent Neural Networks for the Edge with Unstructured Sparsity

    cs.LG 2025-02 conditional novelty 6.0 of 10

    Sparse, 8-bit quantized S5 linear RNNs match dense model audio denoising accuracy with 2x less compute and 36% less memory, and run 42x faster with 149x lower energy on Loihi 2 than a dense FP32 model on Jetson Orin Nano.

Reference graph

Works this paper leans on

3 extracted references · 2 canonical work pages · cited by 1 Pith paper

  1. [1]

    Define tensor x of shape (m, n)

    work page

  2. [2]

    values from a unit normal distribution (b) Launch top- k kernel

    Warmup loop ( 16 iterations): (a) Fill x with i.i.d. values from a unit normal distribution (b) Launch top- k kernel

    work page

  3. [3]

    gpt-fast

    Timing loop ( 512 iterations): (a) Fill x with i.i.d. values from a unit normal distribution (b) Launch sleep kernel (c) Record start CUDA event (d) Launch top- k kernel (e) Record stop CUDA event By launching the sleep kernel within the timing loop we ensure the start/stop events and the top-k kernel have all been queued before the first start event is e...

    2017

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.