Top‑K Methods Comparison

A compact comparison of common top‑k strategies, including whether each method returns the real top‑k set or an approximate top‑k set.

Focus: LLM sampling & sparse attention

At a glance

Method	Real top‑k?	Exact?	Core idea	Best for	Main weakness
Full sort	Real	Yes	Sort everything, then take first `k`	Need full ranking	Does more work than top‑k needs
Heap / partial select	Real	Yes	Keep only the best `k` seen so far	CPU-style selection, small/medium `k`	Comparison-heavy, less GPU friendly
torch.topk	Real	Yes	General-purpose library top‑k	Baseline, compatibility	Not specialized for LLM / sparse attention
Radix top‑k	Real	Yes	Find the pivot byte-by-byte, then collect winners	Large vocab / long score rows	Tie handling and implementation complexity
FlashInfer Multi‑CTA Radix	Real	Yes	Split one long row across CTAs	Very long rows	Needs inter-CTA synchronization
FlashInfer FilteredTopK	Real*	Yes*	Coarse histogram → candidate filter → refine	When candidate bucket stays small	Can overflow on tie-heavy workloads; fallback preserves exactness
FlashInfer clusters exact	Real	Yes	Cluster cooperative-groups exact top‑k	Supported newer GPUs, fast path	Hardware / mode restricted
Fused transform top‑k	Real	Yes	Top‑k + page-table/ragged transform	Sparse attention	Only useful when next stage needs transform

Real top‑k? means the algorithm returns the true top‑k set, not an approximate neighbor/candidate set. Exact? means the returned set is mathematically exact for the input values. Output order and equal-value tie choice may still be non-deterministic unless deterministic tie-breaking is enabled.

Beginner mental model

Full sortRank every student from 1st to last, then keep the top 5. Easy, but wasteful.

HeapCarry a small “best 5 so far” box while scanning all students one by one.

Radix top‑kDo not rank everyone. Just find the cutoff score, then keep everyone above the cutoff plus enough ties.

Fused sparse-attention top‑kAfter choosing the winners, immediately convert them into the exact index format the next kernel needs.

Need top‑k?
You do NOT need the full order.
You only need the pivot = k-th largest value.

Method cards

Full sort

Simple and exact. Best when you truly need all items sorted, not just the top‑k.

Heap / partial select

Classical exact selection strategy. Good on CPU, but not the most natural fit for GPU parallelism.

torch.topk

The easy exact baseline. Great default choice if your shape is small or you want a standard implementation.

Radix top‑k

Exact method: transforms floats to ordered integer keys, then finds the pivot one byte at a time.

FlashInfer Multi‑CTA Radix

Exact long-row radix top‑k using several CTAs for one row.

FlashInfer FilteredTopK

Exact top‑k set with a fast candidate-filter path and fallback for overflow/tie-heavy cases.

FlashInfer clusters exact

Exact specialized fast path using CUDA cluster cooperative groups on supported hardware.

Fused transform top‑k

Exact top‑k plus exact index transform for sparse attention.

Qualitative comparison

Real top‑k set, not approximate top‑k

Full sort

exact

Heap / partial select

exact

torch.topk

exact

Radix top‑k

exact

Multi‑CTA Radix

exact

FilteredTopK

exact*

Clusters exact

exact

Fused transform top‑k

exact

Exactness is about the selected set. Sorted order and equal-value tie order are separate.

GPU friendliness

Full sort

medium

Heap / partial select

lower

torch.topk

baseline

Radix top‑k

high

Multi‑CTA Radix

high

FilteredTopK

very high

Clusters exact

very high

Fused transform top‑k

high+

Implementation simplicity

Full sort

easy

Heap / partial select

easy

torch.topk

easiest

Radix top‑k

harder

Multi‑CTA Radix

complex

FilteredTopK

complex

Clusters exact

specialized

Fused transform top‑k

complex

Recommended choices

Use torch.topk when

You want the simplest exact baseline.
Your tensors are not huge.
Performance is not the main bottleneck.
You want broad framework compatibility.

Use radix top‑k when

Rows are large: big vocab or long context.
You care about GPU efficiency.
You only need top‑k, not full sorting.
Sampling or index selection is performance-sensitive.

Use Multi‑CTA / FilteredTopK when

You are inside a performance-tuned inference stack.
Rows are very long or top‑k is frequent.
You want a specialized FlashInfer path.
You can manage complexity and edge cases.

Use fused transform top‑k when

Top‑k feeds sparse attention directly.
You need page-table or ragged index conversion.
You want fewer kernel launches.
You want to reduce extra memory traffic.

Rule of thumb:
Small / simple case → torch.topk
Large row top‑k → radix top‑k
Long-row specialized path → Multi‑CTA or FilteredTopK
Sparse attention pipeline → fused transform top‑k

One-sentence summary of each method

Method	Real top‑k?	Exact?	One-sentence summary	Verdict
Full sort	Real	Yes	Correct and easy, but does more work than top‑k requires.	usually overkill
Heap / partial select	Real	Yes	A classical exact top‑k method, better matched to CPU-style selection than GPU bulk parallelism.	situational
torch.topk	Real	Yes	The most practical exact default baseline.	great default
Radix top‑k	Real	Yes	Excellent when you want the exact cutoff fast without sorting everything.	LLM-friendly
Multi‑CTA Radix	Real	Yes	The robust exact choice for very long rows when one CTA is not enough.	long-row specialist
FilteredTopK	Real*	Yes*	Exact top‑k with a fast filter/refine path, but more sensitive to ties/overflow.	fast path specialist
Clusters exact	Real	Yes	A hardware-specialized exact path that can be extremely fast on supported GPUs.	platform-dependent
Fused transform top‑k	Real	Yes	Exact top‑k plus transform, best when top‑k feeds sparse-attention indices.	pipeline winner

* FilteredTopK: exact top‑k set when its fallback/overflow logic is used correctly. Tie-heavy or quantized inputs may hurt performance, not intended correctness.

Source-level correctness check

FlashInfer path	Real top‑k?	Why
Radix / Multi‑CTA Radix	Real exact top‑k	Finds the exact k-th pivot, writes all values greater than the pivot, then writes enough values equal to the pivot.
FilteredTopK	Real exact top‑k*	Uses a coarse bucket only as a filter, then refines the threshold. Overflow/tie-heavy cases trigger fallback/rebuild logic for correctness.
Clusters exact	Real exact top‑k	The implementation is explicitly the exact cluster path; final selected indices are written from the exact threshold/refinement result.
Fused page-table / ragged transform	Real exact top‑k + exact transform	The top‑k selection is exact; the epilogue only maps selected indices through a page table or adds a ragged offset.

Important distinction: non-deterministic tie choice is not approximate top‑k. If many values are exactly equal to the pivot, any subset of the tied values can still be a valid exact top‑k set. Deterministic mode only makes the tie choice stable.