---
title: "Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP"
description: "Part 2 of the series examines how PyTorch profiler traces change from a single bias-enabled nn.Linear to a three-layer MLP with an intervening activation, using scripts run on an NVIDIA A100-SXM4-80GB GPU. It explains that nn.Linear represents y = x @ w.T + b, dispatches aten::addmm, and folds bias addition into the cuBLAS GEMM epilogue, while aten::t only rewrites weight strides and launches no GPU kernel. For the MLP, eager execution uses three GEMMs plus separate activation and multiplication kernels, whereas torch.compile fuses the latter operations and keeps the intermediate out of HBM, while leaving the GEMMs unchanged. A hand-written Liger Triton kernel provides the same fusion without Dynamo guards or compile latency, trading Inductor’s shape-specialized speed for generic execution across changing shapes."
---

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

[Hugging Face](https://yomu.fyi/company/hugging-face) · Aritra Roy Gosthipaty, Rémi Ouazan Reboul, Sergio Paniego, Pedro Cuenca, Sayak Paul · Jun 11, 2026

**Type:** Tutorial

## Summary

Part 2 of the series examines how PyTorch profiler traces change from a single bias-enabled nn.Linear to a three-layer MLP with an intervening activation, using scripts run on an NVIDIA A100-SXM4-80GB GPU. It explains that nn.Linear represents y = x @ w.T + b, dispatches aten::addmm, and folds bias addition into the cuBLAS GEMM epilogue, while aten::t only rewrites weight strides and launches no GPU kernel. For the MLP, eager execution uses three GEMMs plus separate activation and multiplication kernels, whereas torch.compile fuses the latter operations and keeps the intermediate out of HBM, while leaving the GEMMs unchanged. A hand-written Liger Triton kernel provides the same fusion without Dynamo guards or compile latency, trading Inductor’s shape-specialized speed for generic execution across changing shapes.

## Context

The post moves from profiling a hand-written matmul-add pair to understanding nn.Linear and then an MLP block. It focuses on identifying CPU dispatch work, GPU kernels, epilogues, tensor views, fusion, and the trade-off between shape-specialized compilation and execution across changing shapes.

## Approach / What changed

It walks through profiler traces from 02\_linear.py, 03\_simple\_mlp.py, and 03\_kernels\_mlp.py, comparing eager and compiled execution on an NVIDIA A100-SXM4-80GB GPU. It examines aten::t, aten::addmm, cuBLAS GEMMs, torch.compile’s fused Triton kernel, and a hand-written Liger Triton kernel.

## Takeaways

- nn.Linear uses aten::addmm with bias folded into the cuBLAS GEMM epilogue, so eager execution does not launch a separate bias-add kernel.
- For the MLP, torch.compile fuses GeLU, multiplication, and reshape into one Triton kernel; the three GEMMs remain unchanged and the intermediate avoids an HBM round trip.
- Liger’s fused kernel runs across changing shapes without recompilation, but its 92.8 µs runtime is slightly slower than Inductor’s 89.4 µs shape-specialized kernel.

**Tags:** [Machine Learning](https://yomu.fyi/topic/machine-learning), [Performance](https://yomu.fyi/topic/performance), [Python](https://yomu.fyi/topic/python)

- Source: [Hugging Face](https://huggingface.co/blog/torch-mlp-fusion)
- Source URL: https://huggingface.co/blog/torch-mlp-fusion
- Ingested by Yomu: 2026-08-29T01:50:20.841Z

[Read original post](https://huggingface.co/blog/torch-mlp-fusion)