# tencent-hunyuan/flashtopk architecture

> Architecture diagram of the GitHub repository tencent-hunyuan/flashtopk, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-11-15). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/tencent-hunyuan/flashtopk
- Repository: https://github.com/tencent-hunyuan/flashtopk

## Overview

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: [`flashtopk/kernels/launch.py`](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/launch.py) (Kernel launcher), [`flashtopk/topk.py`](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/topk.py) (Selector workspace), [`flashtopk/recovery.py`](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/recovery.py) (Dense recovery), [`flashtopk/kernels/scan_sm90.py`](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/scan_sm90.py) (SM90 scan)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_api["Public API"]
  node_select["Select API<br/>[topk.py]"]
  node_torchop["Torch custom op<br/>[ops.py]"]
  node_selector["Selector workspace<br/>[topk.py]"]
end

subgraph group_control["Selection control"]
  node_planner["Adaptive planner<br/>[plan.py]"]
  node_launcher["Kernel launcher<br/>[launch.py]"]
end

subgraph group_gpu["GPU kernels"]
  node_prologue["Gate prologue<br/>[prologue.py]"]
  node_seed["Seed preparation<br/>[seed_prep.py]"]
  node_scan90["SM90 scan<br/>[scan_sm90.py]"]
  node_scan100["SM100 scan<br/>[scan_sm100.py]"]
  node_prune90["SM90 pruning<br/>[prune_sm90.py]"]
  node_prune100["SM100 pruning<br/>[prune_sm100.py]"]
  node_refine90["SM90 refinement<br/>[refine_sm90.py]"]
  node_refine100["SM100 refinement<br/>[refine_sm100.py]"]
  node_selection["Candidate selection<br/>[selection.py]"]
  node_common["Kernel primitives<br/>[common.py]"]
  node_simt["SIMT helpers<br/>[simt.py]"]
end

subgraph group_recovery["Exactness and recovery"]
  node_triage["Recovery triage<br/>[triage.py]"]
  node_recovery["Dense recovery<br/>[recovery.py]"]
end

node_caller(("Library caller"))

node_caller -->|"calls"| node_select
node_caller -->|"invokes"| node_torchop
node_select -->|"runs"| node_selector
node_torchop -->|"dispatches"| node_select
node_selector -->|"plans"| node_planner
node_selector -->|"launches"| node_launcher
node_selector -->|"recovers rows"| node_recovery
node_selector -->|"selects candidates"| node_selection
node_planner -->|"plans kernels"| node_launcher
node_planner -->|"uses"| node_prologue
node_launcher -->|"launches"| node_prologue
node_launcher -->|"launches"| node_seed
node_launcher -->|"dispatches"| node_scan90
node_launcher -->|"dispatches"| node_scan100
node_launcher -->|"dispatches"| node_prune90
node_launcher -->|"dispatches"| node_prune100
node_launcher -->|"dispatches"| node_refine90
node_launcher -->|"dispatches"| node_refine100
node_launcher -->|"launches"| node_selection
node_launcher -->|"launches"| node_triage
node_scan90 -->|"uses"| node_common
node_scan100 -->|"uses"| node_common
node_prune90 -->|"shares primitives"| node_scan90
node_prune100 -->|"shares primitives"| node_scan100
node_refine90 -->|"shares primitives"| node_scan90
node_refine100 -->|"shares primitives"| node_scan100
node_recovery -->|"uses"| node_triage
node_recovery -->|"scores densely"| node_launcher
node_recovery -->|"reselects rows"| node_selection
node_prologue -->|"uses"| node_simt
node_seed -->|"uses"| node_simt
node_selection -->|"uses"| node_simt
node_triage -->|"uses"| node_simt

