如果你跑过基于Transformer的模型推理，或者是接手过线上服务的性能优化，一定对“算子融合”这个词有切身体会。PyTorch作为目前最主流的深度学习框架，用起来确实方便，但默认执行模型有一个天生的短板：一个数学表达式会被拆成一长串独立的GPU kernel，每个kernel各自启动、各自从显存里搬运数据。真正让性能暴跌的，往往不是计算本身，而是这些“各自为战”的小算子在来回搬运数据时浪费掉的带宽。算子融合，就是把多个kernel合并成一个，减少访存和kernel启动开销，这是我在实际项目里见效最快、也最需要抠细节的优化手段。
这篇文章我会结合自己用PyTorch做算子融合的实战经验，把为什么融合能提速、有哪些融合路线、手动写CUDA融合kernel怎么做、Flash Attention式的访存优化怎么落地、torch.compile自动融合什么时候靠谱又什么时候失效，全部梳理一遍。无论你是做推理优化、训练加速，还是想把PyTorch模型塞进边缘设备，这篇文章都值得认真看一遍。
1. 为什么算子融合能带来性能质变：从GPU执行模型说起
1.1 访存瓶颈：一个算子一条命，边界数据来回搬
要理解算子融合的价值，先得理解GPU上的“命门”是什么。GPU算力这几年增长非常快，单块A100的FP16矩阵算力已经到312 TFLOPS，H100更是翻倍。但是显存带宽的增长却慢得多，A100的HBM带宽约2TB/s，H100约3.35TB/s。也就是说，GPU是一个“算得快、搬得慢”的处理器。
对于像y = relu(x + b1) * scale + b2这样的表达式，如果你直接用PyTorch原生的写法：
y = torch.relu(x + b1) * scale + b2PyTorch会按照计算图依次启动kernel：
| 步骤 | 实际执行 | 数据读写 |
|---|---|---|
| 1 | x + b1 -> tmp1 | 读x、读b1、写tmp1 |
| 2 | relu(tmp1) -> tmp2 | 读tmp1、写tmp2 |
| 3 | tmp2 * scale -> tmp3 | 读tmp2、写tmp3 |
| 4 | tmp3 + b2 -> y | 读tmp3、读b2、写y |
这还只是一个很浅的表达式。一个10MB的tensor，每个算子都要把10MB数据从HBM读进寄存器、计算完再写回HBM。一次完整表达式跑下来，访存量是单次操作的好几倍。如果你用torch.profiler去测，会看到4个独立的kernel，其中每一个的耗时都差不多，因为瓶颈不在计算，而在把数据从显存搬进搬出。
1.2 kernel启动开销：别小看那几微秒
除了访存，还有另一个隐藏成本：kernel launch。每次调用一个GPU kernel，CPU都要向GPU提交启动命令，这个过程的延迟大约是3到10微秒，取决于驱动和上下文状态。如果kernel本身执行只需要几十微秒，启动开销就是实打实的损耗。
更要命的是，深度学习模型动辄几十层，每层又拆成多个算子。在一个标准的Transformer推理中，一层可能触发二十多个kernel。模型越大，kernel数量越多，启动开销越积越多。在GPU利用率不高的访存密集场景下，启动开销甚至可以占到总耗时的20%到30%。
我在做移动端和边缘设备优化时感受更强烈。那些设备上的GPU驱动更薄、启动开销占比更高，有时候把50个kernel融合成3个，整个模型的延迟能直接砍掉一半。所以算子融合这条路，本质上是同时解决“访存浪费”和“启动浪费”两个问题。
1.3 什么算子值得融合、什么不值得：Roofline模型的直觉
既然是优化，就要先判断值不值得。业内常用Roofline模型做理论分析，核心指标叫算术强度（Arithmetic Intensity）= 操作数（FLOPs） / 访存量（Bytes），可以理解为“每搬运1字节数据时能做多少次计算”。
- 算术强度很低的算子，比如
relu、add、scale、softmax、LayerNorm、dropout，都属于访存密集算子。每个元素只做几次浮点运算，却要读一遍、写一遍，性能上限完全被带宽卡死。这类算子是最值得融合的。 - 算术强度很高的算子，比如大矩阵乘法、卷积，属于计算密集算子。只要tile切分合理，数据在寄存器/共享内存里反复利用，访存瓶颈已经被压得很低。融合周围的小算子（bias、relu）也能提速，但收益主要来自减少中间写回，而不是“少读了几遍数据”。
我的判断标准很简单：先用torch.profiler跑一遍，看哪些kernel的耗时最长、算术强度最低，然后把它们周围能合并的元素级操作全并进去。对访存密集算子来说，融合两个kernel通常能获得接近一倍的性能提升，因为读写各省了一次。
2. PyTorch算子融合工具箱全景：从手动到自动的三条路线
2.1 三条路线的横向对比
面对一个问题：我想把add -> relu -> scale融合成一个kernel，有哪些手段？我按“控制力”从高到低把常用方案排了一遍：
| 方案 | 易用性 | 性能上限 | 维护成本 | 适用场景 |
|---|---|---|---|---|
| 手动写CUDA Extension | 低 | 极高 | 高 | 固定热点算子，追求极致性能 |
| torch.fx图变换 | 中 | 中高 | 中 | 在Python层做pattern替换，或作为自定义优化pass |
| torch.compile自动融合 | 高 | 高 | 低 | 大多数训练推理场景，一行代码接入 |
| 第三方库（如Triton手写） | 中 | 高 | 中 | 需要灵活控制tile和schedule时 |
这里想多说一句：很多初学者以为有了torch.compile，就完全不需要手动融合了。但在实践中，torch.compile并不能覆盖所有场景。自定义autograd Function、动态shape、复杂的控制流都可能让它“打断图”，退回逐算子执行。而手写CUDA kernel虽然费时，但对于一个已经稳定运行的热点路径，收益是实打实的。后面我会用两个实战案例把这两条路都走一遍。
2.2 torch.fx：在Python层做“手术”
torch.fx是PyTorch官方提供的符号化trace框架。它会把一个nn.Module或函数转换成一张Graph，图里的每个节点是一个算子。有了Graph，你就可以做pattern匹配，把匹配到的多个节点替换成一个融合节点。
举个例子，如果你想实现“把add -> relu -> mul替换成一个自定义算子FusedAddReLUScale”，大概的流程是：
import torch.fx as fx class FuseAddReLUScale(fx.Transformer): def call_function(self, target, args, kwargs): # 自定义匹配逻辑 pass # 用 symbolic_trace 拿到图 graph_module = fx.symbolic_trace(MyModule)torch.fx的优点是纯Python层面、不碰C++，也容易调试。缺点是它本身不执行融合，只是给了你改图的能力；真正的融合kernel还得你自己实现。所以我的经验是：torch.fx适合用来搭建自己的“微型编译器”，但如果你只是想快速提升模型性能，直接上torch.compile更实在。
2.3 torch.compile：编译器接管一切
torch.compile是PyTorch 2.0开始的主推方案，它内部由Dynamo（字节码分析）、AOTAutograd（反向图捕获）、Inductor（代码生成，默认生成Triton kernel）三层组成。
它的使用简直无脑：
compiled_model = torch.compile(model, mode="reduce-overhead")编译后，Inductor会把相邻的元素级算子自动融合成一个Triton kernel。对大多数标准模型来说，torch.compile能直接带来20%到40%的训练吞吐提升，推理场景提升更多。但是，它并不是银弹，后面第五章我会专门讲它什么时候失效，以及怎么排查。
3. 实战一：手写CUDA融合kernel，把ReLU+Add+Scale压缩成一个内核
3.1 基线：原生PyTorch的几个kernel到底有多慢
先定义一个典型的访存密集算子链：
import torch import torch.utils.benchmark as benchmark def baseline(x, b1, scale, b2): tmp = x + b1 tmp = torch.relu(tmp) tmp = tmp * scale return tmp + b2 x = torch.randn(4 * 1024 * 1024, device="cuda") b1 = torch.randn_like(x) scale = 0.5 b2 = torch.randn_like(x)用torch.utils.benchmark.Timer跑100次取均值，在A100上，这个表达式通常耗时在80到120微秒之间（数据量16MB，读写好几遍）。用torch.profiler看，能看到add、relu、mul、add四个kernel依次执行。
对这个简单的表达式，理论最优情况是：读一次x、读一次b1、读一次b2，然后写一次y，总计大约64MB访存。以2TB/s的带宽算，理论上限大约是32微秒。这说明基线版本离理论值还有3倍左右的差距，而这差距就是融合能找回的。
3.2 融合实现：一个kernel搞定所有元素级操作
我写一个简化但可用的CUDA扩展。为了达到更高的访存效率，我用float4向量化读写，一次处理4个float，这会让显存带宽利用率明显提升。
#include <torch/extension.h> #include <cuda_runtime.h> __global__ void fused_add_relu_scale_kernel( const float4* __restrict__ x, const float4* __restrict__ b1, const float4* __restrict__ b2, float4* __restrict__ y, float scale, int n4) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n4) { float4 xv = x[idx]; float4 b1v = b1[idx]; float4 b2v = b2[idx]; float4 yv; yv.x = fmaxf(xv.x + b1v.x, 0.0f) * scale + b2v.x; yv.y = fmaxf(xv.y + b1v.y, 0.0f) * scale + b2v.y; yv.z = fmaxf(xv.z + b1v.z, 0.0f) * scale + b2v.z; yv.w = fmaxf(xv.w + b1v.w, 0.0f) * scale + b2v.w; y[idx] = yv; } } torch::Tensor fused_add_relu_scale( torch::Tensor x, torch::Tensor b1, torch::Tensor b2, double scale) { // 假设都是连续、float32、同shape auto y = torch::empty_like(x); int n = x.numel(); int n4 = n / 4; int threads = 256; int blocks = (n4 + threads - 1) / threads; fused_add_relu_scale_kernel<<<blocks, threads>>>( reinterpret_cast<const float4*>(x.data_ptr<float>()), reinterpret_cast<const float4*>(b1.data_ptr<float>()), reinterpret_cast<const float4*>(b2.data_ptr<float>()), reinterpret_cast<float4*>(y.data_ptr<float>()), scale, n4); return y; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("fused_add_relu_scale", &fused_add_relu_scale); }用torch.utils.cpp_extension.load_inline编译加载即可：
from torch.utils.cpp_extension import load_inline cpp_source = "" # 这里是上面代码的字符串形式 fused_mod = load_inline( name="fused_add_relu_scale", cpp_sources=[cpp_source], functions=["fused_add_relu_scale"], )这个kernel只分配一次输出，读取三个输入各一次，写入一次，访存量从基线的“读多次写多次”降到了“读一次写一次”。
3.3 性能对比与Profiler验证
融合后，我用同样的benchmark跑，实测耗时通常在25到35微秒之间，比基线快3倍左右。这个数字已经非常接近按内存带宽推算的理论极限了。用torch.profiler再看一眼，kernel数量从4个变成了1个，且名字是fused_add_relu_scale_kernel。
这里有几个实操要点：
- 为什么用
float4：因为现代GPU的访存指令以128字节对齐为最优，一次ld.global.v4.f32等于一条指令搬运16字节，能大幅减少指令数量，更容易把带宽打满。 - 为什么不用
grid-stride loop：数据量小的时候，一个线程一个元素正好能铺满；数据量特别大时，才需要用grid-stride loop让每个线程循环处理多个元素，避免线程块太多带来的调度开销。 - 对非
numel % 4 == 0的尾部处理，建议单开一个分支处理剩余元素，或者干脆在Python侧对tensor做pad。在实际项目里，后一种做法更省事。
手动融合的收益非常直观：访存少了、启动少了、带宽利用率反而高了。
4. 实战二：Flash Attention式的访存优化——融合注意力怎么做
4.1 标准Attention为什么慢：中间矩阵是访存黑洞
把视野从元素级算子放大到模块级，最有代表性的融合案例就是Attention。我先说一个扎心的事实：如果在PyTorch里老老实实写：
scores = torch.matmul(q, k.transpose(-2, -1)) / scale probs = torch.softmax(scores, dim=-1) out = torch.matmul(probs, v)那么当seq_len比较大的时候，中间矩阵scores和probs的尺寸是[batch, heads, seq_len, seq_len]。假设batch=1, heads=32, seq_len=4096，那probs就是32 × 4096 × 4096 × 4B ≈ 2GB。这个矩阵被完整地写回显存，又被读出来，再和V做矩阵乘。
这导致一个严重问题：Attention的访存量是O(n²)级别的。而真正的算法逻辑上，我们完全不需要保留完整的scores矩阵——每个query只需要最终输出，Softmax的归一化因子可以通过分块技巧在线更新。这就是Flash Attention的思想：把Attention融合成一个kernel，按block遍历Q、K、V，在寄存器/SRAM级别完成分块矩阵乘和在线Softmax，只把O(n)大小的输出写回HBM。
4.2 在线Softmax的分块更新逻辑
注意力最棘手的地方是Softmax依赖整行的全局最大值。如果我只做一个block的局部Softmax，后面来了更大的值，前面的归一化就全错了。解决办法是用一个可累积的running max：
# 伪代码，展示核心循环结构 m = q.new_full(..., float("-inf")) # running max l = q.new_zeros(...) # running sum of exp acc = q.new_zeros(...) # 分子累加器 for k_chunk in range(0, seq_len, chunk_size): k_chunk_t = k[:, k_chunk:k_chunk+chunk_size] v_chunk = v[:, k_chunk:k_chunk+chunk_size] s = torch.matmul(q, k_chunk_t.transpose(-2, -1)) / scale m_new = torch.maximum(m, s.max(dim=-1, keepdim=True).values) alpha = torch.exp(m - m_new) # 旧数据衰减系数 p = torch.exp(s - m_new) acc = acc * alpha + torch.matmul(p, v_chunk) l = l * alpha + p.sum(dim=-1, keepdim=True) m = m_new out = acc / l这段逻辑如果用原生PyTorch算子写，每次循环仍然会创建临时tensor，但只要把整个循环体放进一个支持融合的编译器里，或者直接用Triton编写，就能避免中间矩阵落回显存。
在实际工程中，我通常直接用Triton写Flash Attention核心循环，每个program处理一个query block，K和V分块加载到SRAM中。分块大小通常设置为64或128，head_dim是64时，整个K/V块都能放进SRAM，计算效率非常高。
4.3 验证正确性：别忽略了“微小差异”
融合Attention的精度验证比普通融合更讲究。由于在线Softmax的数值累积顺序与原生Softmax不同，结果会有微小浮点差异，这属于正常现象。但如果差异超过1e-5量级，就要检查是不是衰减系数alpha算错了。
我习惯用这样的方式验证：
def max_abs_diff(a, b): return (a - b).abs().max().item() # output ~= reference, 允许 1e-5 的误差另外要注意causal mask的融合。如果你对每个query i只允许attend到i之前的key，直接加一个大负数mask就能让Softmax自然变成0，但这样在分块时会有大量无效计算。更高效的做法是在循环里跳过那些完全mask掉的块，或者在分块内做对角线边界判断。这一块非常容易写出bug，建议配合随机mask做单元测试。
5. 实战三：torch.compile自动融合，以及它什么时候罢工
5.1 一段代码看自动融合效果
如果你不想手写kernel，torch.compile是最省力的路径。看一个典型的MLP块：
def mlp_block(x, w1, b1, w2, b2): h = torch.relu(x @ w1 + b1) return h @ w2 + b2 compiled = torch.compile(mlp_block, mode="reduce-overhead")你只需要把函数换成compiled版本，Inductor会做两件事：一是把matmul + bias + relu融合成一个支持bias和relu的gemm kernel（在NVIDIA上通常调用cuBLAS的epilogue功能，或者生成Triton矩阵乘）；二是把第二个matmul + bias也融合。最终你的函数可能只启动1到2个kernel。
在mode="reduce-overhead"下，还会启用CUDA Graph捕获取，进一步削减kernel启动开销。对一个8层的小Transformer，我实测编译后推理延迟能降低35%以上。
5.2 融合失效的常见原因：graph break和它的代价
torch.compile也不是万能的。Dynamo通过字节码分析来捕获Python执行轨迹，只要碰到它无法静态分析的内容，就会产生graph break，通俗说就是“编译到这里断掉了，之后的算子退回Eager模式”。
我踩过的坑大概有几类，整理成表：
| 触发条件 | 典型例子 | 后果 |
|---|---|---|
| 动态shape相关分支 | if x.shape[0] > 8: | graph break，后面的调用走Eager |
| 自定义autograd.Function | 直接调用自定义反向函数 | 无法trace |
| 依赖tensor内容的控制流 | if x.item() > 0: | Dari图中断，且同步CPU/GPU |
| 动态数据结构的操作 | list append后循环 | 生成低效代码 |
| 外部Python对象的复杂交互 | 回调函数、闭包 | 可能直接回退 |
排查方法很直接：用torch._dynamo.config.log_code或者torch._dynamo.explain，再用优化标志把graph break点打印出来：
from torch._dynamo import explain explanation = explain(fn, *args) print(explanation.graph_break_count) print(explanation.break_reasons)如果你发现某个热点区域发生了graph break，我的建议是：把这一小段函数用torch.compile单独编译，并保证函数内部没有分支依赖运行时数据。如果还有动态shape问题，可以尝试用torch._dynamo.mark_dynamic显式标记，或者干脆把该段写成Triton kernel。
5.3 用Inductor生成的Triton代码验证“是否真的融合了”
有时候你以为编译成功了，实际性能却没提升。这时候别急着怀疑硬件，先看看Inductor到底生成了什么kernel。一个非常有效的办法是打开编译产物日志：
import torch._inductor.config as ind_cfg ind_cfg.trace.enabled = True跑一次后，会在当前目录生成torch_compile_debug文件夹，里面能看到Dynamo的graph dump，以及每个Triton kernel的实际源码。在生成的__kernel_name函数里，如果看到类似x + b1、relu*出现在同一个kernel body的多个语句里，说明融合成功；如果发现某个小算子单独生成了kernel，说明它被拆出去或者融合失败了。
另外推荐用torch._inductor.config.debug打印kernel调度顺序。配合torch.profiler看kernel名字，几乎能定位90%的融合失效问题。
6. 避坑指南与性能验证方法论：融合完怎么确认真的变快了
6.1 正确性验证：融合后的数值差异要控制在什么范围
回到最基础的问题：融合算子重排了浮点运算顺序，结果和原生PyTorch对照必然有细微差异。对于a+b+c这类加法，顺序不同可能让U LP差几个ulp；对于Softmax类操作，由于“先exp再除以总和”被换成“衰减累积后除以总和”，差异会稍大一些。
我的验收标准是：
- element-wise类融合，
max_abs_diff不超过1e-5（float32），且整体相对误差在1e-6以内。 - Attention类融合，
max_abs_diff不超过1e-5到1e-4之间可以接受。 - 如果是训练场景，还要跑几个step对比loss曲线，确认收敛行为和基线一致。
一旦差异异常，最先怀疑的不是算法逻辑，而是数据布局问题。比如输入tensor是不是非连续的、dtype是不是被隐式转换了、kernel里有没有把float4强转成char*导致对齐越界。这些bug不会crash，只会在数值上悄悄惩罚你。
6.2 非连续内存和隐式copy：最容易被忽略的隐形杀手
在PyTorch里拿到一个[B, T, C]的tensor，转置成[B, C, T]之后，内存布局就变成了非连续。当你对这种tensor做算子融合时，如果不加.contiguous()，PyTorch会在背后插入一个copy kernel，把数据重新排成连续布局。这个copy kernel的访存开销足以抵消你融合省下来的全部收益。
更隐蔽的情况是：view、narrow、expand创建的张量共享存储，但data_ptr偏移和stride都特殊。手动写CUDA kernel时，如果直接按线性索引访问，数据会完全错位。
我的习惯是：在进入手动融合或torch.compile之前，先统一调用.contiguous()，或者在全链路设计时保证张量始终是连续布局的。如果非连续布局不可避免，那就把融合kernel写成支持stride参数的通用形式，虽然性能会略微下降，但至少正确。
6.3 用Profiler和带宽计算确认真实收益
最后强调一个方法论：性能优化必须靠数据说话，不能靠感觉。我通常是这样一套组合拳：
- 用
torch.profiler跑基线，记录每个kernel的耗时、GPU利用率、带宽估算。 - 用
torch.utils.benchmark.Timer做微基准，多次测量取中位数（不是均值，均值容易被极端值干扰）。 - 算出理论带宽上限，把“优化后的耗时”和“理论耗时”对比，确认融合是否打满了带宽。
- 看整体端到端延迟，而不是只看单个kernel——因为Kernel并发、CUDA Graph capture等都会影响最终效果。
我的一个很深的体会是：很多人做完融合后，发现单看那个kernel确实快了3倍，但整个模型端到端只快了10%。原因往往是模型里另外还有一堆更耗时的算子没处理，或者CPU侧的Python开销、数据加载成了新瓶颈。融合不是孤立技巧，必须放在整个推理/训练链路里看。融合之后用Nsight Systems再看一遍全链路，找出新的热点，继续优化。性能优化是一件反复迭代的事，而不是一次“妙手回春”就能一劳永逸。
从手动写CUDA kernel到用Flash Attention思路重构模块，再到接受torch.compile帮你干活，算子融合这条路我走了很久，也踩过不少坑。最后再分享一个个人经验：融合前一定先用profiler找热点，别凭感觉乱优化；融合后一定要做正确性校验，别只看benchmark数字。只要把这两个习惯养成，你会发现算子融合不仅能让模型跑得更快，还会让你对硬件执行模型的理解深一大截——这种基本功，比任何现成优化库都值钱。
