为什么PyTorch编译如此之快：内核融合

本文是一篇技术原理深度解析文章，其价值不在于宣布新功能，而在于解构torch.compile性能提升背后的底层逻辑。文章通过对比融合前后的代码生成实例，揭示了性能优化的本质。

内核开销与内存瓶颈：被忽视的计算税

现代GPU的强大并行计算能力常受两个隐性成本制约：启动内核的固定开销和中间数据的全局内存往返。文章指出，在未编译的PyTorch中，每个运算符（如乘法、加法）都会触发一次独立的内核启动和完整的内存读写周期。以文中的三层操作为例，未融合时产生了八个内存操作。这些操作并非计算本身的开销，而是由计算范式引发的“税务”开销，极大地稀释了GPU的算力利用率。

垂直融合与计算图优化：从“点”到“线”的跃迁

文章核心介绍的“垂直融合”是针对深度学习计算链最自然的优化策略。它并非简单地将代码合并，而是对计算图进行语义理解后的重构。

• 依赖关系识别：编译器分析出 x * w -> + b -> .sigmoid() 的严格依赖链。
• 操作内化：将中间变量 tmp 的存储从全局内存（DRAM）“提升”到GPU的寄存器（最快的存储层）。寄存器是线程私有的、访问延迟极低，但容量极小。只有融合后的小代码块能充分利用它。
• 生成融合内核：最终生成的单个Triton内核（如文中的triton_poi_fused_add_mul_sigmoid_0）在内部完成所有计算步骤。数据仅在进入内核时从全局内存加载一次，计算结束后写回一次，将内存访问降到了理论最小值。

优化的本质：计算模式从“计算驱动”转向“数据驱动”

文章中未被点明但至关重要的洞察是：融合优化的根本是改变了计算范式。
未融合时，计算由一系列离散的内核调度驱动，内存围绕每个内核被动服务，导致频繁的“加载-计算-存储”循环。
融合后，数据流成为主导。编译器规划出一个高效的数据搬运路径：一次性加载所有输入数据到快速存储器，然后在芯片内完成全部计算，最后一次性写回结果。计算围绕数据的生命周期紧密组织，消除了所有冗余的数据流转。这解释了为何速度提升能达到数量级——它直接干掉了计算图中绝大部分低效的、被隐藏的“搬砖”过程。