★ FEATURED ARTICLE 去年接手一个推理引擎的提速任务整个团队最初的心态是“哪里慢就调哪里”。结果忙了三周改完接口改内存改完内存改并行端到端延迟反而因为引入的各种补丁变得更不稳定。后来我们共同做了一个决定把所有高频小算子从框架里抽出来单独做成一个高性能算子库也就是这篇要说的 ops-math。围绕 ops-math 算子库的高性能实现与调用技巧我想把从设计算账、内核编写、参数调优到调用侧如何不拖后腿的整套经验都摊开讲清楚。如果你也在做推理引擎、离线算子优化或者只是想在报错面前搞清楚算子到底是怎么跑的这篇文章应该能让你少走几步弯路。1. 为什么要单独维护一套算子库而不是继续用现成框架1.1 被框架黑盒卡住的真实场景很多人觉得算子不就一个函数嘛框架里现成何必自己重复造轮子。我们最初也这么想。但当你真的去追性能问题会发现框架里的算子是一个“黑盒”它不暴露内部数据布局不会告诉你中间态是否落过内存也不会为你的业务场景专门调整分块大小。有一次我们优化一个生成模型profiling 显示整体延迟总有一个固定底座去不掉。翻来覆去排查最后发现是框架在 softmax 和后续线性层之间强制生成了一份中间张量这个张量完整写回内存又读出来。要是一层两层无所谓层数一多这种中间态读写就会把内存带宽烧掉大半。也就是从那时起我们决定自己掌控算子内部的访存路径。单一算子自包含意味着你可以针对自己的输入规模、对齐要求、并发模型做定制。这其实和高性能计算领域一贯的做法一致框架负责构图和调度算子库负责把底层数学函数榨干。两者边界清晰各自才能优化到位。1.2 profiling 数据里最咬人的往往不是矩阵乘我们引入自研算子库之前先在一个 16 层的 Transformer 结构上做了完整 profiling。 统计结果大致的耗时分布如下算子类型耗时占比访存性质优化收益GEMM / GEMV48%计算密集高但实现难度大LayerNorm / Softmax 等归一化类17%访存密集很高ElementwiseGELU、Add、Mul8%访存密集中高Transpose / Reshape / Gather6%访存密集高其他调度与框架开销21%不稳定需治理看到矩阵乘只占一半时我们有点惊讶因为绝大多数优化资料都在讲矩阵乘。后来想明白了矩阵乘的单次调用时间确实最长但它一次能处理的数据量大优化空间清楚。反而归一化类和逐元素类算子单个不显眼架不住数量多。模型越深这类小算子越密集它们叠加起来反而成为端到端延迟的主要来源。这引出一个观点算子库不能只做“重计算”更要做“轻访存”。把归一化、逐元素、转置这些看着不起眼的算子优化好往往比把一个矩阵乘再压榨 10% 的收益更大。1.3 ops-math 的定位给上层提供稳定且可扩展的计算基座ops-math 最后定位成很克制的一套库只做算子不做图、不做自动微分、不负责模型结构。它的核心价值有三点提供数学上正确、边界情况有定义的高频算子。同一套接口适配不同 CPU 特性调用方不需要关心底层指令集。对访存路径显式可控支持外部传入 workspace方便上层做内存复用。因为边界清晰我们可以把大量时间花在“访存如何少一次”“分块参数怎么选”这些真正影响性能的问题上而不是被框架的调度机制绑架。2. 设计阶段先算性能账方向错了后面全白忙2.1 算术强度先判断这是访存受限还是计算受限在写任何算子之前我都会先做一道算术题这个算子的计算量除以它必须读写的数据量得到算术强度。算术强度远低于机器平衡点的就是访存受限高于平衡点才值得堆计算优化。以 softmax 为例处理 N 个 float 元素读入 4N 字节写回 4N 字节至少访问 8N 字节。计算上需要做 max、exp、sum、除法大约每元素 4 次浮点运算总计算量约 4N FLOPs。算术强度大约是 4N / 8N 0.5 FLOP/Byte。哪怕一台内存带宽很高、单核算力也不弱的 x86 服务器平衡点通常在 10-20 FLOP/Byte 这个量级。这意味着 softmax 的本质是“搬运数据”而不是“计算”。你优化 exp 函数指令再快收益也有限真正决定性能的是你如何减少对内存的读写次数。我这个习惯是从一个教训里得来的。一开始我们把精力花在优化 exp 的实现上换了一个更快的近似数学库结果整体提速不到 5%。后来改成“只读一次、一并在寄存器里完成行级 max 和 sum”直接缩短了 35% 的耗时。前者是计算优化后者才是访存优化。 以后每次拆算子先算这笔账再决定动作方向。2.2 数据布局和内存对齐看似基础影响却最深远算完算术强度后下一步是确认张量的数据布局。我们统一使用 NCHW 或 NHWC 语义抽象但底层实际存储时都会填充到对齐边界并且显式记录每个维度的 stride。注意这里不是让你把格式改得面目全非。而是在张量描述符里保留通道、高度、宽度语义但允许每一行在末尾多分配几个元素作为 padding确保 SIMD 循环不用处理半个向量。举个例子一个四通道的 float 数据如果每通道偏移是奇数个元素用 256 位向量一次处理 8 个 float 时内存访问会跨 cacheline带宽利用率下降。我们做法是申请内存时按 64 字节对齐通道尾端填到 8 的倍数。多出来的空间换取主循环里完全不判断边界分支少一条循环体更整洁实测带宽吞吐可以提升约 10%-15%。张量描述符大致长这样struct TensorDesc { int64_t dims[4]; int64_t strides[4]; size_t alignment; int dtype; };调用方如果传入未对齐裸指针我们会在调试模式下直接报错而不是偷偷复制一份对齐数据。这个设计看着苛刻实际很有用它逼着上层的内存池在分配时就把对齐要求传下来避免每个算子内部做防御式拷贝。2.3 分派逻辑正确参考实现要留到最后才换掉ops-math 每个算子至少有两套实现一套是绝对正确、但性能平凡的参考实现一套是基于 SIMD、分块、多线程调优的内核。运行时通过分派表决定走哪条路径。分派粒度不是“整个算子”而是“算子的一个内核”。比如 softmax 就要区分小批量行、大批量行、连续多行等场景。小规模时函数启动开销占比高走参考实现反而快规模变大后再切换到向量化内核。代码我们写过类似的结构using SoftmaxKernel void (*)(const float*, float*, int64_t row_size); SoftmaxKernel select_softmax_kernel(const CpuInfo cpu, int64_t batch_size) { if (batch_size 1 cpu.has_avx512()) return softmax_single_row_avx512; if (cpu.has_avx2()) return softmax_multi_row_avx2; return softmax_ref; }这套分派逻辑帮了我们大忙。因为参考实现始终在不管后续怎么调优我们都能用同一组随机输入做一致性校验。很多性能 bug 其实不是数学错了而是向量化后的边界处理和参考实现不一致。分派表加差分测试一起做问题就好定位得多。2.4 融合判断准则不是所有算子都适合粘在一起算子融合是高性能算子库的重要武器但融合不能拍脑袋。我自己总结的判断标准有三条中间张量是否会再次被其他算子使用如果会就要保留中间结果。数学上合并不改变数值稳定性尤其是含有指数、除法和归约的算子。融合后的循环体仍然能保持较整齐的内存访问模式不能为了省一次写而引入随机访存。一个正面例子是 LayerNorm 和后续的线性层融合。LayerNorm 的输出是一个中间向量之后马上被矩阵乘使用。如果中间向量整体写回内存再读出来相当于多了一整轮的 DRAM 往返。融合后线性层直接在寄存器里消费归一化结果省下整块读写。另一个明显的反面例子是把 Softmax 和后面的矩阵乘强行融合。Softmax 输出是稠密向量矩阵乘需要的是完整矩阵或行块融合后需要维护一个巨大的中间窗口缓存根本放不下表面省了一次写实际多了大量随机访问得不偿失。所以融合要具体问题具体分析。3. 高频算子实现拆解从逐元素到矩阵乘3.1 逐元素算子向量化只是起点循环展开和内存事务处理才是关键逐元素算子看着最简单比如 ReLU、Add、Scale但实现时最容易写出“看起来不错、跑起来拉胯”的代码。核心原因是它完全受内存带宽约束向量化只是必要条件。以 ReLU 为例我们用的是 AVX2 内建函数主循环一次处理 8 个 floatvoid relu_avx2(const float* src, float* dst, int64_t n) { const __m256 zero _mm256_setzero_ps(); int64_t i 0; for (; i 31 n; i 32) { __m256 v0 _mm256_loadu_ps(src i); __m256 v1 _mm256_loadu_ps(src i 8); __m256 v2 _mm256_loadu_ps(src i 16); __m256 v3 _mm256_loadu_ps(src i 24); _mm256_storeu_ps(dst i, _mm256_max_ps(v0, zero)); _mm256_storeu_ps(dst i 8, _mm256_max_ps(v1, zero)); _mm256_storeu_ps(dst i 16, _mm256_max_ps(v2, zero)); _mm256_storeu_ps(dst i 24, _mm256_max_ps(v3, zero)); } // 处理尾部标量 for (; i n; i) dst[i] src[i] 0.f ? src[i] : 0.f; }循环展开到 4 个向量连续处理是为了让足够多的内存访问同时处于飞行状态。现代 CPU 有多个内存事务队列如果一次只发一个 load 然后立刻使用内存延迟会把计算时间完全吃掉。展开后多个 load 同时发起等待时间被后续计算覆盖吞吐立刻上来。逐元素算子的经验法则先确保数据连续访问再做向量化最后考虑展开因子。上来就写一堆复杂指令反而可能因为寄存器压力过大而引入额外坑。3.2 归约类算子Softmax 与 LayerNorm 的访存优化Softmax、LayerNorm 这类算子共同点是“每一行内部先归约、再利用归约结果做元素级变换”。我们以 Softmax 为例把单行向量切到多个线程上处理时会遇到一个经典问题每行数据量小但行数很多。我们的策略是让每个线程连续处理多行完整数据而不是让多个线程切分同一行。原因很简单同一行是一段连续内存一个线程顺序读完cache 和预取都能吃满如果多个线程把一行切开线程之间还要做同步归约额外开销可能超过并行收益。归一化的数值稳定性也要特别注意。两遍遍历方案是第一遍求 max第二遍求 exp 之和第三遍写结果。这看着稳定但行数多时会读三遍数据。想要快就得在寄存器里做在线统计一边读一边维护 max 和 sum最后再统一计算并写回。伪代码思路如下// 在线稳定版读一次写一次 float local_max -inf; float local_sum 0.f; for (int64_t i 0; i row_size; i) { float x src[i]; float new_max max(local_max, x); local_sum local_sum * exp(local_max - new_max) exp(x - new_max); local_max new_max; }这种写法在数学上等价于三遍遍历但只读一次数据。最后写回时用 local_max 归一化。对于 LayerNorm均值用标准在线累加方差用 Welford 或两遍法都可以只要注意避免溢出。我们在实际实现中会用 SIMD 版本维护一组 max、sum、sum_of_squares效果非常明显。3.3 矩阵乘寄存器分块和缓存复用才是真正的大头矩阵乘的计算强度远高于逐元素算子优化重点从访存转向“数据复用”。一个朴素三重循环实现内部是每个 C 元素都去读一次 A 的行和 B 的列完全把内存带宽打爆。比较可行的思路是寄存器分块让一个线程在寄存器里维护一个小块的 C例如 8x8然后在这个块的维度上展开 k 循环让 A 和 B 的子块反复从寄存器读取。我们实现 micro-kernel 时选过 6x16 的分块尺寸。选择依据是寄存器数量既要放 C 的 6x16 个累加器又要放足够多的 A 和 B 片段还得留空间给 FMA 指令的临时结果。分块太大寄存器溢出太小又浪费数据复用。不同硬件上最优尺寸不同所以我把这块做成了可配置参数同一套代码编译多份跑基准后选最优。矩阵乘还有个和前面 Softmax 相反之处它适合多线程切分。因为每个线程计算不同区域的 C全局只需要一个屏障没有共享累加冲突。tile 尺寸选在 L2 cache 能装下的范围内外层循环再按 NUMA 节点绑核效果才稳定。3.4 融合算子实例把归一化和线性层粘到一起前面说融合原则这里看一个具体实现。 LayerNorm 输出 y之后跟随一个全连接层权重为 W、偏置为 b。如果分开写写 y 到内存读 y计算 y 乘 W 加 b开发者模式完全可以用一个寄存器窗口去承载归一化输出块然后在同一块内存区域上乘 W 加 b。由于 W 的列数可能很大更实际的融合方式是把线性层也切成与归一化行块匹配的微块保证中间 y 块在 L1/L2 cache 中而不是 DRAM。融合后的算子接口略微复杂因为它要接收归一化参数、权重、偏置三个来源。但它省去的 DRAM traffic 非常可观。在一次对比实验中融合后的端到端耗时下降了约 12%而单个算子本身的计算时间几乎没有变化。这再次印证很多性能收益来自“少移动数据”而不只是“更快地计算”。4. 调用技巧性能库里最容易翻车的都在接口侧4.1 用张量描述符传递参数不要裸指针满天飞很多人调用高性能算子时习惯把数据指针和维度拆分传进去。一旦维度多了很容易在某一层不小心复制出一个非连续视图然后算子默默按照连续内存去读结果就是数据错乱或性能劣化。ops-math 的接口强制要求传TensorDesc内部会校验连续性、对齐、数据类型。连续但未对齐的指针在调试模式下会抛错。虽然这会让刚上手的人觉得麻烦但长期看值得它把“数据是否适合高性能内核执行”的问题提前暴露在编译和调试阶段而不是线上偶发。4.2 Workspace 预分配每个算子都 malloc 一次等于慢性自杀算子内部如果需要临时 buffer不能在每次调用时直接malloc。我们踩过类似的坑Softmax 在批量推理时需要一个大小可变的临时空间存放局部最大值和局部和每次调用分配 8KB看似微不足道但当推理线程以每秒几千次调用运行时分配器的开销会占到总耗时的 5%-8%。解决方法是所有临时空间由调用方通过SetWorkspace传入。算子只负责使用不负责生命周期。上层推理引擎维护一个内存池按线程隔离算子拿到的地址固定且对齐良好。这是高性能库和普通工具库最大的区别之一。4.3 批处理和小算子合并别让线程调度吃掉你的收益推理场景中有大量形状很小的算子例如单行的向量加、单次的 scale。每个算子调用一次线程池线程同步开销可能比计算本身还大。我们的方案是提供“多行批量”接口。同一个算子一次性接收若干独立行内部由调度器自动负载均衡。这样线程只被唤起一次却处理了几十上百行数据。调用方不需要自己管理并行性能库内部完成。这个设计对生产者很有吸引力能显著提升端到端效率。但它也有个代价接口变复杂需要调用方理解“同算子批量”和“不同算子融合”的区别。不过为了性能这个复杂度值得花。4.4 异步执行和预取在 CPU 上同样有效很多人一听到异步就觉得是 GPU 和网络的事情实际上 CPU 算子库也可以做预取。矩阵乘外层循环处理下一个 tile 时提前把数据从 DRAM 预取到 L2 cache可以隐藏内存延迟。实现时可以用__builtin_prefetch或者显式触发非阻塞读取。我们要小心预取距离不要太大太远会导致 cache 污染太近又起不到作用。在一台 x86 服务器上我们测试过预取距离 256 字节到 1KB 之间的差异峰值吞吐相差约 8%。这个需要针对实际内存延迟调整不能照抄。4.5 正确的性能测量姿势很多优化结论都是噪音最后还要强调一下测试方法。我们曾因为测量姿势不对得出过“新内核比旧内核快 30%”的虚假结论后来发现是机器降频。现在固定的测量流程是先跑 20 轮预热把分支预测器和 cache 状态“加热”到真实运行状态。使用固定 CPU 频率关闭动态调频带来的干扰。连续取 100 轮以上去掉最高和最低各 10%用中位数对比。计时采用底层时间戳接口而不是高层的std::chrono::system_clock。这不仅是为了论文式严谨更重要的是算子库的每个内核都会被上层反反复复调用哪怕 1% 的真实差异乘上几百万次调用都会变成不可忽视的端到端差距。5. 实测复盘一组算子的优化前后对比5.1 基准环境说明为了结果可复现这里说明一下测试环境一台双路 x86 服务器64GB 内存编译器使用 g 12开启-O3 -marchnative测试线程数固定为 16CPU 锁定在基准频率。测试的方法是上面说的预热加多轮取中位数。提示以下数据是我们在自己环境下测到的相对结果。硬件不同、编译器版本不同具体数值会有差异但优化趋势和相对收益是可以参考的。5.2 LayerNorm 优化案例从 198 微秒降到 47 微秒一个较大的 LayerNorm 输入形状为 [256, 4096]初始版本用朴素逐行循环实现单次耗时约 198 微秒。第一步优化把内部循环改成 SIMD 向量化不改变算法结构降到 121 微秒。第二步改成在线归约减少一次数据读取降到 86 微秒。第三步加入多线程按行分块负载均衡后降到 47 微秒。回顾每一步收益最大的是“减少数据访问次数”而不是向量化。 向量化本身大概贡献了 40% 的加速但去掉一次读写又贡献了 30%多线程在行数足够多时贡献了剩下比例。这个顺序也印证了第 2 节“先算访存账”的观点。5.3 矩阵乘 micro-kernel 调参分块大小真的需要逐个试矩阵乘 512x512 的矩阵我们先后测试了 4x4、8x8、6x16、8x16 几种寄存器分块。测试下来 4x4 由于数据复用不够速度明显偏慢8x8 中规中矩6x16 在当前机器上最好比 8x8 快约 9%原因是减少了寄存器溢出和内存往返。这给我一个真实经验如果条件允许把 kernel 的可变参数暴露出来用脚本自动跑一遍搜索而不是靠论文里的推荐值一锤定音。 硬件微架构差异会导致最优参数漂移固定参数很容易让你的库只在某一台机器上表现优秀。5.4 常见调优翻车点记录调参过程中踩过的坑也要记录方便以后自查。现象根因解决办法并行后反而变慢行数少线程同步开销超过计算收益小 batch 走单线程内核相同代码在不同批次性能波动大内存分配未对齐访问跨 cacheline强制 64 字节对齐某个形状突然掉点分支预测失效或 tail 处理变成随机访问用 padding 避免边界分支并行线程互相拖慢假共享多个线程写同一 cacheline 的不同字节每个线程独立 buffer按 cacheline 对齐每一条翻车记录后来都转成了代码注释和开发者文档。原因很简单过几个月再看代码的人很可能不是你。6. 踩坑记录与把算子库做成可持续维护的地基6.1 编译器自动向量化不是万能的有一段 Relu 代码编译器自动向量化得很好但同样风格的 GeLU 代码因为包含较复杂的数学函数编译器会退化成标量循环。我们不能假定编译器会把所有循环都向量化而是需要在关键路径上做手工 intrinsic同时保留一份标量 fallback。每次改用 intrinsic 之后还要回到差分测试框架里跑一遍防止精度偏差。6.2 差分测试与数值边界要尽早做算子库最怕“看起来很快结果算错”。我们在 ops-math 内部维护了一套随机测试生成均匀分布、正态分布、极端大数、极端小数、NaN 和 Inf 的组合统一和参考实现比对。对浮点误差设定容差一般按 ulp 上限校验而不是直接相等。这样做有一个意外的收益很多调优过程中发现的边界 bug都不是主路径的问题而是尾部元素处理或 reduce 顺序不同导致的。如果等到集成到上层再暴露排错成本会高很多。6.3 新算子加入的标准流程现在团队内部新增一个算子会固定走下列几个步骤先写参考实现只求正确和明确数值边界。写差分测试覆盖随机输入和边界输入。再写优化内核按“访存账”确定优化方向。跑 benchmark和当前最优开源方案或框架实现对比。记录阈值和适用形状说明进入分派表。最后写一段简短文档说明优化原理防止后人误改。这套流程不复杂但能保证库的质量不会随着算子数量增加而崩盘。算子库的地基属性决定了它必须稳定一个新算子带进一个隐蔽 bug会让上面所有依赖它的模型全军覆没。6.4 面向未来的扩展量化、稀疏和异构ops-math 目前主要面向 CPU 浮点计算但设计时留了扩展余地。量化算子可以复用同一套分派表和 workspace 机制只需要增加新的数据类型枚举稀疏计算可以在矩阵乘内部走稀疏路由异构设备则可以把执行后端抽象成接口CPU、GPU 统一调用入口。不过这都属于下一步的事。在我个人看来先把浮点算子的访存账算透、把分派和测试体系打磨稳比急着铺开一堆新算子更划算。最后分享一个我在这个项目中反复体会到的经验算子库的性能天花板往往不是由某个精妙的指令决定的而是由数据移动次数和缓存行为决定的。 每次你看一个算子觉得“还能再优化”但无从下手时不要盯着计算回头看看它到底读写了几轮内存。把这一点记住很多优化思路会自动浮现出来。
阅读完成 · 觉得有帮助？
