PyTorch剖析显示融合MLP显著加速，torch.compile与手写Triton内核对比

背景介绍

随着大模型在训练和推理阶段对算力的需求日益提升，如何在保持模型精度的前提下压缩算子调度开销，成为业界关注的焦点。本文聚焦PyTorch常用的`nn.Linear`层，进一步通过堆叠三个Linear形成的MLP（GeGLU 变体），对比 torch.compile、Triton 手写内核 与原始 eager 模式的 profiling 数据，揭示了不同优化路径的真实收益。

Linear 层的性能瓶颈

转置元数据：aten::t 仅重写张量的 stride，不涉及 GPU 计算，但在 eager 调度链中会产生额外的 CPU 开销。编译后此步骤被折叠，省去几微秒的调度时间。
Bias 融合：nn.Linear 实际调用的是 aten::addmm，bias 已在 cuBLAS 的 GEMM epilogue 中完成写回，避免了独立的加法 kernel。
GPU Kernel：无论是 eager 还是编译后，实际执行的都是同一条 cuBLAS GEMM kernel（如 cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8），因此单个 Linear 的加速空间有限。

torch.compile 的作用

编译器通过 Inductor 能在编译阶段把视图链（transpose、reshape 等）提前求解，直接生成硬编码 stride 的 addmm 调用。结果是：

CPU 侧调度链缩短，省去约 3‑5µs 的元数据处理。
对于单个 Linear，GPU 端无变化，整体加速不明显。
对于 MLP，编译后会把 GeLU 与乘法以及一次 reshape 融合为一个 Triton pointwise kernel，显著降低了 HBM 读写次数。

手写 Triton 内核的对比

使用 Hugging Face kernels 库提供的 LigerGEGLUMLP，内部实现了一个专门针对 [batch*seq, hidden] 大小的 Triton kernel：_geglu_tanh_forward_kernel。其优势在于：

一次性读取 gate 与 up，在寄存器中完成 GeLU 与乘法，避免了中间 tensor 的全局内存往返。
硬件调参：块大小与线程布局基于列维度自动选择，适配不同的 GPU 架构，无需每次 shape 变化后重新编译。
实测结果显示，手写内核的运行时间约 92.8µs，略慢于 Inductor 编译产生的 89.4µs（后者为特定 shape 专属编译），但在多变 batch/seq 场景下保持稳定，省去了编译成本。

实验结果与启示

方案	GPU Kernel 数	关键点位融合	CPU 调度开销	单次前向耗时
Eager Linear	1	无	包含 transpose/view	0.19ms
Compiled Linear	1	去除 transpose/view	略降	0.19ms
Eager MLP	5 (3 GEMM + GeLU + Mul)	无	3 次 occupancy query	0.55ms
Compiled MLP	5 (3 GEMM + 1 Triton fused)	GeLU+Mul+reshape 合并	去除 pre‑ops	0.51ms
Liger MLP	5 (3 GEMM + 1 Triton fused)	同上且无 compile guard	零	0.53ms

从表中可以看出：

对单个 Linear，torch.compile 只能削减 CPU 元数据调度，提升有限。
对于包含点位运算的 MLP，Fusion（无论是 Inductor 自动生成还是手写 Triton）是主要加速点，能够把一次全局内存往返压缩到寄存器层级。
手写内核在 shape 多变的生产环境下更具鲁棒性，虽然在固定 shape 上略逊于编译专属 kernel，但整体收益更可预测。

小结

本文通过实测阐明了 PyTorch 中 Linear‑MLP 计算路径的内部细节，并对比了 torch.compile 与 Liger Triton 两种加速手段。对研发团队而言，推荐的实践流程是：先使用 torch.compile 检查是否存在明显的调度冗余；若模型包含大量点位运算且对 latency 敏感，可进一步引入经过 CI 编译的 Triton 手写 kernel，以获得更稳定的加速效果。后续我们将继续探索将注意力机制同样进行类似的融合，期待在完整 Transformer 上实现更大幅度的算力压缩。

“猜测 → 验证 → 纠正” 是每一次 profiling 的核心思维，只有在不符合预期时才真正发现优化空间。