PyTorch 性能分析：从 nn.Linear 到融合 MLP

本文是 PyTorch 性能分析系列的第二部分，重点介绍如何通过算子融合优化多层感知机（MLP）。文章首先回顾了 PyTorch 性能分析工具（如 torch.profiler）的使用方法，然后深入分析了 nn.Linear 层的计算瓶颈，包括内存带宽和计算效率。接着，作者展示了如何将多个线性层和激活函数融合为一个自定义的融合 MLP 内核，利用 PyTorch 的 torch.compile 或自定义 CUDA 内核实现。实验结果表明，融合后的 MLP 在推理和训练中均能显著减少内核启动开销和内存访问，提升吞吐量。文章还讨论了融合策略的适用场景和潜在限制，为开发者提供了实用的性能优化指南。