去年接手一个推理引擎的提速任务,整个团队最初的心态是“哪里慢就调哪里”。结果忙了三周,改完接口改内存,改完内存改并行,端到端延迟反而因为引入的各种补丁变得更不稳定。后来我们共同做了一个决定:把所有高频小算子从框架里抽出来,单独做成一个高性能算子库,也就是这篇要说的 ops-math。围绕 ops-math 算子库的高性能实现与调用技巧,我想把从设计算账、内核编写、参数调优,到调用侧如何不拖后腿的整套经验都摊开讲清楚。如果你也在做推理引擎、离线算子优化,或者只是想在报错面前搞清楚算子到底是怎么跑的,这篇文章应该能让你少走几步弯路。
1. 为什么要单独维护一套算子库,而不是继续用现成框架
1.1 被框架黑盒卡住的真实场景
很多人觉得算子不就一个函数嘛,框架里现成,何必自己重复造轮子。我们最初也这么想。但当你真的去追性能问题,会发现框架里的算子是一个“黑盒”:它不暴露内部数据布局,不会告诉你中间态是否落过内存,也不会为你的业务场景专门调整分块大小。
有一次我们优化一个生成模型,profiling 显示整体延迟总有一个固定底座去不掉。翻来覆去排查,最后发现是框架在 softmax 和后续线性层之间强制生成了一份中间张量,这个张量完整写回内存又读出来。要是一层两层无所谓,层数一多,这种中间态读写就会把内存带宽烧掉大半。也就是从那时起,我们决定自己掌控算子内部的访存路径。
单一算子自包含,意味着你可以针对自己的输入规模、对齐要求、并发模型做定制。这其实和高性能计算领域一贯的做法一致:框架负责构图和调度,算子库负责把底层数学函数榨干。两者边界清晰,各自才能优化到位。
1.2 profiling 数据里最咬人的往往不是矩阵乘
我们引入自研算子库之前,先在一个 16 层的 Transformer 结构上做了完整 profiling。 统计结果大致的耗时分布如下:
| 算子类型 | 耗时占比 | 访存性质 | 优化收益 |
|---|---|---|---|
| GEMM / GEMV | 48% | 计算密集 | 高,但实现难度大 |
| LayerNorm / Softmax 等归一化类 | 17% | 访存密集 | 很高 |
| Elementwise(GELU、Add、Mul) | 8% | 访存密集 | 中高 |
| Transpose / Reshape / Gather | 6% | 访存密集 | 高 |
| 其他调度与框架开销 | 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 个 float:
void 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 -march=native,测试线程数固定为 16,CPU 锁定在基准频率。测试的方法是上面说的预热加多轮取中位数。
提示:以下数据是我们在自己环境下测到的相对结果。硬件不同、编译器版本不同,具体数值会有差异,但优化趋势和相对收益是可以参考的。
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 统一调用入口。
不过这都属于下一步的事。在我个人看来,先把浮点算子的访存账算透、把分派和测试体系打磨稳,比急着铺开一堆新算子更划算。
最后分享一个我在这个项目中反复体会到的经验:算子库的性能天花板往往不是由某个精妙的指令决定的,而是由数据移动次数和缓存行为决定的。 每次你看一个算子觉得“还能再优化”但无从下手时,不要盯着计算,回头看看它到底读写了几轮内存。把这一点记住,很多优化思路会自动浮现出来。