训练 vs 推理：AI 芯片算子实现的两条路

副标题同一个算子训练和推理在硬件眼里是两种完全不同的动物。Conv2D、GEMM、Normalization、Attention——每个算子都在训练时和推理时面临不同的 shape、精度、访存和融合需求。本文从算子开发工程师的视角系统性地拆解这些差异。引子一个 Conv2D 算子的两面先从一个最简单的例子说起。Conv2D——深度学习最基础的算子。不管是什么 AI 芯片都要实现它。但同一个 Conv2D在训练和推理时的实现路径可以完全不同训练推理Batch size64-512大1-8小数据类型FP32 / BF16FP16 / INT8 / FP8需要什么Forward Backward ×2Forward 仅核心算法偏好Winograd乘少换快Implicit GEMM无缓存融合模式Conv BNUpdate ReLUConv Bias ReLU前向卷积的数学是一样的——不同在于算子的 shape、精度要求、融合策略和是否需要反向传播把同一个算子的硬件实现拉到了两条不同的优化路径上。这不是 Conv2D 独有的。GEMM、Normalization、Attention 每一个算子都在训练和推理之间展现出了同样深刻的差异。理解这些差异对自研芯片的算子开发有决定性的意义。一、Conv2D四家厂商的路线分歧1.1 四种核心算法Conv2D 的核心挑战是滑动窗口访存不连续。所有芯片厂商的实现都在以下四种算法中选择算法核心思想乘法复杂度显存额外开销适用场景Direct7 层循环硬算O(N·C_out·H_out·W_out·C_in·K²)0几乎不用效率 1%Im2Col GEMM展开成矩阵做 GEMM同上K² 倍输入大小有矩阵乘单元的芯片Implicit GEMMGEMM kernel 内在线算地址同上0不显式展开当前主流Winograd数学变换减少乘法~44%K3 时变换矩阵中间量3×3、大 batchFFTFFT→点乘→IFFTO(N log N)频域缓存大 kernel≥7×71.2 四家厂商的路线选择厂商核心路线为什么走这条路NVIDIA cuDNNImplicit GEMM Tensor CoreTensor Core 只擅长 GEMMImplicit GEMM 把卷积变成 GEMM 同时省掉 im2col 缓存。Find API 在多种算法间自动选最优。AMD MIOpenImplicit GEMM 汇编 WinogradMatrix Core 同样适合 GEMM。Winograd 用纯汇编手写没有 C 版本利用 CDNA 的 64-lane wavefront 做高效 shuffle。Intel oneDNNJIT 编译生成不手写多个 kernel 变体而是运行时按参数生成刚好匹配的 kernel。CPU 上用 AVX-512/AMX 向量化GPU 上用 DPAS。华为昇腾 CANN显式 Im2Col Cube GEMM达芬奇架构的 Cube Unit不能做 Implicit GEMM无法在矩阵乘内部做条件地址计算只能用 Vector Unit 做显式 im2col 展开再喂给 Cube。1.3 训练和推理在 Conv2D 上的差异训练选 Winograd推理选 Implicit GEMM训练N128, 3×3 Conv: Winograd 把乘法从 9 降到 4 N128 分摊了变换矩阵的额外计算 总 FLOPs 减少 ~50% → Winograd 胜出 推理N1, 3×3 Conv: Winograd 的变换步骤A^T·X·A, G·W·G^T多了 3 次小矩阵乘 在 N1 时变换开销可能超过省下的乘法 → Implicit GEMM 胜出无缓存、无变换、1 个 kernel 搞定训练必须 FP32 累加推理可以用 INT8训练时正向的微小误差会被反向传播放大。因此即使输入是 FP16训练卷积的累加器accumulator也必须是 FP32。推理没有反向传播INT8 量化卷积的精度损失可以接受。训练融合 BN update推理吸收 BN训练融合: Conv2D BN(update) ReLU → 融合 kernel 因为 BN 前向要更新 running_mean/var不能直接吸收到 Conv 权重 推理: BN(γ·xβ) 是线性变换直接吸收进 Conv2D 的 W 和 b W γ·W / sqrt(σ²ε) b γ·b / sqrt(σ²ε) β - γ·μ / sqrt(σ²ε) → Conv2D 融合 kernel 不需要 BN 参与更简单训练有反向算子推理没有训练需要额外实现Conv2DBackpropInput计算输入梯度和Conv2DBackpropFilter计算权重梯度。这两个反向算子的计算模式不同BackpropInput等价于转置卷积BackpropFilter需要另一种数据流。Conv2D 的差异在所有算子中是最明显的——因为它的访存模式特殊算法选择多。接下来的 GEMM、Normalization、Attention 的差异更隐蔽但影响同样深远。二、GEMM训练要吞吐推理要延迟GEMMC A B是 AI 推理训练中占比最大的算子LLM 中占 ~65% 算力。但训练和推理场景下的 GEMM 几乎是两个世界。2.1 Shape 差异天壤之别训练 GEMM 的 shapeForward: Y X W^T → M B × L, N d_out, K d_in Backward: dX dY W → M B × L, N d_in, K d_out dW dY^T X → M d_out, N d_in, K B × L 典型值: B64, L2048, d_in4096, d_out4096 → M 131072, N 4096, K 4096 → 巨大的 GEMM算力打满推理 Decode GEMM 的 shapeDecode: Y X W^T → M B, N d_out, K d_in 典型值: B1, d_in7168, d_out4096 → M 1, N 4096, K 7168 → GEMV矩阵×向量算力利用率极低这个差异意味着什么训练 GEMMM131072: Arithmetic Intensity ≈ M·N·K / (M·K N·K M·N) · dtype_size ≈ 131072·4096·4096 / (131072·4096 4096·4096 131072·4096) · 2 ≈ 4.4 TFLOPs / 1.1 GB ~4000 FLOP/Byte → 深度计算密集可以轻松打满 Tensor Core 推理 Decode GEMMM1: Arithmetic Intensity ≈ 1·4096·7168 / (1·7168 4096·7168 1·4096) · 2 ≈ 29M / 29.4MB ~1 FLOP/Byte → 极度访存密集带宽是唯一瓶颈2.2 精度差异训练推理输入精度BF16 / FP16FP16 / INT8 / FP8累加精度必须 FP32FP16 可以接受权重精度FP32 / BF16FP16 / INT4 / FP8为什么训练必须 FP32 累加FP16 的最大表示范围约 65504GEMM 累加时中间结果很容易超过这个范围。一个 4096×4096 的 GEMMK4096 次累加每次累加的是 FP16 乘积值域 ~ ±65504²平均可能达到 10⁹ 的量级——远超 FP16 的表示范围。如果不使用 FP32 累加器精度直接溢出到 NaN梯度全丢。推理中不存在反向传播累加误差最多影响一个输出 token不会被放大。2.3 算法差异训练 GEMMM 1: Tile GEMM: 把大矩阵切成 tile如 128×128每个 tile 对应一个 thread block 在 Tensor Core 上运行共享内存做 double buffering 目标是打满 Tensor Core 利用率H100 可达 ~80% 理论峰值 推理 Decode GEMMM 1: GEMV 优化: M1 时不能用标准的 tile GEMMtile 高度为 128但 M1利用率 1% 有两种方案: A) Split-K: 把 K 维度切分成多份并行计算最后 reduce B) Warp-level GEMV: 一个 warp 处理所有 N 维度逐 K 步进 目标是最大化带宽利用率达到 ~80% HBM 带宽即可对自研芯片的影响如果你的芯片只为训练优化大 tile GEMM推理 decode 的 GEMV 场景下性能会暴跌。必须同时有大 GEMM kernel和小 GEMV kernel或者统一的 kernel 能自适应 M 维度。2.4 一个特殊的精度问题确定性Determinism训练场景有一个推理没有的要求——多批次之间必须 bitwise 一致训练: batch 64 和 batch 128 跑同一个输入 → 前 64 个样本的梯度必须完全一样 否则调 bug 时无法复现分布式训练时更新不一致 → 需要确定性 GEMMsgemm-bi 的做法: 固定累加顺序不用原子操作不同 batch 之间结果 bit-identical 推理: 不需要确定性一次 Forward 就够了 不同的推理步骤可以有不同的精度三、Normalization训练要统计推理要固定3.1 BatchNorm差异最大BatchNorm 是训练和推理行为差异最大的算子训练时: 输入 x (N, C, H, W) ↓ 计算当前 batch 的 μ mean(x), σ² var(x) ↓ y γ · (x - μ) / sqrt(σ² ε) β ↓ 更新 running_mean momentum × running_mean (1-momentum) × μ 更新 running_var momentum × running_var (1-momentum) × σ² → μ 和 σ² 每次都不一样取决于当前 batch 的数据 推理时: 输入 x (N, C, H, W) ↓ y γ · (x - fixed_mean) / sqrt(fixed_var ε) β → 使用训练中积累的固定 running_mean/running_var → μ 和 σ² 是固定的没有数据依赖性对算子实现的直接影响训练 BN: 需要 2 次规约求 mean 和 var 2 次更新running stats 需要存中间输入 x 用于反向传播 → 算子实现更重有额外的统计量计算 推理 BN: 一次逐元素乘加线性变换就够了 → 实践中可以直接吸收进前一个 Conv2D 的权重里 → 根本没有独立的 BN 算子BN 融合进 Conv2D推理-only 优化推理时: Conv2D(x, W, b) → BN(x) γ·x β 合并: Conv2D(x, W, b) W γ·W / sqrt(σ²ε) b γ·(b - μ) / sqrt(σ²ε) β 省: 一次 HBM 读写不需要存 x 中间结果从这个意义上说推理场景的BN根本不是一个算子——它被优化没了。3.2 LayerNorm / RMSNorm训练存中间量推理即算即用相比 BNLayerNorm 和 RMSNorm 的训练/推理差异小得多——因为它们不存在数据依赖的统计量。但差异仍然存在训练 LayerNorm: y γ·(x - μ) / sqrt(σ² ε) β 需要存 x 和 σ² 用于反向传播 → 显存占用: 每层多存 1 个 x 的副本 (d_model,) 推理 LayerNorm: 不需要存任何中间结果 算完 y 直接用x 的内存可以立即释放LLM 中使用的是 RMSNorm比 LayerNorm 少算 μ 和 β。在推理场景下 RMSNorm几乎总是和前面的 GEMM/注意力输出融合——作为 GEMM epilogue 的一部分直接在寄存器中完成归一化不进 HBM。3.3 厂商实现差异厂商训练 Normalization推理 NormalizationNVIDIAcuDNN BN 有训练/推理两个 modeCUDNN_BATCHNORM_TRAINING/CUDNN_BATCHNORM_INFERENCE推理 mode 不更新 running statsIntel oneDNNdnnl::batch_normalization_forward有training和inferenceprop_kind推理跳过 stats 更新华为 CANN训练特有融合BNTrainingUpdate Conv2D BNTrainingReduce推理直接吸收 BN 到 Conv四、Attention计算模式完全不同的两个算子如果 Conv2D 和 GEMM 还是同一个算子、不同参数的话Attention 在训练和推理之间的差异已经到了可以视为两个不同算子的程度。4.1 训练并行处理所有 token训练 Attention以 MHA 为例: Q x W_q # (B×L, h×d_k) ← 所有 token 并行 K x W_k # (B×L, h×d_k) ← 所有 token 并行 V x W_v # (B×L, h×d_k) ← 所有 token 并行 score Q K^T # (B×L, B×L) ← 巨大的注意力矩阵 O(n²) 的核心开销 attn softmax(score / √d) out attn V # (B×L, h×d_v) 反向传播: 需要存 Q、K、V、score、attn、out、x 全部中间量 显存占比: ~70% 的总训练显存花在存 attention 中间激活上关键特征计算密集QK^T 是大 GEMMML, NL, Kd_k打满 Tensor Core显存爆炸L 增大时中间激活 O(L²) 增长反向需要所有中间量催生了 FlashAttention 和 Gradient Checkpointing4.2 推理 Decode逐 token 自回归推理 Decode生成第 t1 个 token: q_t1 x_t1 W_q # (B, h×d_k) ← 只算当前 token k_t1 x_t1 W_k # (B, h×d_k) ← 只算当前 token v_t1 x_t1 W_v # (B, h×d_k) ← 只算当前 token # KV Cache: 把之前所有 token 的 K/V 存起来 K_cache concat(K_cache, k_t1) # (B, L1, h×d_k) V_cache concat(V_cache, v_t1) # (B, L1, h×d_k) # Attention with cache score q_t1 K_cache^T # (B, 1, h) (B, L, h) → (B, 1, L) ← O(n) 不是 O(n²) attn_out score V_cache # (B, 1, L) (B, L, h) → (B, 1, h)关键特征访存密集瓶颈在读取 KV CacheB×L×d_k×2 bytes/步O(n) 不是 O(n²)因为只算了 q_t1 和所有 K 的 score没有算 Q 全矩阵KV Cache 线性增长长度 L 增加每步多存 2×d_k×dtype bytes4.3 一张表看清差异维度训练推理 Decode计算模式所有 token 并行逐 token 自回归核心计算Q K^T (L×L×d_k)q K_cache^T (L×d_k)复杂度O(L²)O(L)瓶颈计算Tensor Core 利用率访存KV Cache 带宽KV Cache不需要必需——每步增长中间激活存 Q/K/V/score/attn 全部不存任何中间量FlashAttention 角色消除 O(L²) 的 HBM 读写优化 KV Cache 加载4.4 FlashAttention 在训练和推理中的不同角色这是“同一个优化技术、面对完全不同的瓶颈”的典型案例训练FlashAttention tiling: 问题: QK^T 的中间结果 (L×L) 太大写 HBM→读 HBM 浪费带宽 解法: 分块 tiling——在共享内存里算完 attention只写回最终结果 收益: I/O 复杂度从 O(L²L) 降到 O(L²/√M)显存省了 10x 效果: 训练速度 2-3x大 batch 下也能跑长序列 推理FlashAttention 解码: 问题: KV Cache 太大每步要加载全部 K/V 到片上 解法: 多个 thread block 并行加载 KV Cache 的不同分块 收益: 饱和 HBM 带宽不浪费 效果: Decode 加速最高 28x 同一个 FlashAttention kernel 训练时解决的瓶颈是中间矩阵写 HBM 太慢 推理时解决的瓶颈是KV Cache 从 HBM 加载太慢。4.5 对自研芯片的意义Attention 在训练和推理之间的差异是所有算子中最大的。如果你的芯片只优化了训练的大 GEMM tiled softmax推理时可能会遇到GEMV 瓶颈q K^T 在 decode 时是 1×d_k d_k×L一个大 GEMV。如果你的芯片没优化 GEMVTensor Core 利用率接近 0%。KV Cache 带宽瓶颈每步要读 L×d_k 的数据。如果 HBM 不够宽decode 延迟直接卡住。PagedAttention 支持推理需要用 PagedAttention 来管理 KV Cache 显存。这不是提速这是能不能跑没有分页64K 上下文的 KV Cache 浪费 60% 显存。五、激活函数差异最小激活函数ReLU、SiLU、GELU、SiTU是训练/推理差异最小的算子。前向计算完全一致。唯一的差异在反向训练: 需要存输入 x → 用于反向计算梯度 ReLU: dL/dx dL/dy · (x 0) ← 需要 x SiLU: dL/dx dL/dy · [sigmoid(x) x·sigmoid(x)·(1-sigmoid(x))] ← 需要 x 如果不存 x就要重新算前向浪费算力 推理: 只算前向: y f(x) 不需要存任何东西 输出可以直接覆盖输入省显存对实现的启示激活函数本身不需要区分训练/推理版本。差异在内存管理策略——训练需要分配额外的 buffer 存输入推理不需要。六、系统级差异贯穿所有算子的主线以上逐算子的分析背后有几条贯穿所有算子的主线6.1 反向传播的存在决定了显存策略训练: Forward 需要保存中间激活 → 反向才能算梯度 通常每层要存 ~4x 的中间量输入、输出、注意力矩阵等 显存大头 中间激活不是权重 推理: Forward 算完即弃 显存大头 权重 KV Cache这对芯片设计的含义——训练芯片需要更大的 HBM 容量存中间激活和计算/显存平衡的设计推理芯片的 HBM 可以主要分配给权重和 KV Cache。6.2 Batch size 决定了计算访存特征算子维度训练推理Batch size64-5121-8M 维度大B×L小B瓶颈计算FLOPs 打满访存HBM 带宽Arithmetic Intensity1000-10000 FLOP/Byte0.5-5 FLOP/Byte训练芯片需要强大的算力TFLOPS推理芯片需要强大的带宽TB/s。这两者是不同的设计目标——H100 有 3.35 TB/s 和 1979 TFLOPSH200 把带宽提到 4.8 TB/s 但算力没变B200 提到 8 TB/s。NVIDIA 一直在加带宽因为推理场景带宽比算力更值钱。6.3 融合策略不同算子训练融合推理融合Conv2DConv BN(update) ReLUConv Bias ReLUGEMM一般不融合GEMM Bias ActivationNorm单独算需要反向吸收进前序 GEMM/ConvAttentionFlashAttentiontile 融合PagedAttention KV Cache推理的融合可以更大胆——因为没有反向传播的限制可以把更多连续操作合并进一个 kernel节省 HBM 读写。6.4 自动调优策略不同训练推理Find API 开销可以接受几十秒 vs 几小时训练不能接受首次加载必须快算法选择大 batch 偏好 Winograd/Tiled GEMM小 batch 偏好 Implicit GEMM/GEMV缓存策略训练过程中 shape 固定缓存一直有效shape 多变可能需要规则引擎确定性需要bitwise 可复现不需要七、实操建议对算子开发团队的参考7.1 如果只能先做一个做推理当前国内 AI 芯片的主要落地场景是推理部署。理由Kernel 数量少推理只需要 Forward不需要 Backward。一个卷积算子是 Conv2D 不是 Conv2D Conv2DBpropInput Conv2DBpropFilter工作量少 2/3。精度要求低INT8/FP16 推理是标准做法不需要 FP32 累加器的大面积设计。更容易做融合不需要考虑反向传播的中间量保留可以大胆做算子融合。市场更成熟云端推理、端侧推理、边缘推理都有明确需求。7.2 推理优先但 kernel 设计要为训练预留虽然优先做推理但 kernel 的接口设计应该为训练预留// 建议的卷积 API训练和推理共用 forward:typedefstruct{intbatch_size;intinput_c,input_h,input_w;intoutput_c,kernel_h,kernel_w;intstride_h,stride_w;intpad_h,pad_w;DataType input_dtype;DataType weight_dtype;DataType accum_dtype;// FP32 for training, FP16 for inferencebool need_activation;// true: fused ReLUbool need_save_intermediate;// true: save input for backward (training)}Conv2DDesc;差别就在accum_dtypeFP32 还是 FP16和need_save_intermediate存不存中间量两个参数上——kernel 的数学逻辑是一样的。7.3 统一算子 两套调度策略不要为训练和推理写两套 kernel。改用一套 kernel 两套调用参数统一 GEMM kernel: gemm(M, N, K, A, B, C, accum_dtype, split_k) 训练调用: gemm(M131072, N4096, K4096, accumfp32, split_k1) 推理调用: gemm(M1, N4096, K7168, accumfp16, split_k8) ↑ 同一套代码用 split_k 处理小 M 场景M 大时自动走 tile GEMM 路径M 小时自动走 Split-K / GEMV 路径accum_dtype控制累加器宽度7.4 精度管理的通用策略场景输入/权重精度累加器精度说明训练 ForwardFP16/BF16FP32防止累加溢出训练 BackwardFP16/BF16FP32梯度精度更敏感推理FP16FP16FP16 或 FP32可选FP32更安全推理INT8INT8INT328 位乘 8 位累加用 32 位推理FP8FP8FP16/FP32FP8 动态范围窄必须高精度累加一条原则累加器 2 × max(dtype bits)。8-bit 输入用 16-bit 累加不安全必须 32-bit。16-bit 输入用 32-bit 累加。32-bit 输入用 32-bit 累加。八、总结核心发现训练和推理的算子差异不在数学定义上而在 shape、精度、反向传播和融合策略上。同样的 Conv2D、GEMM 和 Attention在训练和推理场景下是两套完全不同的优化路径。Attention 的差异最大——训练是 O(L²) 的大 GEMM tiled softmax推理是 O(L) 的 GEMV KV Cache 加载。可以说是同一个名字的两个不同算子。Batch size 是差异的根本来源——训练 MB×L 1所有算子是计算密集的推理 MB通常 1-8所有算子是访存密集的。训练芯片和推理芯片的设计目标不同——训练需要峰值算力TFLOPS推理需要高带宽TB/s。NVIDIA 的 H1003.35 TB/s和 H2004.8 TB/s说明了这个趋势。不要为训练和推理写两套 kernel——用统一 kernel 两套调用参数accum_dtype、split_k 配置。数学逻辑一样调度策略不同。跨算子的差异矩阵算子训练 M 维度推理 M 维度训练精度推理精度训练额外算子推理特有优化Conv2DB×L大B小FP32 累加INT8 可Backprop ×2BN 吸收GEMMB×L大B小FP32 累加INT8 可Backprop ×2Split-K GEMVLayerNormB×L大B小FP32FP16存 x 用于反向融合进 GEMM epilogueBatchNormB大B小FP32FP32running stats 更新吸收进 ConvAttentionL长L长FP32 累加FP16存全中间量KV Cache PagedAttnReLU/SiLUB×LBFP32FP16/INT8存 x 用于反向融合进 GEMM epilogue附录进一步阅读NVIDIA cuDNN 开发指南: docs.nvidia.comAMD MIOpen: github.com/ROCm/MIOpenIntel oneDNN: oneapi-src.github.io/oneDNN华为 CANN 算子开发: hiascend.comFlashAttention: Dao et al. (2022):FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesssgemm-bi: Deterministic, batch-invariant GEMM for training