.png)
Hugging Face published “Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP,” the second post in its profiling series. It moves from a hand-written matmul + add example to an nn.Linear-based MLP and uses profiler traces to show where PyTorch leaves performance on the table, then where torch.compile can recover it.
nn.Linear is the same bottleneck, just packaged differentlyThe post starts by swapping the explicit torch.add(torch.matmul(x, w), b) pattern from Part 1 for nn.Linear(..., bias=True). That matters less for math than for ergonomics: nn.Linear is the primitive most models actually build on, so profiling it gets you much closer to real workloads.
The key point is that nn.Linear does not magically remove the overhead visible in the profiler. It still resolves to a matrix multiply plus bias addition, and the CPU still has to schedule and launch GPU kernels. If the workload is small enough, launch overhead and kernel count dominate; if the workload is large enough, the GEMM dominates and the pointwise work starts to disappear into the noise.
The blog also keeps the same profiling discipline from Part 1: inspect traces, identify kernel boundaries, and reason about whether the model is compute-bound or overhead-bound before changing anything.
The MLP block in the post is deliberately simple: three linear layers with an activation in between. In eager mode, that turns into multiple GPU kernels per forward pass. The post’s central observation is that the intermediate tensors produced by those small pointwise ops are materialized and moved through GPU memory, even when they are structurally part of one fused computation.
That memory traffic is the real tax. It is not just “too many kernels” in the abstract; it is a full round-trip through HBM for intermediate activations that could have stayed on-chip if the compiler had fused the chain more aggressively. On modern accelerators, that distinction matters more than shaving a few microseconds of Python overhead.
This is the kind of case where profiler traces are useful because the bad shape is obvious once you know what to look for: a short sequence of kernels, lots of tiny ops around the matmul-heavy core, and visible gaps between launches. The diagnosis is not “PyTorch is slow”; it is “the execution plan is too fragmented for the hardware.”
torch.compile fuses the pointwise tail, not the GEMMsThe optimization path here is torch.compile, which captures the graph and exposes fusion opportunities around the linear stack. The post is careful about what gets fused and what does not: the large GEMMs remain separate, while nearby pointwise ops and reshapes can be combined into a smaller number of kernels.
That distinction is important for realistic expectations. torch.compile is not turning an MLP into one giant monolithic kernel. It is removing avoidable fragmentation around operations that are cheap individually but expensive when each one forces a launch and a writeback to global memory.
For ML engineers, the practical lesson is to look at the operator mix before assuming compilation will save the day. If your model is already dominated by large GEMMs, fusion may not change much. If your graph has lots of thin pointwise edges around a few heavy ops, compilation can meaningfully reduce launch count and memory traffic.
The post’s example is a good benchmark shape for that reasoning: enough structure to show the compiler’s effect, small enough that the trace remains readable, and close enough to a real MLP block to be useful in practice.
The strongest part of the series is the workflow, not the specific MLP. Start with a trace, identify whether the problem is kernel launch overhead, memory movement, or raw compute, and only then decide whether to rewrite the model or hand control to the compiler.
That approach scales better than micro-optimizing module code by intuition. It also keeps the optimization target honest: if the trace shows the model already saturating the GPU on GEMMs, then fusion is a marginal win. If it shows a trail of tiny ops and materialized intermediates, then fusion is exactly the kind of intervention that pays off.
This is a useful reminder for production inference and training stacks alike: model structure determines how much the compiler can do, and profiler traces tell you whether the structure you wrote is the structure the hardware is actually executing.
