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

Within the first a part of this collection “Profiling in PyTorch”, we used torch.add(torch.matmul(x, w), b) to discover ways to learn PyTorch profiler traces. We additionally mentioned a number of different subjects that got here our means – the CPU dispatch chain, launch overhead, the distinction between an overhead-bound and a compute-bound regime, and a few internals of torch.compile.

Within the second iteration (this weblog submit), we climb one rung up the ladder. We exchange the hand-written matmul-add pair with an nn.Linear (with bias=True). That is the constructing block each deep studying mannequin makes use of. We then stack three of them (particular to our instance), with an activation in between, to kind a Multilayer Perceptron (MLP) block.

The scripts for this weblog submit reside right here: 02_linear.py, 03_simple_mlp.py, and 03_kernels_mlp.py. Like earlier than, it helps to open them in a separate tab and stroll by means of the code as you learn. We use an NVIDIA A100-SXM4-80GB GPU to run the scripts. It’s very easy to arrange a GPU on the Hugging Face infrastructure and experiment with the scripts utilizing Dev Mode with Areas. One may additionally run the scripts with the Hugging Face Jobs pipeline.

Earlier than we start, a fast recap of two concepts we’ll lean on repeatedly:

A GPU kernel is a program that runs in parallel on many threads of the GPU.
The CPU schedules and launches these kernels. A lot of the PyTorch overhead you see in a profiler hint is that this scheduling work.

From matmul-add to Linear

nn.Linear is a module wrapper across the identical matrix multiplication and addition we already profiled in Half 1. The one distinction is that it owns its weight and bias as parameters and exposes a ahead methodology that PyTorch customers have grown aware of.

linear_layer = nn.Linear(in_dim, out_dim, bias=True)
y = linear_layer(x)

The operation at hand might be written as:

y = x @ w.T + b

The place x is the enter, w is the load and b is the bias. Let’s run 02_linear.py and verify the profile.

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

trace-util is a utility that may sync your traces to a Hugging Face bucket after which present the Preffeto URLs in your terminal.

Determine 1 exhibits the profiler hint of a ahead name of the linear layer. We hint the ahead name of the linear layer with an analogous schedule setup because the earlier traces, with wait=1, warmup=1 and energetic=3. This is the reason we see three Profile Steps within the CPU and GPU lanes.

What’s the transpose doing?

If we zoom into the profiler hint, as we do in Determine 2, we discover an aten::t (transpose) op earlier than the aten::addmm (multiplication and addition) op. We will already work out that nn.Linear transposes the load parameter after which multiplies it with the enter. That is the rationale we see an aten::t op.

An necessary factor to note is that aten::t does probably not copy or reorganize information: it solely rewrites tensor metadata (form and stride) on the CPU to characterize the transposed matrix. It doesn’t launch a kernel on the GPU. One can confirm this two methods: by trying on the GPU lane within the hint, or by checking the aten::t row within the profiler desk and the time it took on CUDA.

Why are there no separate mul and add kernels?

There isn’t a aten::add (the bias addition) within the dispatch chain of the linear layer, as seen in Determine 3. It’s because the bias addition has been folded into the matrix multiplication kernel, utilizing what is known as an epilogue.

An epilogue is a small computation {that a} GEMM (GEneral Matrix Multiply) kernel does on the very finish, simply earlier than it writes its outcome again to HBM (Excessive Bandwidth Reminiscence, the GPU’s predominant reminiscence). Including a bias, making use of an activation, or scaling by a relentless are all traditional epilogues. The purpose of an epilogue is to keep away from loading or writing to HBM a second time, since reminiscence visitors makes an operation costly.

nn.Linear calls torch.nn.purposeful.linear, which, in flip, calls aten::linear. aten::linear appears on the inputs, notices {that a} bias was handed, and dispatches aten::addmm(bias, x, weight) as an alternative of doing a matmul and an add individually. addmm computes:

out = x @ weight.T + bias

The cuBLAS GEMM kernel that runs on the GPU has a bias-add variant inbuilt, and that is the kernel aten::addmm picks. The add by no means seems as a separate kernel as a result of it’s a part of the matmul kernel’s writeback, which is precisely what an epilogue is.

That is the second to note one thing refined. The kernel you noticed in Half 1 underneath –compile (addmm) is the kernel that keen nn.Linear already makes use of. There’s nothing left for torch.compile to fuse right here, which is the following factor we’ll confirm.

