GPU基数排序：从CUB工程实现到性能调优

写 GPU 排序，我很少一上来就吹某种算法，但基数排序是个例外。对于整数、定长字节这类键，GPU 基数排序确实是最能接近“性能最好”这个说法的通用算法。它不是靠比较来排序，而是靠“按位计数 + 分段搬运”，一趟就是一次连续读写，天然适合 GPU 那种高带宽、多线程的架构。这篇文章我会把一份能直接跑的代码和详细注释给你，再把背后的 histogram、scan、scatter 三步拆开讲清楚，顺便聊一聊你真正调参时会遇到的那些坑。

1. 先下结论：基数排序凭什么算“性能最好的通用算法”

先说清楚，我这里说的“通用”不是指所有数据类型，而是指固定长度、能编码成整数序列的键，比如uint32_t、uint64_t、int32_t、float，以及 pack 成 uint64 的多字段键。这种场景下，基数排序基本是 GPU 上的首选，理由很直接：

• 比较排序的下界是O(n log n)，但 GPU 上“比较”并不便宜。分支、乱序访存、比较器开销都会把常数放大，尤其是随机数据，分支预测基本失效。
• 基数排序一次 pass 只做三件事：统计 digit 出现次数、做前缀和、按偏移把元素写回新位置。没有value < pivot这类比较，也没有跨 block 的归并等待。
• 数据访问模式非常规整。按 8-bit 一个 digit 来处理，一趟 pass 就是对整个数组做一次读 + 一次写，GPU 的 memory controller 和 cache 最喜欢这种 sequence。

如果只拿“32 位无符号整数”做基准，8-bit 一次 pass 就是 256 个桶，32 位总共 4 次 pass。每次 pass 的访存量是n * (4 字节读 + 4 字节写)，4 次 pass 就是32n字节。你可以用自己显卡的内存带宽粗算一下：比如带宽 1TB/s，1 亿个 uint32 的理论下限就是32 * 1e8 / 1e12 = 3.2ms左右，加上临时空间和 scan 开销，实际能跑到 4-6ms 已经很不错。这个量级是快排、归并那些方案很难碰到的。

拿常见的几类排序在 GPU 上的表现做个对比，可能更直观：

这里有个很重要的限定条件：不要说“任何情况都最快”。如果你要排的是变长字符串，或者只有几千个元素，基数排序的优势基本没了。上万元素以下，直接cudaMemcpy到 CPU 用标准库排可能更快；变长字符串还是要上比较排序。基数的“通用”，是“在 GPU 排序定长键”这个范围内的通用。

2. 一份能直接跑的代码：CUB 封装后面的 4 趟 LSD 基数排序

你完全可以手写 GPU 基数排序，但工程上我更建议先跑通 CUB。CUB 的DeviceRadixSort已经是高度优化的 LSD 基数排序实现，内部把访存合并、block rank、scan、共享内存冲突这些问题都处理好了。自己写的第一版性能大概率打不过它。先把生产代码跑起来，再拆原理。

下面这段是完整的最小可运行思路，我用两个 buffer 交替作为输入输出：

#include <cub/cub.cuh> #include <cuda_runtime.h> #include <cstdint> // 将 d_keys 中的 n 个 uint32 排成升序。 // 调用前需要两个设备端 buffer： // d_keys：输入数据 // d_out ：输出数据/临时目标 // d_temp：CUB 工作空间，第一次调用后按需分配 // 如果你想原地排序，可以把 d_keys 和 d_out 传同一个指针， // CUB 内部会借助临时空间完成 in-place 操作。 void RunGpuRadixSort( uint32_t* d_keys, uint32_t* d_out, uint32_t n, uint32_t*& d_temp, size_t& temp_bytes) { // 第一次调用只传 nullptr 工作空间，CUB 会返回需要的字节数。 // 这是 CUB 的标准用法，几乎所有 Device 算法都这样。 size_t bytes = 0; cudaError_t err = cub::DeviceRadixSort::SortKeys( nullptr, // d_temp_storage bytes, // temp_storage_bytes（输出） d_keys, // d_keys_in d_out, // d_keys_out (int)n, // num_items 0, // begin_bit 32 // end_bit，闭开区间 [0, 32) ); if (err != cudaSuccess) { // 错误处理省略，生产环境记得判断 return; } // 如果已有临时空间不够，就重新分配 if (d_temp == nullptr || temp_bytes < bytes) { if (d_temp) cudaFree(d_temp); cudaMalloc(&d_temp, bytes); } temp_bytes = bytes; // 真正执行排序。 // 这里一次性排完整个 32 位。 // 理解上可以想象成内部做了 4 趟 8-bit 的 LSD pass： // 第 0 趟按 bit[0,7) 排序 // 第 1 趟按 bit[8,15) 排序 // 第 2 趟按 bit[16,23) 排序 // 第 3 趟按 bit[24,32) 排序 err = cub::DeviceRadixSort::SortKeys( d_temp, temp_bytes, d_keys, d_out, (int)n, 0, 32 ); // 如果 n 超过 int 范围，需要看平台是否支持 int64 版本接口。 // 另一端的数据已经写到 d_out，需要的话用它覆盖 d_keys。 }

