top-k, sort, and the ops that are not GEMM

Because GPUs are throughput machines for regular work, and every property of selection is irregular.

• No arithmetic intensity. You read $n$ values and perform comparisons. There is nothing to amortise the memory traffic against, so you are pinned at the memory roof from the start — 1 MB of logits at 3 TB/s is ~0.3 µs of pure bandwidth, and any real implementation is far above that floor, which tells you the cost is structure, not data movement.
• A serial dependency. Selection is fundamentally a reduction, and a reduction’s depth is at best $\log n$ — 17 dependent levels for 128k. Heap-based partial sorts are worse: inherently serial, branch-heavy, and a poor fit for SIMT execution.
• Data-dependent control flow. Whether an element survives depends on its value, so lanes in a warp diverge. Compaction requires a prefix sum and a scatter, which is another reduction plus irregular writes.
• Small output, large input. Most of the machine finishes early and idles while the final reduction levels play out with almost no parallelism left.

The consequence to internalise: FLOP counting is useless here. A top-$k$ over 128k values does essentially no arithmetic and can cost more wall-clock than a matrix multiply doing billions of operations, because its cost is set by dependency depth and irregularity. Any performance model built on FLOPs will mispredict it by orders of magnitude, and that is precisely why it was overlooked until it started showing up in profiles.

Radix select finds the $k$-th largest value by examining the bits of the representation from most significant to least, and it never orders anything it does not have to.

Take 8 bits at a time. Histogram all $n$ values by their top 8 bits — 256 buckets, one atomic-add pass. Scan the histogram from the high bucket down, accumulating counts, until you find the bucket containing the $k$-th element. Every value in higher buckets is definitely in the top $k$; everything in lower buckets is definitely out. Recurse into the boundary bucket with the next 8 bits, now over a much smaller candidate set.

Why it fits the machine:

• Each pass is a histogram, which is a well-understood, coalesced, high-throughput GPU primitive.
• Passes are bounded. 32-bit floats at 8 bits per pass is at most 4 passes, and typically 2 in practice because the candidate set collapses immediately.
• No data movement. You are computing a threshold, not rearranging elements. The final step is one predicated pass extracting values above it.
• Order-preserving keys. Floats can be monotonically mapped to unsigned integers with a couple of bit operations, so radix ordering matches numeric ordering exactly.

Against $O(n \log n)$ sorting: sorting computes the full permutation of 128,000 elements, moving data $\log n$ times, to answer a question about one boundary value. Radix select does a handful of histogram passes and one filter, and it is the difference between ~100–500 µs and ~20–60 µs.

The transferable principle is the one this module is really about: solve the problem you have. Sorting is a strictly harder problem that happens to contain the answer, and reaching for it is a habit from a serial world where the constant factors made it not matter.

By reframing it as finding a threshold, then using rejection sampling — because the sorted order was never actually part of the answer.

The definition of top-$p$ is: the smallest set of tokens whose cumulative probability exceeds $p$. That reads like it needs a sorted cumulative sum, but what it determines is a probability cutoff $\tau$ such that the tokens with $p_i \ge \tau$ are exactly that set. Given $\tau$, the sorted order is irrelevant — you sample from the renormalised distribution restricted to those tokens.

So the loop is:

1. Guess $\tau$, from the max and a running estimate.
2. Two reductions compute the mass above it and the count, both fully parallel.
3. Adjust by bisection or a fixed-point step.
4. Sample from the masked distribution by inverse-CDF over a prefix sum, or by rejection.

A handful of parallel reductions, no ordering, no data movement. This is the essence of FlashInfer’s sorting-free sampling, and it composes: top-$k$ and top-$p$ together become a joint threshold search, and rejection sampling lets you draw the token and verify the constraint in the same kernel.

Two properties worth appreciating. It is exact, not an approximation — the rejection step guarantees the sample comes from the correct distribution, so this is not a quality trade. And it is shape-static: the work does not depend on how many tokens pass the filter, so unlike a compaction approach it stays CUDA-graphable (Module 9), which matters enormously since sampling sits inside the decode loop you spent Part L making capturable.

Everywhere the workload stops looking like dense linear algebra, and the pattern is always the same: an op with low intensity, irregular access, or data-dependent output size sitting on a critical path that everything else has been optimised off.

Instances worth recognising:

• MoE routing (Module 11) is a top-$k$ over experts per token, plus a scatter — the same op, and it feeds a data-dependent shape.
• Speculative decoding compares draft and target distributions per position; done naively it materialises and sorts more than it needs.
• Beam search is top-$k$ over beams × vocabulary — a much larger selection, done every step.
• Sparse and long-context attention selects top-$k$ blocks or keys, so the selection sits inside the attention kernel.
• Embedding gather/scatter in recommenders is irregular memory access with no arithmetic at all, and is usually the dominant cost of those models.
• Sequence packing and unpadding are prefix sums and scatters — cheap in FLOPs, awkward on a GPU.

The general rule, and it is the useful takeaway: once you have optimised the GEMMs, your bottleneck moves to the ops that are not GEMMs, and none of your GEMM intuitions transfer. FLOP counts do not predict their cost, arithmetic-intensity reasoning gives a floor that is far below reality, and the compiler will not fix them — Inductor fuses elementwise chains well and has essentially nothing to say about a selection algorithm.

Which is why these are the ops that end up hand-written in Triton or CUTLASS (Module 7), and why libraries like FlashInfer exist at all. They are the residue that the general machinery cannot reach.