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

Source: https://huggingface.co/blog/torch-mlp-fusion Back to Articles

• From matmul-add to Linear- What is the transpose doing? - Why are there no separatemulandaddkernels? - Can --compile help a single Linear? - Where did the transpose go? Kernel layouts and pre-ops
• Stacking three Linears: the MLP- Why are there two types of GEMM kernels? - What doestorch\.compiledo? - The fused Triton kernel
• Let’s use hand tuned kernels- Why use the kernels library - Why tuned kernels are better
• Conclusion

In thefirst part of this series “Profiling in PyTorch”, we usedtorch\.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 oftorch\.compile.

In the second iteration (this blog post), we climb one rung up the ladder. We replace the hand-written matmul-add pair with annn\.Linear(withbias=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,03\_simple\_mlp\.py, and03\_kernels\_mlp\.py. Like before, it helps to open them in a separate tab and walk through the code as you read. We use anNVIDIA A100\-SXM4\-80GBGPU to run the scripts. It is really easy to set up a GPU on the Hugging Face infrastructure and experiment with the scripts usingDev Mode with Spaces. One could also run the scripts with theHugging Face Jobs pipeline.

Before we begin, a quick recap of two ideas we will lean on repeatedly:

1. A GPUkernelis a program that runs in parallel on many threads of the GPU.
2. The CPUschedules and launchesthese kernels. Most of the PyTorch overhead you see in a profiler trace is this scheduling work.

https://huggingface.co/blog/torch-mlp-fusion#from-matmul-add-to-linearFrom matmul-add to Linear

nn\.Linearis a module wrapper around the same matrix multiplication and addition we already profiled inPart 1. The only difference is that it owns its weight and bias as parameters and exposes aforwardmethod 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

Wherexis the input,wis the weight andbis the bias. Let’s run02\_linear\.pyand check the profile.

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64
uvx trace-util traces -b traces

trace\-utilis a utility that will sync your traces to aHugging Face bucketand then provide thePreffeto URLson your terminal.

Figure 1: Profiler trace ofnn\.Linear Figure 1 shows the profiler trace of a forward call of the linear layer. We trace theforwardcall of the linear layer with a similarschedulesetup as the previous traces, withwait=1,warmup=1andactive=3. This is why we see three Profile Steps in the CPU and GPU lanes.

https://huggingface.co/blog/torch-mlp-fusion#what-is-the-transpose-doingWhat is the transpose doing?

Figure 2: The transpose CPU row If we zoom into the profiler trace, as we do in Figure 2, we notice anaten::t(transpose) op before theaten::addmm(multiplication and addition) op. We can already figure out thatnn\.Lineartransposes the weight parameter and then multiplies it with the input. This is the reason we see anaten::top.

An important thing to notice is thataten::tdoes 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 theaten::trow in the profiler table and the time it took on CUDA.

https://huggingface.co/blog/torch-mlp-fusion#why-are-there-no-separate-mul-and-add-kernelsWhy are there no separatemulandaddkernels?

Figure 3: Noaten::addin the profile of a linear layer There is noaten::add(the bias addition) in the dispatch chain of the linear layer, as seen in Figure 3. This is because the bias addition has beenfoldedinto the matrix multiplication kernel, using what is called anepilogue.

Anepilogueis 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\.Linearcallstorch\.nn\.functional\.linear, which, in turn, callsaten::linear.aten::linearlooks at the inputs, notices that a bias was passed, and dispatchesaten::addmm\(bias, x, weight\)instead of doing a matmul and an add separately.addmmcomputes:

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 kernelaten::addmmpicks. The add never appears as a separate kernel because it ispart 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 inPart 1 under\-\-compile(addmm) is the kernel that eagernn\.Linearalready uses. There is nothing left fortorch\.compileto fuse here, which is the next thing we will verify.

https://huggingface.co/blog/torch-mlp-fusion#can—compile-help-a-single-linearCan --compile help a single Linear?

Let’s compile the forward call and look at the profiler trace. (The profiler trace is visualized in thenext section)

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 singlenn\.Linear’sforward, you will find:

• The same cuBLAS GEMM kernel on the GPU.
• The sameaten::addmmop on the CPU.
• A few extra rows on the CPU lane unique to compile.

This is worth internalizing. A common reflex is to reach fortorch\.compilewhenever 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 bylooking at an MLP.

https://huggingface.co/blog/torch-mlp-fusion#where-did-the-transpose-go-kernel-layouts-and-pre-opsWhere 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 whereaten::linearwalks throughaten::t(transpose) and thenaten::addmm Figure 5: Compiled dispatch chain whereaten::addmmis called directly, with no transpose The eager CPU dispatch chain insideaten::linearisaten::tfollowed byaten::addmm(Figure 4). To understand whataten::tactually does, we need a quick detour intostridesandviews.

A tensor stores its data as one flat, contiguous run of numbers in memory. Theshapeandstrideare metadata that sit on top of that run and tell PyTorch how to walk it: a stride of\(s0, s1\)means “steps0elements to move one row, steps1to move one column”. Change the metadata and you get a differentviewof thesameraw 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 buffer0, 1, 2, 3, 4, 5in transposed order. The underlying data is identical; only the metadata differs.

This is exactly whataten::tdoes inside the linear layer: it does not allocate a new tensor or copy any data, it produces aviewof the weight with rewritten strides.

As we can see in Figure 5, compile did not remove a GPU kernel: it removed theCPU overheadof dispatching that view. Inductor traced through the view chain at compile time, computed the resulting strides once, and emitted a directaten::addmmcall 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 thesamekernel 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
                                                          ^^

Thattnis the layout descriptor. cuBLAS and CUTLASS precompile aseparate kernel binaryfor each combination of input layouts.

n(non-transposed) andt(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\_,bf16vsfp16, ors16816gemmvss161616gemm) 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.

https://huggingface.co/blog/torch-mlp-fusion#stacking-three-linears-the-mlpStacking 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).

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. 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. Theforwardfunction does a fair amount of computation, but most of it is already familiar to us.

We should expect threeaten::lineardispatches, one for eachnn\.Linearlayer. 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 toconfirm or breaka guess, not to form one from scratch.

Figure 7: The profiler trace for a GeGLU MLP Figure 8: The occupancy queries highlighted in the linear projection CPU lane From Figure 7 we can pat ourselves on the back, as our intuition was correct. Per forward pass (onemlp\_fwd), the GPU runs exactly 5 kernels. Figure 8 highlights the “occupancy query” as seen in the CPU lane for the linear projection layers.

OpCPU opGPU kernellaunchesgate\_proj``aten::linear``ampere\_bf16\_s16816gemm\_bf16\_128x128\_\.\.\.occupancy query + cudaLaunchKernelup\_proj``aten::linear``ampere\_bf16\_s16816gemm\_bf16\_128x128\_\.\.\.occupancy query + cudaLaunchKernelgelu``aten::gelu``vectorized\_elementwise\_kernel<4, GeluCUDAKernelImpl\.\.\.\>cudaLaunchKernelh \* u``aten::mul``vectorized\_elementwise\_kernel<4, \.\.\.MulFunctor\.\.\.\>cudaLaunchKerneldown\_proj``aten::linear``ampere\_bf16\_s16816gemm\_bf16\_128x256\_\.\.\.occupancy query + cudaLaunchKernel The three GEMMs each do an extracudaOccupancyMaxActiveBlocksPerMultiprocessorcall before the launch. We have a separate section on this in Part 1,you can find it here. 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.

Figure 9: The table shows that some ops launch zero kernels Theaten::t,aten::transpose,aten::reshape,aten::view,aten::as\_strided, andaten::\_unsafe\_viewops launch zero kernels. They show0\.000usof 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.

https://huggingface.co/blog/torch-mlp-fusion#why-are-there-two-types-of-gemm-kernelsWhy 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 forbatchand 128 forseq, so that’s where the8192(batch \* seq = 64 \* 128) below comes from.

From the trace:

Linearaten::mminput dimsM·K·NcuBLAS kernelavg CUDAgate\_proj``\[8192,768\] x \[768,3072\]``8192·768·3072``…128x128…stages\_32x5\_tn0.19msup\_proj``\[8192,768\] x \[768,3072\]``8192·768·3072``…128x128…stages\_32x5\_tn0.19msdown\_proj``\[8192,3072\] x \[3072,768\]``8192·3072·768``…128x256…stages\_64x3\_tn0.17ms All three GEMMs have the same FLOP count,2·8192·768·3072 ≈ 38\.7 GFLOPeach, yetdown\_projis about10%faster. Same work, different shape (N=768instead of3072), so cuBLAS picks a different tile (128×256, with a deeperstages\_64x3pipeline) that gets better reuse for that shape.

If you want to learn more about tiling in depth,here is a great resourceto get started with.

This is exactly why the table had two GEMM rows (Figure 9): the128x128row is gate+up and the128x256row is down.

https://huggingface.co/blog/torch-mlp-fusion#what-does-torchcompile-doWhat doestorch\.compiledo?

Before compiling theforwardmethod 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

Figure 10: The profiler trace for the compiled GeGLU MLP In eager mode, eachnn\.Linearwas 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\.compileremoves 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 howmmis called. We can see three bareaten::mmexternal 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\_tnfor gate and up, and\.\.\.128x256\.\.\.stages\_64x3\_tnfor down.

https://huggingface.co/blog/torch-mlp-fusion#the-fused-triton-kernelThe fused Triton kernel

Figure 11: 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 kernelspoi, reductionsred, and persistent reductionsper).
• 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 intermediateh = 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 readsganduonce, computesgelu\(g\) \* u, and writes the result once. One whole round trip of the intermediate through global memory is gone.

