Radix Top‑K: find the boundary, not the full order

FlashInfer-style radix top-k converts scores into order-preserving integer keys, finds the k-th largest pivot byte by byte, then collects all values above the pivot plus enough ties.

Example: float32 scores · K = 5

Big idea for beginners

Top‑K is like choosing finalists, not ranking everyone.

If you only need the best 5 students, you do not need to sort all 1,000 students from best to worst. You only need to know the cutoff score.

The pivot is the cutoff line.

For top‑5, the pivot is the 5th largest score. Everyone above it is definitely selected. People exactly on the cutoff are ties.

Radix means “look at digits from left to right.”

Instead of comparing full numbers again and again, radix top‑k asks: which high digit contains the cutoff? Then which next digit? Then the next.

Histogram means “count, do not sort.”

A 256-bin histogram counts how many values fall into each byte bucket. Counting is usually cheaper than sorting.

The insight: sorting answers too much. Top‑K only needs the boundary between selected and discarded items.

Toy example: school scores

Find top‑3 scores without fully sorting everyone.

A
99
B
92
C
91
D
88
E
76
F
65
G
40
Tens digit
90s: 3 scores 80s: 1 70s: 1 60s: 1 40s: 1
Ones digit
99 92 91 ← pivot
K = 3
pivot = 91
selected = scores > 91 plus enough scores = 91

Vocabulary before reading the CUDA version

BucketA group of values that share the same current byte/digit.
PrefixThe high bytes already chosen. It narrows the candidate set.
remaining_kHow many more winners we still need inside the current bucket.
PivotThe final cutoff value: the k‑th largest score.
In FlashInfer, the “digits” are not decimal digits like 90s or 80s. They are bytes of an ordered integer key. For float32, there are 4 byte rounds. For fp16/bf16, there are only 2 byte rounds.

Algorithm mental model

Read this as a “cutoff finder.” The algorithm keeps asking narrower questions until it knows the exact cutoff value.

1
Map float → ordered integerAfter mapping, larger score means larger unsigned key. For float32 this gives a 32-bit key; for fp16/bf16 it is 16-bit.
2
Radix-select the pivotBuild 256-bin histograms over one byte at a time. Use suffix sums to locate the bucket that contains the k-th largest element.
3
Narrow by prefixAfter each round, keep only values matching the selected high-bit prefix. Update remaining_k by subtracting values already greater than the selected bucket.
4
Collect resultsAll keys greater than pivot are definitely selected. Keys equal to pivot fill the remaining slots. Keys below pivot are discarded.
count_ge(bucket) = #values with byte ≥ bucket
count_gt(bucket) = #values with byte > bucket
choose bucket where: count_ge ≥ remaining_k and count_gt < remaining_k

This is why radix top-k can avoid sorting. It only needs the threshold value and a final collection pass.

Real example: top‑5 from 12 scores

> pivot: selected = pivot: tie region < pivot: discard
Id Score Float bits Ordered key Status

Radix rounds: 4 bytes for float32

Each round finds one byte of the final pivot key. The selected bucket is the one containing the current k-th largest candidate.

Final pivot and collection

After the pivot is known, the rest is simple: take everything above the cutoff, then fill the remaining slots from the tie group.

Output policy for ties

When multiple values equal the pivot, the fast path may fill the remaining slot with any equal-pivot element. Deterministic tie-break makes that choice stable.