torch.compile

torch.compile (PyTorch 2.x) traces your model into a graph, optimizes it, and generates device code. TorchInductor often emits Triton kernels for fusions, and delegates big GEMMs to cuBLAS-like libraries. You gain automatic fusion and less Python in the hot loop. You pay compile time and you fight dynamic shapes.

The big idea

torch.compile is a chef who preps the mise en place once, then cooks fast. The first dinner can take minutes. The next ones are quick, until someone orders a wildly different plate and prep starts over.

The picture breaks for multi-request serving. Engines like vLLM have their own schedulers. torch.compile is one tool in the stack, not a serving solution.

Walking through it

import torch
import torch.nn as nn

class Tiny(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(8, 8)

    def forward(self, x):
        return torch.relu(self.lin(x))

model = torch.compile(Tiny(), mode="reduce-overhead")
x = torch.randn(4, 8)
print(tuple(model(x).shape))

Output:

(4, 8)

mode="reduce-overhead" leans on CUDA graphs for further launch reduction. max-autotune spends more compile time searching kernel configs.

What bites you:

• Compile time. First run, and cache misses, can take seconds to minutes.
• Dynamic shapes. Varying batch or sequence length can cause recompiles. Serving engines often pad or bucket to stabilize shapes.
• Graph breaks. Data-dependent Python drops you back to eager.

Compile and autotune caches must travel with the deployment. Cold pods that recompile on first request create multi-second TTFT cliffs. Bake caches into images or warm pods before taking traffic. Treat compile as part of the release artifact.

Serving engines mix these: library attention, compiled MLP glue, graphs around the step. Know which layer owns a hotspot before rewriting anything.

Common mistakes

• Compiling at request time on a cold replica. That is a TTFT outage, not a speedup.
• Letting every sequence length through unpadded. You will recapture and recompile under load.
• Treating torch.compile as a replacement for a serving engine.

Check yourself

1. Why do serving stacks pad or bucket lengths under torch.compile? Answer: To keep shapes inside the compiled guards so you do not recompile every request.
2. What does mode="reduce-overhead" lean on besides fusion? Answer: CUDA graphs, to cut launch tax after the graph is built.

Key takeaways

• torch.compile fuses the glue and leaves big GEMMs to libraries.
• First compile is slow. Ship the cache with the replica.
• Dynamic shapes and graph breaks are the usual regressions.
• It is one layer in a serving stack, not the stack.

Further learning

Back to the kernels, fusion and compilation hub.