https://huggingface.co/blog/torch-mlp-fusion#lets-use-hand-tuned-kernelsLet’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 theLigerGEGLUMLPlayer, that we can easily fetch from theHugging Face Hubwith thekernelslibrary.

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.

uv run 03_kernels_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072
uvx trace-util traces -b traces

Figure 12: The profiler trace for theLigerGEGLUMLPlayer Figure 12 shows the profile for theLigerGEGLUMLPlayer using the Liger kernels from the Hub.

https://huggingface.co/blog/torch-mlp-fusion#why-use-the-kernels-libraryWhy use the kernels library

Writing kernels in Triton or CUDA is one problem andshippingthem 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”, missingnvcc, wrong Triton version).

Thekernelslibrary moves that build step off your machine.get\_kernel\("kernels\-community/liger\-kernels", version=1\)downloads apre-built, version-pinnedkernel 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=1pins 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\.layersattribute with drop-innn\.Modules (likeLigerGEGLUMLP). You swap your module for theirs and nothing else in your model changes.

https://huggingface.co/blog/torch-mlp-fusion#why-tuned-kernels-are-betterWhy tuned kernels are better

When we say “tuned”, we mean two concrete things, and both are visible in the trace.

Figure 13: The compiled run pays for pre-ops (Dynamo, guards, prologue) before any GEMM runs Figure 14: The Liger kernel has no pre-ops — the box where they would be is empty