Can –compile assist a single Linear?

Let’s compile the ahead name and take a look at the profiler hint. (The profiler hint is visualized within the subsequent part)

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

In the event you evaluate the keen and compiled traces for a single nn.Linear’s ahead, you will see:

The identical cuBLAS GEMM kernel on the GPU.
The identical aten::addmm op on the CPU.
A couple of further rows on the CPU lane distinctive to compile.

That is price internalizing. A typical reflex is to succeed in for torch.compile at any time when a mannequin feels sluggish. For a single GEMM-with-bias, compile has little or no to do. This isn’t a bug, that is simply that compile wants multiple operation to presumably do any fusing. Let’s show that by taking a look at an MLP.

The place did the transpose go? Kernel layouts and pre-ops

A cautious reader of the 2 traces (keen vs compile) will discover that the keen CPU dispatch chain has extra in it than the compiled one.

The keen CPU dispatch chain inside aten::linear is aten::t adopted by aten::addmm (Determine 4). To know what aten::t truly does, we want a fast detour into strides and views.

A tensor shops its information as one flat, contiguous run of numbers in reminiscence. The form and stride are metadata that sit on high of that run and inform PyTorch the best way to stroll it: a stride of (s0, s1) means “step s0 parts to maneuver one row, step s1 to maneuver one column”. Change the metadata and also you get a distinct view of the identical uncooked information, with no copy:

>>> M = torch.tensor([[0, 1],
… [2, 3],
… [4, 5]])
>>> M.form, M.stride()
(torch.Measurement([3, 2]), (2, 1))

>>> T = M.t()
>>> T.form, T.stride()
(torch.Measurement([2, 3]), (1, 2))
>>> T
tensor([[0, 2, 4],
[1, 3, 5]])
>>> T.flatten()
tensor([0, 2, 4, 1, 3, 5])

M.t() didn’t transfer a single quantity. It returned a brand new view whose strides are swapped, so studying it row-by-row now walks the unique buffer 0, 1, 2, 3, 4, 5 in transposed order. The underlying information is an identical; solely the metadata differs.

That is precisely what aten::t does contained in the linear layer: it doesn’t allocate a brand new tensor or copy any information, it produces a view of the load with rewritten strides.

As we are able to see in Determine 5, compile didn’t take away a GPU kernel: it eliminated the CPU overhead of dispatching that view. Inductor traced by means of the view chain at compile time, computed the ensuing strides as soon as, and emitted a direct aten::addmm name with these strides hard-coded. A couple of microseconds of CPU work disappear whereas the GPU does an identical math.

As one would anticipate, when the enter information violates the strides precomputed by the compiler, it is going to throw an error.

In the event you take a look at the GPU lane in each traces, there’s precisely one kernel per ahead, and it’s the identical kernel each occasions:

cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8

If no transpose kernel ran, who taught the GEMM to learn the load matrix in transposed order? The reply is within the kernel’s identify. Have a look at the suffix:

cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8
^^

That tn is the structure descriptor. cuBLAS and CUTLASS precompile a separate kernel binary for every mixture of enter layouts.

n (non-transposed) and t (transposed) describe how a kernel walks its enter throughout the interior loop. The dispatcher’s job is to have a look at the enter strides, resolve which suffix mixture matches, and decide the fitting precompiled kernel.

The kernel identify in a profiler hint is a hash dump of the kernel’s identification. If two runs present the identical kernel identify, the GPU is doing the identical work. In the event that they differ (e.g., _tn_ vs _nn_, bf16 vs fp16, or s16816gemm vs s161616gemm) then the GPU is doing totally different work, and the dispatcher took a distinct department. Studying to learn this identify is among the most helpful habits when evaluating traces.

Stacking three Linears: the MLP

On this part, we’ll profile a Multilayer Perceptron (MLP). To make this extra attention-grabbing, we’ll profile a feed-forward community with the GeGLU activation variant (which is sort of closely utilized in apply). That is additionally our means of paying tribute to one of many best strains ever written within the historical past of deep studying analysis (Determine 6).

class SimpleGeGLUMLP(nn.Module):
def __init__(self, dim, hidden):
tremendous().__init__()
self.gate_proj = nn.Linear(dim, hidden, bias=False)
self.up_proj = nn.Linear(dim, hidden, bias=False)
self.down_proj = nn.Linear(hidden, dim, bias=False)

