How tencent-hunyuan/flashtopk works
FlashTopK is a GPU kernel library for exact, deterministic top-k selection over per-query key windows in sparse attention. The primary workflow is a caller invoking the selection API, which plans and launches architecture-specific scoring; optional adaptive pruning bounds keys before refinement, and candidate selection produces the result. Recovery rescans rows that overflow candidates or return too few keys. The README also documents a PyTorch custom-op entry point and packed-sequence window helper. Kernel internals are only partially sampled, so edges involving them rely on indexed references or documented behavior.
Main entry points: launch.py Kernel launcher · topk.py Selector workspace · recovery.py Dense recovery · scan_sm90.py SM90 scan
Components 19
Public API
- Select API
flashtopk/topk.py - Torch custom op
flashtopk/ops.py - Selector workspace
flashtopk/topk.py
Selection control
- Adaptive planner
flashtopk/plan.py - Kernel launcher
flashtopk/kernels/launch.py
GPU kernels
- Gate prologue
flashtopk/kernels/prologue.py - Seed preparation
flashtopk/kernels/seed_prep.py - SM90 scan
flashtopk/kernels/scan_sm90.py - SM100 scan
flashtopk/kernels/scan_sm100.py - SM90 pruning
flashtopk/kernels/prune_sm90.py - SM100 pruning
flashtopk/kernels/prune_sm100.py - SM90 refinement
flashtopk/kernels/refine_sm90.py - SM100 refinement
flashtopk/kernels/refine_sm100.py - Candidate selection
flashtopk/kernels/selection.py - Kernel primitives
flashtopk/kernels/common.py - SIMT helpers
flashtopk/kernels/simt.py
Exactness and recovery
- Recovery triage
flashtopk/kernels/triage.py - Dense recovery
flashtopk/recovery.py
Other components
- Library caller
Connections 33
- Library caller → Select API: calls
README.md - Library caller → Torch custom op: invokes
README.md - Select API → Selector workspace: runs
flashtopk/topk.py - Torch custom op → Select API: dispatches
flashtopk/ops.py - Selector workspace → Adaptive planner: plans
flashtopk/topk.py - Selector workspace → Kernel launcher: launches
flashtopk/topk.py - Selector workspace → Dense recovery: recovers rows
flashtopk/topk.py - Selector workspace → Candidate selection: selects candidates
flashtopk/topk.py - Adaptive planner → Kernel launcher: plans kernels
flashtopk/plan.py - Adaptive planner → Gate prologue: uses
flashtopk/plan.py - Kernel launcher → Gate prologue: launches
flashtopk/kernels/launch.py - Kernel launcher → Seed preparation: launches
flashtopk/kernels/launch.py - Kernel launcher → SM90 scan: dispatches
flashtopk/kernels/launch.py - Kernel launcher → SM100 scan: dispatches
flashtopk/kernels/launch.py - Kernel launcher → SM90 pruning: dispatches
flashtopk/kernels/launch.py - Kernel launcher → SM100 pruning: dispatches
flashtopk/kernels/launch.py - Kernel launcher → SM90 refinement: dispatches
flashtopk/kernels/launch.py - Kernel launcher → SM100 refinement: dispatches
flashtopk/kernels/launch.py - Kernel launcher → Candidate selection: launches
flashtopk/kernels/launch.py - Kernel launcher → Recovery triage: launches
flashtopk/kernels/launch.py - SM90 scan → Kernel primitives: uses
flashtopk/kernels/scan_sm90.py - SM100 scan → Kernel primitives: uses
flashtopk/kernels/scan_sm100.py - SM90 pruning → SM90 scan: shares primitives
flashtopk/kernels/prune_sm90.py - SM100 pruning → SM100 scan: shares primitives
flashtopk/kernels/prune_sm100.py - SM90 refinement → SM90 scan: shares primitives
flashtopk/kernels/refine_sm90.py - SM100 refinement → SM100 scan: shares primitives
flashtopk/kernels/refine_sm100.py - Dense recovery → Recovery triage: uses
flashtopk/recovery.py - Dense recovery → Kernel launcher: scores densely
flashtopk/recovery.py - Dense recovery → Candidate selection: reselects rows
flashtopk/recovery.py - Gate prologue → SIMT helpers: uses
flashtopk/kernels/prologue.py - Seed preparation → SIMT helpers: uses
flashtopk/kernels/seed_prep.py - Candidate selection → SIMT helpers: uses
flashtopk/kernels/selection.py - Recovery triage → SIMT helpers: uses
flashtopk/kernels/triage.py