(20) CBR融合算子：Conv + BatchNorm + ReLU的三位一体

——“一个人用AI如何写出比PyTorch更快的自研深度学习框架”系列文章之二十

在几乎任何现代卷积神经网络里，你都能找到这样一种结构：一个卷积层后面紧跟批归一化（BatchNorm），再紧跟 ReLU 激活。ResNet、VGG、MobileNet、EfficientNet 都是如此。在 PyTorch 风格的代码里，它通常写成三行：

nn.Conv2d(...),
nn.BatchNorm2d(...),
nn.ReLU(...)

这三个算子在数学上各司其职——卷积提取空间特征，BatchNorm 稳定训练分布，ReLU 引入非线性——但在硬件执行上它们却紧紧相邻。正是这种相邻关系，给了框架一个巨大的优化机会：把它们融合成一个算子来执行。

这就是 Tech-Renaissance 中的 CBR 融合算子。本文要讲清楚三件事：CBR 融合为什么能快、为什么训练场景比推理场景更难做、以及 Tech-Renaissance 是如何在 AMP 路径上把 CBR 打造成一等算子的。

一、分立算子的开销：带宽才是真正的瓶颈

在没有融合的情况下，一个 CBR 块在前向传播中至少要经历三次独立的 CUDA kernel 提交：

Conv：读取输入特征图 X 和卷积核 W，计算卷积结果，把中间结果 conv_output 写回全局显存；
BatchNorm：读取 conv_output，计算或套用均值、方差、缩放、偏移，把 bn_output 写回显存；
ReLU：读取 bn_output，做 max(0, x)，把最终结果写回显存。

三次 launch、三次全局内存事务、两份完整大小的中间激活张量。对于 ResNet-50 这种有 50 多个卷积层的网络，这种开销会一层一层累加。

更重要的是，现代 CNN 训练在 A100、RTX 5090 这类高端 GPU 上往往是显存带宽受限而非算力受限。以 A100 为例，FP16 Tensor Core 的（稠密）峰值算力约为 312 TFLOPS；A100 40GB 的 HBM 带宽约为 1.55 TB/s，80GB 版约为 2.0 TB/s。按照 roofline 模型的口径，若按 40GB 版计算，运算强度需要达到约 200 FLOPs/byte 才能不被带宽拖慢；若按 80GB 版计算则约为 156 FLOPs/byte。卷积本身的计算密度较高，通常能达到这个门槛；但 BN 和 ReLU 都是逐元素操作，计算量极小、访存量极大，是典型的”内存墙”操作。如果能让这些逐元素操作紧挨着卷积执行，把中间结果留在寄存器或共享内存里，而不是来回搬运到全局显存，就能显著减少带宽压力。

所以 CBR 融合最根本的收益不是”算得更快”，而是少读少写中间结果。这也是后面所有设计决策的出发点。

举个具体例子。假设一个中间层输出尺寸为 N=256, H=W=56, C=64，使用 FP16，那么一张完整特征图占用：

bytes = N * H * W * C * sizeof(half);   // sizeof(half) == 2
      = 256 * 56 * 56 * 64 * 2
      ≈ 102 MB

分立实现下，Conv 写一次 conv_output、BN 读它再写 bn_output、ReLU 读 bn_output 再写最终结果，仅这一层就多产生约 200 MB 的显存读写。ResNet-50 有 50 多层卷积，而 VGG16BN 虽然只有 13 个卷积层，但特征图分辨率更高、通道数更大，单层特征图的数据量更大。CBR 融合把其中一次完整特征图的写和一次读省掉，收益就这样一层层放大。

二、训练比推理更难融合：为什么不能直接折叠 BN

在推理场景，CBR 融合有一个非常漂亮的数学捷径：把 BatchNorm 的参数折叠进卷积的权重里。

假设卷积输出为：

y = conv(x, W, b);

BN 的变换为：

bn(y) = gamma * (y - mean) / sqrt(var + eps) + beta;

把两式合并，可得到一个等价的卷积：

W_fold = W * gamma / sqrt(var + eps);
b_fold = (b - mean) * gamma / sqrt(var + eps) + beta;
out    = conv(x, W_fold, b_fold);
out    = relu(out);

