Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP
02:00 · June 11, 2026 · Hugging Face Blog
Summary
The second installment in this PyTorch profiling series moves from raw matrix multiplication to the nn.Linear module that underpins most transformer blocks. The author shows that nn.Linear is essentially a thin wrapper around the same GEMM-plus-bias pattern examined in Part 1, but now expressed through aten::addmm so that the bias addition occurs inside the cuBLAS kernel’s epilogue rather than as a separate write to HBM. Because the epilogue already fuses the bias, torch.compile has little additional work to perform on a lone linear layer; the only measurable change is the removal of a few microseconds of CPU-side view and stride bookkeeping.
Stacking three such layers with a GeGLU activation produces a representative MLP block. In eager mode the profiler records five distinct GPU kernels per forward pass—three GEMMs plus separate pointwise launches for the GELU and the subsequent multiplication—plus repeated cudaOccupancyMaxActiveBlocksPerMultiprocessor queries on the linear paths. The intermediate activation tensor between the gate and up projections must travel through HBM twice, adding roughly 50 MB of memory traffic for the chosen batch and sequence lengths.
When torch.compile is applied, the two pointwise operations and the intervening reshape collapse into a single Triton kernel that keeps the intermediate values in registers. The three GEMM kernels remain byte-for-byte identical to their eager counterparts, confirming that the compiler’s benefit lies in eliminating the CPU dispatch chain and the extra HBM round-trip rather than in altering the matrix-multiplication code itself. Layout descriptors embedded in the kernel names (for example, _tn_ versus _nn_) reveal that cuBLAS selects different tiling strategies for the down-projection, which explains the observed 10 % runtime difference despite identical FLOP counts.
The post closes by replacing the compiled MLP with an expert-written Triton kernel taken from the Hugging Face Hub. This substitution lets the reader compare hand-tuned memory access patterns against both the eager and compiled baselines on an A100, reinforcing how profiler traces can guide the decision of whether further fusion or a custom kernel is worthwhile.
Why it matters
Directly addresses production-level PyTorch optimization, kernel fusion, and profiling for ML Engineers building or tuning models, with actionable scripts and trace interpretation that Dutch teams can apply immediately.
