Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP
Hugging Face Blog
PyTorch's profiler traces show that nn.Linear uses an optimized cuBLAS kernel (addmm) that folds bias addition into the matrix multiplication via an epilogue, eliminating the need for a separate add kernel. When torch.compile optimizes a single Linear layer, it removes only CPU dispatch overhead by pre-computing strides at compile time—the GPU kernel remains identical. In a three-layer MLP with GeGLU activation, the forward pass launches exactly five GPU kernels: three GEMMs for the linear projections plus one each for GeLU and element-wise multiplication.