只要推理时使用固定的 running_mean 和 running_var，这个折叠是数学精确的。TensorRT、ONNX Runtime、TensorFlow Lite 等推理引擎普遍采用这种策略，收益非常直接。

但训练场景不行。训练时 BN 的均值和方差来自当前 batch，还要更新 running statistics，并且反向传播需要保存 saved_mean 和 saved_inv_var 来计算 d_scale 和 d_bias。你不能简单地把 BN 折叠进卷积权重，因为每次迭代的统计量都在变。因此，训练阶段的 CBR 融合必须保留 BN 的完整语义，同时尽量把卷积、统计量生成、归一化、激活在物理上靠得更近。

训练时 BN 的反向传播涉及三个梯度：上游梯度 dY、输入 X 的梯度 dX、以及参数 gamma/beta 的梯度。要算 dX，必须知道当前 batch 的 saved_mean 和 saved_inv_var；要更新 running statistics，又需要 batch 的 mean 和 var。这些量都无法像推理那样预先折叠进卷积权重。更进一步，BN 的反向通常需要先做一次 dY * mask 来处理 ReLU 的死亡神经元，然后才能进入 batchnorm_backward。所以训练路径的 CBR 融合不是消去 BN，而是把”ReLU 掩码应用 + BN 反向 + 卷积梯度”的流水线尽量紧凑地组织起来。

这也是为什么主流训练框架对 CBR 融合相对保守。PyTorch 的 torch.compile / TorchInductor 可以在 FX 图级别做算子融合，但效果受动态图、shape 变化、CUDA Graph 稳定性等因素影响；TensorFlow 的 XLA 也取决于图能否被完整捕获和编译。Tech-Renaissance 的选择是：在 AMP 混合精度训练路径上，把 CBR 作为一个一等算子完整实现，包括前向、反向、首层特化和推理路径。

三、从 BluePrint 到计算图：让框架自动认出 CBR

在 Tech-Renaissance 的高层 DSL BluePrint 里，用户可以直接写：

auto layer = cbr(64, 3, 1, 1);        // out_ch=64, kernel=3, stride=1, padding=1
// 或者等价写法
auto layer = conv_bn_relu(64, 3, 1, 1);

这会生成一个 NodeKind::CBR 节点。在 ArchPlan::step9_merge_triple() 中，当 AMP 开启时，框架会自动扫描连续的 Conv → Bn2d → ReLU 三层，把它们合并成一个 LayerKind::CBR：

merge_pattern_triple(LayerKind::Conv, LayerKind::Bn2d, LayerKind::ReLU, LayerKind::CBR, build_cbr);

也就是说，即使用户写的是三个独立层，只要开启了 AMP，编译器也会尝试把它们熔铸成一个融合算子。这种”图级模式匹配 + 算子级特化实现”的分层设计，是静态图框架的优势：我们在编译期就知道整个网络长什么样，可以大胆地做全局重写，而不必在运行时动态捕获图。

在 Compiler 构建计算图时，CBR 层被分配完整的参数组：卷积权重、BN 的 gamma/beta、BN 的 running mean/var，以及用于反向的梯度占位。src/graph/layer_descriptor_registry.cpp 中为 CBR 定义了 23 个张量（Conv 部分 8 个、BN 部分 13 个、ReLU 部分 2 个），这些张量的布局、数据类型、内存区域在 MemoryPlan 阶段就全部确定下来。运行时，前向算子 CBR_AMP_FWD 的实际输入输出为：

输入（8 个）：X（FP16 特征图）、amp_w（FP16 卷积核）、bn_w（FP32 gamma）、bn_b（FP32 beta）、prev_mean（FP32 running mean）、prev_var（FP32 running var）、eps（FP32 标量）、mom（FP32 标量）；
输出（10 个）：conv_output、bn_sum、bn_sq_sum、bn_output、saved_mean、saved_inv_var、relu_output、relu_mask、next_mean、next_var。

其中 X 由编译器自动前置，eps 和 mom 是全局标量张量 ID；bn_output 在训练前向中并不存放完整特征图，而是被复用来放置 eq_scale/eq_bias 这两个小向量。这种显存复用只有在 MemoryPlan 的静态规划下才是安全的。

四、前向：三段式 cuDNN Graph，中间结果能不写就不写

