Skip to content

Instantly share code, notes, and snippets.

@zwx2238
zwx2238 / topk_bit_select_numpy.md
Created May 26, 2026 08:12
Top-k bit selection with NumPy pseudocode

Top-k Bit Selection with NumPy Pseudocode

下面用 bit 二分,也就是 radix=2 的特例,说明 CUDA top-k selection 的基本思想。

目标是找出 top-k largest。算法分两步:

  1. 先找第 k 大元素的阈值 kth_value
  2. 再收集所有 > kth_value 的元素,并用 == kth_value 的元素补满 k 个。

这里为了突出算法本身,假设输入已经是 unsigned sortable key。对于 uint32 数据可以直接用;对于 float,真实 CUDA 实现会先把 float 转成保持排序关系的 unsigned bit key。