去年帮一位朋友调一个AIGC模型的推理延迟，我把整整两天耗在一个看起来很小的点上：把残差连接和LayerNorm从两个算子合并成一个。模型跑在昇腾平台上，算子库里的现成方案不算少，但那个场景的batch load特别重，每个step都要搬运大量中间张量，访存开销高得离谱。真正把延迟压下去的那一刻，我突然意识到，AIGC时代性能的杠杆，很多时候根本不在模型结构，而在最底层的自定义算子。这篇文章就把那段时间积累的东西拆开聊聊——用CANN ops-nn做AIGC自定义算子，以及围绕它做性能优化的完整思路。
如果你正在做Transformer、Diffusion这类模型的移植或加速，又不满足于算子库现有接口的性能表现，这篇内容应该对你有用。我会从路径选型、算子实现、性能调优讲到真实踩坑，尽量把每个"为什么这么做"都说清楚。
1. AIGC结构迭代太快，算子库开始跟不上模型的"新形状"
1.1 模型每升级一次，底层算子形态就变一次
AIGC模型迭代速度有多快，做过部署的人都清楚。注意力机制从原始多头改到Grouped Query Attention，激活函数从ReLU换成SwiGLU，归一化从LayerNorm换到RMSNorm，位置编码从绝对位置换成RoPE。每一次结构变化，在计算图上都是一组新的算子组合。
标准算子库覆盖的是"成熟且被广泛使用"的算子。一个算子要沉淀到库里面，需要经过需求收集、设计、开发、验证、发布，周期并不短。可AIGC模型的实验节奏是按周甚至按天走的，两者天然存在时间差。这个时间差造成的缝隙，就得靠自定义算子补上。某个开发者经常跟我说一句话：跑通一个最新结构的模型，往往不取决于模型代码写得多好，而取决于你能否把新结构里的关键算子快速落地成高效实现。
1.2 融合算子的意义，不只是省一次调度开销
自定义算子里很大一类是融合算子，比如把QK^T、softmax、V加权融合成一个Attention算子，把残差加法、归一化、Dropout融合成一个Block算子。很多人以为融合只是减少kernel launch次数，这个理解太浅了。
kernel launch节省的那点CPU开销，在推理场景通常是微秒量级，真正的大头是融合消除了中间张量的读写。以Attention为例，如果拆开计算，QK^T的结果是一个[batch, head, seq_len, seq_len]的矩阵，这个中间量要写回显存，softmax之后还要再读出来，内存带宽就被白白消耗了两轮。AIGC模型里seq_len普遍在2048以上，这个中间矩阵的体积相当可观。融合之后，整块数据留在片上寄存器或高速缓存里，省掉的访存时间往往能把算子端到端耗时减掉一半以上。
1.3 "异构计算缝隙"这个概念值得多说两句
拿昇腾这类异构计算平台举例，它的芯片上除了通用CPU核，还有专门做矩阵运算和向量运算的AI Core。跑一个算子时，数据要从外部内存搬进AI Core的片上缓存，算完再搬出去。每个算子都重复这个搬运过程。算子拆得越碎，搬运次数越多，性能损耗越大。
融合算子本质上是把"多次搬运"压缩成"一次搬运"，把多个计算密集操作放在同一批已加载的数据上完成。这个概念不仅适用于Attention，也适用于Diffusion里的U-Net、LLM里的FeedForward。理解了这个缝隙，就能理解为什么在AIGC场景做自定义算子不是锦上添花，而是绕不开的步骤。
2. 开发路径选型：DSL脚本还是原生算子，别急着写代码
2.1 CANN ops-nn里两条路的基本面
CANN的算子开发框架里，大致有两条技术路线。一条是基于DSL的脚本式开发，写起来像是在写数学公式，描述输入输出张量和中间计算的逻辑关系，框架负责把它编译成芯片能执行的指令；另一条是更贴近底层的原生C/C++开发，直接操作缓冲区、指针和指令流，灵活度更高，代价是要自己管理的东西也更多。
两条路线没有绝对的好坏，只看项目处在哪个阶段。我见过不少团队一上来就选原生路线，结果写了一个月还在调整边界逻辑。也见过全程用DSL实现，结果性能始终压不下去，最后不得不回炉重写。选错路线的代价，比写错代码的代价大得多。
2.2 我的选型判定标准：访存密度和算子生命周期
我的判断标准有两个：访存密度和算子生命周期。
如果算子逻辑以向量化逐元素操作为主，数据搬运模式简单，比如只是把两个向量相加再做归一化，那DSL路线足够，开发效率高，代码也容易维护。这种算子即便性能差一些，在一个完整模型里占的时间比例也有限。
如果算子核心是复杂的数据分块和循环重排，比如Attention里需要同时在行和列两个方向切分数据，或者算子要获得极致性能就得手动安排多级缓存的流转，那我会直接选原生路线。这类算子一旦进入模型主循环，生命周期很长，值得投入精力打磨。AIGC场景里绝大多数值得自定义的算子，都属于这一类。
2.3 环境准备中最容易忽略的一步：确认版本和算子注册表
不管是哪条路线，环境准备里最容易被忽略的不是编译链，而是算子的注册和调度配置。CANN平台上自定义算子要能被模型调用，需要在算子的描述文件中声明输入输出的数据格式、精度类型、分块策略等元信息。漏掉一个layout字段，编译时可能不报错，跑到模型加载阶段才暴露出算子找不到的错。
我的建议是，动手写计算逻辑之前，先用一个最简单的"拷贝算子"或"加和算子"走通全流程：编写、编译、注册、在模型侧调用、跑通、打印结果。全链路通了，后面所有复杂逻辑都是在一条已知可行的路径上叠加。很多入门者跳过了这一步，直接写复杂算子，出了问题根本分不清是代码逻辑问题还是框架集成问题。
3. 按"数据流向"拆解多头注意力融合算子的实现
3.1 算子规格设计：把计算图先翻译成数据流
拿一个典型的多头注意力融合算子举例。输入通常有Q、K、V三个张量，形状类似[batch, head, seq_len, head_dim]，还可能带一个attention mask。输出是[batch, head, seq_len, head_dim]的加权结果。算法逻辑很简单：对每个位置的query，计算它与所有key的点积，除以缩放因子，做softmax，再对V加权求和。
落到算子设计上，第一步不是写计算逻辑，而是设计数据如何流动。芯片上的计算核心有片上缓存，但容量有限，seq_len=2048个位置的QK^T结果不可能一次性全放进去。所以必须做分块：把seq_len方向切成若干小块，每一块完整走完"乘Q、算softmax、加权V"的过程，再滑到下一块。这个切块大小既要匹配片上缓存容量，又要保证单块的计算密度足够高。
3.2 核心计算循环的两种组织方式
这里有个关键设计决策：外层循环按"行"走还是按"列"走。按行（query维度）切分，意味着一组query对应的所有key都要遍历一遍，谁先谁后由数据缓存策略决定；按列（key/value维度）切分，则是对所有query同时处理，但每次只加载一部分key和value，做部分累积结果。
两种方式我都测过。按query行切分时，softmax的计算是完整一次性的，不需要跨块做归一化校正，逻辑简洁，但每个query要重复读到所有key，访存量偏大。按key列切分时，访存更友好，因为key和value是流式读取的，天然适合流水线，但softmax的分母需要在多块之间做跨块累加，需要用online softmax的技巧维护一个局部最大修正项。
AIGC的主流做法是采用后者。访存上的优势在长序列场景非常明显，多出来的那点逻辑修正复杂度，相比省下的带宽开销，完全值得。
3.3 一段能表达思路的伪代码
用伪代码描述核心逻辑，能更清楚地看出数据流结构。以下是我那个算子某个版本的简化骨架，略去内存管理细节：
# 伪代码：FusedAttention one-block computation # block_seq表示一个分块内的序列长度 for q_block in range(seq_len // block_seq): load_q(q_block) # 加载一组query到片上缓存 for kv_block in range(seq_len // block_seq): load_k(kv_block) load_v(kv_block) for i in range(block_seq): for j in range(block_seq): score[i][j] = dot(q_block[i], k_block[j]) * scale if mask is not None: score[i][j] = where(mask[i][j] == 0, -inf, score[i][j]) local_max = row_max(score) local_sum = row_sum(exp(score - local_max)) # online softmax 修正跨kv_block的归一化 update_global_max_and_sum(global_max, global_sum, local_max, local_sum) acc += exp(score - global_max) @ v_block里面每个循环的i和j最终都会被向量化成底层批量指令，但逻辑结构保持这样，读代码的人一眼就能看出哪些数据被复用了、哪些步骤存在跨块依赖。
3.4 跑通顺序：先不优化，先确保数值可对齐
写完之后不要急着调性能。我会先用一个很小的序列长度，比如seq_len=64，跟框架自带的原生Attention结果做数值对齐，误差阈值设在1e-3以内。这个阶段出现精度偏大，通常不是计算逻辑有问题，而是softmax的修正项写错了。等数值对齐了，再上板测时间。
4. 性能优化纵深：带宽、并发、指令三个层面的博弈
4.1 访存优化：把搬运次数逼近理论下限
性能优化的第一刀，几乎永远是访存。一个算子实际能跑多快，很大程度上由它从外部内存搬了多少数据决定。昇腾平台上AI Core有独立的高速缓存空间，数据从外部内存搬进缓存的速度远低于计算速度，所以计算单元的利用率往往取决于能不能把数据访问隐藏好。
实践里最有效的手段是double buffer。简单说，就是计算当前一块数据的同时，预取下一块数据到备用缓存，计算和搬运在时间上重叠。这样AI Core不会空等数据到达，带宽被持续占满。我见过不少性能不达标的算子，问题不是算法慢，而是没开双缓冲，AI Core大部分时间在等待数据就位。
另一个方法是控制分块的大小。块太小，搬运次数多，每次搬运的固定开销占比高；块太大，缓存放不下，反而触发额外的换入换出。这个尺寸没有通解，得结合序列长度和缓存容量实测，通常用二分的方式找到拐点。
4.2 多核并发：分块粒度决定AI Core吃饱的程度
单核优化做到一定程度后，真正的提速来自多核并行。芯片上有多个AI Core，如果算子只用一个核，性能天花板就是单核能力；如果切给多个核，理论上可以线性翻倍。
但多核并发有个前置条件：分块要互相独立。Attention按head维切分时逻辑最干净，因为不同head的QKV互不相干，计算过程中的部分和也不需要跨核合并。这是我推荐的首选切分方向。如果head数量不足以喂饱所有核，再叠加序列方向的切分，但这时要注意部分和归约的开销，跨核通信和归约如果处理不好，收益会被延迟吃回去。
我实际项目里的经验是：多核分配的比例按处理时延来动态调整，谁先算完谁多领下一块，比静态均分好不少。调度逻辑多几十行代码，换来的是整卡吞吐十几个百分点的提升。
4.3 指令级优化：向量化、循环展开与缓存冲突
更往下走，就是指令层面的功夫。AIGC算子装满了矩阵乘和逐元素操作，恰好都是芯片向量单元的强项。要让向量单元不空转，循环一定要写成可以被编译器批量向量化的形式。相比让编译器去解复杂的数据依赖，更靠谱的做法是算子代码里直接按向量处理的节奏来组织数据布局。
循环展开也值得做。展开可以减少循环控制指令的占比，让向量指令流水更紧凑。但展开因子不是越大越好，它会增加寄存器压力，寄存器溢出后效率反而下降。我通常从4开始试，观察时间曲线选拐点。
缓存冲突是个容易忽略的坑。当多个并行任务访问片上缓存时，如果地址恰好在同一组索引位置，会产生冲突互斥，访问被迫串行。调整数据在缓冲区里的摆放偏移量能显著改善这个现象，这个操作听起来玄学，实测往往能带来10%以上的提升。
4.4 一个验证过的调优工作节奏
分享一个我自己习惯的节奏。先用profiling工具确认当前瓶颈是访存还是计算，这个判断决定后续方向。然后按"分块尺寸 -> 双缓冲 -> 多核切分 -> 指令展开"的顺序依次动手，每改一步就全量验证正确性，再记录耗时。
这样做的原因是每一层都有耦合，一次性全改动，出问题根本定位不了。分步调整虽然慢，但每一步都清楚知道收益来源。最终性能曲线往往不是匀速下降，而是平台期和跳变交替，跳变通常对应某个资源瓶颈刚被解除的时刻。
5. 实测数据复盘：一个融合算子从慢到快的全过程
5.1 基线：把功能跑对，性能难看没关系
当时测试的是seq_len=2048、head_dim=128、单batch、8个头的Attention融合算子。第一版逻辑正确，没做任何优化，跑出来的耗时大概是12.6毫秒。用profiling一测，AI Core利用率只有20%，大部分时间都在搬数据。
这个基线数据其实是良好的开始。它至少证明逻辑没问题，后面所有优化都有参照系。我见过有人第一版就跑出不错的数据，结果反而不知道该往哪个方向投入精力，基线差一点不是坏事。
5.2 第一刀：调整分块尺寸
把分块尺寸从512一路往下试，分别测了256、128、64。曲线的趋势是块越小越好，但到64之后变化不再明显。最终选了128作为分块尺寸，理由是它在缓存占用和搬运次数上比较均衡。这一轮之后耗时降到7.8毫秒，吞吐提升主要来自缓存命中率改善和搬运次数收敛。
5.3 第二刀：双缓冲加多核切分
接着开双缓冲，把搬运和计算重叠起来，耗时从7.8毫秒降到5.1毫秒。然后是8个head切到8个核并行，因为每核独立处理一个head，不需要跨核归约。这次降到了1.6毫秒，效果是断崖式的。做AIGC算子时，head维度的天然并行一定要优先吃掉。
5.4 第三刀：省掉一次多余的中间写回
最后压时间来自一个细节。初版代码在算完softmax后会把归一化的概率写了一次缓存再读回来做V加权，这个中间步骤其实可以省略——直接在一个累加器里完成概率计算和加权两个操作。去掉这次存读后，耗时降到1.2毫秒左右，相比基线提升超过10倍。
5.5 和未融合版本做个对照
同样配置下，用三个独立算子完成同样的Attention计算，总耗时约3.4毫秒。也就是说，融合加优化的版本比直接拆开算子快将近三倍。这组数据也侧面回答了一个问题：为什么费力气做算子融合——它不是锦上添花，而是直接把一个关键路径的延迟压到了三分之一。
6. 踩过的几个坑：从数值偏差到偶发挂死的排查链路
6.1 输出对不上：先怀疑精度，再怀疑逻辑
第一次跑融合算子时，输出和参考实现差了2%左右。起初怀疑是softmax修正项写错了，反复看了好几遍逻辑都没发现问题。后来换了更高精度的中间累加类型，偏差直接消失了。原因是softmax里指数操作对输入的动态范围非常敏感，某些序列位置上中间结果的绝对值偏大，单精度累加的舍入误差被放大了。
这个经历换来一条经验：做融合算子时优先使用高精度的内部累加器，只在最终输出时截断到目标精度，性能损失很小，正确性却稳得多。
6.2 偶发挂死：先看日志和core dump，不要瞎试
有段时间算子在长序列下会偶发挂死，没什么规律，复现概率很低。一开始我以为是显存不足，试了各种配置都没用。后来仔细翻了系统的日志和内存转储文件，定位到是分块计算时某个边界索引越界，在多数情况下恰好访问到无害的地址，只在某些数据长度组合下踩到了非法区域。
这个排查链路值得记下来：偶发挂死的优先怀疑对象，不是随机错误，而是未被完整测试的边界条件。写算子时，边界索引有条件尽量写成断言形式，宁可主动报错，也好过偶发挂死后的几小时排查。
6.3 多核结果不一致：问题往往出在归约顺序
多核并行版本出来后，发现单batch结果一致，但多batch偶尔对不上。排查到最后是部分和归约顺序的问题。浮点加法不满足结合律，不同归约顺序会产生微小差异。模型推理对微小的数值差异通常不敏感，但如果你在做面向精度边界测试，就要注意这个问题。
6.4 一组适用的调试清单
最后整理一个调试清单，是我每次写自定义算子都会过的：
- 边界索引是否全覆盖，有没有写成断言
- 内部累加器精度是否高于输出精度
- 分块是否能被缓存完整容纳
- 多核切分是否有跨核依赖
- 中间结果的存读是否真的必要
- 每个改动点是否单独做了数值比对
7. 写在最后的几条体会
这几条是我反复栽跟头后才固化的做法，分享出来希望你能少走弯路。
第一，AIGC模型迭代快，但算子开发其实是可以沉淀的。一旦你为一个模型写好了一套融合算子，下一个模型大概率能复用其中大半。真正通用的不是某个算子的代码，而是你拆解计算图、设计数据流、做性能剖析的那套方法。
第二，不要迷信理论上的计算密度，实际运行时间永远以剖析数据为准。有些算子按计算量看应该是访存瓶颈，实测却是计算瓶颈，原因可能是向量指令没有被有效利用。纸上算得多仔细，都不如跑一遍profiling来得准。
第三，性能优化的每一步改动，都要能独立验证并保留历史版本。你当时以为没用的参数，可能在其他配置下才是关键。留好改动记录，方便回头横向对比，也比凭记忆乱试更高效。
第四，团队里如果有不同技术栈背景的人一起做算子开发，一定要统一中间数据摆放的约定和命名习惯。这个听起来虚，但在多模块协作时，它节省的沟通时间可能比写算子本身还多。
自定义算子在AIGC推理优化里确实是个投入产出比很高的方向。只要把数据流、缓存、并行这几个核心问题想清楚，性能提升的空间往往会超出你的预期。