CBR 前向算子 CBR_AMP_FWD 的实现位于 src/backend/ops/dtensor/cbr_op.cpp。它的核心不是写一个巨大的自定义 CUDA kernel，而是把问题拆成三张 cuDNN Frontend Graph，分别跑在三条计算流上，用 CUDA event 做跨流同步。

4.1 Conv + GenStats（COMP_1）

第一张图只做两件事：用 conv_fprop 计算卷积输出；用 genstats 从卷积输出直接生成通道级的 sum 和 sq_sum。

auto conv_out = graph->conv_fprop(X, W, conv_opts);
auto [sum, sq_sum] = graph->genstats(conv_out, genstats_opts);

这一步的关键在于 genstats。在单独实现 BN 时，框架需要先把 conv_output 写回全局显存，再读取一次来计算均值和方差；而 cuDNN 的 genstats 可以在卷积 forward 的同时，从寄存器或共享内存里直接归约出统计量。这意味着我们省掉了”为了算 BN 而再读一遍整张特征图”的带宽开销。对于 256×56×56×64 的 FP16 特征图，这一次就避免了约 102 MB 的全局显存读取。

4.2 BN Finalize（COMP_2）

拿到 sum 和 sq_sum 后，第二张图 bn_finalize 在纯 FP32 精度下完成：

计算 saved_mean、saved_inv_var；
更新 next_mean、next_var；
计算本次前向真正用于缩放偏移的 eq_scale 和 eq_bias。

它们的计算公式为：

N_ac        = N * H * W;                              // 每个通道累积的样本数
mean        = sum   / N_ac;
var         = sq_sum / N_ac - mean * mean;
inv_std     = 1.0f / sqrtf(var + eps);
saved_mean  = mean;                                   // 保存给反向
saved_inv_var = inv_std;                              // 保存给反向
eq_scale    = gamma * inv_std;                        // 用于 Apply 的等价缩放
eq_bias     = beta - gamma * mean * inv_std;          // 用于 Apply 的等价偏移
next_mean   = (1 - momentum) * prev_mean + momentum * mean;
next_var    = (1 - momentum) * prev_var  + momentum * var;

这里 sum、sq_sum 是两个长度为 C 的 FP32 小向量，计算量极小，几乎可以完全隐藏在 Conv+GenStats 的执行过程中。eq_scale 和 eq_bias 被写入 bn_output 缓冲区的前 2 * K * sizeof(float) 字节，saved_mean、saved_inv_var、next_mean、next_var 则各自写入独立的 FP32 缓冲区。其中 momentum 是 cuDNN 风格的 exponential average factor（新 batch 统计量的权重，默认 0.1）。

4.3 BN Apply + ReLU（COMP_3）

第三张图接收卷积输出、eq_scale/eq_bias，把 BN 仿射变换和 ReLU 激活融合在一起：

auto scaled  = graph->pointwise(conv_out, eq_scale_ta, MUL);
auto shifted = graph->pointwise(scaled,  eq_bias_ta,  ADD);
auto relu    = graph->pointwise(shifted, RELU_FWD);
auto mask    = graph->pointwise(shifted, zero, CMP_GT);   // shifted > 0

同时用 CMP_GT 生成一个位压缩的布尔掩码（cuDNN 的 BOOLEAN 类型，每个元素 1 bit），保存下来给反向传播用。这与分立 ReLU 算子常用的每元素 1 字节 INT8 mask 不同：CBR 内部使用位压缩格式来节省显存，测试代码里专门做了两种 mask 格式之间的转换，以保证等价性校验通过。

4.4 多流同步

三个阶段分别跑在 COMP_1、COMP_2、COMP_3 上。MultiStreamCaptureState 负责注册流并在阶段之间插入 cudaStreamWaitEvent：Finalize 等 Conv+GenStats 完成，Apply+ReLU 等 Finalize 完成。这套机制与我们在第 17 篇 CUDA Graph 全捕获中讲的多流依赖管理是同一套基础设施，保证了即使算子内部拆成多个流，对外仍然呈现为一个原子的 CBR_AMP_FWD 节点。

