企业官网建设流程全解析

很多写CUDA算法的人第一次碰CUB，都是从一个排序或者一个大数组归约开始的。但你会在官网示例里看到一段“标准三步走”：先调一次API拿临时显存大小，再cudaMalloc，再调一次API真正干活。第一次看会觉得“这不是脱裤子放屁吗”，后来写多了才发现，这套设计的背后是设备端内存分配的硬约束。不过问题也随之而来：每次写这种代码都要重复三遍，实在啰嗦。CUB单次调用API的定位就是把这一段过程收进口袋里，让开发者在写归约、扫描、排序这类基础原语时，只需要关心“数据”和“算法”，不用再纠结“临时存储”这几个字。

这篇文章适合三类人看：正在用CUB写算子但被temp_storage反复折磨的CUDA工程师；想把CUB里面的DeviceReduce、DeviceScan、DeviceSegmentedRadixSort真正用起来的学生；还有那些想了解“库接口为什么要这么设计”而不是只抄代码的读者。我会把单次调用API到底改了什么、内部是怎么做的、性能和手动分配到底差多少、有哪些坑，一次讲清楚。

1. 为什么CUB的经典API写起来这么累：临时存储的前世今生

1.1 一段“标准三步走”代码带来的精神污染

先来看最经典的数组求和。用CUB的DeviceReduce，不用任何封装，最原始的写法是这样的：

#include <cub/cub.cuh> #include <cuda_runtime.h> #include <cstdio> int main() { const int n = 1 << 20; float *h_in = new float[n]; float h_out = 0.0f; for (int i = 0; i < n; ++i) h_in[i] = 1.0f; float *d_in = nullptr; float *d_out = nullptr; cudaMalloc(&d_in, n * sizeof(float)); cudaMalloc(&d_out, sizeof(float)); cudaMemcpy(d_in, h_in, n * sizeof(float), cudaMemcpyHostToDevice); // 第一步：查询临时存储大小 void *d_temp = nullptr; size_t temp_bytes = 0; cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, n); // 第二步：分配临时存储 cudaMalloc(&d_temp, temp_bytes); // 第三步：正式开始归约 cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, n); cudaMemcpy(&h_out, d_out, sizeof(float), cudaMemcpyDeviceToHost); printf("sum = %f\n", h_out); cudaFree(d_temp); cudaFree(d_in); cudaFree(d_out); delete[] h_in; return 0; }

这段代码本身没啥毛病，但每次写都像在重复填表。更麻烦的是，如果你在同一个函数里又想要一个前缀和，又想要一个排序，就会出现好几个d_temp、好几个temp_bytes，分配完还要记得释放。我的一个习惯是写一个局部结构体把临时存储包起来，但不同算法的对齐要求不一样，光对齐这个事就足以让人想放弃CUB转投Thrust。

其实这一段“三步走”不是CUB故意找存在感，这里面有非常实际的原因。

1.2 GPU设备端为什么不让你直接裸奔分配临时内存

CPU上写并行库，算法内部临时开个vector，是件很自然的事。GPU设备端虽然支持malloc和new，但代价和限制都很大：设备端malloc本质上是在设备堆上分配，慢，而且需要额外的显存池管理，还只能在kernel内部用，不能在host端随便触发。更重要的是，一个device kernel内部做动态分配，会严重影响占用率，因为每个线程都可能持有堆指针，根本无法预估显存占用。

CUB这种底层原语库，面向的是极致性能，它的算法内部需要一大块scratch memory做block级别的中间结果交换、双buffer切换、位图排序的计数数组等等。所以它明面上选择了一个最保守也最通用的策略：让调用者准备一块临时存储，你准备好我再干活。这就是为什么会有d_temp_storage和temp_storage_bytes这两个参数。

