很多写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::vectorfloat 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_ptrvoid, 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建议从这个角度去重新审视一遍自己手头的老代码。 
长河君 · 2026/10/21 8:37:44 · 阅读 1.2k CUB单次调用API：告别临时存储三步走，聚焦算法本身
很多写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::vectorfloat 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_ptrvoid, 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建议从这个角度去重新审视一遍自己手头的老代码。