你可能会问：为什么不直接用一张巨大的图把 Conv+BN+ReLU 全包进去？原因在于 cuDNN 当前对训练路径的 Conv+GenStats+BN finalize+Apply+ReLU 这五连操作，并不总能成功实例化为一张可执行图；而拆成三段后，每一段都更小、更稳定，同时通过多流重叠让三段之间的等待时间被其他计算隐藏。另外，分段也让我们能在 BN finalize 阶段复用 bn_output 缓冲区放 eq_scale/eq_bias，这对显存管理很关键。

每张图首次构建后都会被缓存到以 cuDNN handle、形状、pad/stride 为键的 unordered_map 中，后续同形状 batch 直接复用，不需要重复 build。

五、反向：把 BN 和 ReLU 的梯度也算进一张图

CBR 的反向 CBR_AMP_BWD 同样不是简单地”分别调用 Conv BWD、BN BWD、ReLU BWD”。它的执行也分布在三条流上：

COMP_1：执行 BN+ReLU BWD 融合子图；
COMP_3：执行 WGrad（权重梯度）；
COMP_2：执行 DGrad（数据梯度）。

5.1 BN + ReLU 反向融合（COMP_1）

先用 MUL(dY, mask) 把上游梯度与 ReLU 前向掩码相乘，屏蔽掉 ReLU 关闭的神经元；然后直接调用 batchnorm_backward，输出三样东西：传到卷积的梯度 dL/d(conv_out)、BN gamma 的梯度 d_scale、BN beta 的梯度 d_bias。

auto dy_masked = graph->pointwise(dY_ta, mask_ta, MUL);
auto [dx_bn, dscale, dbias] = graph->batchnorm_backward(dy_masked, x_ta, scale_ta, bn_opts);

这一步被封装在一张 cuDNN Frontend Graph 里，dy_masked 不需要写回 HBM。输出 dL/d(conv_out) 被写进 bn_output 缓冲区——也就是前向里那个只存了 eq_scale/eq_bias 的缓冲区——因为进入反向之后，前向的 eq_scale/eq_bias 已经不再需要，这块内存可以安全复用。

5.2 WGrad（COMP_3）与 DGrad（COMP_2）

bn_output 此时保存的是 dL/d(conv_out)。WGrad 用它和原始输入 X 计算卷积权重的梯度 dW。DGrad 用 dL/d(conv_out) 和卷积权重计算输入 X 的梯度 dX。

源码里明确注释了 DGrad 为什么要等 WGrad 完成：

“等待 WGrad，确保 WGrad 已读完原始 X 后 DGrad 才覆盖 X。”

这个同步不是数据依赖——DGrad 不依赖 WGrad 的输出——而是因为两者都读取原始输入 X，而 DGrad 的输出 dX 会原地覆盖 X 的存储位置。如果 DGrad 在 WGrad 之前启动，它可能会在 WGrad 读取 X 之前就覆盖了 X 的内容，导致 WGrad 计算出错误的梯度。这是融合算子中需要特别小心处理的读写冲突。

5.3 首层特化

对于网络的第一层，CBR_AMP_BWD_FIRST_LAYER 跳过了 DGrad——因为首层的输入是原始训练数据，不是中间特征图，不需要计算输入数据的梯度。这进一步节省了首层的计算开销。在编译器层面，首层和非首层的 BWD 算子通过 to_first_layer_bwd_op 函数进行自动转换，用户无需关心这个差异。

六、推理路径：能压成一图就压成一图

训练路径需要生成 GenStats 和中间统计量，因此拆成了三张子图。但推理路径不需要这些——它只需要 Conv → BN(INF) → ReLU 的纯前向计算，既不需要本批次的统计量，也不需要更新 running statistics。

因此 CBR_AMP_INF 首先尝试把整个 Conv → MUL(eq_scale) → ADD(eq_bias) → ReLU 表达成一张 cuDNN Frontend Graph：

auto Y      = graph->conv_fprop(X, W, conv_opts);
auto scaled = graph->pointwise(Y, eq_scale_ta, MUL);
auto shifted= graph->pointwise(scaled, eq_bias_ta, ADD);
auto relu   = graph->pointwise(shifted, RELU_FWD);

如果 cuDNN 接受这个图并成功实例化，就只需一次 execute 调用；如果因为某些形状、精度或引擎限制导致单图构建失败，框架会静默回退到三段式执行：先 Conv INF，再自定义 BN INF kernel，最后 ReLU INF kernel。这种”先尝试最优，再保证正确”的策略在不牺牲健壮性的前提下，尽可能拿到融合收益。