1. The fusion is baked in.TheLigerGEGLUMLPforward isdown\_proj\(LigerGELUMulFunction\.apply\(gate\_proj\(x\), up\_proj\(x\)\)\). TheLigerGELUMulFunctionruns a single Triton kernel,\_geglu\_tanh\_forward\_kernel, that computesgelu\(gate\) \* upin one pass. This is exactly what we saw fromtorch\.compile, where the intermediate never makes a round-trip through HBM. We get it herewithout the compiler, as shown in Figures 13 and 14 (no Dynamo guards, no compile latency, no recompilation risk).
2. **The launch parameters were chosen for the hardware.**The kernel does not guess its block size at random. Liger’scalculate\_settingspicks them from the column count.

It is worth being honest about the trade-off here, because the raw numbers can be misleading. The Liger kernel runs in92.8 µs, while Inductor’s fused kernel from the compile run was89.4 µs. At first glance the hand-written kernel looks slightly slower, but that comparison hides the cost that makes it worthwhile.

torch\.compilespecializes for astatic shape. Inductor’s89\.4 µskernel is fast precisely because it was generated forthis 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 isa fast generic kernel vs a kernel specialized for one particular input shape. The Liger kernel takes one set of launch parameters and runs them foranyshape with no recompilation. It gives up the last few microseconds that per-shape specialization would buy, in exchange for being robust to changing shapes.

https://huggingface.co/blog/torch-mlp-fusion#conclusionConclusion

The table below collects what each step changed on the GPU and what it left untouched.

SetupWhat changedWhat stayed the sameEagernn\.LinearBaseline: bias add is already folded into the GEMM epilogue (addmm), so it isonecuBLAS kernel, not a matmul plus an add—Compilednn\.LinearA few CPU dispatch ops (theaten::tview bookkeeping) disappearSame single cuBLAS GEMM kernel, byte-for-byte. Compile has nothing to fuseEager MLP5 GPU kernels: 3 GEMMs + a GeLU + a mul. The\[8192, 3072\]intermediate makes a full round-trip through HBMEach GEMM is still the same bias-free cuBLAS kernel as a standalone linearCompiled MLPGeLU + mul + reshape collapse intoonefused Triton kernel; the intermediate stays in registers. Pays compile pre-ops (Dynamo, guards)The 3 GEMMs are untouched with identical cuBLAS kernel namesLiger MLPSame fusion, but baked into a hand-written Triton kernel with hardware-tuned launch params withnoDynamo, guards, or compile latencyThe 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 theProfiling in PyTorchseries. 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 toNoe FlandreandPedro Gabriel Gengo Lourençofor their reviews on the early draft of the post!