Profiling PyTorch from a Linear layer up to a fused MLP

Aritra Roy Gosthipaty and a team at Hugging Face published a follow-up to their PyTorch profiler tutorial, walking from a single nn.Linear up to a full MLP and showing where the easy performance wins actually live. The first finding is that nn.Linear with bias does not need a separate add kernel. PyTorch folds the bias into the GEMM epilogue, so the addition happens in the same kernel just before the result hits memory. The GEMM dispatcher then picks among precompiled kernel binaries based on tensor layout, with names like the tn suffix encoding stride patterns. The same logical computation can land on different kernels depending on input shape, and the only way to see this is in the trace.

The MLP example, a feed-forward block with GeGLU activation, runs as five GPU kernels in eager mode: three matmuls plus a separate GeLU and an elementwise multiply. Calling torch.compile collapses the GeLU and the multiply into one kernel, keeps intermediate tensors in registers, and saves about 50 MB of memory traffic per forward pass. The post also benchmarks against the Liger kernel library, which ships hand-tuned, version-pinned kernels. On static shapes Liger comes in at 92.8 µs versus 89.4 µs for the compiled version, slightly slower, but it avoids the per-shape recompilation when input dimensions change. The writeup is disciplined about methodology in a way most profiler tutorials are not: form a prediction first, then read the trace, then dig only where prediction and trace disagree.

Why it matters

If you train or serve transformers and have never opened a profiler trace, this is the cleanest introduction to where time and memory actually go inside an MLP block, with concrete numbers attached. The torch.compile versus Liger comparison also doubles as a deployment decision: if your serving shapes are static, compile is faster; if shapes vary, prebuilt kernels save you the recompile bill.