click node_select "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/topk.py"
click node_torchop "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/ops.py"
click node_selector "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/topk.py"
click node_planner "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/plan.py"
click node_launcher "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/launch.py"
click node_prologue "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/prologue.py"
click node_seed "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/seed_prep.py"
click node_scan90 "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/scan_sm90.py"
click node_scan100 "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/scan_sm100.py"
click node_prune90 "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/prune_sm90.py"
click node_prune100 "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/prune_sm100.py"
click node_refine90 "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/refine_sm90.py"
click node_refine100 "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/refine_sm100.py"
click node_selection "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/selection.py"
click node_triage "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/triage.py"
click node_recovery "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/recovery.py"
click node_common "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/common.py"
click node_simt "https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/simt.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_select,node_torchop,node_selector toneBlue
class node_planner,node_launcher toneAmber
class node_prologue,node_seed,node_scan90,node_scan100,node_prune90,node_prune100,node_refine90,node_refine100,node_selection,node_common,node_simt toneMint
class node_triage,node_recovery toneRose
class node_caller toneIndigo
```

## Components

### Public API

- [Select API](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/topk.py): `flashtopk/topk.py`
- [Torch custom op](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/ops.py): `flashtopk/ops.py`
- [Selector workspace](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/topk.py): `flashtopk/topk.py`

### Selection control

- [Adaptive planner](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/plan.py): `flashtopk/plan.py`
- [Kernel launcher](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/launch.py): `flashtopk/kernels/launch.py`

### GPU kernels

- [Gate prologue](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/prologue.py): `flashtopk/kernels/prologue.py`
- [Seed preparation](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/seed_prep.py): `flashtopk/kernels/seed_prep.py`
- [SM90 scan](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/scan_sm90.py): `flashtopk/kernels/scan_sm90.py`
- [SM100 scan](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/scan_sm100.py): `flashtopk/kernels/scan_sm100.py`
- [SM90 pruning](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/prune_sm90.py): `flashtopk/kernels/prune_sm90.py`
- [SM100 pruning](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/prune_sm100.py): `flashtopk/kernels/prune_sm100.py`
- [SM90 refinement](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/refine_sm90.py): `flashtopk/kernels/refine_sm90.py`
- [SM100 refinement](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/refine_sm100.py): `flashtopk/kernels/refine_sm100.py`
- [Candidate selection](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/selection.py): `flashtopk/kernels/selection.py`
- [Kernel primitives](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/common.py): `flashtopk/kernels/common.py`
- [SIMT helpers](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/simt.py): `flashtopk/kernels/simt.py`

### Exactness and recovery

- [Recovery triage](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/kernels/triage.py): `flashtopk/kernels/triage.py`
- [Dense recovery](https://github.com/tencent-hunyuan/flashtopk/blob/main/flashtopk/recovery.py): `flashtopk/recovery.py`

### Other components

- **Library caller**

## Connections

- Library caller → Select API: calls (evidence: `README.md`)
- Library caller → Torch custom op: invokes (evidence: `README.md`)
- Select API → Selector workspace: runs (evidence: `flashtopk/topk.py`)
- Torch custom op → Select API: dispatches (evidence: `flashtopk/ops.py`)
- Selector workspace → Adaptive planner: plans (evidence: `flashtopk/topk.py`)
- Selector workspace → Kernel launcher: launches (evidence: `flashtopk/topk.py`)
- Selector workspace → Dense recovery: recovers rows (evidence: `flashtopk/topk.py`)
- Selector workspace → Candidate selection: selects candidates (evidence: `flashtopk/topk.py`)
- Adaptive planner → Kernel launcher: plans kernels (evidence: `flashtopk/plan.py`)
- Adaptive planner → Gate prologue: uses (evidence: `flashtopk/plan.py`)
- Kernel launcher → Gate prologue: launches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → Seed preparation: launches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → SM90 scan: dispatches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → SM100 scan: dispatches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → SM90 pruning: dispatches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → SM100 pruning: dispatches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → SM90 refinement: dispatches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → SM100 refinement: dispatches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → Candidate selection: launches (evidence: `flashtopk/kernels/launch.py`)
- Kernel launcher → Recovery triage: launches (evidence: `flashtopk/kernels/launch.py`)
- SM90 scan → Kernel primitives: uses (evidence: `flashtopk/kernels/scan_sm90.py`)
- SM100 scan → Kernel primitives: uses (evidence: `flashtopk/kernels/scan_sm100.py`)
- SM90 pruning → SM90 scan: shares primitives (evidence: `flashtopk/kernels/prune_sm90.py`)
- SM100 pruning → SM100 scan: shares primitives (evidence: `flashtopk/kernels/prune_sm100.py`)
- SM90 refinement → SM90 scan: shares primitives (evidence: `flashtopk/kernels/refine_sm90.py`)
- SM100 refinement → SM100 scan: shares primitives (evidence: `flashtopk/kernels/refine_sm100.py`)
- Dense recovery → Recovery triage: uses (evidence: `flashtopk/recovery.py`)
- Dense recovery → Kernel launcher: scores densely (evidence: `flashtopk/recovery.py`)
- Dense recovery → Candidate selection: reselects rows (evidence: `flashtopk/recovery.py`)
- Gate prologue → SIMT helpers: uses (evidence: `flashtopk/kernels/prologue.py`)
- Seed preparation → SIMT helpers: uses (evidence: `flashtopk/kernels/seed_prep.py`)
- Candidate selection → SIMT helpers: uses (evidence: `flashtopk/kernels/selection.py`)
- Recovery triage → SIMT helpers: uses (evidence: `flashtopk/kernels/triage.py`)