PyTorch's nn.Linear already runs a fused kernel—here's why torch.compile adds nothing for one layer
Curated by the Inblix editorial team
If you’ve been sprinkling torch.compile on every layer hoping for free speed, the profiler has some humbling news. In the second installment of our deep dive into PyTorch profiling, we swapped a hand-rolled matmul-and-add for a standard nn.Linear layer, the literal building block of modern deep learning. The trace revealed that eager-mode PyTorch isn’t as naive as you might think.
When you call a linear layer with bias=True, PyTorch doesn’t launch separate matrix multiplication and addition kernels. It dispatches to aten::addmm, which calls a cuBLAS GEMM kernel that fuses the bias addition directly into its writeback stage, a technique called an epilogue. This avoids an extra round trip to high-bandwidth memory, which is exactly the kind of optimization you’d hope a compiler would perform. The only other operation we spotted, an aten::t transpose on the weight matrix, is a metadata-only sleight of hand on the CPU. No data moves; no GPU kernel launches.
This means the kernel you see under eager execution is already the optimized, fused kernel. So when we applied torch.compile to a single nn.Linear layer, the profiler confirmed there was nothing left to fuse. The compiler’s magic is real, but it needs a chain of operations to work with. A lone linear layer is already as fast as it gets out of the box.
The practical upshot? Don’t waste your mental energy micro-optimizing individual layers. The real wins from torch.compile, as we’ll see when we stack three of these into an MLP, come from fusing the computational graph between layers, like merging the activation function with the preceding matrix multiply. The profiler isn’t just a debugging tool; it’s a bullshit detector for optimization claims.
💡 Key Takeaways
- Eager-mode nn.Linear already dispatches to a fused cuBLAS kernel that bakes bias addition into the matmul's writeback, eliminating a separate memory round-trip.
- The aten::t transpose you see in the trace is a pure metadata operation on the CPU that doesn't copy data or launch a GPU kernel.
- torch.compile provides no benefit for a single linear layer because there are no separate operations to fuse—the kernel is already optimal.
- Profiler traces reveal that optimization headroom lies in fusing operations *between* layers, not within already-optimized primitives like nn.Linear.
Keep reading: See related articles below for more coverage on this topic.
Get smarter about AI
The sharpest AI news, curated daily. Delivered free to your inbox.