def ahead(self, x):
g = self.gate_proj(x)
u = self.up_proj(x)
h = F.gelu(g, approximate=“tanh”)
m = h * u
y = self.down_proj(m)
return y

One can find your complete script right here: 03_simple_mlp.py. Execute it like so:

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

Earlier than we open the hint, let’s assume collectively about what we should always anticipate to see. The ahead perform does a good quantity of computation, however most of it’s already acquainted to us.

We must always anticipate three aten::linear dispatches, one for every nn.Linear layer. We must also anticipate two pointwise kernel launches, one for the GeLU and one for the multiplication. Forming this expectation earlier than trying is the one most helpful behavior within the profiling journey: you learn the hint to substantiate or break a guess, to not kind one from scratch.

From Determine 7 we are able to pat ourselves on the again, as our instinct was right. Per ahead go (one mlp_fwd), the GPU runs precisely 5 kernels. Determine 8 highlights the “occupancy question” as seen within the CPU lane for the linear projection layers.

The three GEMMs every do an additional cudaOccupancyMaxActiveBlocksPerMultiprocessor name earlier than the launch. We now have a separate part on this in Half 1, yow will discover it right here. That’s cuBLAS sizing the grid. The pointwise ops (GeLU and mul) launch instantly, with no occupancy question. So “a linear” is definitely question + launch, whereas “a pointwise op” is simply launch.

The aten::t, aten::transpose, aten::reshape, aten::view, aten::as_strided, and aten::_unsafe_view ops launch zero kernels. They present 0.000us of CUDA time within the desk (Determine 9) as a result of they solely rewrite tensor metadata (form and stride) on the CPU. A reader scanning the desk sees round six op names per linear, however solely certainly one of them (mm) ever reaches the GPU.

Why are there two varieties of GEMM kernels?

The MLP flattens [batch, seq, dim] to [batch * seq, dim] for the matmul. In our command-line invocation we used 64 for batch and 128 for seq, in order that’s the place the 8192 (batch * seq = 64 * 128) under comes from.

From the hint:

All three GEMMs have the identical FLOP depend, 2·8192·768·3072 ≈ 38.7 GFLOP every, but down_proj is about 10% sooner. Similar work, totally different form (N=768 as an alternative of 3072), so cuBLAS picks a distinct tile (128×256, with a deeper stages_64x3 pipeline) that will get higher reuse for that form.

If you wish to study extra about tiling in depth, right here is a superb useful resource to get began with.

That is precisely why the desk had two GEMM rows (Determine 9): the 128×128 row is gate+up and the 128×256 row is down.

What does torch.compile do?

Earlier than compiling the ahead methodology and visualizing it, let’s do the psychological train once more of asking ourselves what we anticipate to see within the hint. It is a enjoyable experiment, and an necessary one to repeat each time you profile one thing your self. All the time construct in your instinct, and the second one thing doesn’t match, cease and work out why.

uv run 03_simple_mlp.py –batch 64 —seq 128 –dim 768 –hidden 3072 –compile
uvx trace-util traces -b traces

In keen mode, every nn.Linear was expanded into a series of dispatcher ops (aten::linear → aten::t → aten::transpose → aten::matmul → aten::reshape → aten::mm). These are the high-level wrappers that ATen walks by means of earlier than reaching the actual GEMM. torch.compile removes that chain.

By the point the compiled graph runs, there isn’t a linear, no matmul, no transpose or reshape and people metadata ops have been folded into how mm is known as. We will see three naked aten::mm exterior calls (Determine 10). The proof that it’s the identical GEMM is that the kernel names are byte-for-byte an identical to keen: …128×128…stages_32x5_tn for gate and up, and …128×256…stages_64x3_tn for down.

The fused Triton kernel

That is the headline of the entire compile lesson. The 2 keen pointwise kernels (GeLU and mul) plus a reshape collapsed into one kernel, triton_poi_fused__unsafe_view_gelu_mul_0 (Determine 11). Let’s decode the identify:

triton: generated by Inductor’s Triton backend (not cuBLAS, not ATen).
poi: pointwise (Inductor tags pointwise kernels poi, reductions pink, and protracted reductions per).
fused__unsafe_view_gelu_mul: the ops it merged: the _unsafe_view (reshape), the GeLU, and the mul.
0: the distinctive id throughout the graph.