七、Mode C 经验搜索：为每一层匹配最优引擎

cuDNN 为同一个卷积操作提供了多种引擎实现，不同的引擎在不同的张量形状、内存布局和硬件架构下，性能差异可能非常显著。常规做法是使用 cuDNN 的启发式引擎选择，让 cuDNN 根据少数几个参数做一个粗略估计，选一个”应该还行”的引擎。但启发式算法无法保证选到最优。

Tech-Renaissance 的 CBR 融合算子引入了 Mode C 经验搜索机制：针对 A100 和 RTX 5090 等目标 GPU，预先对所有可能的形状组合进行穷举式基准测试，将最优引擎的 tag 记录为经验数据，以 C++ constexpr 数组的形式编译进框架二进制中。以 A100 FP16 经验表为例，它包含 92 条记录，覆盖了 ResNet-50 和 VGG16BN 等经典模型的卷积层形状组合，以及 conv_fprop、conv_genstats、conv_wgrad、conv_dgrad 等操作类型。

每条经验记录大致包含：

shape_key：查询键，精确到 GPU 型号、CUDA/cuDNN 版本、算子类型、张量各维度、stride、padding、布局等；
winner_tag：基准测试中最优的引擎标签；
backup1_tag / backup2_tag：两个备选标签，用于最优引擎不可用时回退；
workspace_bytes / benchmark_time_ms：工作空间需求和实测耗时。

运行时，框架根据当前张量形状构造查询键，在 constexpr 表中二分查找匹配的记录，然后按三级优先级依次尝试构建对应引擎：winner → backup1 → backup2。如果全部失败，则优雅回退到标准的启发式选择。

这种设计的优势在于零运行时搜索开销：经验表是编译期常量，二分查找在纳秒级完成，匹配到标签后直接调用 build_plan_at_index，不需要遍历和评估所有候选引擎。这对 Conv+GenStats、WGrad、DGrad 三个子图都独立适用。

经验数据通过独立的 Python 脚本生成：脚本遍历目标模型的所有卷积层，提取形状参数，用 Mode A 枚举所有可用引擎，逐个 warmup 后计时，选出最快的引擎作为 winner，再选两个次优作为 backup，最终写入 include/generated/cbr_experience_*.hpp。目前框架为 A100 和 RTX 5090 分别生成了 FP16 经验表。

八、约束、限制与正确性验证

CBR 融合算子目前有一些明确的约束，它们大多来自底层库能力边界或工程取舍：

仅 CUDA 路径：CPU 路径直接抛出 NotImplementedError。这不是设计缺陷，而是 CBR 的收益主要来自 GPU 张量核心和全局显存带宽，CPU 端做融合意义有限。
仅 AMP/FP16 模式：cuDNN 的 Conv + GenStats 组合目前对 FP16 支持最完整；FP32 模式下即使开启融合，带宽节省带来的收益也不足以抵消工程复杂度，且整体速度往往仍不及 AMP 路径。因此 Tech-Renaissance 的 FP32 路径仍使用独立的 Conv、BN2D、ReLU 算子。
输出通道需为 8 的倍数：这是 cuDNN TensorCore BN 内部约束。如果用户模型不满足，编译器会报错并提示插入 channel_padding 或调整通道数。
统一使用 NHWC 特征图布局和 KRSC 卷积核布局：这与 PyTorch 默认的 NCHW 不同，但和 cuDNN 内部的偏好一致，能让卷积和 BN 更容易被融合执行。框架在编译期就完成布局推导，运行时不需要额外的 transpose。

为了确保融合不变成”偷换数学”，我们在 tests/op/ 中为 CBR 写了四组等价性测试：

test_cbr_amp_fwd.cpp
test_cbr_amp_bwd.cpp
test_cbr_amp_bwd_first_layer.cpp
test_cbr_amp_inf.cpp

每组都把 CBR 路径和独立的 Conv+BN2D+ReLU 路径做数值对比。由于 CBR 使用 cuDNN 的位压缩 BOOLEAN 掩码，而独立 ReLU 使用每元素 INT8 掩码，测试里还专门做了 mask 转码对齐。

