{
  "@context": "https://schema.org",
  "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-06-11-profiling-in-pytorch-part-2-from-nn-linear-to-a-fused-mlp/#webpage",
  "@type": [
    "WebPage",
    "ArchiveComponent"
  ],
  "articleBody": "[![Thumbnail of the blog post](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/thumbnail.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/thumbnail.png)\n\nIn 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`.\n\nIn 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.\n\n> The scripts for this blog post live here: [`02_linear.py`](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/02_linear.py), [`03_simple_mlp.py`](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_simple_mlp.py), and [`03_kernels_mlp.py`](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_kernels_mlp.py). Like before, it helps to open them in a separate tab and walk through the code as you read. We use an `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](https://huggingface.co/docs/hub/spaces-dev-mode). One could also run the scripts with the [Hugging Face Jobs pipeline](https://huggingface.co/docs/huggingface_hub/en/guides/jobs).\n\nBefore we begin, a quick recap of two ideas we will lean on repeatedly:\n\n1. A GPU **kernel** is a program that runs in parallel on many threads of the GPU.\n2. The CPU **schedules and launches** these kernels. Most of the PyTorch overhead you see in a profiler trace is this scheduling work.\n\n## From matmul-add to Linear\n\n`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.\n\n```py\n# bias=True would truly emulate the multiplication and addition\n# operations we have seen in part 1 of the series\nlinear_layer = nn.Linear(in_dim, out_dim, bias=True)\ny = linear_layer(x)\n```\n\nThe operation at hand can be written as:\n\n```\ny = x @ w.T + b\n```\n\nWhere `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.\n\n```bash\nuv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\n> [`trace-util`](https://x.com/ariG23498/status/2054811716727517374) is a utility that will sync your traces to a [Hugging Face bucket](https://huggingface.co/storage) and then provide the [Preffeto URLs](https://perfetto.dev/) on your terminal.\n\n| [![PyTorch profiler trace of an `nn.Linear` forward pass: three short Profile Steps and `linear_fwd` annotations on the CPU lane, a tiny kernel on the GPU lane, and a long `cudaDeviceSynchronize` bar at the end](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/linear-profile-trace.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/linear-profile-trace.png) |\n| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 1: Profiler trace of `nn.Linear`                                                                                                                                                                                                                                                                                                                                                                                                                                           |\n\nFigure 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.\n\n### What is the transpose doing?\n\n| [![Zoomed in CPU dispatch chain showing the aten::t transpose op nested before aten::addmm inside aten::linear, with no matching activity on the GPU lane](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/transpose-cpu-dispatch.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/transpose-cpu-dispatch.png) |\n| ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 2: The transpose CPU row                                                                                                                                                                                                                                                                                                                                                                                              |\n\nIf 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.\n\nAn 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.\n\n### Why are there no separate `mul` and `add` kernels?\n\n| [![Profiler trace of the linear layer with the dispatch chain highlighted, showing aten::linear, aten::t and aten::addmm but no separate aten::add op](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-aten-add.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-aten-add.png) |\n| -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 3: No `aten::add` in the profile of a linear layer                                                                                                                                                                                                                                                                                                                                          |\n\nThere 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**.\n\nAn **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.\n\n`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:\n\n```\nout = x @ weight.T + bias\n```\n\nThe 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.\n\nThis 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.\n\n### Can --compile help a single Linear?\n\nLet's compile the forward call and look at the profiler trace. (The profiler trace is visualized in the [next section](https://huggingface.co/blog/torch-mlp-fusion#where-did-the-transpose-go-kernel-layouts-and-pre-ops))\n\n```bash\nuv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 --compile\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\nIf you compare the eager and compiled traces for a single `nn.Linear`'s `forward`, you will find:\n\n- The same cuBLAS GEMM kernel on the GPU.\n- The same `aten::addmm` op on the CPU.\n- A few extra rows on the CPU lane unique to compile.\n\nThis 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](https://huggingface.co/blog/torch-mlp-fusion#stacking-two-linears-the-mlp).\n\n### Where did the transpose go? Kernel layouts and pre-ops\n\nA 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.\n\n| [![Eager CPU dispatch chain with the aten::t transpose and aten::addmm boxed separately under aten::linear](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/eager.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/eager.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 4: Eager dispatch chain where `aten::linear` walks through `aten::t` (transpose) and then `aten::addmm`                                                                                                                                                                                                                              |\n\n| [![Compiled CPU dispatch chain showing a Torch-Compiled Region and a single aten::addmm call, with no transpose op](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 5: Compiled dispatch chain where `aten::addmm` is called directly, with no transpose                                                                                                                                                                                                                                                             |\n\nThe 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*.\n\nA 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:\n\n```py\n>>> M = torch.tensor([[0, 1],\n...                   [2, 3],\n...                   [4, 5]])\n>>> M.shape, M.stride()\n(torch.Size([3, 2]), (2, 1))   # two steps per row, one step per column\n\n>>> T = M.t()                  # transpose\n>>> T.shape, T.stride()\n(torch.Size([2, 3]), (1, 2))   # shape and stride swapped, data untouched\n>>> T\ntensor([[0, 2, 4],\n        [1, 3, 5]])\n>>> T.flatten()                # forced to materialize, so the data is reordered\ntensor([0, 2, 4, 1, 3, 5])\n```\n\n`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.\n\nThis 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.\n\nAs 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.\n\nAs one would expect, when the input data violates the strides precomputed by the compiler, it will throw an error.\n\nIf you look at the GPU lane in both traces, there is exactly one kernel per forward, and it is the *same* kernel both times:\n\n```\ncutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8\n```\n\nIf 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:\n\n```\ncutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8\n                                                          ^^\n```\n\nThat `tn` is the layout descriptor. cuBLAS and CUTLASS precompile a *separate kernel binary* for each combination of input layouts.\n\n`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.\n\n> 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.\n\n## Stacking three Linears: the MLP\n\nIn 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).\n\n| [![Conclusions section of the GLU Variants Improve Transformer paper, with the closing sentence attributing the architectures' success to divine benevolence highlighted](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/geglu-paper.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/geglu-paper.png) |\n| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 6: The conclusion section of the [GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202) paper.                                                                                                                                                                                                                                                                                                  |\n\n```py\nclass SimpleGeGLUMLP(nn.Module):\n    def __init__(self, dim, hidden):\n        super().__init__()\n        self.gate_proj = nn.Linear(dim, hidden, bias=False)\n        self.up_proj = nn.Linear(dim, hidden, bias=False)\n        self.down_proj = nn.Linear(hidden, dim, bias=False)\n\n    def forward(self, x):\n        g = self.gate_proj(x)\n        u = self.up_proj(x)\n        h = F.gelu(g, approximate=\"tanh\")\n        m = h * u\n        y = self.down_proj(m)\n        return y\n```\n\nYou 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:\n\n```bash\nuv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\nBefore 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.\n\nWe 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.\n\n| [![Profiler trace of the GeGLU MLP forward pass, with five boxed groups on the CPU lane labelled linear, linear, gelu, mul, linear](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-eager.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-eager.png) |\n| ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 7: The profiler trace for a GeGLU MLP                                                                                                                                                                                                                                                                                                                                              |\n\n| [![Occupancy Queries highlighted in the linear projection traces](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/occupancy-queries.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/occupancy-queries.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 8: The occupancy queries highlighted in the linear projection CPU lane                                                                                                                                                                                                                                             |\n\nFrom 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.\n\n| Op          | CPU op         | GPU kernel                                                | launches                           |\n| ----------- | -------------- | --------------------------------------------------------- | ---------------------------------- |\n| `gate_proj` | `aten::linear` | `ampere_bf16_s16816gemm_bf16_128x128_...`                 | occupancy query + cudaLaunchKernel |\n| `up_proj`   | `aten::linear` | `ampere_bf16_s16816gemm_bf16_128x128_...`                 | occupancy query + cudaLaunchKernel |\n| `gelu`      | `aten::gelu`   | `vectorized_elementwise_kernel<4, GeluCUDAKernelImpl...>` | cudaLaunchKernel                   |\n| `h * u`     | `aten::mul`    | `vectorized_elementwise_kernel<4, ...MulFunctor...>`      | cudaLaunchKernel                   |\n| `down_proj` | `aten::linear` | `ampere_bf16_s16816gemm_bf16_128x256_...`                 | occupancy query + cudaLaunchKernel |\n\nThe 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.\n\n| [![Profiler table for the GeGLU MLP listing op names and their CUDA times, where metadata ops like aten::transpose and aten::as_strided show 0.000us of CUDA time](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-table.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-table.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |\n| Figure 9: The table shows that some ops launch zero kernels                                                                                                                                                                                                                                                                                                                                                              |\n\nThe `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.\n\n### Why are there two types of GEMM kernels?\n\nThe 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.\n\nFrom the trace:\n\n| Linear      | `aten::mm` input dims      | M·K·N           | cuBLAS kernel             | avg CUDA |\n| ----------- | -------------------------- | --------------- | ------------------------- | -------- |\n| `gate_proj` | `[8192,768] x [768,3072]`  | `8192·768·3072` | `…128x128…stages_32x5_tn` | 0.19ms   |\n| `up_proj`   | `[8192,768] x [768,3072]`  | `8192·768·3072` | `…128x128…stages_32x5_tn` | 0.19ms   |\n| `down_proj` | `[8192,3072] x [3072,768]` | `8192·3072·768` | `…128x256…stages_64x3_tn` | 0.17ms   |\n\nAll 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.\n\n> If you want to learn more about tiling in depth, [here is a great resource](https://alvinwan.com/how-to-tile-matrix-multiplication/) to get started with.\n\nThis is exactly why the table had two GEMM rows (Figure 9): the `128x128` row is gate+up and the `128x256` row is down.\n\n### What does torch.compile do?\n\nBefore 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.\n\n```bash\nuv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 --compile\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\n| [![Profiler trace of the compiled GeGLU MLP showing three aten::mm calls and one fused triton kernel on the CPU lane, labelled mm, mm, fused, mm](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-compile-trace.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-compile-trace.png) |\n| ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 10: The profiler trace for the compiled GeGLU MLP                                                                                                                                                                                                                                                                                                                                                                |\n\nIn 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.\n\nBy 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.\n\n### The fused Triton kernel\n\n| [![Compiled MLP trace with the triton_poi_fused__unsafe_view_gelu_mul_0 kernel boxed on the CPU lane, replacing the separate gelu and mul kernels from the eager run](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/fused.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/fused.png) |\n| ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 11: The fused Triton kernel                                                                                                                                                                                                                                                                                                                                                                    |\n\nThis 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:\n\n- `triton`: generated by Inductor's Triton backend (not cuBLAS, not ATen).\n- `poi`: pointwise (Inductor tags pointwise kernels `poi`, reductions `red`, and persistent reductions `per`).\n- `fused__unsafe_view_gelu_mul`: the ops it merged: the `_unsafe_view` (reshape), the GeLU, and the mul.\n- `0`: the unique id within the graph.\n\nWhy 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.\n\n## Let's use hand tuned kernels\n\nSo 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.\n\n```python\nfrom kernels import get_kernel\n\nkernels_layers = get_kernel(\"kernels-community/liger-kernels\", version=1).layers\nkernels_geglu_mlp = kernels_layers.LigerGEGLUMLP(Config()).to(device, dtype=torch.bfloat16).eval()\n```\n\nThe full script is here: [`03_kernels_mlp.py`](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_kernels_mlp.py).\n\n```bash\nuv run 03_kernels_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\n| [![Profiler trace of the LigerGEGLUMLP forward pass showing three aten::linear groups and a single LigerGELUMulFunction group on the CPU lane](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/kernels-profile.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/kernels-profile.png) |\n| -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 12: The profiler trace for the `LigerGEGLUMLP` layer                                                                                                                                                                                                                                                                                                                                        |\n\nFigure 12 shows the profile for the `LigerGEGLUMLP` layer using the Liger kernels from the Hub.\n\n### Why use the kernels library\n\nWriting 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).\n\nThe [`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:\n\n- The kernels are compiled once, in CI, for many architectures and version combinations. You download the right binary instead of compiling it yourself.\n- `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\".\n- 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.\n\n### Why tuned kernels are better\n\nWhen we say \"tuned\", we mean two concrete things, and both are visible in the trace.\n\n| [![Compiled MLP trace with the TorchDynamo, prologue and guard pre-ops boxed on the CPU lane before the compiled graph runs](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile-preops.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile-preops.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |\n| Figure 13: The compiled run pays for pre-ops (Dynamo, guards, prologue) before any GEMM runs                                                                                                                                                                                                                                                                                   |\n\n| [![LigerGEGLUMLP trace with an empty box where the compile pre-ops would be, showing the hand-written kernel has no Dynamo or guard overhead](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-preops.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-preops.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 14: The Liger kernel has no pre-ops — the box where they would be is empty                                                                                                                                                                                                                                                                                                     |\n\n1. **The fusion is baked in.** The [`LigerGEGLUMLP`](https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/layers.py#L307) forward is `down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x)))`. The [`LigerGELUMulFunction`](https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L130) runs a single Triton kernel, [`_geglu_tanh_forward_kernel`](https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L97), that computes `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).\n\n2. **The launch parameters were chosen for the hardware.** The kernel does not guess its block size at random. Liger's [`calculate_settings`](https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L95) picks them from the column count.\n\nIt 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.\n\n`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.\n\nSo 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.\n\n## Conclusion\n\nThe table below collects what each step changed on the GPU and what it left untouched.\n\n| Setup                | What changed                                                                                                                               | What stayed the same                                                       |\n| -------------------- | ------------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------- |\n| 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              | —                                                                          |\n| 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 |\n| 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 |\n| 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               |\n| 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                              |\n\nIf 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.\n\nThis 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.\n\nThanks 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!\n\n> The blog post was polished using an LLM. This in no way means that we have let an agent run in the background and let it generate the blog. Some of us in the team are non-english speakers and think LLMs (which are mostly trained in the English Language) can rectify silly grammar mistakes or rephrase sentences that sound less intimidating and cleaner. Hope this helps with the idea of \"why should I read, if this was LLM generated\". 🤗\n",
  "authors": [],
  "captured_at": "2026-10-19T01:26:53.998049939+00:00",
  "category": "ai",
  "content_html": "<p><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/thumbnail.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/thumbnail.png\" alt=\"Thumbnail of the blog post\" loading=\"eager\" decoding=\"async\" fetchpriority=\"high\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></p>\n<p>In the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/torch-profiler\">first part of this series &quot;Profiling in PyTorch&quot;</a>, we used <code>torch.add(torch.matmul(x, w), b)</code> 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 <code>torch.compile</code>.</p>\n<p>In the second iteration (this blog post), we climb one rung up the ladder. We replace the hand-written matmul-add pair with an <code>nn.Linear</code> (with <code>bias=True</code>). 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.</p>\n<blockquote>\n<p>The scripts for this blog post live here: <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/02_linear.py\"><code>02_linear.py</code></a>, <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_simple_mlp.py\"><code>03_simple_mlp.py</code></a>, and <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_kernels_mlp.py\"><code>03_kernels_mlp.py</code></a>. Like before, it helps to open them in a separate tab and walk through the code as you read. We use an <code>NVIDIA A100-SXM4-80GB</code> 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 <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/docs/hub/spaces-dev-mode\">Dev Mode with Spaces</a>. One could also run the scripts with the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/docs/huggingface_hub/en/guides/jobs\">Hugging Face Jobs pipeline</a>.</p>\n</blockquote>\n<p>Before we begin, a quick recap of two ideas we will lean on repeatedly:</p>\n<ol>\n<li>A GPU <strong>kernel</strong> is a program that runs in parallel on many threads of the GPU.</li>\n<li>The CPU <strong>schedules and launches</strong> these kernels. Most of the PyTorch overhead you see in a profiler trace is this scheduling work.</li>\n</ol>\n<h2>From matmul-add to Linear</h2>\n<p><code>nn.Linear</code> is a module wrapper around the same matrix multiplication and addition we already profiled in <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/torch-profiler\">Part 1</a>. The only difference is that it owns its weight and bias as parameters and exposes a <code>forward</code> method that PyTorch users have grown familiar with.</p>\n<pre><code class=\"language-py\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> bias=True would truly emulate the multiplication and addition\n</span><span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> operations we have seen in part 1 of the series\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">linear_layer</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">nn</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">Linear</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">in_dim</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">out_dim</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">bias</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">y</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">linear_layer</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">x</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<p>The operation at hand can be written as:</p>\n<pre><code><span class=\"code-snippet\">y = x @ w.T + b\n</span></code></pre>\n<p>Where <code>x</code> is the input, <code>w</code> is the weight and <code>b</code> is the bias. Let's run <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/02_linear.py\"><code>02_linear.py</code></a> and check the profile.</p>\n<pre><code class=\"language-bash\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uv</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> run 02_linear.py<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>batch</span> 1024<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>in_dim</span> 32<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>out_dim</span> 64</span>\n<span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uvx</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> trace-util<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>f</span> traces<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>b</span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&lt;</span>hf_uname<span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&gt;</span>/traces</span>\n</span></span></code></pre>\n<blockquote>\n<p><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://x.com/ariG23498/status/2054811716727517374\"><code>trace-util</code></a> is a utility that will sync your traces to a <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/storage\">Hugging Face bucket</a> and then provide the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://perfetto.dev/\">Preffeto URLs</a> on your terminal.</p>\n</blockquote>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/linear-profile-trace.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/linear-profile-trace.png\" alt=\"PyTorch profiler trace of an nn.Linear forward pass: three short Profile Steps and linear_fwd annotations on the CPU lane, a tiny kernel on the GPU lane, and a long cudaDeviceSynchronize bar at the end\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 1: Profiler trace of <code>nn.Linear</code></td>\n</tr>\n</tbody>\n</table>\n<p>Figure 1 shows the profiler trace of a forward call of the linear layer. We trace the <code>forward</code> call of the linear layer with a similar <code>schedule</code> setup as the previous traces, with <code>wait=1</code>, <code>warmup=1</code> and <code>active=3</code>. This is why we see three Profile Steps in the CPU and GPU lanes.</p>\n<h3>What is the transpose doing?</h3>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/transpose-cpu-dispatch.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/transpose-cpu-dispatch.png\" alt=\"Zoomed in CPU dispatch chain showing the aten::t transpose op nested before aten::addmm inside aten::linear, with no matching activity on the GPU lane\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 2: The transpose CPU row</td>\n</tr>\n</tbody>\n</table>\n<p>If we zoom into the profiler trace, as we do in Figure 2, we notice an <code>aten::t</code> (transpose) op before the <code>aten::addmm</code> (multiplication and addition) op. We can already figure out that <code>nn.Linear</code> transposes the weight parameter and then multiplies it with the input. This is the reason we see an <code>aten::t</code> op.</p>\n<p>An important thing to notice is that <code>aten::t</code> 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 <code>aten::t</code> row in the profiler table and the time it took on CUDA.</p>\n<h3>Why are there no separate <code>mul</code> and <code>add</code> kernels?</h3>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-aten-add.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-aten-add.png\" alt=\"Profiler trace of the linear layer with the dispatch chain highlighted, showing aten::linear, aten::t and aten::addmm but no separate aten::add op\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 3: No <code>aten::add</code> in the profile of a linear layer</td>\n</tr>\n</tbody>\n</table>\n<p>There is no <code>aten::add</code> (the bias addition) in the dispatch chain of the linear layer, as seen in Figure 3. This is because the bias addition has been <em>folded</em> into the matrix multiplication kernel, using what is called an <strong>epilogue</strong>.</p>\n<p>An <strong>epilogue</strong> 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.</p>\n<p><code>nn.Linear</code> calls <code>torch.nn.functional.linear</code>, which, in turn, calls <code>aten::linear</code>. <code>aten::linear</code> looks at the inputs, notices that a bias was passed, and dispatches <code>aten::addmm(bias, x, weight)</code> instead of doing a matmul and an add separately. <code>addmm</code> computes:</p>\n<pre><code><span class=\"code-snippet\">out = x @ weight.T + bias\n</span></code></pre>\n<p>The cuBLAS GEMM kernel that runs on the GPU has a bias-add variant built in, and that's the kernel <code>aten::addmm</code> picks. The add never appears as a separate kernel because it is <strong>part of the matmul kernel's writeback</strong>, which is exactly what an epilogue is.</p>\n<p>This is the moment to notice something subtle. The kernel you saw in <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/torch-profiler#did-we-fuse-the-matmul-and-add-kernels-into-one\">Part 1 under <code>--compile</code></a> (<code>addmm</code>) is the kernel that eager <code>nn.Linear</code> already uses. There is nothing left for <code>torch.compile</code> to fuse here, which is the next thing we will verify.</p>\n<h3>Can --compile help a single Linear?</h3>\n<p>Let's compile the forward call and look at the profiler trace. (The profiler trace is visualized in the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/torch-mlp-fusion#where-did-the-transpose-go-kernel-layouts-and-pre-ops\">next section</a>)</p>\n<pre><code class=\"language-bash\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uv</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> run 02_linear.py<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>batch</span> 1024<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>in_dim</span> 32<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>out_dim</span> 64<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>compile</span></span>\n<span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uvx</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> trace-util<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>f</span> traces<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>b</span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&lt;</span>hf_uname<span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&gt;</span>/traces</span>\n</span></span></code></pre>\n<p>If you compare the eager and compiled traces for a single <code>nn.Linear</code>'s <code>forward</code>, you will find:</p>\n<ul>\n<li>The same cuBLAS GEMM kernel on the GPU.</li>\n<li>The same <code>aten::addmm</code> op on the CPU.</li>\n<li>A few extra rows on the CPU lane unique to compile.</li>\n</ul>\n<p>This is worth internalizing. A common reflex is to reach for <code>torch.compile</code> 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 <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/torch-mlp-fusion#stacking-two-linears-the-mlp\">looking at an MLP</a>.</p>\n<h3>Where did the transpose go? Kernel layouts and pre-ops</h3>\n<p>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.</p>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/eager.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/eager.png\" alt=\"Eager CPU dispatch chain with the aten::t transpose and aten::addmm boxed separately under aten::linear\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 4: Eager dispatch chain where <code>aten::linear</code> walks through <code>aten::t</code> (transpose) and then <code>aten::addmm</code></td>\n</tr>\n</tbody>\n</table>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile.png\" alt=\"Compiled CPU dispatch chain showing a Torch-Compiled Region and a single aten::addmm call, with no transpose op\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 5: Compiled dispatch chain where <code>aten::addmm</code> is called directly, with no transpose</td>\n</tr>\n</tbody>\n</table>\n<p>The eager CPU dispatch chain inside <code>aten::linear</code> is <code>aten::t</code> followed by <code>aten::addmm</code> (Figure 4). To understand what <code>aten::t</code> actually does, we need a quick detour into <em>strides</em> and <em>views</em>.</p>\n<p>A tensor stores its data as one flat, contiguous run of numbers in memory. The <code>shape</code> and <code>stride</code> are metadata that sit on top of that run and tell PyTorch how to walk it: a stride of <code>(s0, s1)</code> means &quot;step <code>s0</code> elements to move one row, step <code>s1</code> to move one column&quot;. Change the metadata and you get a different <em>view</em> of the <em>same</em> raw data, with no copy:</p>\n<pre><code class=\"language-py\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">M</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">torch</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">tensor</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">0</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span>\n<span class=\"syntax-constant syntax-language syntax-python\">...</span>                   <span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">3</span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span>\n<span class=\"syntax-constant syntax-language syntax-python\">...</span>                   <span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">4</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">5</span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">M</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">shape</span></span>, <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">M</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">stride</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-group syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-group syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">torch</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">Size</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">3</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-separator syntax-tuple syntax-python\">,</span> <span class=\"syntax-meta syntax-group syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-group syntax-begin syntax-python\">(</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-separator syntax-tuple syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-section syntax-group syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-section syntax-group syntax-end syntax-python\">)</span></span>   <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> two steps per row, one step per column\n</span>\n<span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">T</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">M</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">t</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>                  <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> transpose\n</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">T</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">shape</span></span>, <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">T</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">stride</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-group syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-group syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">torch</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">Size</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">3</span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-separator syntax-tuple syntax-python\">,</span> <span class=\"syntax-meta syntax-group syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-group syntax-begin syntax-python\">(</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-separator syntax-tuple syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-section syntax-group syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-section syntax-group syntax-end syntax-python\">)</span></span>   <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> shape and stride swapped, data untouched\n</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">T</span></span>\n<span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">tensor</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">0</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">4</span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span>\n        <span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">3</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">5</span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span><span class=\"syntax-keyword syntax-operator syntax-comparison syntax-python\">&gt;</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">T</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">flatten</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>                <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> forced to materialize, so the data is reordered\n</span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">tensor</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">0</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">2</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">4</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">3</span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">5</span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<p><code>M.t()</code> 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 <code>0, 1, 2, 3, 4, 5</code> in transposed order. The underlying data is identical; only the metadata differs.</p>\n<p>This is exactly what <code>aten::t</code> does inside the linear layer: it does not allocate a new tensor or copy any data, it produces a <em>view</em> of the weight with rewritten strides.</p>\n<p>As we can see in Figure 5, compile did not remove a GPU kernel: it removed the <em>CPU overhead</em> of dispatching that view. Inductor traced through the view chain at compile time, computed the resulting strides once, and emitted a direct <code>aten::addmm</code> call with those strides hard-coded. A few microseconds of CPU work disappear while the GPU does identical math.</p>\n<p>As one would expect, when the input data violates the strides precomputed by the compiler, it will throw an error.</p>\n<p>If you look at the GPU lane in both traces, there is exactly one kernel per forward, and it is the <em>same</em> kernel both times:</p>\n<pre><code><span class=\"code-snippet\">cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8\n</span></code></pre>\n<p>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:</p>\n<pre><code><span class=\"code-snippet\">cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8\n                                                          ^^\n</span></code></pre>\n<p>That <code>tn</code> is the layout descriptor. cuBLAS and CUTLASS precompile a <em>separate kernel binary</em> for each combination of input layouts.</p>\n<p><code>n</code> (non-transposed) and <code>t</code> (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.</p>\n<blockquote>\n<p>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., <code>_tn_</code> vs <code>_nn_</code>, <code>bf16</code> vs <code>fp16</code>, or <code>s16816gemm</code> vs <code>s161616gemm</code>) 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.</p>\n</blockquote>\n<h2>Stacking three Linears: the MLP</h2>\n<p>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).</p>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/geglu-paper.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/geglu-paper.png\" alt=\"Conclusions section of the GLU Variants Improve Transformer paper, with the closing sentence attributing the architectures&#39; success to divine benevolence highlighted\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 6: The conclusion section of the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://arxiv.org/abs/2002.05202\">GLU Variants Improve Transformer</a> paper.</td>\n</tr>\n</tbody>\n</table>\n<pre><code class=\"language-py\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-class syntax-python\"><span class=\"syntax-storage syntax-type syntax-class syntax-python\">class</span> <span class=\"syntax-entity syntax-name syntax-class syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">SimpleGeGLUMLP</span></span><span class=\"syntax-meta syntax-class syntax-inheritance syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-inheritance syntax-begin syntax-python\">(</span></span></span><span class=\"syntax-meta syntax-class syntax-inheritance syntax-python\"><span class=\"syntax-entity syntax-other syntax-inherited-class syntax-python\">nn<span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span>Module</span><span class=\"syntax-punctuation syntax-section syntax-inheritance syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-class syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-class syntax-begin syntax-python\">:</span></span>\n<span class=\"syntax-meta syntax-function syntax-python\">    <span class=\"syntax-storage syntax-type syntax-function syntax-python\">def</span> <span class=\"syntax-entity syntax-name syntax-function syntax-python\"><span class=\"syntax-support syntax-function syntax-magic syntax-python\">__init__</span></span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-parameters syntax-begin syntax-python\">(</span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-python\"><span class=\"syntax-variable syntax-parameter syntax-python\">self</span><span class=\"syntax-punctuation syntax-separator syntax-parameters syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">dim</span><span class=\"syntax-punctuation syntax-separator syntax-parameters syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">hidden</span><span class=\"syntax-punctuation syntax-section syntax-parameters syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-function syntax-begin syntax-python\">:</span></span>\n        <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-support syntax-function syntax-builtin syntax-python\">super</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-support syntax-function syntax-magic syntax-python\">__init__</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-language syntax-python\">self</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">gate_proj</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">nn</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">Linear</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">dim</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">hidden</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">bias</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">False</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-language syntax-python\">self</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">up_proj</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">nn</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">Linear</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">dim</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">hidden</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">bias</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">False</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-language syntax-python\">self</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">down_proj</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">nn</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">Linear</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">hidden</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">dim</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">bias</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">False</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n<span class=\"syntax-meta syntax-function syntax-python\">    <span class=\"syntax-storage syntax-type syntax-function syntax-python\">def</span> <span class=\"syntax-entity syntax-name syntax-function syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">forward</span></span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-parameters syntax-begin syntax-python\">(</span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-python\"><span class=\"syntax-variable syntax-parameter syntax-python\">self</span><span class=\"syntax-punctuation syntax-separator syntax-parameters syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">x</span><span class=\"syntax-punctuation syntax-section syntax-parameters syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-function syntax-begin syntax-python\">:</span></span>\n        <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">g</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-language syntax-python\">self</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">gate_proj</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">x</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">u</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-language syntax-python\">self</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">up_proj</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">x</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">h</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">F</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">gelu</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">g</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">approximate</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">tanh<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">m</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">h</span></span> <span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">*</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">u</span></span>\n        <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">y</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-language syntax-python\">self</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">down_proj</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">m</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-keyword syntax-control syntax-flow syntax-return syntax-python\">return</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">y</span></span>\n</span></span></code></pre>\n<p>You will find the entire script here: <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_simple_mlp.py\"><code>03_simple_mlp.py</code></a>. Execute it like so:</p>\n<pre><code class=\"language-bash\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uv</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> run 03_simple_mlp.py<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>batch</span> 64<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>seq</span> 128<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>dim</span> 768<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>hidden</span> 3072</span>\n<span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uvx</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> trace-util<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>f</span> traces<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>b</span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&lt;</span>hf_uname<span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&gt;</span>/traces</span>\n</span></span></code></pre>\n<p>Before we open the trace, let's think together about what we should expect to see. The <code>forward</code> function does a fair amount of computation, but most of it is already familiar to us.</p>\n<p>We should expect three <code>aten::linear</code> dispatches, one for each <code>nn.Linear</code> 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 <em>confirm or break</em> a guess, not to form one from scratch.</p>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-eager.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-eager.png\" alt=\"Profiler trace of the GeGLU MLP forward pass, with five boxed groups on the CPU lane labelled linear, linear, gelu, mul, linear\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 7: The profiler trace for a GeGLU MLP</td>\n</tr>\n</tbody>\n</table>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/occupancy-queries.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/occupancy-queries.png\" alt=\"Occupancy Queries highlighted in the linear projection traces\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 8: The occupancy queries highlighted in the linear projection CPU lane</td>\n</tr>\n</tbody>\n</table>\n<p>From Figure 7 we can pat ourselves on the back, as our intuition was correct. Per forward pass (one <code>mlp_fwd</code>), the GPU runs exactly 5 kernels. Figure 8 highlights the &quot;occupancy query&quot; as seen in the CPU lane for the linear projection layers.</p>\n<table>\n<thead>\n<tr>\n<th>Op</th>\n<th>CPU op</th>\n<th>GPU kernel</th>\n<th>launches</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><code>gate_proj</code></td>\n<td><code>aten::linear</code></td>\n<td><code>ampere_bf16_s16816gemm_bf16_128x128_...</code></td>\n<td>occupancy query + cudaLaunchKernel</td>\n</tr>\n<tr>\n<td><code>up_proj</code></td>\n<td><code>aten::linear</code></td>\n<td><code>ampere_bf16_s16816gemm_bf16_128x128_...</code></td>\n<td>occupancy query + cudaLaunchKernel</td>\n</tr>\n<tr>\n<td><code>gelu</code></td>\n<td><code>aten::gelu</code></td>\n<td><code>vectorized_elementwise_kernel&lt;4, GeluCUDAKernelImpl...&gt;</code></td>\n<td>cudaLaunchKernel</td>\n</tr>\n<tr>\n<td><code>h * u</code></td>\n<td><code>aten::mul</code></td>\n<td><code>vectorized_elementwise_kernel&lt;4, ...MulFunctor...&gt;</code></td>\n<td>cudaLaunchKernel</td>\n</tr>\n<tr>\n<td><code>down_proj</code></td>\n<td><code>aten::linear</code></td>\n<td><code>ampere_bf16_s16816gemm_bf16_128x256_...</code></td>\n<td>occupancy query + cudaLaunchKernel</td>\n</tr>\n</tbody>\n</table>\n<p>The three GEMMs each do an extra <code>cudaOccupancyMaxActiveBlocksPerMultiprocessor</code> call before the launch. We have a separate section on this in Part 1, <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/blog/torch-profiler#why-does-matmul-have-an-extra-cuda-runtime-call\">you can find it here</a>. That is cuBLAS sizing the grid. The pointwise ops (GeLU and mul) launch directly, with no occupancy query. So &quot;a linear&quot; is actually query + launch, while &quot;a pointwise op&quot; is just launch.</p>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-table.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-table.png\" alt=\"Profiler table for the GeGLU MLP listing op names and their CUDA times, where metadata ops like aten::transpose and aten::as_strided show 0.000us of CUDA time\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 9: The table shows that some ops launch zero kernels</td>\n</tr>\n</tbody>\n</table>\n<p>The <code>aten::t</code>, <code>aten::transpose</code>, <code>aten::reshape</code>, <code>aten::view</code>, <code>aten::as_strided</code>, and <code>aten::_unsafe_view</code> ops launch zero kernels. They show <code>0.000us</code> 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 (<code>mm</code>) ever reaches the GPU.</p>\n<h3>Why are there two types of GEMM kernels?</h3>\n<p>The MLP flattens <code>[batch, seq, dim]</code> to <code>[batch * seq, dim]</code> for the matmul. In our command-line invocation we used 64 for <code>batch</code> and 128 for <code>seq</code>, so that's where the <code>8192</code> (<code>batch * seq = 64 * 128</code>) below comes from.</p>\n<p>From the trace:</p>\n<table>\n<thead>\n<tr>\n<th>Linear</th>\n<th><code>aten::mm</code> input dims</th>\n<th>M·K·N</th>\n<th>cuBLAS kernel</th>\n<th>avg CUDA</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><code>gate_proj</code></td>\n<td><code>[8192,768] x [768,3072]</code></td>\n<td><code>8192·768·3072</code></td>\n<td><code>…128x128…stages_32x5_tn</code></td>\n<td>0.19ms</td>\n</tr>\n<tr>\n<td><code>up_proj</code></td>\n<td><code>[8192,768] x [768,3072]</code></td>\n<td><code>8192·768·3072</code></td>\n<td><code>…128x128…stages_32x5_tn</code></td>\n<td>0.19ms</td>\n</tr>\n<tr>\n<td><code>down_proj</code></td>\n<td><code>[8192,3072] x [3072,768]</code></td>\n<td><code>8192·3072·768</code></td>\n<td><code>…128x256…stages_64x3_tn</code></td>\n<td>0.17ms</td>\n</tr>\n</tbody>\n</table>\n<p>All three GEMMs have the same FLOP count, <code>2·8192·768·3072 ≈ 38.7 GFLOP</code> each, yet <code>down_proj</code> is about <code>10%</code> faster. Same work, different shape (<code>N=768</code> instead of <code>3072</code>), so cuBLAS picks a different tile (<code>128×256</code>, with a deeper <code>stages_64x3</code> pipeline) that gets better reuse for that shape.</p>\n<blockquote>\n<p>If you want to learn more about tiling in depth, <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://alvinwan.com/how-to-tile-matrix-multiplication/\">here is a great resource</a> to get started with.</p>\n</blockquote>\n<p>This is exactly why the table had two GEMM rows (Figure 9): the <code>128x128</code> row is gate+up and the <code>128x256</code> row is down.</p>\n<h3>What does torch.compile do?</h3>\n<p>Before compiling the <code>forward</code> 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.</p>\n<pre><code class=\"language-bash\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uv</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> run 03_simple_mlp.py<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>batch</span> 64<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>seq</span> 128<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>dim</span> 768<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>hidden</span> 3072<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>compile</span></span>\n<span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uvx</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> trace-util<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>f</span> traces<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>b</span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&lt;</span>hf_uname<span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&gt;</span>/traces</span>\n</span></span></code></pre>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-compile-trace.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-compile-trace.png\" alt=\"Profiler trace of the compiled GeGLU MLP showing three aten::mm calls and one fused triton kernel on the CPU lane, labelled mm, mm, fused, mm\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 10: The profiler trace for the compiled GeGLU MLP</td>\n</tr>\n</tbody>\n</table>\n<p>In eager mode, each <code>nn.Linear</code> was expanded into a chain of dispatcher ops (<code>aten::linear</code> → <code>aten::t</code> → <code>aten::transpose</code> → <code>aten::matmul</code> → <code>aten::reshape</code> → <code>aten::mm</code>). Those are the high-level wrappers that ATen walks through before reaching the real GEMM. <code>torch.compile</code> removes that chain.</p>\n<p>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 <code>mm</code> is called. We can see three bare <code>aten::mm</code> external calls (Figure 10). The proof that it is the same GEMM is that the kernel names are byte-for-byte identical to eager: <code>...128x128...stages_32x5_tn</code> for gate and up, and <code>...128x256...stages_64x3_tn</code> for down.</p>\n<h3>The fused Triton kernel</h3>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/fused.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/fused.png\" alt=\"Compiled MLP trace with the triton_poi_fused__unsafe_view_gelu_mul_0 kernel boxed on the CPU lane, replacing the separate gelu and mul kernels from the eager run\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 11: The fused Triton kernel</td>\n</tr>\n</tbody>\n</table>\n<p>This is the headline of the whole compile lesson. The two eager pointwise kernels (GeLU and mul) plus a reshape collapsed into one kernel, <code>triton_poi_fused__unsafe_view_gelu_mul_0</code> (Figure 11). Let's decode the name:</p>\n<ul>\n<li><code>triton</code>: generated by Inductor's Triton backend (not cuBLAS, not ATen).</li>\n<li><code>poi</code>: pointwise (Inductor tags pointwise kernels <code>poi</code>, reductions <code>red</code>, and persistent reductions <code>per</code>).</li>\n<li><code>fused__unsafe_view_gelu_mul</code>: the ops it merged: the <code>_unsafe_view</code> (reshape), the GeLU, and the mul.</li>\n<li><code>0</code>: the unique id within the graph.</li>\n</ul>\n<p>Why is this a win? In eager mode, the intermediate <code>h = gelu(g)</code> is a full <code>[8192, 3072]</code> 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 <code>g</code> and <code>u</code> once, computes <code>gelu(g) * u</code>, and writes the result once. One whole round trip of the intermediate through global memory is gone.</p>\n<h2>Let's use hand tuned kernels</h2>\n<p>So far we have let PyTorch (eager) and the compiler (<code>torch.compile</code>) pick our kernels. Now we plug in a kernel that a human expert wrote and tuned by hand. We use the <code>LigerGEGLUMLP</code> layer, that we can easily fetch from the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/kernels/kernels-community/liger-kernels\">Hugging Face Hub</a> with the <code>kernels</code> library.</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">kernels</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">get_kernel</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">kernels_layers</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">get_kernel</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">kernels-community/liger-kernels<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">version</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">layers</span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">kernels_geglu_mlp</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">kernels_layers</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">LigerGEGLUMLP</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">Config</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">to</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">device</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">dtype</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">torch</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">bfloat16</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">eval</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<p>The full script is here: <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_kernels_mlp.py\"><code>03_kernels_mlp.py</code></a>.</p>\n<pre><code class=\"language-bash\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uv</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> run 03_kernels_mlp.py<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>batch</span> 64<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>seq</span> 128<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>dim</span> 768<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>hidden</span> 3072</span>\n<span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uvx</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> trace-util<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>f</span> traces<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>b</span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&lt;</span>hf_uname<span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&gt;</span>/traces</span>\n</span></span></code></pre>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/kernels-profile.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/kernels-profile.png\" alt=\"Profiler trace of the LigerGEGLUMLP forward pass showing three aten::linear groups and a single LigerGELUMulFunction group on the CPU lane\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 12: The profiler trace for the <code>LigerGEGLUMLP</code> layer</td>\n</tr>\n</tbody>\n</table>\n<p>Figure 12 shows the profile for the <code>LigerGEGLUMLP</code> layer using the Liger kernels from the Hub.</p>\n<h3>Why use the kernels library</h3>\n<p>Writing kernels in Triton or CUDA is one problem and <em>shipping</em> 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 (&quot;works on my machine&quot;, missing <code>nvcc</code>, wrong Triton version).</p>\n<p>The <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/huggingface/kernels\"><code>kernels</code></a> library moves that build step off your machine. <code>get_kernel(&quot;kernels-community/liger-kernels&quot;, version=1)</code> downloads a <strong>pre-built, version-pinned</strong> kernel package from the Hugging Face Hub and caches it locally (here under <code>~/.cache/...kernels-community--liger-kernels</code>). The benefits are:</p>\n<ul>\n<li>The kernels are compiled once, in CI, for many architectures and version combinations. You download the right binary instead of compiling it yourself.</li>\n<li><code>version=1</code> pins the exact build, so everyone running your script gets the same kernel. There is no &quot;it got slower after I updated a package&quot;.</li>\n<li>The package exposes a <code>.layers</code> attribute with drop-in <code>nn.Module</code>s (like <code>LigerGEGLUMLP</code>). You swap your module for theirs and nothing else in your model changes.</li>\n</ul>\n<h3>Why tuned kernels are better</h3>\n<p>When we say &quot;tuned&quot;, we mean two concrete things, and both are visible in the trace.</p>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile-preops.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile-preops.png\" alt=\"Compiled MLP trace with the TorchDynamo, prologue and guard pre-ops boxed on the CPU lane before the compiled graph runs\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 13: The compiled run pays for pre-ops (Dynamo, guards, prologue) before any GEMM runs</td>\n</tr>\n</tbody>\n</table>\n<table>\n<thead>\n<tr>\n<th><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-preops.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-preops.png\" alt=\"LigerGEGLUMLP trace with an empty box where the compile pre-ops would be, showing the hand-written kernel has no Dynamo or guard overhead\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Figure 14: The Liger kernel has no pre-ops — the box where they would be is empty</td>\n</tr>\n</tbody>\n</table>\n<ol>\n<li>\n<p><strong>The fusion is baked in.</strong> The <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/layers.py#L307\"><code>LigerGEGLUMLP</code></a> forward is <code>down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x)))</code>. The <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L130\"><code>LigerGELUMulFunction</code></a> runs a single Triton kernel, <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L97\"><code>_geglu_tanh_forward_kernel</code></a>, that computes <code>gelu(gate) * up</code> in one pass. This is exactly what we saw from <code>torch.compile</code>, where the intermediate never makes a round-trip through HBM. We get it here <strong>without the compiler</strong>, as shown in Figures 13 and 14 (no Dynamo guards, no compile latency, no recompilation risk).</p>\n</li>\n<li>\n<p><strong>The launch parameters were chosen for the hardware.</strong> The kernel does not guess its block size at random. Liger's <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L95\"><code>calculate_settings</code></a> picks them from the column count.</p>\n</li>\n</ol>\n<p>It is worth being honest about the trade-off here, because the raw numbers can be misleading. The Liger kernel runs in <strong>92.8 µs</strong>, while Inductor's fused kernel from the compile run was <strong>89.4 µs</strong>. At first glance the hand-written kernel looks slightly slower, but that comparison hides the cost that makes it worthwhile.</p>\n<p><code>torch.compile</code> specializes for a <strong>static shape</strong>. Inductor's <code>89.4 µs</code> kernel is fast precisely because it was generated for <em>this exact</em> <code>[8192, 3072]</code> 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.</p>\n<p>So the real choice is not &quot;slow human kernel vs fast compiled kernel&quot;. It is <strong>a fast generic kernel vs a kernel specialized for one particular input shape</strong>. The Liger kernel takes one set of launch parameters and runs them for <em>any</em> 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.</p>\n<h2>Conclusion</h2>\n<p>The table below collects what each step changed on the GPU and what it left untouched.</p>\n<table>\n<thead>\n<tr>\n<th>Setup</th>\n<th>What changed</th>\n<th>What stayed the same</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Eager <code>nn.Linear</code></td>\n<td>Baseline: bias add is already folded into the GEMM epilogue (<code>addmm</code>), so it is <em>one</em> cuBLAS kernel, not a matmul plus an add</td>\n<td>—</td>\n</tr>\n<tr>\n<td>Compiled <code>nn.Linear</code></td>\n<td>A few CPU dispatch ops (the <code>aten::t</code> view bookkeeping) disappear</td>\n<td>Same single cuBLAS GEMM kernel, byte-for-byte. Compile has nothing to fuse</td>\n</tr>\n<tr>\n<td>Eager MLP</td>\n<td>5 GPU kernels: 3 GEMMs + a GeLU + a mul. The <code>[8192, 3072]</code> intermediate makes a full round-trip through HBM</td>\n<td>Each GEMM is still the same bias-free cuBLAS kernel as a standalone linear</td>\n</tr>\n<tr>\n<td>Compiled MLP</td>\n<td>GeLU + mul + reshape collapse into <strong>one</strong> fused Triton kernel; the intermediate stays in registers. Pays compile pre-ops (Dynamo, guards)</td>\n<td>The 3 GEMMs are untouched with identical cuBLAS kernel names</td>\n</tr>\n<tr>\n<td>Liger MLP</td>\n<td>Same fusion, but baked into a hand-written Triton kernel with hardware-tuned launch params with <strong>no</strong> Dynamo, guards, or compile latency</td>\n<td>The 3 GEMMs are still the same cuBLAS kernels</td>\n</tr>\n</tbody>\n</table>\n<p>If there is one habit to carry forward, it is the one we practiced before every trace: <strong>guess first, then look.</strong> State what you expect the trace to contain, open it, and treat any mismatch as the most interesting thing on the screen.</p>\n<p>This was the second stop in the <strong>Profiling in PyTorch</strong> 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.</p>\n<p>Thanks to <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/NoeFlandre\">Noe Flandre</a> and <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/pedrogengo\">Pedro Gabriel Gengo Lourenço</a> for their reviews on the early draft of the post!</p>\n<blockquote>\n<p>The blog post was polished using an LLM. This in no way means that we have let an agent run in the background and let it generate the blog. Some of us in the team are non-english speakers and think LLMs (which are mostly trained in the English Language) can rectify silly grammar mistakes or rephrase sentences that sound less intimidating and cleaner. Hope this helps with the idea of &quot;why should I read, if this was LLM generated&quot;. 🤗</p>\n</blockquote>\n",
  "content_markdown": "[![Thumbnail of the blog post](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/thumbnail.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/thumbnail.png)\n\nIn 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`.\n\nIn 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.\n\n> The scripts for this blog post live here: [`02_linear.py`](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/02_linear.py), [`03_simple_mlp.py`](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_simple_mlp.py), and [`03_kernels_mlp.py`](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_kernels_mlp.py). Like before, it helps to open them in a separate tab and walk through the code as you read. We use an `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](https://huggingface.co/docs/hub/spaces-dev-mode). One could also run the scripts with the [Hugging Face Jobs pipeline](https://huggingface.co/docs/huggingface_hub/en/guides/jobs).\n\nBefore we begin, a quick recap of two ideas we will lean on repeatedly:\n\n1. A GPU **kernel** is a program that runs in parallel on many threads of the GPU.\n2. The CPU **schedules and launches** these kernels. Most of the PyTorch overhead you see in a profiler trace is this scheduling work.\n\n## From matmul-add to Linear\n\n`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.\n\n```py\n# bias=True would truly emulate the multiplication and addition\n# operations we have seen in part 1 of the series\nlinear_layer = nn.Linear(in_dim, out_dim, bias=True)\ny = linear_layer(x)\n```\n\nThe operation at hand can be written as:\n\n```\ny = x @ w.T + b\n```\n\nWhere `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.\n\n```bash\nuv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\n> [`trace-util`](https://x.com/ariG23498/status/2054811716727517374) is a utility that will sync your traces to a [Hugging Face bucket](https://huggingface.co/storage) and then provide the [Preffeto URLs](https://perfetto.dev/) on your terminal.\n\n| [![PyTorch profiler trace of an `nn.Linear` forward pass: three short Profile Steps and `linear_fwd` annotations on the CPU lane, a tiny kernel on the GPU lane, and a long `cudaDeviceSynchronize` bar at the end](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/linear-profile-trace.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/linear-profile-trace.png) |\n| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 1: Profiler trace of `nn.Linear`                                                                                                                                                                                                                                                                                                                                                                                                                                           |\n\nFigure 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.\n\n### What is the transpose doing?\n\n| [![Zoomed in CPU dispatch chain showing the aten::t transpose op nested before aten::addmm inside aten::linear, with no matching activity on the GPU lane](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/transpose-cpu-dispatch.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/transpose-cpu-dispatch.png) |\n| ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 2: The transpose CPU row                                                                                                                                                                                                                                                                                                                                                                                              |\n\nIf 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.\n\nAn 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.\n\n### Why are there no separate `mul` and `add` kernels?\n\n| [![Profiler trace of the linear layer with the dispatch chain highlighted, showing aten::linear, aten::t and aten::addmm but no separate aten::add op](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-aten-add.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-aten-add.png) |\n| -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 3: No `aten::add` in the profile of a linear layer                                                                                                                                                                                                                                                                                                                                          |\n\nThere 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**.\n\nAn **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.\n\n`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:\n\n```\nout = x @ weight.T + bias\n```\n\nThe 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.\n\nThis 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.\n\n### Can --compile help a single Linear?\n\nLet's compile the forward call and look at the profiler trace. (The profiler trace is visualized in the [next section](https://huggingface.co/blog/torch-mlp-fusion#where-did-the-transpose-go-kernel-layouts-and-pre-ops))\n\n```bash\nuv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 --compile\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\nIf you compare the eager and compiled traces for a single `nn.Linear`'s `forward`, you will find:\n\n- The same cuBLAS GEMM kernel on the GPU.\n- The same `aten::addmm` op on the CPU.\n- A few extra rows on the CPU lane unique to compile.\n\nThis 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](https://huggingface.co/blog/torch-mlp-fusion#stacking-two-linears-the-mlp).\n\n### Where did the transpose go? Kernel layouts and pre-ops\n\nA 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.\n\n| [![Eager CPU dispatch chain with the aten::t transpose and aten::addmm boxed separately under aten::linear](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/eager.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/eager.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 4: Eager dispatch chain where `aten::linear` walks through `aten::t` (transpose) and then `aten::addmm`                                                                                                                                                                                                                              |\n\n| [![Compiled CPU dispatch chain showing a Torch-Compiled Region and a single aten::addmm call, with no transpose op](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 5: Compiled dispatch chain where `aten::addmm` is called directly, with no transpose                                                                                                                                                                                                                                                             |\n\nThe 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*.\n\nA 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:\n\n```py\n>>> M = torch.tensor([[0, 1],\n...                   [2, 3],\n...                   [4, 5]])\n>>> M.shape, M.stride()\n(torch.Size([3, 2]), (2, 1))   # two steps per row, one step per column\n\n>>> T = M.t()                  # transpose\n>>> T.shape, T.stride()\n(torch.Size([2, 3]), (1, 2))   # shape and stride swapped, data untouched\n>>> T\ntensor([[0, 2, 4],\n        [1, 3, 5]])\n>>> T.flatten()                # forced to materialize, so the data is reordered\ntensor([0, 2, 4, 1, 3, 5])\n```\n\n`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.\n\nThis 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.\n\nAs 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.\n\nAs one would expect, when the input data violates the strides precomputed by the compiler, it will throw an error.\n\nIf you look at the GPU lane in both traces, there is exactly one kernel per forward, and it is the *same* kernel both times:\n\n```\ncutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8\n```\n\nIf 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:\n\n```\ncutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8\n                                                          ^^\n```\n\nThat `tn` is the layout descriptor. cuBLAS and CUTLASS precompile a *separate kernel binary* for each combination of input layouts.\n\n`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.\n\n> 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.\n\n## Stacking three Linears: the MLP\n\nIn 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).\n\n| [![Conclusions section of the GLU Variants Improve Transformer paper, with the closing sentence attributing the architectures' success to divine benevolence highlighted](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/geglu-paper.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/geglu-paper.png) |\n| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 6: The conclusion section of the [GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202) paper.                                                                                                                                                                                                                                                                                                  |\n\n```py\nclass SimpleGeGLUMLP(nn.Module):\n    def __init__(self, dim, hidden):\n        super().__init__()\n        self.gate_proj = nn.Linear(dim, hidden, bias=False)\n        self.up_proj = nn.Linear(dim, hidden, bias=False)\n        self.down_proj = nn.Linear(hidden, dim, bias=False)\n\n    def forward(self, x):\n        g = self.gate_proj(x)\n        u = self.up_proj(x)\n        h = F.gelu(g, approximate=\"tanh\")\n        m = h * u\n        y = self.down_proj(m)\n        return y\n```\n\nYou 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:\n\n```bash\nuv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\nBefore 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.\n\nWe 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.\n\n| [![Profiler trace of the GeGLU MLP forward pass, with five boxed groups on the CPU lane labelled linear, linear, gelu, mul, linear](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-eager.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-eager.png) |\n| ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 7: The profiler trace for a GeGLU MLP                                                                                                                                                                                                                                                                                                                                              |\n\n| [![Occupancy Queries highlighted in the linear projection traces](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/occupancy-queries.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/occupancy-queries.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 8: The occupancy queries highlighted in the linear projection CPU lane                                                                                                                                                                                                                                             |\n\nFrom 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.\n\n| Op          | CPU op         | GPU kernel                                                | launches                           |\n| ----------- | -------------- | --------------------------------------------------------- | ---------------------------------- |\n| `gate_proj` | `aten::linear` | `ampere_bf16_s16816gemm_bf16_128x128_...`                 | occupancy query + cudaLaunchKernel |\n| `up_proj`   | `aten::linear` | `ampere_bf16_s16816gemm_bf16_128x128_...`                 | occupancy query + cudaLaunchKernel |\n| `gelu`      | `aten::gelu`   | `vectorized_elementwise_kernel<4, GeluCUDAKernelImpl...>` | cudaLaunchKernel                   |\n| `h * u`     | `aten::mul`    | `vectorized_elementwise_kernel<4, ...MulFunctor...>`      | cudaLaunchKernel                   |\n| `down_proj` | `aten::linear` | `ampere_bf16_s16816gemm_bf16_128x256_...`                 | occupancy query + cudaLaunchKernel |\n\nThe 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.\n\n| [![Profiler table for the GeGLU MLP listing op names and their CUDA times, where metadata ops like aten::transpose and aten::as_strided show 0.000us of CUDA time](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-table.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-table.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |\n| Figure 9: The table shows that some ops launch zero kernels                                                                                                                                                                                                                                                                                                                                                              |\n\nThe `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.\n\n### Why are there two types of GEMM kernels?\n\nThe 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.\n\nFrom the trace:\n\n| Linear      | `aten::mm` input dims      | M·K·N           | cuBLAS kernel             | avg CUDA |\n| ----------- | -------------------------- | --------------- | ------------------------- | -------- |\n| `gate_proj` | `[8192,768] x [768,3072]`  | `8192·768·3072` | `…128x128…stages_32x5_tn` | 0.19ms   |\n| `up_proj`   | `[8192,768] x [768,3072]`  | `8192·768·3072` | `…128x128…stages_32x5_tn` | 0.19ms   |\n| `down_proj` | `[8192,3072] x [3072,768]` | `8192·3072·768` | `…128x256…stages_64x3_tn` | 0.17ms   |\n\nAll 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.\n\n> If you want to learn more about tiling in depth, [here is a great resource](https://alvinwan.com/how-to-tile-matrix-multiplication/) to get started with.\n\nThis is exactly why the table had two GEMM rows (Figure 9): the `128x128` row is gate+up and the `128x256` row is down.\n\n### What does torch.compile do?\n\nBefore 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.\n\n```bash\nuv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 --compile\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\n| [![Profiler trace of the compiled GeGLU MLP showing three aten::mm calls and one fused triton kernel on the CPU lane, labelled mm, mm, fused, mm](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-compile-trace.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/simple-mlp-compile-trace.png) |\n| ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 10: The profiler trace for the compiled GeGLU MLP                                                                                                                                                                                                                                                                                                                                                                |\n\nIn 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.\n\nBy 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.\n\n### The fused Triton kernel\n\n| [![Compiled MLP trace with the triton_poi_fused__unsafe_view_gelu_mul_0 kernel boxed on the CPU lane, replacing the separate gelu and mul kernels from the eager run](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/fused.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/fused.png) |\n| ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 11: The fused Triton kernel                                                                                                                                                                                                                                                                                                                                                                    |\n\nThis 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:\n\n- `triton`: generated by Inductor's Triton backend (not cuBLAS, not ATen).\n- `poi`: pointwise (Inductor tags pointwise kernels `poi`, reductions `red`, and persistent reductions `per`).\n- `fused__unsafe_view_gelu_mul`: the ops it merged: the `_unsafe_view` (reshape), the GeLU, and the mul.\n- `0`: the unique id within the graph.\n\nWhy 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.\n\n## Let's use hand tuned kernels\n\nSo 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.\n\n```python\nfrom kernels import get_kernel\n\nkernels_layers = get_kernel(\"kernels-community/liger-kernels\", version=1).layers\nkernels_geglu_mlp = kernels_layers.LigerGEGLUMLP(Config()).to(device, dtype=torch.bfloat16).eval()\n```\n\nThe full script is here: [`03_kernels_mlp.py`](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_kernels_mlp.py).\n\n```bash\nuv run 03_kernels_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072\nuvx trace-util -f traces -b <hf_uname>/traces\n```\n\n| [![Profiler trace of the LigerGEGLUMLP forward pass showing three aten::linear groups and a single LigerGELUMulFunction group on the CPU lane](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/kernels-profile.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/kernels-profile.png) |\n| -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 12: The profiler trace for the `LigerGEGLUMLP` layer                                                                                                                                                                                                                                                                                                                                        |\n\nFigure 12 shows the profile for the `LigerGEGLUMLP` layer using the Liger kernels from the Hub.\n\n### Why use the kernels library\n\nWriting 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).\n\nThe [`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:\n\n- The kernels are compiled once, in CI, for many architectures and version combinations. You download the right binary instead of compiling it yourself.\n- `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\".\n- 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.\n\n### Why tuned kernels are better\n\nWhen we say \"tuned\", we mean two concrete things, and both are visible in the trace.\n\n| [![Compiled MLP trace with the TorchDynamo, prologue and guard pre-ops boxed on the CPU lane before the compiled graph runs](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile-preops.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/compile-preops.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |\n| Figure 13: The compiled run pays for pre-ops (Dynamo, guards, prologue) before any GEMM runs                                                                                                                                                                                                                                                                                   |\n\n| [![LigerGEGLUMLP trace with an empty box where the compile pre-ops would be, showing the hand-written kernel has no Dynamo or guard overhead](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-preops.png)](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-mlp-fusion/no-preops.png) |\n| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| Figure 14: The Liger kernel has no pre-ops — the box where they would be is empty                                                                                                                                                                                                                                                                                                     |\n\n1. **The fusion is baked in.** The [`LigerGEGLUMLP`](https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/layers.py#L307) forward is `down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x)))`. The [`LigerGELUMulFunction`](https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L130) runs a single Triton kernel, [`_geglu_tanh_forward_kernel`](https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L97), that computes `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).\n\n2. **The launch parameters were chosen for the hardware.** The kernel does not guess its block size at random. Liger's [`calculate_settings`](https://huggingface.co/kernels/kernels-community/liger-kernels/blob/v1/build/torch-cuda/geglu.py#L95) picks them from the column count.\n\nIt 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.\n\n`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.\n\nSo 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.\n\n## Conclusion\n\nThe table below collects what each step changed on the GPU and what it left untouched.\n\n| Setup                | What changed                                                                                                                               | What stayed the same                                                       |\n| -------------------- | ------------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------- |\n| 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              | —                                                                          |\n| 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 |\n| 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 |\n| 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               |\n| 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                              |\n\nIf 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.\n\nThis 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.\n\nThanks 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!\n\n> The blog post was polished using an LLM. This in no way means that we have let an agent run in the background and let it generate the blog. Some of us in the team are non-english speakers and think LLMs (which are mostly trained in the English Language) can rectify silly grammar mistakes or rephrase sentences that sound less intimidating and cleaner. Hope this helps with the idea of \"why should I read, if this was LLM generated\". 🤗\n",
  "dateCreated": "2026-10-19T01:26:53.998049939+00:00",
  "date_created": "2026-06-11T00:00:00+00:00",
  "date_modified": "2026-06-11T00:00:00+00:00",
  "date_published": "2026-06-11T00:00:00+00:00",
  "description": "In the first part of this series \"Profiling in PyTorch\", 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…",
  "external_url": "https://huggingface.co/blog/torch-mlp-fusion",
  "git": {
    "commit": "ebb0941c5c06a2ab2377342d607535262a3e969e",
    "path": "items/huggingface-co/2026/06/2026-06-11-profiling-in-pytorch-part-2-from-nn-linear-to-a-fused-mlp.md",
    "permalink": "https://github.com/aymericbeaumet/aggr-instance/blob/ebb0941c5c06a2ab2377342d607535262a3e969e/items/huggingface-co/2026/06/2026-06-11-profiling-in-pytorch-part-2-from-nn-linear-to-a-fused-mlp.md"
  },
  "id": "urn:aggr:item:c5c5b5d56052ceee86d477b6274ad8b089ac5b5c",
  "isBasedOn": {
    "@id": "https://huggingface.co/blog/torch-mlp-fusion"
  },
  "mainEntity": {
    "@id": "https://huggingface.co/blog/torch-mlp-fusion",
    "@type": "CreativeWork",
    "archivedAt": {
      "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-06-11-profiling-in-pytorch-part-2-from-nn-linear-to-a-fused-mlp/#webpage"
    },
    "dateModified": "2026-06-11T00:00:00+00:00",
    "datePublished": "2026-06-11T00:00:00+00:00",
    "headline": "Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP",
    "image": {
      "@type": "ImageObject",
      "caption": null,
      "height": 90,
      "url": "https://huggingface.co/blog/assets/torch-mlp-fusion/thumbnail.png",
      "width": 160
    },
    "timeRequired": "PT17M",
    "url": "https://huggingface.co/blog/torch-mlp-fusion",
    "wordCount": 3609
  },
  "name": "Archived snapshot: Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP",
  "preview": {
    "alt": null,
    "color": null,
    "height": 90,
    "placeholder": {
      "data_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=",
      "hash": ""
    },
    "url": "https://huggingface.co/blog/assets/torch-mlp-fusion/thumbnail.png",
    "width": 160
  },
  "reading_minutes": 17,
  "replicated_at": null,
  "source": "huggingface-co",
  "source_name": "Hugging Face - Blog",
  "summary": null,
  "tags": [],
  "temporalCoverage": "2026-10-19T01:26:53.998049939+00:00",
  "title": "Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP",
  "url": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-06-11-profiling-in-pytorch-part-2-from-nn-linear-to-a-fused-mlp/",
  "word_count": 3609
}