B200-specific highlights in fast_topk_clusters_exact
This version separates what is truly B200/Blackwell-specific in current FlashInfer dispatch from what is a generic radix/top-k optimization or a general SM90+ cluster capability.
B200 / Blackwell-specific in current dispatch
SM90+ cluster feature
Generic exact top-k optimization
Current dispatch gate: compute capability major == 10
Pipeline with B200-specific parts highlighted
1
Blackwell-gated pathPython chooses clusters exact only when compute capability major is 10, non-deterministic, and graph-safe mode is off.B200-specific dispatch2
Cluster count policyBatch-size based cluster count: 8, 4, 2, or 1 clusters per row. Disabled to 1 cluster for short rows < 8192.B200-tuned policy3
Thread-block clustersUses CUDA cluster cooperative groups: multiple CTAs cooperate on one row.SM90+ capability4
Distributed shared memoryUses cluster shared-memory mapping to combine CTA histograms without global-memory reduction.SM90+ capability5
Boundary cache refineFull scan once, cache threshold-bin candidates, then refine only cached candidates byte by byte.generic algorithmic win6
Exact outputWrite all values above pivot plus enough equal-pivot ties. This is exact top-k, not approximate top-k.exact top-k
Most of the math is still exact radix selection. The B200-specific part is mainly the current Python dispatch gate and the tuning/policy around the cluster exact fast path.
What is B200-specific?
| Component | Classification | Why |
|---|
cap[0] == 10 dispatch | B200 / Blackwell-specific | The Python API only auto-selects clusters exact when compute capability major is 10. |
get_fast_topk_clusters() | B200-tuned policy | Allocates more clusters for low batch size to increase per-row parallelism; short rows force 1 cluster. |
| Thread-block clusters | SM90+ feature | CUDA cluster cooperative groups require SM90+, so not exclusively B200, but used by the B200 path. |
| Distributed shared memory | SM90+ feature | cluster.map_shared_rank lets CTAs read peer CTA shared memory inside the cluster. |
| Large shared-memory carveout | hardware-sensitive | The kernel trades shared memory for fewer global-memory accesses and larger candidate cache. |
| Radix pivot selection | Generic | Same exact top-k idea: find the pivot and collect winners. |
| Boundary candidate cache | Generic algorithmic optimization | Full-row scan once, then refine only threshold-bin candidates. |
| Page-table/ragged fused transform | Pipeline optimization | Useful for sparse attention regardless of whether the top-k kernel is B200-gated. |
Dispatch logic to remember
use clusters exact top-k only if:
not dsa_graph_safe
not deterministic
FLASHINFER_TOPK_ALGO is unset or "clusters"
compute capability major == 10
So for a slide, the safest wording is: “FlashInfer has a Blackwell/B200-gated clusters exact top-k fast path.” Avoid saying the generic radix top-k kernel itself is B200-specific.
How to explain it in one slide
Do not say
- “Radix top-k is a B200-specific algorithm.”
- “The result is approximate but faster.”
- “All cluster optimizations are Blackwell-only.”
Better wording
- “The public FlashInfer top-k API can select a Blackwell/B200-gated clusters exact fast path.”
- “The fast path uses SM90+ cluster cooperative groups and distributed shared memory.”
- “It remains exact top-k; tie choice may be non-deterministic.”