九、CBR 与分立算子的根本差异

下表把 CBR 与独立 Conv+BN2D+ReLU 在实现层面做一个对比：

维度	分立 Conv+BN2D+ReLU	CBR 融合算子
计算图节点	3 个节点	1 个 CBR 节点
默认调度流	各按 `op_stream_policy` 分配	入口在 `COMP_1`，内部 fan out 到 `COMP_1/2/3`
前向 kernel/图提交	Conv、BN、ReLU 各一次	3 张 cuDNN FE Graph，内部尽可能融合
BN 统计量	BN 单独读取 `conv_output` 算 mean/var	Conv 同时生成 `sum`/`sq_sum`
中间激活	`conv_output`、`bn_output` 都作为完整张量	`bn_output` 复用为 `eq_scale`/`eq_bias` 或反向 `dL/d(conv_out)`
ReLU 掩码	每元素 INT8	cuDNN `BOOLEAN` 位压缩
通道约束	一般无特殊要求	输出通道需为 8 的倍数
精度支持	FP32 / AMP	仅 AMP

从显存事务角度看，前向传播中分立路径的 conv_output 要经历 1 次写入和 2 次读取（统计量计算 + BN Apply），而 CBR 融合中 conv_output 只有 1 次写入和 1 次读取（BN Apply），且这次读取的结果在寄存器中直接流向 ReLU，不再产生 bn_output 的中间写入。反向传播中，分立路径的 ReLU_BWD 和 BN_BWD 各需要一次 HBM 写入和读取来传递中间梯度；CBR 融合中这两个操作被合并，dy_masked 直接通过寄存器传递。

十、CBR 在整个系统里的位置

单独看一个 CBR 算子，它节省的不过是几次中间特征图的读写；但把它放在整个网络里看，收益会被层数放大。在 README 报告的 VGG16BN A100×8 训练中，Tech-Renaissance 达到了 9,310.13 images/sec，比 PyTorch torch.compile 的 7,351.20 images/sec 高出 26.65%。这个差距不是某一个单点造成的，而是静态图编译、CUDA Graph 全捕获、MemoryPlan、多流调度和 CBR 这类融合算子叠加起来的结果。

CBR 在这个系统里的角色，是把”高层网络中随处可见的 Conv+BN+ReLU”这个模式压榨到极致。它让框架在不做任何用户侵入式改造的前提下，就能更有效地利用 GPU 带宽。当然，我们不能把 26.65% 的加速全部归功于 CBR。它是一个乘数因子：CUDA Graph 消除了 launch overhead，MemoryPlan 保证了地址稳定，多流调度隐藏了同步等待，CBR 则把最热的算子模式压榨到更接近硬件极限。它们彼此依赖，缺少任何一个，其他优化都难以充分发挥。

十一、小结

CBR 融合算子是 Tech-Renaissance 算子层优化的一个缩影。它的设计思路可以概括为：

图级自动识别：在 ArchPlan 里把 Conv+BN+ReLU 自动合并成 CBR；
底层 cuDNN Graph：用 Frontend Graph 把卷积、统计量生成、BN 归一化、ReLU 尽可能串成最少的执行单元；
多流重叠：把不同阶段拆到 COMP_1/2/3，用 event 保证依赖，提高并行度；
内存复用：bn_output 在前向和反向扮演不同角色，减少显存占用；
训练与推理兼顾：前向保留完整 BN 统计更新，推理则尝试单图融合并带 fallback；
引擎经验搜索：用预编译的 Mode C 经验表跳过运行时启发式搜索；
严格等价性测试：确保融合后的结果与分立算子一致。

它不是把三段代码粗暴地粘在一起，而是在保证数学正确的前提下，把”该省的内存读写、该合并的 kernel launch、该并行的计算阶段”全部做到极致。图级自动合并让用户无感知，底层 cuDNN Graph 和多流调度让硬件满负荷，静态内存规划让复用安全可控，等价性测试让速度不以牺牲正确性为代价——这正是 Tech-Renaissance 做算子融合的基本方法论。

下一篇，我们将继续沿着”融合”这条主线，进入优化器层面——看看 SGD、AdamW、LARS 的权重更新又是如何被批量融合到整区内存上的。