1. 项目背景LSS方案里的cumsum到底卡在哪做BEV感知部署尤其是走LSSLift-Splat-Shoot这条技术路线的模型真正费时间的地方往往不在主干网络而在Splat这一步。论文里轻描淡写一句“把提升出来的特征分配到栅格上再做累积求和”到了真机部署时就是一大串排序、前缀和、索引回填的活。PyTorch里几行代码跑得很快因为框架把内存管理、算子调度都包了真到自己动手做工程化就必须把这些操作拆成一个一个可以受控的算子。这篇文章是我在LSS-BEV部署系列里的第4篇主题是一个部署插件QuickCumsum。简单说它把“体素内重复点归并”这一长串操作做成了一个高性能的CUDA插件形态面向的恰好是LSS路线最容易被忽略的cumsum环节。无论是正在折腾LSS量化的还是已经尝过自定义算子苦头的应该都能从这篇里拿到一套可以直接复用的思路。这里先把结论放在前面cumsum本身不难真正难的是它和排序、散射、段边界识别纠缠在一起。QuickCumsum的价值不在于提出了一种新的前缀和数学理论而是把这条链路在GPU上重新编排了一遍让内存读写次数和内核启动次数同时降下来最终能换来接近一个数量级的耗时代差。1.1 Splat步骤为什么要做累积求和我希望先把“为什么需要cumsum”讲透因为它决定了后续所有优化方向。Lift阶段会把图像特征和深度分布组合起来生成一个很大的“点集”每个点都带着一组特征向量还有一个三维坐标比如BEV格子里的x、y和深度z。Splat阶段则要把这些点按坐标落到鸟瞰视角划好的网格里并对特征做累加。问题在于不同相机的同一物理位置、同一相机的相邻像素、同一个像素的不同深度估计最终都会落进同一个体素中。如果硬枚举每个体素包含哪些点再逐个累加逻辑上当然可行但实现效率惨不忍睹。带LSS风格的参考代码几乎都使用同一个组合拳先按坐标排序让同一个格子里的点变得连续接着对排序后的特征做cumsum最后再根据排序索引和体素边界做一次差分把重复点的贡献分摊回正确位置。这一步如果从头写复杂度很容易失控而用前缀和实现时可以借GPU的并行扫描能力把整个归并过程变成几次大规模的内存搬运。关键点在于排序后的序列中每个体素对应的多个点会连续排列。只要知道这段连续区间的起点和终点就能用“前缀和最后一项减前缀和前一项”得到整个体素的特征和。这就是所有LSS实现里那句“cumsum trick”的来历。它本质上是一种segment sum只是用前缀和来等价实现从而避开了逐段串行累加。1.2 朴素实现推理时的性能痛处我先说一组实践中看到的量级一个常见的多相机LSS模型假设6路相机、深度离散成几十个bin单帧的采样点数往往在千万级别。每个点又带着几十维特征。也就是说Splat的输入是一个千万行以上的特征矩阵要按索引完成排序、前缀和、回填。PyTorch版朴素计算大概是这个形状feats feats.flatten(0, 1) coords coords.flatten(0, 1) rank torch.argsort(coords, descendingFalse, stableTrue) sorted_coords coords[rank] sorted_feats feats[rank] prefix torch.cumsum(sorted_feats, dim0) boundary sorted_coords[1:] ! sorted_coords[:-1]这一段在训练时完全没问题推理时却有几个具体瓶颈。第一个瓶颈是多次内核启动。排序是一个内核gather是一个内核cumsum是一个内核逆索引回填又是一个内核。每次都把整个特征矩阵从全局内存读一遍、写一遍。千万行乘以64维特征单次搬运已经是百MB级别连续四五次操作下来内存带宽被反复消耗真正的计算反而没有多少。第二个瓶颈是前缀和本身。cumsum天然有串行依赖PyTorch底层虽然做了并行scan但它是按整段顺序处理的没有利用“同一个体素的点只在排序序列中占一小段”这个特性。实际LSS场景里一个体素内的点往往是几个到几十个极少出现把整行都吞掉的情况。逐行扫描意味着大量线程在等待前一段的计算结果利用率上不去。第三个瓶颈在边界识别。朴素实现通常用“当前坐标与上一个坐标是否相同”来标记段起点这个比较本身很简单但为了做这个比较必须额外保存一份排序后的坐标索引有时候还会产生一个布尔掩码把这些中间结果再扫一遍。这些看起来不起眼的操作叠加起来就成了端到端延时里不容忽视的一块。在我测过的几个项目里Splat这段整体占到推理延时的12%到20%。更尴尬的是模型精度没有任何变化纯粹是工程实现拖了后腿。所以我的建议很直接只要准备把LSS模型真正拿去做部署这部分的加速就值得单独立项千万别指望框架自动帮你优化掉。提示QuickCumsum是一个插件不改变模型参数不改变训练逻辑只在推理侧替换原生的排序加前缀和链路。2. 设计思路QuickCumsum的核心思想2.1 天真扫描和并行scan的差距串行前缀和的原生实现可以理解为“每个元素都要等前一个元素算完”。GPU上有上千个线程不可能让它们互相等待于是就有了并行scan。经典的并行scan结构包括Hillis-Steele算法和Blelloch算法思路都是用迭代方式两两相加利用偏移翻倍来把前缀和摊开算完整体时间复杂度从O(N)降到O(N log N)但好处是N个线程能同时干活。如果你有一万多行的输入串行扫描对单核来说确实是线性时间并行scan的价值就在于N越大线程越多时间摊得越薄。尤其当特征通道数C等于64甚至128时扫描过程可以直接对整行做向量化操作进一步摊薄开销。QuickCumsum没有发明一个新的scan理论它做的是把scan和边界识别融合到一起。此前缀和的不是整段从头扫到尾而是只关心“段结束的那个位置”上的前缀和值以及“段开始前一个位置”上的前缀和值。中间所有元素的前缀和对最终结果来说都不需要显式存储。顺着这个思路设计目标就很清楚把排序之后的特征扫描限制在一个个短段内完成段与段之间互不阻塞不再出现“更长的一段拖垮整体等待时间”的问题。2.2 排序、识别段边界、回填一次完成假设已经拿到排好序的坐标key这个key把三维坐标压缩成一维例如x 40 | y 20 | z和对应的特征。我们需要的是每个段里所有特征之和并且把这个和写回到原始token位置上。QuickCumsum的关键观察是不需要对整段做全局前缀和。段的边界其实是“当前元素key和上一个key相同下一个key就变了”。一旦知道段起点和终点段和就等于“前缀和[终点] - 前缀和[起点-1]”。因此算法顺序可以这样安排对key做一次排序得到rank与inverse_rank。对排序后的特征做一次并行scan得到前缀和。同时扫描key的相邻差异产出段边界表。用段边界表做一次向量化gather用“前缀和[end] - 前缀和[start-1]”得到每段和。把每段和用inverse_rank回填到原始特征顺序。分开写需要启动四到六个内核QuickCumsum的做法是把第2、3、4作为同一内核中的前后阶段第1步和第5步再用高效的排序库和索引交换实现。内核数量少了读写放大也随之降低。这里有一个很实际的细节排序后的key数组就已经包含了段边界的信息所以第3步其实不需要单独分配一个布尔掩码数组可以在scan内核里通过读取相邻线程自己的key来判断。这样省掉一块中间内存也省掉一次全局内存遍历。2.3 为什么选warp shuffle在并行scan里区分“段内”和“段间”非常重要。真实LSS模型里容易落在同一个体素里的点往往是几个到几十个很少达到成千上万。大多数字段都能塞进一个warp32线程里。因此让一个warp各自处理相邻的段就能用__shfl_down_sync做五次蝶形交换完成32长度内的前缀和无需访问共享内存。这个选择的原因很朴素共享内存有bank冲突和同步开销而warp shuffle直接在寄存器之间交换数据延迟只有几个周期。如果段再长一点就先用warp内scan得到部分和再交给共享内存做第二级扫描最后把结果广播回去。这其实是两级scan架构在可并行scan的实现里属于性价比很高的做法。顺带一提很多实现里纠结的“多个体素点凑不满一个warp”的问题在LSS场景里其实并不可怕。因为一个warp处理的不只一个体素而是把多个相邻体素的短段拼接起来一起处理只要段与段之间有清晰的终止标记warp内部完全可以通过key比较来复位累加器并不需要每个体素独占一个完整warp。3. 插件实现一个可落地的CUDA版本3.1 数据结构与内存布局写插件之前先把内存布局定好。输入特征在推理引擎里一般是一个连续tensorshape是[S, C]S是排序前所有的采样点总数C是特征维度。坐标key则是一个[S]的一维数组。如果模型里有多个批量叠加通常把batch维打平到S里用key的最高几位区分不同batch。输出特征和输入特征形状一致语义是“每个原始位置拿到的体素聚合和”。这一步必须保持shape不变否则下游的BEV格子化操作没法对接。我建议用结构体把插件内部状态收拢起来struct QuickCumsumMeta { int num_points; int feat_dim; int batch_size; int grid_x, grid_y, grid_z; int64_t* key_buffer; int* sort_order; int* inverse_order; float* prefix_buffer; float* output_buffer; };其中sort_order保存排序后的原始索引inverse_order回填时使用。prefix_buffer是扫描过程的临时空间。这些buffer在插件初始化时一次性申请不要在每次推理时重复分配。3.2 排序key的编码坐标key建议直接用64位整数。LSS的体素坐标不一定是定点数它可能是浮点栅格坐标量化后的整数。量化公式往往是int64_t key (int64_t)batch_id 48 | (int64_t)x 32 | (int64_t)y 16 | (int64_t)z;x、y和z的取值范围需要提前确认。比如BEV网格是128x128乘以64个深度bin那x和y各占16位、z占16位就够。如果网格规模更大比如超过65535那这个编码就会溢出碰撞。更稳妥的做法是先用网格形状算一遍最大位宽再动态决定移位位数。注意编码碰撞是排序里最隐蔽的错误。一旦两个不同体素的key相同cumsum会把它们当成同一个段结果就会少算一截。调精度时这种错误还很难定位因为大部分点是对的只有边界上的点错了。3.3 内核主体分段并行扫描这里只展示核心逻辑完整的排序和scatter步骤用库调用配合。假设排序已经完成sorted_key是排序后的key数组sorted_feat是按排序顺序摆放的特征。每个线程负责一行特征也就是一个token。__global__ void fused_segment_cumsum_kernel( const float* sorted_feat, // [S, C] const int64_t* sorted_key, // [S] float* prefix, // [S, C] int S, int C) { int row blockIdx.x * blockDim.x threadIdx.x; if (row S) return; int lane threadIdx.x 31; int warp_in_row row / 32; // 先把自身特征读进寄存器 float val sorted_feat[row]; int64_t key sorted_key[row]; // warp内蝶形scan #pragma unroll for (int off 1; off 32; off 1) { float nbr __shfl_up_sync(0xffffffff, val, off); int64_t nbr_key __shfl_up_sync(0xffffffff, key, off); if (lane off nbr_key key) { val nbr; } } // 如果当前已是warp中同key的最后一段把结果写进prefix // 注意这里还需要处理跨warp的段实际工程里要用shared memory做二级scan prefix[row] val; }真实工程里比这个复杂一些因为一个段可能横跨多个warp。我的做法是先做warp内扫描把每个warp最后一个元素的部分和写到共享内存数组再做一次针对这些“段尾部分和”的并行scan最后把上级scan结果广播回来。这样无论段有多长都能正确算出每个位置的前缀和同时大部分短段都不需要经过第二级。3.4 回填与收尾拿到前缀和后最后一步是把每段和写回原始token位置。排序阶段会输出一个sort_order数组它记录了“排序后第i个位置对应原始第几个token”。逆索引inverse_order可以把结果反推回去。__global__ void scatter_back_kernel( const float* prefix, const int64_t* sorted_key, const int* sort_order, float* output, int S, int C) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx S) return; int64_t key sorted_key[idx]; bool is_end (idx S - 1) || (sorted_key[idx 1] ! key); if (is_end) { int start idx; while (start 0 sorted_key[start] key) --start; start; // 找到段起点 float seg_sum 0; if (start 0) { seg_sum prefix[idx]; } else { seg_sum prefix[idx] - prefix[start - 1]; } // 写回本段所有原始位置 for (int p start; p idx; p) { int orig sort_order[p]; for (int c 0; c C; c) { output[orig * C c] seg_sum; } } } }这段代码里有个明显问题while循环在极端情况下会退化尤其是当你有一个特别大的体素段时。实际上我用的是排序阶段预计算的段起始表而不是在回填时线性扫描。更好的做法是在扫描边界时直接产出segment_start和segment_end两张表回填时只做索引查表避免while循环。3.5 和训练侧结果的精度对齐插件写完后的第一件事是写一个数值对比脚本。关键点并行归约改变了浮点累加顺序结果不会和PyTorch的cumsum严格相等。但插件内部保持确定性即相同输入一定得到相同输出。对比策略如下。以fp32对比时最大绝对误差控制在1e-4以内比较稳妥。以fp16对比时并行scan改变了累加顺序误差级别可能在1e-3上下。对检测任务来说这种误差不会明显影响结果但如果你把插件插进一个对数值高度敏感的分支比如在线后处理里做阈值判断那最好单独验证。一个容易踩的坑排序库输出的是“逻辑索引”而实际特征可能是上一个阶段产出的连续tensor。如果前面有transpose或reshape只替换插件而忽略内存顺序变化结果会整体错位。避免方法是把整个前向链路用同样的一组原语跑一遍把tensor的stride信息与训练脚本反复核对。4. 部署集成把QuickCumsum接进推理引擎4.1 插件化的基本流程推理引擎的插件机制大同小异。一般流程是把模型导出成中间格式然后在图上把cumsum节点替换成自定义插件节点或者在引擎构建阶段直接插入插件层。我的做法是在训练框架里用一个自定义PyTorch函数把整段排序加前缀和包起来导出时标记为自定义算子推理引擎侧注册同名插件插件内部把排序、scan、scatter三个子步骤全部实现。这样网络结构图里只会多出一个节点其它算子一个都不用动。工程代码里一个重要步骤是插件注册。这里有一个简化版的结构示意class QuickCumsumPlugin : public IPluginV2DynamicExt { public: DimsExprs getOutputDimensions(...) override { // 输出和输入shape保持一致 } size_t getWorkspaceSize(...) const override { // 临时空间 key缓冲 排序辅助空间 scan缓冲 } int enqueue(...) override { // 调用内部排序、scan、散射kernel } };在enqueue里尽量把多个内部kernel放进同一个CUDA stream。如果有多路相机并行也可以用多stream进一步压吞吐但要注意key缓冲和输出缓冲之间的依赖同步避免两个stream同时写同一块workspace。4.2 动态Batch、动态空间的处理LSS系列部署时batch并不是固定值。很多时候为了吞吐把多个相机的特征拼成一个batchbatch大小可能随相机数变化网格大小也可能随原始输入分辨率变化。插件如果要支持动态shape需要注意三个地方。第一workspace按最大batch申请而不是按第一次出现的shape申请。动态shape环境下第一次推理可能恰好是小batch如果workspace随之缩水第二次大batch进来就会越界。这个错误往往不会立刻崩而是隔几帧才在排序阶段报非法地址排查起来非常难受。第二排序临时buffer大小随N线性增长不能有隐式固定数组。很多排序库接口都需要单独传入临时空间大小这个大小要用保守公式估算比如temp_bytes N * 16 1024之类的。第三所有中间状态必须留在device端host侧不能直接访问key缓冲。尤其是在调试时如果你把key拷回host打印可能顺手把stream同步点带进推理链路性能会掉一截。4.3 实测数据我在某个项目中记录了评测结果。测试场景是单卡推理、FP16特征、特征维度64、千万级token。把Splat里的cumsum链路分别用四种实现来量得到下面的相对耗时。实现相对耗时说明PyTorch cumsum直接跑1.0多次kernel启动加大量中间读写共享内存扫描未融合边界0.42减少部分读写仍有多余内核局部warp scan加合并边界0.21大部分段在32线程内完成QuickCumsum含排序与散射融合0.09所有步骤集中为一遍内存搬运绝对时间会随显卡波动但相对关系很稳定。如果cumsum链路在端到端延时里占15%换掉之后整体大概能少13%左右。这个收益并不来自某一个魔法内核而是把“无谓的全局内存反复读写”给去掉了。5. 常见问题与排查记录5.1 key编码溢出导致不同体素被合并症状很隐蔽精度测试时输出只在某些边界位置有误差整体误差不大但单独检查某个体素时发现它的特征和比预期大了一倍。这种大概率就是key编码碰撞。排查方法把输入缩到极小的网格比如4x4x4让所有坐标都能被编码再放大网格复测。如果你发现小网格完美、大网格出错基本可以锁定是移位不够。修正方式是先统计当前batch中所有x、y、z的最大值动态计算位宽而不是拍脑袋定死48位、32位、16位。5.2 排序不稳定导致逆索引错乱如果使用非稳定排序同一个key内部的token顺序可能被打乱同段的和仍然正确但scatter回填时容易出问题。尤其是当你同时做数值对比发现“前几个通道正确后面通道全错”时多半是秩数组和特征数组没有同步。解决方式是显式传一个token_id参与排序也就是排序的key是(key, token_id)的复合结构保证排序结果唯一。这样inverse_order必然是一一映射不存在二义性。5.3 空段与越界访问体素网格并不是每个格子都有采样点。如果排序后的key出现跳变那么两个相邻段之间会有空隙。此时如果段起始表用“后一个段的起点减1”来推导上一个段的终点就会算出一个空段start和end相等甚至end小于start。在散射内核里必须加判断seg_end - seg_start 1直接跳过不要做减法运算。越界访问通常表现为偶发的cuda error很难稳定复现最好在host侧加一层检查逻辑打印每个段的长度分布及时发现异常。5.4 推理结果偏差较大但逐段测试正常这种问题一般不是内核本身的毛病而是上游数据布局不一致。我吃过一次亏训练侧的特征排序是针对[B, C, D, H, W]展开的部署侧因为中间换过layout变成了[B, D, H, W, C]导致按行读取时的通道维度错位。调试时我会做一个最小复现用例构造4个体素、少量token分别用原脚本和插件跑一遍打印每一步的中间tensor。只要定位到排序后的特征和key数组是否对上了问题就能快速浮出水面。6. 一点工程体会这个插件做完之后我最大的感受是LSS部署的加速空间往往不在卷积核里而在这些容易被当成“框架已优化”的基础操作里。cumsum只是一个前缀和但因为它的上下链路涉及排序、gather、scatter复杂度一下子起来了。如果你手头的模型也有类似的splat步骤我的建议是先量化它占推理延时的比例如果超过10%就值得动手。优化的顺序也很重要先做内核融合减少读放大再考虑warp级并行不要一上来就追求写出极致的并行scan那样调试成本太高收益却不一定成正比。我自己后面还在尝试的一个方向是把QuickCumsum和BEV池化算子合并成一个插件让排序、前缀和、散射以及后续的二维柱体池化只经过一次全局内存往返。代价是临时缓冲要按体素数量重新规划边界情况更多但一旦跑通端到端的收益还能再上一个台阶。这个等稳定了再单独写一篇。
工具包)
