Why is pytorch compile so fast?

PyTorch's Inductor compiler automatically groups dependent operations together into single, efficient Triton kernels. This keeps data in faster memory close to the register and cuts down on kernel overhead. In this article, we'll look at an example of fusion, and you'll see exactly how torch.compile transforms your PyTorch operations into optimized GPU code. When you use PyTorch's compiler, your model runs up to 10x faster, but what's actually happening? Without compilation, the GPU runs a kernel (a function on the GPU) for each torch operation in your code. This causes things to slow down for two reasons: Time is spent moving data in memory, and there's overhead of starting each new kernel. Every time the GPU launches a kernel, it pays an overhead cost, and every intermediate result means writing to and reading from memory. This is where vertical fusion comes in, and why you need to understand how to use it. To get the most out of this article, you need basic familiarity with PyTorch, and a general understanding of GPU programming concepts. What is Vertical Fusion? Think of vertical fusion as a way to link steps, so the output of one goes straight into the next. It's called "vertical" because if you picture the computation graph, these operations stack vertically, with each one dependent on the result of the previous step. This is the most common fusion pattern in deep learning because neural networks are chains of operations: Normalization, then linear layers, then activation functions, and so on. The big win is eliminating intermediate results. Those temporary tensors never need to be written to or read from global memory. They stay in fast registers where the GPU can reach them more quickly. Let's dive into an example of vertical fusion, specifically pointwise fusion. Pointwise fusion example Pointwise operations are simple math kernels that work on each element: Addition, multiplication, activation functions, and more. Let's look at a pattern you might see in a neural network layer: import torch def pointwise_example(x, w, b): # Multiple element-wise operations tmp = x * w # multiply tmp = tmp + b # add tmp = tmp.sigmoid() # sigmoid activation return tmp Unfused: Three separate kernels Without fusion, Inductor creates three separate Triton kernels. Don't worry if the Triton syntax looks intimidating. The important part isn't memorizing the syntax, but understanding the pattern. Each kernel loads data, does one operation, and writes the result. Kernel 1: Multiply @triton.jit def mul_kernel(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr): xoffset = tl.program_id(0) * XBLOCK xindex = xoffset + tl.arange(0, XBLOCK)[:] xmask = xindex < xnumel x0 = xindex tmp0 = tl.load(in_ptr0 + x0, xmask) tmp1 = tl.load(in_ptr1 + x0, xmask) tmp2 = tmp0 * tmp1 tl.store(out_ptr0 + x0, tmp2, xmask) For succinctness, I've included just the signatures of the next kernels, because they're nearly identical. See my Git repository for the full source code. Kernel 2: Add @triton.jit def add_kernel(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr) Kernel 3: Sigmoid @triton.jit def sigmoid_kernel(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr) Across the three kernels, you're performing eight memory operations: Reading inputs twice for multiply, reading multiply's result and the bias for add, reading add's result for sigmoid, and writing all three results. That's a lot of memory traffic. Fused: One kernel With fusion, torch.compile creates a single kernel: Kernel 4: Fused @triton.jit def triton_poi_fused_add_mul_sigmoid_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, XBLOCK: tl.constexpr): xoffset = tl.program_id(0) * XBLOCK xindex = xoffset + tl.arange(0, XBLOCK)[:] xmask = xindex < xnumel x0 = xindex # Load all inputs once tmp0 = tl.load(in_ptr0 + (x0), xmask) tmp1 = tl.load(in_ptr1 + (x0), xmask) tmp3 = tl.load(in_ptr2 + (x0), xmask) # Fused pointwise operations: mul -> add -> sigmoid tmp2 = tmp0 * tmp1 tmp4 = tmp2 + tmp3 tmp5 = tl.sigmoid(tmp4) # Store final result only tl.store(out_ptr0 + (x0), tmp5, xmask) Notice the difference: We load all inputs once, do all three operations in a row, and store only the final result. The intermediate values (tmp2 and tmp4 ) stay in registers (the fastest memory on the GPU). They never touch the slower global memory. Benefits of funsion - Kernel launches: Three reduced to one. - Intermediate buffers: Two eliminated (multiply result and add result). - Memory bandwidth: Reading five full tensors and writing three full tensors (eight memory operations) reduced to reading three tensors and writing one (four memory operations). That's a 50% reduction in memory traffic. Other fusion types Pointwise fusion is just one type of vertical fusion. Inductor uses other forms of vertical fusion to keep your GPU efficient: - Reduction fusion: Combines reducing operations like max, mean, or sum, with the operations that happen before and after them. This is critical for operations like batch normalization. - GEMM + Epilogue fusion: Attaches simple math to the end of heavy matrix calculations. Instead of doing a matrix multiply, writing the result to memory, then reading it back to add bias and apply ReLU, the bias and activation happen right after the multiply in the same kernel. - Prologue fusion: The opposite of epilogue. Preprocessing happens as data loads. For instance, normalizing input before matrix multiplication can happen even as the data comes in. In addition to vertical fusion (the most prominent type of fusion), Inductor also uses horizontal fusion. - Horizontal fusion: Runs multiple independent operations on the same input at once. For example, computing both sin(x) andcos(x) in a single kernel, loadingx only once instead of twice. Get started: See fusion in your own code Here's a complete example using a reduction pattern. Step 1: Create a simple reduction example Create a file called fusion_example.py : import torch def reduction_example(x): # Pointwise operation followed by reduction tmp = x * 2.0 result = tmp.sum(dim=-1) result = result + 1.0 return result # Create test input x = torch.randn(1024, 1024, device='cuda') compiled_fn = torch.compile(reduction_example) result_fused = compiled_fn(x) Step 2: View the generated code Run your script with the TORCH_LOGS environment variable to see what Inductor generated: TORCH_LOGS="output_code" python fusion_example.py This outputs the generated Triton kernels to your terminal. Look for a kernel named something like triton_per_fused_add_mul_sum_0 . In this context, per in the kernel name indicates that it's a "per-reduction" kernel. The rest of the name tells you that add, mul, and sum were all fused together. Conclusion Fusion is one of the most important optimizations that torch.compile does. By linking dependent operations into single kernels, it cuts down memory traffic and kernel overhead, often the main cause of slow GPU work. Try accelerating your own code with torch compile . As you've seen, there's no need to change your implementation. Just add a torch compiler decorator, and let the compiler do the work. Learn more - PyTorch documentation has complete guides on compilation and optimization strategies. - Reference my Git repository for the full source code from this article.