PyTorch：torch.compile通过内核融合实现最高10倍加速

PyTorch编译器torch.compile无需对模型代码进行任何修改即可在GPU上实现最高10倍加速。这种加速背后的关键机制是内核融合——将多个操作合并为一个GPU内核以消除向全局内存写入中间值的昂贵操作。

为何GPU内存会成为模型训练的瓶颈？

GPU计算的不同操作成本并不相同。启动每个GPU内核都有固定开销，而每次向全局内存写入中间值比处理器本身寄存器中的操作慢数百倍。没有编译器时，PyTorch将每个操作（multiply、add、sigmoid）作为独立内核执行，每个内核都有自己的全局内存读写操作。

具体来说，三个连续的逐元素操作（x * w、+ b、.sigmoid()）不经融合时需要5次读和3次写——共计8次完整的张量I/O操作。内核融合后：3次读 + 1次写 = 4次操作，即内存流量减少50%。

torch.compile如何自动合并内核？

torch.compile使用Inductor编译器分析计算图并应用多种类型的融合：

• 垂直融合（最常见）：合并依赖操作使中间值保留在GPU寄存器中。消除了所有否则需要进入全局内存的临时张量。
• 逐元素融合：将逐元素操作（multiply → add → sigmoid）合并为一个Triton内核。
• 归约融合：将归约操作（如求和）与周围操作结合，对批归一化至关重要。
• GEMM + Epilogue融合：将偏置和激活函数直接添加到矩阵乘法末尾。
• 水平融合：对相同输入的独立操作（如sin(x)和cos(x)）在一个内核中并行执行。

实践中内核会发生什么？

融合生成的Triton内核将所有中间值保留在GPU寄存器中，最后才写入结果：

加载x、加载w、加载b
→ tmp = x * w       （保留在寄存器中）
→ tmp = tmp + b     （保留在寄存器中）
→ tmp = sigmoid(tmp)（保留在寄存器中）
写入结果             （唯一写入全局内存的操作）

Inductor自动使用命名约定triton_poi_fused_*（逐元素）和triton_per_fused_*（归约）命名融合内核，通过设置TORCH_LOGS="output_code"变量可以看到这些。

结果与应用

PyTorch社区记录了由于融合、消除内核启动开销和减少内存流量而带来的实际模型高达10倍加速。所需代码只有一行：

model = torch.compile(model)

Inductor编译器随后独立构建最优Triton或CUDA内核，无需手动编写GPU代码。这种自动化使torch.compile成为PyTorch生态系统中生产机器学习训练和推理的最重要工具之一。