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

PyTorch's `nn.Linear` module transposes its weight tensor before performing matrix multiplication and addition, as revealed by profiler traces showing an `aten::t` operation that only modifies tensor metadata on the CPU without launching a GPU kernel. The bias addition is folded into a single `aten::addmm` kernel rather than appearing as separate multiplication and addition operations, demonstrating how PyTorch optimizes the linear layer's forward pass. This profiling analysis builds on Part 1's examination of manual matmul-add pairs to show how the framework's module-level abstractions handle these operations internally.

Updated • 62 • 1 Profiling in PyTorch Part 2 : From nn.Linear to a Fused MLP Update on GitHub https://github.com/huggingface/blog/blob/main/torch-mlp-fusion.md In the first part of this series "Profiling in PyTorch" https://huggingface.co/blog/torch-profiler , we used torch.add torch.matmul x, w , b to learn how to read PyTorch profiler traces. We also discussed several other topics that came our way - the CPU dispatch chain, launch overhead, the difference between an overhead-bound and a compute-bound regime, and some internals of torch.compile . In the second iteration this blog post , we climb one rung up the ladder. We replace the hand-written matmul-add pair with an nn.Linear with bias=True . This is the building block every deep learning model uses. We then stack three of them specific to our example , with an activation in between, to form a Multilayer Perceptron MLP block. The scripts for this blog post live here: , 02 linear.py , and 03 simple mlp.py . Like before, it helps to open them in a separate tab and walk through the code as you read. We use an 03 kernels mlp.py NVIDIA A100-SXM4-80GB GPU to run the scripts. It is really easy to set up a GPU on the Hugging Face infrastructure and experiment with the scripts using Dev Mode with Spaces . One could also run the scripts with the Hugging Face Jobs pipeline . Before we begin, a quick recap of two ideas we will lean on repeatedly: - A GPU kernel is a program that runs in parallel on many threads of the GPU. - The CPU schedules and launches these kernels. Most of the PyTorch overhead you see in a profiler trace is this scheduling work. From matmul-add to Linear nn.Linear is a module wrapper around the same matrix multiplication and addition we already profiled in Part 1 https://huggingface.co/blog/torch-profiler . The only difference is that it owns its weight and bias as parameters and exposes a forward method that PyTorch users have grown familiar with. bias=True would truly emulate the multiplication and addition operations we have seen in part 1 of the series linear layer = nn.Linear in dim, out dim, bias=True y = linear layer x The operation at hand can be written as: y = x @ w.T + b Where x is the input, w is the weight and b is the bias. Let's run 02 linear.py https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/02 linear.py and check the profile. uv run 02 linear.py --batch 1024 --in dim 32 --out dim 64 uvx trace-util traces -b traces is a utility that will sync your traces to a trace-util Hugging Face bucket and then provide the Preffeto URLs on your terminal. Figure 1 shows the profiler trace of a forward call of the linear layer. We trace the forward call of the linear layer with a similar schedule setup as the previous traces, with wait=1 , warmup=1 and active=3 . This is why we see three Profile Steps in the CPU and GPU lanes. What is the transpose doing? If we zoom into the profiler trace, as we do in Figure 2, we notice an aten::t transpose op before the aten::addmm multiplication and addition op. We can already figure out that nn.Linear transposes the weight parameter and then multiplies it with the input. This is the reason we see an aten::t op. An important thing to notice is that aten::t does not really copy or reorganize data: it only rewrites tensor metadata shape and stride on the CPU to represent the transposed matrix. It does not launch a kernel on the GPU. One can verify this two ways: by looking at the GPU lane in the trace, or by checking the aten::t row in the profiler table and the time it took on CUDA. Why are there no separate mul and add kernels? There is no aten::add the bias addition in the dispatch chain of the linear layer, as seen in Figure 3. This is because the bias addition has been folded into the matrix multiplication kernel, using what is called an epilogue . An epilogue is a small computation that a GEMM GEneral Matrix Multiply kernel does at the very end, just before it writes its result back to HBM High Bandwidth Memory, the GPU's main memory . Adding a bias, applying an activation, or scaling by a constant are all classic epilogues. The point of an epilogue is to avoid loading or writing to HBM a second time, since memory traffic makes an operation expensive. nn.Linear calls torch.nn.functional.linear , which, in turn, calls aten::linear . aten::linear looks at the inputs, notices that a bias was passed, and dispatches aten::addmm bias, x, weight instead of doing a matmul and an add separately. addmm computes: out = x @ weight.T + bias The cuBLAS GEMM kernel that runs on the GPU has a bias-add variant built in, and that's the kernel aten::addmm picks. The add never appears as a separate kernel because it is part of the matmul kernel's writeback , which is exactly what an epilogue is. This is the moment to notice something subtle. The kernel you saw in Part 1 under --compile https://huggingface.co/blog/torch-profiler did-we-fuse-the-matmul-and-add-kernels-into-one addmm is the kernel that eager nn.Linear already uses. There is nothing left for torch.compile to fuse here, which is the next thing we will verify. Can --compile help a single Linear? Let's compile the forward call and look at the profiler trace. The profiler trace is visualized in the next section where-did-the-transpose-go-kernel-layouts-and-pre-ops uv run 02 linear.py --batch 1024 --in dim 32 --out dim 64 --compile uvx trace-util traces -b traces If you compare the eager and compiled traces for a single nn.Linear 's forward , you will find: - The same cuBLAS GEMM kernel on the GPU. - The same aten::addmm op on the CPU. - A few extra rows on the CPU lane unique to compile. This is worth internalizing. A common reflex is to reach for torch.compile whenever a model feels slow. For a single GEMM-with-bias, compile has very little to do. This is not a bug, this is just that compile needs more than one operation to possibly do any fusing. Let's prove that by looking at an MLP stacking-two-linears-the-mlp . Where did the transpose go? Kernel layouts and pre-ops A careful reader of the two traces eager vs compile will notice that the eager CPU dispatch chain has more in it than the compiled one. Figure 4: Eager dispatch chain where aten::linear walks through aten::t transpose and then aten::addmm | The eager CPU dispatch chain inside aten::linear is aten::t followed by aten::addmm Figure 4 . To understand what aten::t actually does, we need a quick detour into strides and views . A tensor stores its data as one flat, contiguous run of numbers in memory. The shape and stride are metadata that sit on top of that run and tell PyTorch how to walk it: a stride of s0, s1 means "step s0 elements to move one row, step s1 to move one column". Change the metadata and you get a different view of the same raw data, with no copy: M = torch.tensor 0, 1 , ... 2, 3 , ... 4, 5 M.shape, M.stride torch.Size 3, 2 , 2, 1 two steps per row, one step per column T = M.t transpose T.shape, T.stride torch.Size 2, 3 , 1, 2 shape and stride swapped, data untouched T tensor 0, 2, 4 , 1, 3, 5 T.flatten forced to materialize, so the data is reordered tensor 0, 2, 4, 1, 3, 5 M.t did not move a single number. It returned a new view whose strides are swapped, so reading it row-by-row now walks the original buffer 0, 1, 2, 3, 4, 5 in transposed order. The underlying data is identical; only the metadata differs. This is exactly what aten::t does inside the linear layer: it does not allocate a new tensor or copy any data, it produces a view of the weight with rewritten strides. As we can see in Figure 5, compile did not remove a GPU kernel: it removed the CPU overhead of dispatching that view. Inductor traced through the view chain at compile time, computed the resulting strides once, and emitted a direct aten::addmm call with those strides hard-coded. A few microseconds of CPU work disappear while the GPU does identical math. As one would expect, when the input data violates the strides precomputed by the compiler, it will throw an error. If you look at the GPU lane in both traces, there is exactly one kernel per forward, and it is the same kernel both times: cutlass 80 wmma tensorop bf16 s161616gemm bf16 32x32 32x1 tn align8 If no transpose kernel ran, who taught the GEMM to read the weight matrix in transposed order? The answer is in the kernel's name. Look at the suffix: cutlass 80 wmma tensorop bf16 s161616gemm bf16 32x32 32x1 tn align8 ^^ That tn is the layout descriptor. cuBLAS and CUTLASS precompile a separate kernel binary for each combination of input layouts. n non-transposed and t transposed describe how a kernel walks its input during the inner loop. The dispatcher's job is to look at the input strides, decide which suffix combination matches, and pick the right precompiled kernel. The kernel name in a profiler trace is a hash dump of the kernel's identity. If two runs show the same kernel name, the GPU is doing the same work. If they differ e.g., tn vs nn , bf16 vs fp16 , or s16816gemm vs s161616gemm then the GPU is doing different work, and the dispatcher took a different branch. Learning to read this name is one of the most useful habits when comparing traces. Stacking three Linears: the MLP In this section, we will profile a Multilayer Perceptron MLP . To make this more interesting, we will profile a feed-forward network with the GeGLU activation variant which is quite heavily used in practice . This is also our way of paying tribute to one of the greatest lines ever written in the history of deep learning research Figure 6 . | Figure 6: The conclusion section of the | python class SimpleGeGLUMLP nn.Module : def init self, dim, hidden : super . init self.gate proj = nn.Linear dim, hidden, bias=False self.up proj = nn.Linear dim, hidden, bias=False self.down proj = nn.Linear hidden, dim, bias=False def forward self, x : g = self.gate proj x u = self.up proj x h = F.gelu g, approximate="tanh" m = h u y = self.down proj m return y You will find the entire script here: 03 simple mlp.py https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03 simple mlp.py . Execute it like so: uv run 03 simple mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 uvx trace-util traces -b traces Before we open the trace, let's think together about what we should expect to see. The forward function does a fair amount of computation, but most of it is already familiar to us. We should expect three aten::linear dispatches, one for each nn.Linear layer. We should also expect two pointwise kernel launches, one for the GeLU and one for the multiplication. Forming this expectation before looking is the single most useful habit in the profiling journey: you read the trace to confirm or break a guess, not to form one from scratch. From Figure 7 we can pat ourselves on the back, as our intuition was correct. Per forward pass one mlp fwd , the GPU runs exactly 5 kernels. Figure 8 highlights the "occupancy query" as seen in the CPU lane for the linear projection layers. | Op | CPU op | GPU kernel | launches | |---|---|---|---| gate proj | aten::linear | ampere bf16 s16816gemm bf16 128x128 ... | occupancy query + cudaLaunchKernel | up proj | aten::linear | ampere bf16 s16816gemm bf16 128x128 ... | occupancy query + cudaLaunchKernel | gelu | aten::gelu | vectorized elementwise kernel<4, GeluCUDAKernelImpl... | cudaLaunchKernel | h u | aten::mul | vectorized elementwise kernel<4, ...MulFunctor... | cudaLaunchKernel | down proj | aten::linear | ampere bf16 s16816gemm bf16 128x256 ... | occupancy query + cudaLaunchKernel | The three GEMMs each do an extra cudaOccupancyMaxActiveBlocksPerMultiprocessor call before the launch. We have a separate section on this in Part 1, you can find it here https://huggingface.co/blog/torch-profiler why-does-matmul-have-an-extra-cuda-runtime-call . That is cuBLAS sizing the grid. The pointwise ops GeLU and mul launch directly, with no occupancy query. So "a linear" is actually query + launch, while "a pointwise op" is just launch. The aten::t , aten::transpose , aten::reshape , aten::view , aten::as strided , and aten:: unsafe view ops launch zero kernels. They show 0.000us of CUDA time in the table Figure 9 because they only rewrite tensor metadata shape and stride on the CPU. A reader scanning the table sees around six op names per linear, but only one of them mm ever reaches the GPU. Why are there two types of GEMM kernels? The MLP flattens batch, seq, dim to batch seq, dim for the matmul. In our command-line invocation we used 64 for batch and 128 for seq , so that's where the 8192 batch seq = 64 128 below comes from. From the trace: | Linear | aten::mm input dims | M·K·N | cuBLAS kernel | avg CUDA | |---|---|---|---|---| gate proj | 8192,768 x 768,3072 | 8192·768·3072 | …128x128…stages 32x5 tn | 0.19ms | up proj | 8192,768 x 768,3072 | 8192·768·3072 | …128x128…stages 32x5 tn | 0.19ms | down proj | 8192,3072 x 3072,768 | 8192·3072·768 | …128x256…stages 64x3 tn | 0.17ms | All three GEMMs have the same FLOP count, 2·8192·768·3072 ≈ 38.7 GFLOP each, yet down proj is about 10% faster. Same work, different shape N=768 instead of 3072 , so cuBLAS picks a different tile 128×256 , with a deeper stages 64x3 pipeline that gets better reuse for that shape. If you want to learn more about tiling in depth, here is a great resource to get started with. This is exactly why the table had two GEMM rows Figure 9 : the 128x128 row is gate+up and the 128x256 row is down. What does torch.compile do? Before compiling the forward method and visualizing it, let's do the mental exercise again of asking ourselves what we expect to see in the trace. This is a fun experiment, and an important one to repeat every time you profile something yourself. Always build on your intuition, and the moment something does not match, stop and figure out why. uv run 03 simple mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 --compile uvx trace-util traces -b traces In eager mode, each nn.Linear was expanded into a chain of dispatcher ops aten::linear → aten::t → aten::transpose → aten::matmul → aten::reshape → aten::mm . Those are the high-level wrappers that ATen walks through before reaching the real GEMM. torch.compile removes that chain. By the time the compiled graph runs, there is no linear, no matmul, no transpose or reshape and those metadata ops were folded into how mm is called. We can see three bare aten::mm external calls Figure 10 . The proof that it is the same GEMM is that the kernel names are byte-for-byte identical to eager: ...128x128...stages 32x5 tn for gate and up, and ...128x256...stages 64x3 tn for down. The fused Triton kernel This is the headline of the whole compile lesson. The two eager pointwise kernels GeLU and mul plus a reshape collapsed into one kernel, triton poi fused unsafe view gelu mul 0 Figure 11 . Let's decode the name: triton : generated by Inductor's Triton backend not cuBLAS, not ATen . poi : pointwise Inductor tags pointwise kernels poi , reductions red , and persistent reductions per . fused unsafe view gelu mul : the ops it merged: the unsafe view reshape , the GeLU, and the mul. 0 : the unique id within the graph. Why is this a win? In eager mode, the intermediate h = gelu g is a full 8192, 3072 bf16 tensor around 50 MB that the GeLU kernel writes to HBM and the mul kernel immediately reads back. Fusion keeps it in registers memory that resides inside the chip and are closer than the HBM . The Triton kernel reads g and u once, computes gelu g u , and writes the result once. One whole round trip of the intermediate through global memory is gone. Let's use hand tuned kernels So far we have let PyTorch eager and the compiler torch.compile pick our kernels. Now we plug in a kernel that a human expert wrote and tuned by hand. We use the LigerGEGLUMLP layer, that we can easily fetch from the Hugging Face Hub https://huggingface.co/kernels/kernels-community/liger-kernels with the kernels library. python from kernels import get kernel kernels layers = get kernel "kernels-community/liger-kernels", version=1 .layers kernels geglu mlp = kernels layers.LigerGEGLUMLP Config .to device, dtype=torch.bfloat16 .eval The full script is here: 03 kernels mlp.py https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03 kernels mlp.py . uv run 03 kernels mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 uvx trace-util traces -b traces Figure 12 shows the profile for the LigerGEGLUMLP layer using the Liger kernels from the Hub. Why use the kernels library Writing kernels in Triton or CUDA is one problem and shipping them is another. The kernel has to be compiled for your exact combination of GPU architecture, CUDA version, and PyTorch version. This is the step that usually breaks "works on my machine", missing nvcc , wrong Triton version . The kernels https://github.com/huggingface/kernels library moves that build step off your machine. get kernel "kernels-community/liger-kernels", version=1 downloads a pre-built, version-pinned kernel package from the Hugging Face Hub and caches it locally here under ~/.cache/...kernels-community--liger-kernels . The benefits are:- The kernels are compiled once, in CI, for many architectures and version combinations. You download the right binary instead of compiling it yourself. version=1 pins the exact build, so everyone running your script gets the same kernel. There is no "it got slower after I updated a package".- The package exposes a .layers attribute with drop-in nn.Module s like LigerGEGLUMLP . You swap your module for theirs and nothing else in your model changes. Why tuned kernels are better When we say "tuned", we mean two concrete things, and both are visible in the trace. The fusion is baked in. Theforward is LigerGEGLUMLP down proj LigerGELUMulFunction.apply gate proj x , up proj x . Theruns a single Triton kernel, LigerGELUMulFunction , that computes geglu tanh forward kernel gelu gate up in one pass. This is exactly what we saw from torch.compile , where the intermediate never makes a round-trip through HBM. We get it here without the compiler , as shown in Figures 13 and 14 no Dynamo guards, no compile latency, no recompilation risk . The launch parameters were chosen for the hardware. The kernel does not guess its block size at random. Liger'spicks them from the column count. calculate settings It is worth being honest about the trade-off here, because the raw numbers can be misleading. The Liger kernel runs in 92.8 µs , while Inductor's fused kernel from the compile run was 89.4 µs . At first glance the hand-written kernel looks slightly slower, but that comparison hides the cost that makes it worthwhile. torch.compile specializes for a static shape . Inductor's 89.4 µs kernel is fast precisely because it was generated for this exact 8192, 3072 problem. Change the batch size, the sequence length, or the hidden dimension, Dynamo re-traces, and you pay the compile cost all over again to get a new specialized kernel. So the real choice is not "slow human kernel vs fast compiled kernel". It is a fast generic kernel vs a kernel specialized for one particular input shape . The Liger kernel takes one set of launch parameters and runs them for any shape with no recompilation. It gives up the last few microseconds that per-shape specialization would buy, in exchange for being robust to changing shapes. Conclusion The table below collects what each step changed on the GPU and what it left untouched. | Setup | What changed | What stayed the same | |---|---|---| Eager nn.Linear | Baseline: bias add is already folded into the GEMM epilogue addmm , so it is one cuBLAS kernel, not a matmul plus an add | — | Compiled nn.Linear | A few CPU dispatch ops the aten::t view bookkeeping disappear | Same single cuBLAS GEMM kernel, byte-for-byte. Compile has nothing to fuse | | Eager MLP | 5 GPU kernels: 3 GEMMs + a GeLU + a mul. The 8192, 3072 intermediate makes a full round-trip through HBM | Each GEMM is still the same bias-free cuBLAS kernel as a standalone linear | | Compiled MLP | GeLU + mul + reshape collapse into one fused Triton kernel; the intermediate stays in registers. Pays compile pre-ops Dynamo, guards | The 3 GEMMs are untouched with identical cuBLAS kernel names | | Liger MLP | Same fusion, but baked into a hand-written Triton kernel with hardware-tuned launch params with no Dynamo, guards, or compile latency | The 3 GEMMs are still the same cuBLAS kernels | If there is one habit to carry forward, it is the one we practiced before every trace: guess first, then look. State what you expect the trace to contain, open it, and treat any mismatch as the most interesting thing on the screen. This was the second stop in the Profiling in PyTorch series. In the next post we will keep climbing the ladder, moving from this MLP block towards the attention block and, eventually, a full model. Thanks to Noe Flandre https://huggingface.co/NoeFlandre and Pedro Gabriel Gengo Lourenço https://huggingface.co/pedrogengo for their reviews on the early draft of the post