做算子优化这件事最怕的不是算法不复杂而是看数据一个一个跳出来的时候你还得憋得住火。前几天我接到一个任务在DCU上把Softmax算子性能提上来。当时我还想着这个简单无非是求个最大值、求个和、再做个除法三行公式的事。结果一看Profiling报告这个看似人畜无害的算子竟然在端到端推理里占了将近8%的时间。于是就有了今天这篇实战记录。如果你接下来要在DCU这类高性能加速卡上手写算子或者只是想把GPU/DCU上的基础算子优化思路理顺这篇文章里应该能找到几个能直接用的招。DCU的编程模型和CUDA很像所以里面的思路搬到其他加速卡上也基本通用。我会把踩坑过程、关键代码、优化前后数据全部摊开来说。1. 项目背景与优化思路1.1 为什么Softmax值得单独做一次优化Softmax在深度学习里几乎无处不在从Transformer的注意力分数归一化到大模型输出层的概率分布转换再到多分类任务的输出端都要靠它把一组logits变成总和为1的概率。它的数学形式很简单p_i exp(x_i - max(x)) / sum_j exp(x_j - max(x))但简单不意味着好算。里面有两个全局归约操作一个全局max一个全局sum。在GPU/DCU这类并行架构上归约操作会引入跨线程的数据依赖把访存和同步变成主要开销。我在实际Profile里看到一个看起来只占整体计算0.1%的Softmax在端到端模型里常常能吃掉3%到8%的时间尤其在推理batch较大、vocab或者序列长度很长的时候。这不是算子本身有多复杂而是它把大量数据来回读了好几遍访存带宽就浪费在这上面。针对这个算子做专项优化投资收益非常明确代码逻辑不算长优化空间却很大而且做完之后整个网络的耗时能实打实地降下来。同时Softmax是很多模型的公共子模块优化一次可以被多个模型复用属于那种“花半天时间做后面一直受益”的工作。所以接到需求后我并没有急着去翻软件说明书或者套现成实现而是先把这个算子在不同shape下的访存特征梳理了一遍再决定从哪一步动手。1.2 DCU上写算子先想清楚这三点先说下DCU本身。DCU这类加速卡在硬件设计上更接近AMD的GCN架构编程模型主要走HIP代码风格和CUDA非常接近很多习惯从CUDA迁移过来是能直接用的。但有三点和写普通CUDA程序时需要额外注意。第一DCU的显存带宽很高但隐藏延迟依赖足够多的并行线程你的kernel如果某个串行部分卡住了延迟就会直接暴露出来所以访存要尽量一次拿全减少反复访问。第二DCU的调度单元是wavefront一条wavefront通常包含64条work item这和NVIDIA的warp32线程不同所有按32线程假设写死的归约逻辑都要重新验证。第三编译器和驱动对某些优化指令的支持没有想象中那么激进比如自动向量化经常不靠谱最好在代码里显式用float4/float2的方式来引导。在写算子之前还要把数据布局想清楚。这次优化里输入是二维矩阵行为样本列为softmax维度。这种布局下每一行是一个独立的softmax单元行与行之间没有任何依赖正好适合交给不同block并行处理block内部再做线程间归约。这个思路本身不稀奇但具体到访存模式、寄存器容量、归约方式就有很多地方需要抠。我习惯先把数据流画出来哪里读、哪里写、哪里有同步再动手。后面的所有改动本质上都是在优化这张数据流图。1.3 本次优化的目标场景与约束我把范围设定在一个典型推理场景输入规模是[batch8, seq_len1024, vocab2048]的float32张量最后一维做softmax等价于一个8192行乘以2048列的矩阵每行做一个softmax。为什么选这个规模因为这是LLM推理时非常常见的形状vocab2048不算大但seq_len拉长之后行数会很多矩阵总元素达到1600万以上足够暴露访存问题。约束条件是不改模型结构、不改数据精度前向仍然用fp32只改算子实现。目标有三个第一是把多个kernel合并成一个第二是让每个全局内存元素尽量只被读一次、写一次第三是把块内归约的同步开销降到可接受范围。后面所有优化都围绕这三条展开。2. 第一版朴素实现能跑但很慢2.1 教科书式Softmax的实现与问题第一版我几乎是把伪代码直接翻译成HIP的思路是一个block处理一行block里只用一个线程算其他线程都闲着。这种实现正确性没问题但慢得离谱。当然第一版更多是为了验证接口和结果不能当作优化起点。真正想当基准看的是一个“看起来并行其实不太并行”的版本一个block处理一行所有线程都参与但需要先全局归约max再全局归约sum最后再读一遍数据计算输出。这样每个元素至少要读两次甚至会读三次再加上每个归约阶段都要启动一次kernel或者用原子操作来跨block归约启动和同步的开销也很可观。这里有个典型反模式为了省事把整个矩阵看成一个大数组然后调用现成的并行reduce函数分别算max和sum。逻辑上没错但Softmax的归约是逐行进行的在大数组上做全局归约等于把行与行之间的独立性全丢了不但多算了很多跨block同步还让计算量加倍。正确的姿势是行内归约行间完全并行。从这个问题出发我们发现所谓优化其实是从数据并行模型上做减法去掉多余的同步去掉重复访存。2.2 基线性能数据与瓶颈定位先跑了一版“每个block处理一行block内256线程每个线程用for循环读取该行的一段连续数据”的实现这个版本虽然每个线程会读一段连续数据但同一时刻相邻线程访问地址不连续也不是好写法在8192x2048规模上耗时大约278微秒。后来又改成更接近正确写法的版本每个线程按跨步方式读耗时降到198微秒。但两次kernel仍然要等所有block归约完才能算sum访存了两遍Profiling显示全局读取量大约是理论最小值的1.9倍。用工具看这个kernel的实际带宽利用率只有30%左右。原因很简单数据读了两遍且第二遍读取一般都会触发缓存miss等于白白浪费了带宽。另一个问题是两轮kernel之间有一次隐式的全局同步模型里其他kernel只能等着进一步拖慢端到端时间。这两个问题就是后续优化的直接靶子。我当时的判断是这个算子不能通过简单调整访存方式蒙混过关必须从结构上把重复读去掉才能拿到质的提升。3. 融合kernel把三次访存变成一次3.1 优化思路一次读取寄存器暂存核心改动是消除第二次全局读。既然每行的数据在一开始就能拿到为什么不让每个线程把自己负责的那部分元素读到寄存器里等max和sum归约完成后再用Softmax计算本来就需要两遍归约但数据不需要跟着读两遍。每个线程处理固定几个元素第一步读进寄存器数组顺便算局部max做完block内max归约后直接在寄存器数组上算exp并累加局部sum做完sum归约后再直接写回结果。这样一来每个全局内存元素只从显存读一次、写一次理论上可以把访存量减半。当然这里面要处理一个问题block内归约max和sum都需要所有线程的数据但寄存器是线程私有的别人拿不到。所以我们在归约时借助共享内存做跨线程聚合但共享内存中只放每个线程的中间归约结果比如局部max、局部sum而不是把整行数据都放进去。整行数据还是留在寄存器里这样共享内存占用很小block可以开得足够大寄存器压力也可控。3.2 核心代码实现下面这段是我在实际项目里调通的主干逻辑。每个block负责一行block内每个线程按跨步方式从全局内存读元素存到寄存器数组vals里。注意这里读的时候是按tid j * nthreads取下标所以每轮循环内相邻线程访问的内存地址是连续的访存是合并的这一点很关键。__global__ void softmax_fused_kernel( const float* __restrict__ input, float* __restrict__ output, int rows, int cols) { extern __shared__ float sdata[]; int row blockIdx.x; if (row rows) return; const float* x input (size_t)row * cols; float* y output (size_t)row * cols; int tid threadIdx.x; int nthreads blockDim.x; // 演示版本每个线程最多处理 8 个元素调用时保证 cols nthreads * 8 const int MAX_VALS 8; float vals[MAX_VALS]; int cnt 0; for (int idx tid; idx cols; idx nthreads) { vals[cnt] x[idx]; } // 线程内局部 max float lmax -FLT_MAX; for (int j 0; j cnt; j) { lmax fmaxf(lmax, vals[j]); } // block 内归约 max sdata[tid] lmax; __syncthreads(); for (int s nthreads / 2; s 0; s 1) { if (tid s) { sdata[tid] fmaxf(sdata[tid], sdata[tid s]); } __syncthreads(); } float max_val sdata[0]; // 用 max_val 计算 exp同时求局部 sum float lsum 0.0f; for (int j 0; j cnt; j) { lsum expf(vals[j] - max_val); } // block 内归约 sum sdata[tid] lsum; __syncthreads(); for (int s nthreads / 2; s 0; s 1) { if (tid s) { sdata[tid] sdata[tid] sdata[tid s]; } __syncthreads(); } float sum_val sdata[0]; // 写回 for (int j 0; j cnt; j) { int idx tid j * nthreads; y[idx] expf(vals[j] - max_val) / sum_val; } }解释几个关键点。vals数组保存当前线程读到的所有元素后续两步都用它不用再访问全局内存。sdata是动态共享内存调用kernel时需要传入字节数。这里故意复用同一块共享内存做两次归约第一次放max第二次放sum省共享内存。两次归约之间用__syncthreads()保证顺序归约循环里同步在if外面确保所有线程执行相同步数不会出现死锁。3.3 访存模式与边界处理为什么用for (int idx tid; idx cols; idx nthreads)而不是让每个线程处理一段连续数据原因很简单合并访问。在任意一次循环迭代里tid0的线程访问x[row*cols0]tid1访问x[row*cols1]相邻线程访问相邻地址硬件会把这一组访问合并成少数几次大的内存事务。如果反过来每个线程处理一段连续数据虽然单线程内部连续但同一条内存请求里放的是不同线程离得很远的地址合并效率就低了。边界情况如果cols不是nthreads的整数倍会有部分线程少读一个元素。上面代码用cnt计数只在有效范围内计算所以不会有越界。但如果你想追求性能可以把cols对齐到nthreads的整数倍或者用向量化读取把cols处理成4的倍数。我的做法是保持通用循环然后在更激进的版本里用float4读取那个版本要求cols是4的倍数我们这次没有遇到非倍数场景所以先用当前版本保证正确性。如果你拿到手的模型正好有一个奇怪的维度最简单的方式还是先跑通用版再针对热点做特化别一上来就写死。4. 进阶优化向量化与归约策略4.1 向量化不只是凑代码是对齐带宽v2把访存减半后耗时从198微秒降到96微秒。再往下压就得抠指令效率和访问粒度了。DCU这类加速卡对128bit即float4级别的访存很友好一次能搬4个float比四个单独的32bit load要少很多指令和地址计算。于是我在v3里把跨步读取改成按float4处理每个线程每次读一个float4相当于一次处理4个元素。整体逻辑不变每个线程持有若干个float4先算局部max再做归约。核心代码片段大概是这样const float4* x4 reinterpret_castconst float4*(x row * cols); float4 v x4[tid j * nthreads]; float v0 v.x, v1 v.y, v2 v.z, v3 v.w;这里有个前提每行的起始地址必须16字节对齐并且cols必须是4的倍数。在我们场景里hipMalloc分配的内存通常256字节对齐2048也正好是4的倍数所以可以直接转。如果你的维度不是4倍数就得单独处理尾部几个元素或者先把输入pad一下。用向量化之后全局读取指令数变成原来的四分之一实测在DCU上带宽利用率能提升10到15个百分点。4.2 warp shuffle的诱惑与DCU的差别优化到v3后block内归约用的是共享内存加二叉树归约。有经验的朋友可能会问为什么不用warp shuffleNVIDIA GPU上warp shuffle能在warp内部通过寄存器换数据不经过共享内存非常快。但我在DCU上试过之后发现DCU的wavefront是64条lane和NVIDIA的32条lane不同很多写死在32上的shfl归约代码直接搬过来是有问题的。另外一些HIP版本的shuffle接口在掩码位数和width参数上的处理并不完全一致稍微不注意就只归约到半个wavefront然后拿到一个错误的结果。我的建议是如果你能确认驱动和运行时对shuffle的支持情况可以用否则共享内存归约的稳定性已经足够性能差距在Softmax这种访存密集型kernel里并不会特别夸张。这里踩一次坑的记录是因为我直接把CUDA的32线程shfl归约改成64线程但mask还写0xffffffff结果只有前32个lane参与归约最大值和总和算错而且由于结果只体现为概率不对debug的时候很容易忽略建议在测试代码里加入针对小输入的数值校验。4.3 block size与元素负载怎么配v3的另一个调节旋钮是block大小和每个线程处理元素数。block太小时共享内存归约深度小但并行度不够典型配置是128block太大时归约串行深度增加而且动态共享内存占用也会限制同卡并发block数。我扫了128、256、512三档结果256是当前规模的最优解每行2048个元素256线程每个线程恰好处理8个元素没有空转同时256线程的2棵归约树深度合适延迟可接受。如果一行元素特别多比如4096或者8192我建议把一行拆给多个block每个block做部分归约后使用原子操作或者二次kernel完成跨block合并。虽然引入了额外同步但总比一个block里让每个线程串行处理上百个元素要好。如果一行很短比如只有32个元素那就不值得开这么多block直接用warp或者一个block处理多行来摊薄开销。整体来说block大小和每个线程处理元素数是两个需要一起调的参数先用计算量大的shape跑一遍扫描选稳定值。5. 实测性能对比与调参实录5.1 四个版本的耗时对比现在我们直接看表。这是相同环境、相同数据规模下的实测数据取100次平均。naive_v0是每个block一行但只用一个线程算的版本multi-pass是两次kernel分别归约max和sum的版本fused_v2是本文第3节的融合单kernel版本vec_v3是在v2基础上做float4向量化的版本。耗时越低越好。版本实现方式平均耗时相对加速全局读次数naive_v0每行单线程1120 us1x每元素3次multi-pass两次归约kernel198 us5.7x每元素2次fused_v2融合单kernel寄存器暂存96 us11.7x每元素1次vec_v3float4向量化融合58 us19.3x每元素1次从数据看真正拉开差距的是从multi-pass到fused_v2这一步访存量直接减半kernel启动次数从2次变1次时间砍掉一半以上。后一步vec_v3虽然访存次数没变但指令数和访存效率进一步优化又砍掉近40%。这说明在访存密集型算子里消除重复读永远是第一优先级的优化手段向量化是第二优先级。5.2 用Profile定位吞吐瓶颈优化期间我用rocprof看指标。第一版multi-pass里L2 cache命中率低带宽利用率30%原因是重复读数据fused_v2中全局读取量下降带宽利用率到60%vec_v3继续提升到75%左右。需要注意Softmax这种小算子很难把带宽打到90%以上因为还有exp的运算时间和归约的同步不能只看带宽。当看到某个版本耗时不再下降时建议先看是不是同步和尾效应而不是继续加向量化。尾效应表现在block之间执行时间不均衡可以用更多block来缓解。还有一点只看kernel耗时容易忽略启动开销。multi-pass版本有两个kernelPCIe/驱动层面的启动延迟即使只有几微秒在短kernel里也占不小比例。合并成一个kernel后这部分基本消失了。所以我在汇报优化进展时喜欢把端到端profiling一起拿出去而不是只汇报单个kernel的耗时这样更有说服力。5.3 换形状换环境后的稳定性同样的优化代码换到[batch16, seq_len128, vocab65536]这种大vocab、少行的场景性能趋势不一样。行数少但每行很长单个block处理一行会遇到并行度不够此时更适合每行用多个block分段归约。我加了一层配置当cols blockDim.x * 8时自动切分行每段用一个block做局部max和sum然后用global memory中预先分配的空间做跨block归约。这个分支在现有场景下没启用但保证扩展性。另外换不同型号DCU之后warpSize可能保持一致但主频和显存带宽不同最优block大小也许会漂移建议每次换硬件后重新跑一遍block扫描。最优解不是一劳永逸而是要在目标型号上重新验一遍。6. 常见问题与排查技巧6.1 浮点精度差点把结果搞崩代码写完后我第一件事是拿小矩阵和CPU的double结果对比。在fp32下2048个exp累加如果使用__expf这种快速近似函数max和sum稍有问题就会导致概率无法严格归一。我最后在归一化前保留fp32的expf没有开-use-fast-math。如果只是推理你也许可以考虑快速exp但差个千分之一在分类场景可能看不出来在训练场景反向传播时Softmax的梯度公式是p_i * (delta_ij - p_j)前向概率的误差会直接放大到梯度导致训练不稳定。所以我的建议是训练模型用精确expf纯推理且时间压力大的场景再换__expf。顺带提一句如果做的是训练Kernel别只看前向。前向里把max、sum归约顺序调整后反向的p值可能变化一点点你可能需要在反向算子测试里重新对齐误差范围。这个坑我在之前另一个优化项目里踩过一次前向精度验收全过了结果训练两天后loss异常最后定位到反向梯度对前向概率的误差太敏感。6.2 同步、边界、共享内存的几个坑同步坑主要出现在共享内存归约循环。我的归约循环写法是if (tid s) { sdata[tid] ...; } __syncthreads();同步在if外面保证所有线程执行同一轮同步。如果你手滑把同步写进if分支里block内不同线程经过的同步次数不一致程序会直接挂起。检查方法很简单在循环里数每个线程执行了多少次__syncthreads必须完全相同。边界坑如果你用了float4向量化版本必须确认每行起始地址是16字节对齐。有时候行内offset是row * cols而cols不是4的倍数那么从第二行开始就错位程序不报错但数值乱掉debug成本很高。我建议在kernel里加一个assert(cols % 4 0)发布版本再关掉。共享内存坑动态共享内存的大小必须由调用方传入。我最初用的是静态共享内存__shared__ float smem[256]block一改成512就爆了。后来改成extern dynamic同时把调用参数写成bytes blockDim.x * sizeof(float)block大小随便调再没遇到越界。这个细节看起来小但如果你打算做block size扫描静态版本就是又一个定时炸弹。6.3 编译选项和快速数学库的取舍hipcc编译时我用-O3没有用--use-fast-math。DCU驱动对标准cmath浮点函数的支持比较稳定expf编译后会自动映射到硬件指令性能已经够用。如果你想再压一点可以单独用__expf替换热点但必须做精度回归。还有一个容易被忽略的选项是-ffast-math会影响编译器对浮点运算做重关联可能会改变求和顺序。Softmax的sum是正数累加重关联后误差通常不大但在严格归一化校验里可能差2-3个ULP。我的经验是算子类代码保持默认浮点语义如果确实要冲性能把-ffast-math编译选项放到最后再试同时在验收测试里绑定一个精度上限。这次优化做下来最大的体会不是某个花活而是把访存次数从3次减到1次的过程中每一步都能在Profiler里看到对应的变化。DCU上的算子优化没有太多玄学先保证数据只读一次再考虑向量化和归约结构性能基本不会差。如果你下一步要优化别的算子这套思路也能复用先画数据流把重复访存和多余同步找出来然后才是写代码。