这种做法某种程度上有点类似C语言里的malloc(size)之前先要malloc_size那种“咨询”机制，也有一点像你在组装家具时需要先比划一下柜子尺寸。两次调用确实麻烦，但它保证了一件事：无论调用者是谁、无论显存多大、无论是哪个GPU架构，CUB都不会在内部偷偷分配显存，也就不会在kernel满天飞的时候突然触发一个显存分配失败。

1.3 两步API的真实代价：不只是多写三行代码

很多人觉得，多写三步就多写三步，反正也就几行。真正让这东西变成负担的，是它在工程里的“传染性”。

第一，你写一个函数用CUB，这个函数就必须把临时存储参数暴露出去，或者自己内部管掉。内部管掉就是三步，暴露出去则是把心智负担传播给上层调用者。

第二，查询大小的那一步，虽然看起来只是同步地算一个size，但你要是忘记处理temp_bytes == 0的边界，就会踩出莫名其妙的bug。

第三，临时存储释放的时机也是一个隐患。在多流并发场景下，你在A流上调用CUB算法，如果A流还没执行完就把临时存储free了，另外一条流把这块内存复用给别的kernel，数据就乱了。

所以单次调用API解决的不是“少写三行代码”的问题，它把“临时存储生命周期”这个变量从你的开发流程里直接拿掉了。你不需要再问：这块内存谁分配、谁释放、什么时候释放。这对写原型、写教学demo、写一次性的数据处理脚本，价值非常大。

2. 单次调用API到底改了什么：把封装边界拆开看

2.1 所谓单次调用，内部其实是“三步并一步”

CUB里的单次调用API，字面意思就是你不用再自己先query再调用了。但拆开来看，它内部做的事情跟三步走是一模一样的：还是先查询临时存储大小，还是要在设备端分配一块scratch空间，最后启动kernel。只不过这一步从“调用者手动完成”变成了“库内部自动完成”。

我以某个较新版本CUB里DeviceReduce的简化签名来举例：

cub::DeviceReduce::Sum(d_in, d_out, n);

你不需要传临时存储指针，也不用传字节数。如果你需要指定流：

cub::DeviceReduce::Sum(d_in, d_out, n, stream);

单次调用API通常会在内部使用一个缓存分配器，第一次调用时真正向CUDA运行时申请显存，用完以后不是马上释放，而是把这块显存给缓存起来；下一次调用如果size匹配或者能复用，就直接从缓存里拿，省掉cudaMalloc的开销。

如果你是手动三步流，就没有这个缓存，每次调用完就cudaFree了。所以单次调用API在“重复调用同一个原语”的场景下，性能并不会差，甚至因为省掉了重复QR阶段的API往返，整体还能快一点。

CUB里cub::CachingDeviceAllocator就是干这个的。它本质上维护了一个空闲显存块列表，分配时先看有没有大小合适的空闲块，没有才真正走cudaMalloc。单次调用API的默认后端基本就是这路子。理解这一点之后，你就知道单次调用API不是“黑魔法”，它只是把你本来该做的脏活累活，收编到了库的封装层里面。

2.2 三种常见的形式：无参重载、分配器注入、host端helper封装

不同版本CUB，甚至不同device原语之间，单次调用API的公开形态并不完全一样。我把它分成三种理解，这样你去查文档、看源码的时候，不会一头雾水。

第一种是官方直接提供免临时存储参数的重载。就像上面写的那样，参数列表里压根没有d_temp_storage，这是最理想的情况。

第二种是允许你往里注入一个分配器策略。你传一个自定义的分配器对象，库在需要临时存储时调用它去拿内存。这种做法更偏底层，适合跟你自己项目的显存池做对接。

第三种是host端helper封装。这通常是社区里或者你自己项目里写的工具函数，把三步走包成一个lambda调用或者一个宏。我自己至少写过两个版本，一个基于std::function，一个基于模板套娃。模板版本大概长这样：

