signal.

Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP

Hugging Face published the second part of its series on profiling in PyTorch, exploring the transition from individual nn.Linear layers to a fused Multilayer Perceptron. The post details how PyTorch handles matrix multiplication, weight transposition via tensor metadata views, bias folding through epilogues, and the specific behaviors of eager versus compiled execution on an NVIDIA GPU.

Hugging Face released the second installment of its profiling in PyTorch series, focusing on moving from a hand-written matrix multiplication and addition pair to an nn.Linear module and subsequently a Multilayer Perceptron block. Utilizing an NVIDIA A100 GPU, the blog post examines the mechanics of linear layers, tensor views, and kernel fusion. When analyzing the profiler traces, the article highlights that the aten::t operation does not copy data or launch a GPU kernel, but instead rewrites tensor metadata and strides on the CPU to create a transposed view. Furthermore, bias addition does not appear as a separate kernel because it is folded into the matrix multiplication kernel as an epilogue using cuBLAS. The text contrasts eager and compiled execution for a single nn.Linear module, demonstrating that torch.compile removes CPU dispatch overhead rather than changing the underlying GPU kernel, since a single GEMM-with-bias lacks multiple operations for compile to fuse. The analysis shows that layout descriptors in kernel names, such as the tn suffix, indicate how a kernel walks its input during the inner loop. Finally, the post introduces the profiling of a Multilayer Perceptron incorporating the GeGLU activation variant, providing practical insights into multi-operation neural network blocks.

Why it matters

Understanding how PyTorch dispatches operations, handles tensor views, and utilizes kernel epilogues helps developers optimize deep learning models and avoid common misconceptions about when torch.compile provides performance gains.

Topic dossiers NVIDIA Hugging Face