Why is that this a win? In keen mode, the intermediate h = gelu(g) is a full [8192, 3072] bf16 tensor (round 50 MB) that the GeLU kernel writes to HBM and the mul kernel instantly reads again. Fusion retains it in registers (reminiscence that resides contained in the chip and are nearer than the HBM). The Triton kernel reads g and u as soon as, computes gelu(g) * u, and writes the outcome as soon as. One entire spherical journey of the intermediate by means of world reminiscence is gone.

Let’s use hand tuned kernels

To this point we now have let PyTorch (keen) and the compiler (torch.compile) decide our kernels. Now we plug in a kernel {that a} human skilled wrote and tuned by hand. We use the LigerGEGLUMLP layer, that we are able to simply fetch from the Hugging Face Hub with the kernels library.

from kernels import get_kernel

kernels_layers = get_kernel(“kernels-community/liger-kernels”, model=1).layers
kernels_geglu_mlp = kernels_layers.LigerGEGLUMLP(Config()).to(machine, dtype=torch.bfloat16).eval()

The complete script is right here: 03_kernels_mlp.py.

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

Determine 12 exhibits the profile for the LigerGEGLUMLP layer utilizing the Liger kernels from the Hub.

Why use the kernels library

Writing kernels in Triton or CUDA is one downside and transport them is one other. The kernel needs to be compiled in your precise mixture of GPU structure, CUDA model, and PyTorch model. That is the step that often breaks (“works on my machine”, lacking nvcc, mistaken Triton model).

The kernels library strikes that construct step off your machine. get_kernel(“kernels-community/liger-kernels”, model=1) downloads a pre-built, version-pinned kernel package deal from the Hugging Face Hub and caches it regionally (right here underneath ~/.cache/…kernels-community–liger-kernels). The advantages are:

The kernels are compiled as soon as, in CI, for a lot of architectures and model combos. You obtain the fitting binary as an alternative of compiling it your self.
model=1 pins the precise construct, so everybody operating your script will get the identical kernel. There isn’t a “it obtained slower after I up to date a package deal”.
The package deal exposes a .layers attribute with drop-in nn.Modules (like LigerGEGLUMLP). You swap your module for theirs and nothing else in your mannequin modifications.

Why tuned kernels are higher

Once we say “tuned”, we imply two concrete issues, and each are seen within the hint.

The fusion is baked in. The LigerGEGLUMLP ahead is down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x))). The LigerGELUMulFunction runs a single Triton kernel, _geglu_tanh_forward_kernel, that computes gelu(gate) * up in a single go. That is precisely what we noticed from torch.compile, the place the intermediate by no means makes a round-trip by means of HBM. We get it right here with out the compiler, as proven in Figures 13 and 14 (no Dynamo guards, no compile latency, no recompilation threat).

The launch parameters have been chosen for the {hardware}. The kernel doesn’t guess its block measurement at random. Liger’s calculate_settings picks them from the column depend.

It’s price being trustworthy concerning the trade-off right here, as a result of the uncooked numbers might be deceptive. The Liger kernel runs in 92.8 µs, whereas Inductor’s fused kernel from the compile run was 89.4 µs. At first look the hand-written kernel appears barely slower, however that comparability hides the associated fee that makes it worthwhile.

torch.compile specializes for a static form. Inductor’s 89.4 µs kernel is quick exactly as a result of it was generated for this precise [8192, 3072] downside. Change the batch measurement, the sequence size, or the hidden dimension, Dynamo re-traces, and also you pay the compile price yet again to get a brand new specialised kernel.

So the actual alternative just isn’t “sluggish human kernel vs quick compiled kernel”. It’s a quick generic kernel vs a kernel specialised for one explicit enter form. The Liger kernel takes one set of launch parameters and runs them for any form with no recompilation. It offers up the previous few microseconds that per-shape specialization would purchase, in alternate for being strong to altering shapes.

Conclusion

The desk under collects what every step modified on the GPU and what it left untouched.

If there’s one behavior to hold ahead, it’s the one we practiced earlier than each hint: guess first, then look. State what you anticipate the hint to comprise, open it, and deal with any mismatch as probably the most attention-grabbing factor on the display screen.

This was the second cease within the Profiling in PyTorch collection. Within the subsequent submit we’ll maintain climbing the ladder, transferring from this MLP block in the direction of the eye block and, finally, a full mannequin.

Due to Noe Flandre and Pedro Gabriel Gengo Lourenço for his or her opinions on the early draft of the submit!