template <typename KernelInvoker> void cub_once(KernelInvoker &&invoke) { void *d_temp = nullptr; size_t temp_bytes = 0; invoke(d_temp, temp_bytes); // 第一次调用，传入 empty 指针，查大小 if (temp_bytes > 0) cudaMalloc(&d_temp, temp_bytes); invoke(d_temp, temp_bytes); // 第二次调用，正式执行 if (d_temp) cudaFree(d_temp); }

用法很直观：

cub_once([&](void *d_temp, size_t temp_bytes) { cub::DeviceScan::ExclusiveSum(d_temp, temp_bytes, d_in, d_out, n); });

当然这种helper是同步的，如果需要异步，你可以在invoke里带上流，并且在释放前追加cudaStreamSynchronize或者干脆让分配器的生命周期长过流。这也是我自己早期没有注意到的地方。

2.3 哪些原语适合单次调用，哪些最好别用

CUB的Device级原语挺多，包括DeviceReduce、DeviceScan、DeviceHistogram、DeviceSegmentedRadixSort、DeviceSelect、DevicePartition、DeviceMergeSort等。从工程角度看，不是所有原语都适合无脑用单次调用。

适合的有这么几种：一次性的数据处理流程、调试用的临时算子、输入规模不大而且对首次延迟无所谓的场景。比如说我在某个模拟项目里，需要对一组几百万粒子的坐标做一次分段归约，统计各区间的总能量，这种活我直接用单次调用API，心智负担最低。

不适合的也有两类。第一类是热路径上的原语，比如在迭代求解器里，每一轮都要跑一次reduce，而且要的是确定性延迟，那建议手动分配一次temp，然后循环里复用。第二类是输入规模极大、比如超过可用显存一半的输入。自动分配器并不了解你整个进程的显存占用情况，它只会一看到就分配，不会替你精打细算。这种场景手动管理更稳妥。

3. 实操：用单次调用API重写一个排序-扫描组合管线

3.1 环境准备：版本、头文件、编译参数

开始动手之前，先说环境。CUB头文件库，传统上可以从GitHub拉源码，也可以直接通过系统包管理器装。较新的CUDA Toolkit实际上已经直接把CUB集成进去了，不再需要单独下载。这一点让很多人困惑，因为老教程还在教你git clone cub再指定-I路径。

现在比较省心的做法是：

• 安装CUDA Toolkit 11.1以上，大部分版本自带CUB头文件。
• include路径里找cub/cub.cuh，如果找不到，再去找CUB源码。
• 编译的时候，注意加上--expt-relaxed-constexpr，否则某些复杂模板实例化会报错；这个flag在较新CUDA里虽然不是必须，但加上能省很多奇奇怪怪的编译错误。

一个最简单的CMakeLists.txt大概是这样：

cmake_minimum_required(VERSION 3.18) project(cub_single_pass LANGUAGES CUDA CXX) find_package(CUDAToolkit REQUIRED) add_executable(single_pass_demo main.cu) target_link_libraries(single_pass_demo PRIVATE CUDA::cudart)

头文件包含就是这个：

#include <cub/cub.cuh>

如果你是老版本CUB用户，建议先看device_reduce.h头文件里有没有不带d_temp_storage参数的接口。没有的话，用我上面那个cub_once的模板helper，效果也大差不差。

3.2 上手：单次调用的归约和扫描

先写一个完整的单次调用归约demo，假设我的CUB版本支持免临时存储参数的重载：

#include <cub/cub.cuh> #include <cuda_runtime.h> #include <cstdio> #include <vector> int main() { int n = 1 << 22; std::vector<float> h_in(n, 1.0f); float *d_in = nullptr; float *d_out = nullptr; cudaMalloc(&d_in, n * sizeof(float)); cudaMalloc(&d_out, sizeof(float)); cudaMemcpy(d_in, h_in.data(), n * sizeof(float), cudaMemcpyHostToDevice); cudaStream_t stream; cudaStreamCreate(&stream); // 单次调用：不需要临时存储，指定流 cub::DeviceReduce::Sum(d_in, d_out, n, stream); float h_out = 0.0f; cudaMemcpyAsync(&h_out, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream); cudaStreamSynchronize(stream); printf("sum = %f, expect = %f\n", h_out, (float)n); cudaStreamDestroy(stream); cudaFree(d_in); cudaFree(d_out); return 0; }

这里有一个很容易忽略的点：单次调用API在内部如果用了缓存分配器，它会去分配显存，然后把kernel提交到指定的stream上。提交是异步的，所以你在后面cudaMemcpyAsync之前不需要额外同步，但一定不能在使用完d_out之前提前结束进程或者释放输入缓冲区。

扫描和归约是同一个套路。CUB里DeviceScan::ExclusiveSum和DeviceScan::InclusiveSum是常用的前缀和原语，单次调用版本几乎是复制粘贴：

cub::DeviceScan::InclusiveSum(d_in, d_out, n, stream);

自定义二元操作符的归约也可以这样用，比如求最大值：

cub::DeviceReduce::Reduce(d_in, d_out, n, cub::Max(), 0.0f, stream);

到这里你会注意到，代码里完全看不见“临时存储”这四个字了。视觉上清爽非常多。

3.3 组合起来：归一化加上前缀和，再做一个分段统计

单独写一个reduce和scan还看不出单次调用API的真正威力。我现在给一个更有代表性的场景：我需要把一组非负数值做归一化，算累积分布函数，然后再按区间分段统计数量。这就是一个典型的“排序-扫描-选择”组合管线。

第一步，用DeviceReduce求最大值，这里用单次调用：

float max_val = 0.0f; cub::DeviceReduce::Max(d_in, &max_val, n, stream);

第二步，把所有元素除以最大值，并同时计算前缀和。这里用一个自定义算子做“归一化后累加”：

struct NormalizeAdd { float divisor; __device__ float operator()(float a, float b) const { return (a / divisor) + b; } };

然后调用InclusiveScan，注入这个fancy算子：

cub::DeviceScan::InclusiveScan(d_in, d_out, n, NormalizeAdd{max_val}, 0.0f, stream);

第三步，如果还要按区间统计，就用DeviceSelect里的Flag或者SegmentedRadixSort辅助。整个过程中，每一个CUB原语的调用都只需要一行，没有一处去手动处理临时存储。

这种写法的好处，对做原型验证的人来说是质变。过去你光是排布临时存储的分配和释放，就需要写一坨和算法本身无关的胶水代码。现在算法主线和内存管理完全解耦了，你可以先快速验证一个想法对不对，等确认要做性能优化，再把热路径上的原语替换成手动管理临时存储的版本。

3.4 老版本怎么办：自己封装一个helper顶上去

绝大多数人其实卡在一个不新不旧的CUDA版本上，官方单次调用API不一定有。这时候直接用我前面那个cub_once模板就行，但它也有个问题：如果原语内部的不同实例需要的临时存储大小不一样，同一个helper没问题；但如果原语本身是异步的，cub_once里直接在invoke之后cudaFree，这就会出问题。

所以自己写helper要额外注意同步。严谨一点可以这样：

template <typename KernelInvoker> void cub_once_sync(KernelInvoker &&invoke) { void *d_temp = nullptr; size_t temp_bytes = 0; invoke(d_temp, temp_bytes, 0); // 第一次调用，默认流 if (temp_bytes > 0) cudaMalloc(&d_temp, temp_bytes); invoke(d_temp, temp_bytes, 0); // 第二次调用 cudaStreamSynchronize(0); // 确保kernel完成再free if (d_temp) cudaFree(d_temp); }

更推荐的做法是封装一个“临时存储所有者（temp_storage owner）”，把分配、释放和析构绑定在一起：

class TempBuffer { public: template <typename Invoker> void invoke(Invoker &&invoke) { void *d_temp = nullptr; size_t temp_bytes = 0; invoke(d_temp, temp_bytes); if (temp_bytes > 0) { cudaMalloc(&d_temp, temp_bytes); } invoke(d_temp, temp_bytes); if (d_temp) { // 车辆入库：不是立刻释放，而是留到析构时 data_.reset(d_temp); } } private: std::unique_ptr<void, CudaFreeDeleter> data_; };

这样，临时存储的生命周期至少能延续到封装对象析构那一刻，比裸写法安全得多。

4. 性能与显存开销：单次调用是免费午餐吗

4.1 分配开销的实测思路：从微秒到毫秒的跨档

很多人在乎单次调用API是不是比手动三步流慢，这里我不给死板结论，先说测试思路。

你要关心的开销主要来自两部分。第一是第一次调用时内部真去cudaMalloc，这个调用的延迟，实测下来大概在几微秒量级，取决于驱动和显存占用状态。第二是内部分配器维护缓存，这个开销几乎可以忽略，就是一次哈希查找或者线性扫描。

所以对于输入规模在几十万到几百万元素的reduce、scan这类原语，kernel本身要跑几十微秒甚至几百微秒，用单次调用API多出来的那几微秒分配开销完全可以忽略。但如果你的输入数组只有一千个元素，kernel本身几微秒，这时候再叠加一次cudaMalloc，就显得不划算了。

我自己的测试习惯是：把同一个原语循环调用1000次，手动管理temp和单次调用各测一遍，取中间值而不是平均值。这里有一个很容易踩的坑，就是CachingDeviceAllocator缓存命中后，第二次调用几乎不分配，性能会跟手动管理非常接近，但不能拿这个去代表首次调用的开销，一定要把第一次和后续分开统计。

另外一个观察是，单次调用API如果默认使用CachingDeviceAllocator，它在重复调用时的开销几乎为零。真正让人感到“慢”的，通常不是分配本身，而是某些实现为了做到“跨流安全”，不得不在内部做流同步。如果你本来在一个stream里连续调用多个CUB原语，而API内部每次都要等待上一个原语结束，流水线就废了。

4.2 什么时候该回到手动管理临时存储

我个人划了一条线，这条线很粗但很有用：如果同一块输入数据在同一个函数里会连续调用三次以上CUB原语，我就手动分配一块足够大的临时存储，然后复用它。

为什么是三次？因为第一次调用要分配，第二次调用分配器缓存命中，第三次调用其实已经进入稳定态，手动和自动没啥区别。但在三次调用之前，代码里如果每个原语都要求独立临时存储，自动分配器会让三块内存同时在显存里活着，显存占用会比复用方案高三倍。这在显存紧张的项目里就是一个真实的危机。

那手动复用的标准写法是：

void *d_temp = nullptr; size_t temp_bytes = 0; cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, n); cudaMalloc(&d_temp, temp_bytes); // 三个原语共用一块临时存储 cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, n); cub::DeviceScan::InclusiveSum(d_temp, temp_bytes, d_in, d_scan, n); cub::DeviceReduce::Min(d_temp, temp_bytes, d_scan, d_min, n);

注意这里的CUB原语实际需要的临时存储大小可能不一样，比如排序通常比归约需要更大的scratch，所以你取所有原语里最大的那个temp_bytes即可。如果不够，直接连续分配三次，也没有错，只是浪费。手动模式下，这些都由你掌控。

4.3 preclude：显存也是一等公民

再提一个很容易被忽视的事情。单次调用API背后的分配器，通常是“要多少给多少”，不会主动做显存池上限控制。如果你的程序已经占用了不少显存，又突然来一个非常大的单次调用操作，极有可能触发cudaMalloc失败，返回cudaErrorMemoryAllocation。这时候你以为问题出在算法上，其实只是显存不够了。

我的建议是，在任何用到单次调用API的原型项目里，显存剩余量低于200MB就不要再做大型CUB操作。200MB是个经验值，不同平台不一样，但它的意义在于：给CUDA上下文、驱动预留、还有你眼前看不到的内存碎片留一点缓冲。

如果在生产环境里，单次调用API的显存策略不符合你的预期，建议还是自己写一个显存池，或者给CUB的分配器注入一个带配额限制的自定义allocator。这会多一点代码，但能避免半夜线上任务因为显存不足全挂。

5. 几个容易踩的坑，以及我现在写CUB的个人习惯

5.1 零字节临时存储：一个几乎一定踩到的边界

手动三步走时，你可能会遇到一种情况：CUB原语经查询后temp_bytes == 0。这个状态在早期CUB版本里是合法状态，意思是这个原语在当前输入规模下不需要额外scratch。这时候如果代码无脑cudaMalloc(&d_temp, 0)，返回的指针可能是非空也可能是空，行为在不同驱动上并不一致。

单次调用API虽然在内部处理了这块逻辑，但如果你在封装helper时偷懒，忘了判断temp_bytes > 0，就会踩到。我见过的实际案例是：输入一个空数组或者长度为1的数组，reduce查询返回0，然后cudaMalloc(0)，接着kernel启动时直接报invalid argument。

处理方式就一句话：分配前判断，先处理temp_bytes为0的情况，否则直接用原指针nullptr也能跑很多原语。

5.2 多流并发：自动分配会不会串流

多流并发是另一个高频翻车点。假设你写了两个stream，同时跑两个不同的CUB算法，两个都用了单次调用API，内部都使用同一个默认分配器。如果分配器内部做了流同步，你的两个stream就被强制串行化了；如果没做同步，分配器分配出来的显存块又有可能被另一个stream拿到。

最稳妥的做法分两级。第一级，如果只是测试或者原型，直接全部放到同一个stream里，不要搞花活。第二级，如果确实要多流并发，并且强调性能，就给每条流单独一个临时存储buffer，不要依赖自动分配。我就是这么干的：每个stream持有一个TempBuffer对象，所有CUB调用都走这个对象，绝不在两个流之间共享。

5.3 版本差异：接口不是一直一致的

CUB的接口在不同版本之间变动很大。有些版本支持无temp参数的DeviceReduce，有些版本不支持；有些版本在cub/cub.cuh里全部暴露，有些版本需要单独包含cub/device/device_reduce.cuh。我遇到过最头疼的情况，是同一个项目在不同服务器上因为CUDA Toolkit版本不同，导致其中一个编译不过。

所以我现在的习惯是：写任何CUB代码，都记一个最小兼容版本号，放在注释里。比如“Requires CUB >= 1.16 or CUDA 11.1”。真要在老版本上编译，就直接用helper封装，把单次调用API写在新代码层，底层还是经典的temp_storage接口。这样既不牺牲兼容性，又能让业务代码保持简洁。

5.4 我的模板化“一行流”习惯

最后分享一个我现在写CUB的个人习惯。我基本上会把所有单次调用原语再包一层模板函数，名字叫cub_call，接受Lambda、输入输出指针和大小，内部处理临时存储并且做错误检查。这样业务代码就真的变成流水账了，逻辑主线非常清楚：

CUB_CALL(cub::DeviceReduce::Sum(d_in, d_out, n, stream)); CUB_CALL(cub::DeviceScan::ExclusiveSum(d_out, d_scan, n, stream)); CUB_CALL(cub::DeviceSelect::If(d_scan, d_selected, d_num_selected, n, Predicate(), stream));

每行一个原语，临时存储、错误检查、显存管理全部藏在宏或者封装的背后。刚开始可能不适应，但写过几个算子之后就会觉得，这才是CUB打开的正确方式。它让GPU并行算法开发真正回到了“算法”本身，而不是被“并行库怎么分配内存”这件事反复绑架。如果你也经常写CUB，建议从这个角度去重新审视一遍自己手头的老代码。