推理延迟居高不下是模型上线时最常遇到的问题之一。很多团队习惯先怀疑模型太大，于是压缩通道、蒸馏或换小模型，但有时精度已经明显下降，延迟却只改善了一点点。问题很可能不在参数量，而在推理图里存在大量细碎算子：一个卷积后面跟着批归一化、激活、池化，每个算子都单独启动一次 GPU 内核，中间结果反复在全局内存中读写。算子融合正是从执行机制上消除这类开销，让一次内核完成更多计算，是部署优化中性价比最高的手段之一。
一、推理慢的根因：不是算得慢，而是访存和启动太频繁
在 GPU 或专用加速器上，一次推理由成百上千个内核依次执行完成。每个内核的启动都需要 CPU 向 GPU 提交命令、调度线程块、分配共享内存等。单个启动成本通常只有几微秒到几十微秒，但当网络中有几百个小算子时，启动开销叠加起来就可能达到毫秒级。对于延迟敏感的服务，这已经非常可观。
更隐蔽的问题是中间张量的读写。假设卷积输出一个 56×56×64 的特征图，以 float32 存储，单次写入约 800KB，后续批归一化再读取并写回，激活函数又读一次写一次。如果一个卷积块包含五个逐元素算子，中间数据可能被读写多轮，访存时间远超实际矩阵乘加的时间。内存带宽有限，例如常见推理卡带宽约 200GB/s 到 900GB/s，当访存量过大，计算单元只能等待。
因此推理优化不能只看 FLOPs。用 Roofline 模型分析，很多逐元素算子位于内存带宽受限区，计算强度很低。算子融合可以把多个访存密集型操作合并，把中间数据留在寄存器或共享内存中，从而显著降低全局内存流量。
二、算子融合的核心原理与典型模式
算子融合的本质是把原本需要多次内核调用、多次全局内存访问的计算序列，合并到单个内核中。融合后，数据从全局内存载入寄存器或共享内存，完成多个阶段的数学运算后再写回，中间结果不再回到全局内存。融合还减少了内核启动和同步开销，并给编译器更多指令级并行机会。
常见的融合模式可以按数据流划分。纵向融合是指上下游算子共享同一数据流，例如卷积、批归一化、ReLU、池化依次处理同一张特征图，可以合并为一个内核。横向融合则针对相互独立的算子，例如多个形状相同的逐元素加法或缩放，可以合并到一个内核中以增加并行度。实际工程里纵向融合使用更多，收益也更直接。
以卷积块为例，训练阶段通常保留独立的批归一化参数，推理时可以把 BN 的缩放、平移参数折叠进卷积权重和偏置，再将 ReLU 直接放在卷积内核后面。公式上，卷积输出 y，BN 计算为 gamma*(y-mean)/sqrt(var+eps)+beta，可以重写为 y_new = W_fold*x + b_fold，其中 W_fold=gamma/sqrt(var+eps)*W，b_fold=gamma/sqrt(var+eps)*(b-mean)+beta。激活函数则在同一次计算完成后立即对结果应用，不需要单独访存。
# 未融合：三次内核调用，三遍全局内存读写
def forward_unfused(x, weight, bias, bn_gamma, bn_beta, bn_mean, bn_var, eps=1e-5):
y = conv2d(x, weight, bias) # 第一次写回全局内存
y = batch_norm(y, bn_gamma, bn_beta, bn_mean, bn_var, eps) # 第二次读写
y = relu(y) # 第三次读写
return y
# 融合后：一次内核完成，中间结果保留在局部存储
def forward_fused(x, weight_folded, bias_folded):
y = conv2d(x, weight_folded, bias_folded)
y = relu(y) # 在同一个内核内完成
return y
另一个典型模式是矩阵乘法加偏置和激活。未融合时，矩阵乘法输出写入全局内存，随后偏置加法和激活又分别启动内核读取和写回。融合内核可以在累加结束后直接加上偏置并判断是否大于零，结果只写回一次。注意力机制中的缩放、mask 加法、softmax 与 dropout 也属于高频融合对象，尤其在 Transformer 类模型中，这些操作中间矩阵很大，融合后显存占用和延迟都能明显下降。
// 未融合：每个线程块单独处理一个操作
__global__ void scale_kernel(float* x, float a, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) x[i] = x[i] * a;
}
__global__ void relu_kernel(float* x, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) x[i] = fmaxf(x[i], 0.0f);
}
// 融合后：同一个内核完成 scale + relu
__global__ void scale_relu_fused_kernel(float* x, float a, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float v = x[i] * a;
x[i] = v > 0.0f ? v : 0.0f;
}
}
三、主流推理框架中的融合机制
手动识别融合模式并修改模型结构虽然直观，但在复杂网络里难以穷举。主流推理引擎都提供了自动融合能力，核心是计算图上的模式匹配与重写。它们先遍历计算图，找到符合规则的子图，然后替换为融合节点或自定义内核，最后进行代码生成。
TensorRT 是最典型的代表。它会在网络解析和优化阶段执行垂直融合、水平融合、Conv+BN+ReLU 合并、矩阵乘法与激活合并等。对于不支持的内核，它还会生成专门的 CUDA kernel。用户通常只需要开启 FP16 或 INT8 精度并设置合理的工作流，无需手动改模型。ONNX Runtime 则通过 GraphTransformer 和 ExecutionProvider 实现类似能力，CPU 与 GPU 后端都支持算子融合。
TVM 的融合思路更开放。它在 Relay 层提供 FuseOps Pass，按照一定规则将可融合的算子合并为子图，再由 AutoTVM/Ansor 对子图或整图搜索最优内核实现。TVM 的优势是可以针对新的硬件后端扩展融合规则，适合自研加速器或特殊算子。XLA 则面向 TensorFlow/JAX，把计算图聚合成 HLO 模块，通过融合流水线减少内核数量。PyTorch 2.0 引入的 torch.compile 也使用 Dynamo 捕获图，再由 Inductor 做融合和内核生成。
这些工具的共同点是：用户不应只依赖默认导出后的原始图。对比 benchmark 时，应确认引擎真正完成融合。例如某些动态形状或控制流会阻止融合，需要将输入形状固定或简化控制逻辑。
四、手写融合内核的实践与注意事项
当自动融合无法覆盖自定义层或特殊数据布局时，手写融合内核仍然必要。例如在 Transformer 推理中，注意力分数计算通常包含缩放、mask 加法、softmax 和 dropout，这些操作如果由分别的内核完成，中间矩阵会反复读写。将它们写成一个 CUDA kernel，可以显著降低显存占用与延迟。
下面给一个简化的矩阵乘法加偏置和 ReLU 的融合 CUDA 内核。为了展示结构，这里假设使用 1D 线程块和简化索引，真实实现还需要考虑共享内存分块、向量化访存和 bank conflict。融合后的内核在一次遍历中完成矩阵乘累加、偏置加法和激活，中间结果保存在寄存器中。
// A: M x K, B: K x N, C: M x N + bias + relu
__global__ void matmul_bias_relu_kernel(
const float* A, const float* B, const float* bias,
float* C, int M, int N, int K) {
int row = blockIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= M || col >= N) return;
float acc = 0.0f;
for (int k = 0; k < K; ++k) {
acc += A[row * K + k] * B[k * N + col];
}
acc += bias[col];
C[row * N + col] = acc > 0.0f ? acc : 0.0f;
}
手写内核的收益取决于是否真正减少了全局内存流量。如果融合后内核内部出现过多分支、寄存器溢出或负载不均衡，性能可能反而不如编译器的自动版本。开发时建议先用 NVIDIA Nsight Compute 或类似工具查看访存与占用率，再决定是否继续优化。
另一个常见误区是只融合计算，不调整内存布局。例如将通道优先改为块优先的 NHWC 或 NCHW8c，可以让相邻线程访问连续地址，提升带宽利用率。融合和布局优化通常要一起做。量化、FP16 与 INT8 也能与融合叠加，但需要注意数值范围与溢出风险。
五、从系统视角继续压缩延迟
算子融合解决的是算子粒度和访存问题，但推理延迟还受其他因素影响。完整的部署优化流程通常先做图级简化：常量折叠、删除无用的形状变换、合并重复节点、用更高效的实现替换低效算子。例如将多个 reshape 和 transpose 合并，或把 gather 替换为更快的索引实现。之后再进行算子融合。
内存管理同样关键。推理过程中频繁分配和释放 GPU 显存会导致同步与碎片。提前规划输入输出张量并复用内存池，可以消除分配开销。推理服务还可以使用多流并行，将预处理、H2D 传输、计算、D2H 传输重叠起来。对于小模型，CPU 与 GPU 的通信往往占主导，此时可以考虑将部分算子留在 CPU 执行。
动态 shape 是融合的大敌。很多框架在输入形状固定时才能进行激进融合和内存规划。如果模型需要支持可变长度，可以为常见长度分别编译计算图，或者通过 padding 到固定长度换取更好的内核性能。批量推理时，动态 batch 与自动 padding 能提高吞吐，但会增加延迟，需要根据 SLA 权衡。
最终评估时应使用与线上一致的吞吐和延迟指标，避免只看单个内核的加速比。融合前后的完整推理耗时、P99 延迟、显存占用、首包时间等都要记录，确保优化没有牺牲可维护性和数值稳定性。