这段代码的注释已经把重点写出来了：CUB 会自己处理 temp storage，你不要在排序循环里反复cudaMalloc。临时空间分配一次，后面所有 pass 复用。之所以强调这一点，是因为cudaMalloc一次可能几十微秒到上百微秒，对于一次几毫秒的排序来说，反而不算小开销了。

如果你想把 LSD 的“一趟一字节”看得很明确，也可以手动拆成四次调用：

uint32_t* src = d_keys; uint32_t* dst = d_out; for (int pass = 0; pass < 4; ++pass) { int begin_bit = pass * 8; int end_bit = begin_bit + 8; cub::DeviceRadixSort::SortKeys( d_temp, temp_bytes, src, dst, (int)n, begin_bit, end_bit ); // 下一趟的输入是这一趟的输出，所以交换 src/dst std::swap(src, dst); }

手动拆四趟的性能通常会比一次调用整段 32 位略差，因为每次SortKeys都有点额外启动开销，但逻辑上更接近“LSD 基数排序”的教学模型。生产环境直接一次调用[0, 32)就行。

3. 把黑盒拆开：8-bit 一趟的 histogram、scan、scatter 三段论

CUB 好用，但你如果不知道它在干什么，碰到性能异常和稳定性问题时会很慌。一次 8-bit 基数排序 pass，本质上是三步：

1. histogram：统计每个 block 里 256 个 digit 各出现多少次。
2. scan：对 block 间的统计做前缀和，算出每个 digit 在全局输出里的起点。
3. scatter：按起点把元素写回新数组。

下面我用“教学版”内核演示这三步。注意，这段代码不是冲着极限性能去的，也不是直接可以替代 CUB 的生产代码，但它能让你看清 pass 的结构。

第一步，每个 block 统计自己负责那段元素里 256 个 digit 的直方图：

#define BLOCK_THREADS 256 #define ITEMS_PER_THREAD 4 // 每个 block 处理 BLOCK_THREADS * ITEMS_PER_THREAD 个连续元素。 // 统计结果写入 histBlocks[blockIdx.x * 256 + d]。 template<int BT, int IPT> __global__ void Histogram256( const uint32_t* keys, uint32_t* histBlocks, int n, int shift) { __shared__ uint32_t sHist[256]; // 初始化共享内存中的 256 个桶 for (int d = threadIdx.x; d < 256; d += BT) sHist[d] = 0; __syncthreads(); int tileStart = blockIdx.x * (BT * IPT); // 每个线程顺序处理 IPT 个元素 for (int item = 0; item < IPT; ++item) { int idx = tileStart + item * BT + threadIdx.x; if (idx < n) { uint32_t digit = (keys[idx] >> shift) & 0xFFu; atomicAdd(&sHist[digit], 1u); } } __syncthreads(); // 把本 block 的 256 个桶写回全局内存 for (int d = threadIdx.x; d < 256; d += BT) histBlocks[blockIdx.x * 256 + d] = sHist[d]; }

这里的atomicAdd是往共享内存上做原子加，而不是全局内存原子加，速度还算能接受。但极端情况下，如果所有线程都命中同一个 digit，原子操作还是会串行化。CUB 内部会做多级 histogram 来分散原子冲突，这也是专业实现和教学实现的差距之一。

第二步，对所有 block 的直方图做一次“跨 block + 跨 digit”的前缀和：

// 计算每个 block、每个 digit 之前的元素总量。 // blockOffsets[b * 256 + d] = 前 b 个 block 中 digit d 的累计个数。 // digitStart[d] = 全局范围内 digit d 的起始偏移。 // 这个 kernel 只启动一个 block，线程数 256，一个线程负责一个 digit。 __global__ void ExclusiveScanBlockOffsets( const uint32_t* histBlocks, // [numBlocks * 256] uint32_t* blockOffsets, // [numBlocks * 256] uint32_t* digitStart, // [256] int numBlocks) { __shared__ uint32_t sTotals[256]; int d = threadIdx.x; // 0..255 uint32_t sum = 0; // 串行遍历所有 block，为当前 digit 做 block 级前缀和 for (int b = 0; b < numBlocks; ++b) { blockOffsets[b * 256 + d] = sum; sum += histBlocks[b * 256 + d]; } // sTotals[d] 就是全局 digit d 的总个数 __syncthreads(); sTotals[d] = sum; __syncthreads(); // 再做一次 256 个桶的串行前缀和，得到每个 digit 的起始位置 uint32_t start = 0; for (int j = 0; j < d; ++j) start += sTotals[j]; digitStart[d] = start; }

第三步，按计算好的偏移把元素写回输出：

// 教学版 scatter：只保证每个 digit 的元素落在正确区间， // 不保证同一个 digit 内的元素保持原顺序。 // 因为 LSD 多趟排序需要“稳定”，真正生产环境不能直接用这个版本。 template<int BT, int IPT> __global__ void Scatter256TeachingDemo( const uint32_t* keys, uint32_t* out, const uint32_t* blockOffsets, const uint32_t* digitStart, int n, int shift) { __shared__ uint32_t sCursor[256]; // 本 block 的可写起始位置 = 全局 digit 起始 // + 之前所有 block 中本 digit 的数量 for (int d = threadIdx.x; d < 256; d += BT) sCursor[d] = digitStart[d] + blockOffsets[blockIdx.x * 256 + d]; __syncthreads(); int tileStart = blockIdx.x * (BT * IPT); for (int item = 0; item < IPT; ++item) { int idx = tileStart + item * BT + threadIdx.x; if (idx < n) { uint32_t digit = (keys[idx] >> shift) & 0xFFu; // 取出当前元素在输出数组中的位置，然后把游标 +1 uint32_t pos = atomicAdd(&sCursor[digit], 1u); out[pos] = keys[idx]; } } }

看到这里你可能会想：这个 scatter 不是挺直接的么？问题就出在atomicAdd上。GPU 多个线程并发抢同一个 digit 的游标时，谁先拿到位置是调度器决定的，不是按原数组顺序决定的。对于只排序数值本身，可能结果还是对的；但对 LSD 多趟排序来说，如果相同 digit 的内部顺序被打乱，下一趟按更高位排序时，低位信息就丢了，最终结果可能是错的。

所以专业实现里，scatter 前必须算出每个元素在block 内的稳定 rank，比如用类似BlockScan的算法统计“前面有多少个和我 digit 相同的元素”，再把这个 rank 加到全局偏移上。这才是 GPU 基数排序里最容易写错、也最值钱的部分。CUB 已经把这些细节封装好了，你自己实现时千万不要只拿原子操作糊弄过去。

4. 调参不是拍脑袋：radix bit 数、块大小、数据宽度怎么选

很多文章只告诉你“用 8-bit 就行”，但没说为什么。这个选择背后是数学和硬件限制的折中。

如果一趟处理R个 bit，桶数量就是2^R，32 位整数的 pass 次数是ceil(32 / R)：

从 pass 次数看，R 越大越好，因为访存次数少了。但 R 变大后，histogram 的原子冲突、scan 的开销、共享内存压力都会上升。8-bit 在 GPU 上是个甜点：256 个桶只需要 1KB 共享内存，4 个 pass 刚好覆盖 32 位，而且(key >> shift) & 0xFF的位运算一个周期就能出结果。

块大小方面，我常用的起点是 256 线程、每个线程 4 元素。这样一个 block 一次处理 1024 个元素。ITEMS_PER_THREAD太大会导致寄存器压力上升，太小又不够掩盖内存延迟。你可以先跑一版，再用 CUDA 自带的 occupancy 工具看寄存器占用和 occupancy 的平衡。

还有一个很容易被忽略的点：访存一定要合并。教学代码里tileStart + item * BT + threadIdx.x的写法，就是为了让连续线程访问连续地址。读uint32_t时，尽量保证一个 warp 访问的地址在同一段连续内存里。现代 GPU 对 128-bit 的 vector load 也有优化，如果你压榨到最后，可以试试float4/uint4方式加载数据，一次读 16 字节，内存指令数会明显下降。

数据宽度也要算清楚。排uint32_t和排uint64_t不是一个量级：

• uint32_t：8-bit pass 4 趟，访存量32n字节。
• uint64_t：8-bit pass 8 趟，访存量64n字节。
• float：先做位变换映射成uint32_t，仍然是 4 趟。
• key-value对：如果 value 也是 4 字节，每趟读 8 字节、写 8 字节，访存量直接翻倍，排序时间也基本翻倍。

所以，不要拿“只排 key”的基准去估算“排 key-value”的耗时。排序的瓶颈很多时候不是计算，而是内存搬移。

5. 从“能跑”到“跑得快”：unsigned、int、float、多个 key 的工程化处理

这一步才是真正从 demo 走向生产的关键：数据编码。

5.1 无符号整数

最直接，uint32_t就是基数排序的原生对象。不需要额外处理，直接排。

5.2 有符号整数

int32_t直接当uint32_t排，顺序就错了，因为负数在无符号表示里是很大的正数。标准做法是把符号位翻转：

// 假设 x 是 int32_t uint32_t u = (uint32_t)x ^ 0x80000000u;

这个变换会把-2147483648映射成0，把2147483647映射成0xFFFFFFFF，从而保持字典序。排完之后，再u ^ 0x80000000u转回 int 即可。

5.3 浮点数

IEEE 754 的float也能映射成可排序的无符号整数，但要稍微复杂一点：

uint32_t bits = ...; // 从 float 直接位拷贝出来 // bits 符号位为 0：原浮点数为正数，翻转符号位即可 // bits 符号位为 1：原浮点数为负数，全部 bit 取反 if (bits & 0x80000000u) bits ^= 0xFFFFFFFFu; else bits ^= 0x80000000u;

如果是double，把同样的处理扩展到 64 位即可。但浮点的NaN排序行为在不同实现里可能不一样，你要根据业务决定是先剔除还是把 NaN 排到固定位置。千万不要默认 NaN 的顺序和 CPU 一致，这在工程上会埋雷。

5.4 多个关键字段

如果是一个复合 key，比如结构体里(a, b)，你需要先把它编码成单个整数：

• a、b都非负且各占 16 bit：直接key = ((uint64_t)a << 32) | b，排uint64_t。
• 有符号字段：先把每个字段做符号位翻转，再按位拼起来。
• 字段太多放不进 64 bit：不要硬拼，改用cub::DeviceRadixSort::SortPairs，或者自己写多 pass 的定长结构体排序。

这里有个原则：能编码成整数就编码成整数，编码后的排序最简单。如果键已经不是定长二进制序列，比如字符串，那就不适合直接上基数排序了。对字符串可以先按长度分组，或者用公共前缀，复杂度会高不少。

5.5 key-value 对怎么办

需要同时移动 value 的，直接用cub::DeviceRadixSort::SortPairs：

// 这个方法除了传 key，还要传 value cub::DeviceRadixSort::SortPairs( d_temp, temp_bytes, d_keys_in, d_values_in, d_keys_out, d_values_out, n, 0, 32 );

不要自己排完 key 再手动按 key 的最终位置去挪 value，那样会多一次额外的重建索引和访存，而且容易把 CUDA 的并行效率打没。

6. 我踩过的坑和最后的建议

最后聊几个实际项目里容易踩的坑，都是经验之谈。

第一个坑是temp storage 没按两次调用分配。CUB 的标准用法是先用nullptr查大小，再分配，再真正执行。如果你省掉第一步，直接塞一个很小的栈上数组进去，轻则报错，重则显存越界，非常难查。我习惯把所有 CUB 调用都封装成一个 helper，内部判断temp_bytes是否够用。

第二个坑是小数据量也硬上 GPU。几千个元素排序，启动一个 kernel 的开销可能比排序本身还大。在我这边，通常n < 50k左右，CPU 的std::sort已经很快了；硬用 GPU 只是给自己找麻烦。当然阈值取决于你的硬件和调用频率，最好自己做一次交叉测试。

第三个坑是in-place 理解错误。DeviceRadixSort::SortKeys支持传入同一个输入输出指针，但并不是说不需要额外空间，它只是把临时结果放在临时空间里。你依然要给d_temp分配足够的 workspace。如果你把两个 buffer 交替使用，还要注意最后一趟结果在哪个 buffer，别在后续处理时拿错了。

第四个坑是直方图阶段的原子冲突。如果你手写 histogram，全部线程用同一个atomicAdd(&sHist[digit], 1)，在高位集中分布的数据上会非常惨。比如排序一个标志位字段，绝大多数元素都落在 digit 0 和 digit 1，256 个线程疯狂抢两个地址，性能会下降一个数量级。CUB 内部的做法会更聪明，它能通过多副本 histogram 把冲突摊开。你自研时如果发现性能比 CUB 差很多，多半就是这一段的问题。

第五个坑是稳定性被忽略。很多人写 GPU 基数排序第一版只求“看起来排对了”，但数据里有大量重复 key 时，LSD 多趟排序必须保证稳定。一旦某趟同 digit 内顺序打乱，最终结果可能只在某些测试数据上正确，换个数据集就崩溃。你要是只排唯一键，稳定性可以不管；但只要涉及 key-value 对、复合 key，或者只是重复值很多，稳定性就是硬要求。

如果让我给一个落地方案，我会这样选：普通业务里排数字，直接用 CUB 的DeviceRadixSort，别手写；如果你是为了学习或者特殊场景要自研，先从单 block、小数组的稳定版本写起，再扩展到多 block。性能优化的优先级是：先保证访存合并，再做多级 histogram 分散冲突，最后才考虑 vector load 和汇编级优化。基数排序的源码看起来简单，真正的复杂度都藏在“稳定 rank”和“并行 scan”里。把这层窗户纸捅破，你才算真正掌握 GPU 基数排序。