{"slug": "how-torch-compile-actually-works", "title": "How torch.compile Actually Works", "summary": "Torch.compile is not a traditional compiler but a system that intercepts Python bytecode at runtime, extracts compilable regions through a multi-stage pipeline, and stitches them back with eager Python fallbacks, with boundaries determining speedup. The pipeline consists of four stages: TorchDynamo for graph capture, AOTAutograd for ahead-of-time backward tracing, Inductor for fusion and delegation, and fallback for unsupported operations.", "body_md": "# How torch.compile Actually Works\n\nMost PyTorch users think `torch.compile`\n\nis a compiler. It is not — at least not in the way `gcc`\n\nor `nvcc`\n\nis a compiler. It is a system that intercepts Python bytecode at runtime, extracts the parts it can handle, compiles those parts through a multi-stage pipeline, and stitches the compiled regions back together with eager Python fallbacks for everything it cannot handle.\n\nThe interesting engineering is not in the compilation itself. It is in the *boundaries* — where the compiler gives up and hands control back to Python. Those boundaries determine whether you get a 2x speedup or no speedup at all.\n\nThis post walks through the four stages of the `torch.compile`\n\npipeline: how Dynamo captures a graph from live Python, how AOTAutograd traces the backward pass ahead of time, how Inductor decides what to fuse and what to delegate to cuBLAS, and where the whole thing falls apart.\n\n## The Pipeline\n\n``` php\ngraph LR\n    PY[\"Python bytecode\"] --> DY[\"TorchDynamo<br/>graph capture\"]\n    DY --> FX[\"FX Graph<br/>(ATen ops)\"]\n    FX --> AOT[\"AOTAutograd<br/>fwd + bwd tracing\"]\n    AOT --> FWD[\"Forward graph\"]\n    AOT --> BWD[\"Backward graph\"]\n    FWD --> IND[\"TorchInductor\"]\n    BWD --> IND\n    IND --> TRI[\"Triton kernels<br/>(GPU)\"]\n    IND --> CPP[\"C++ code<br/>(CPU)\"]\n    IND --> LIB[\"cuBLAS / cuDNN<br/>(vendor libs)\"]\n\n    style PY fill:#fce4ec,color:#1a1a1a\n    style DY fill:#e8eaf6,color:#1a1a1a\n    style FX fill:#e0f2f1,color:#1a1a1a\n    style AOT fill:#fff3e0,color:#1a1a1a\n    style FWD fill:#f3e5f5,color:#1a1a1a\n    style BWD fill:#f3e5f5,color:#1a1a1a\n    style IND fill:#e8f5e9,color:#1a1a1a\n    style TRI fill:#fce4ec,color:#1a1a1a\n    style CPP fill:#e0f2f1,color:#1a1a1a\n    style LIB fill:#fff3e0,color:#1a1a1a\n```\n\nWhen you call `torch.compile(model)`\n\n, nothing compiles yet. The compilation happens lazily on the first forward pass. Each stage feeds the next, and the final output — Triton kernels, C++ code, or cuBLAS calls — gets cached on disk so the second run skips the whole pipeline.\n\n## Stage 1: TorchDynamo — Capturing a Graph from Live Python\n\nDynamo is the most unusual part of the stack. It does not parse Python source code. It does not require type annotations or a restricted language subset. Instead, it hooks into CPython’s frame evaluation API (PEP 523) and intercepts Python bytecode *as it executes*.[1](#fn:1)\n\nHere is what that means concretely. When Python is about to execute a frame (a function call), Dynamo steps in before CPython’s default evaluator runs. It walks the bytecode instructions one by one, symbolically executing them. When it sees a PyTorch operation — a `torch.add`\n\n, a `nn.Linear`\n\n, a tensor indexing op — it records it into an FX graph. When it sees plain Python (a list append, a `print`\n\n, a call into a C extension), it has a choice: trace through it if it can, or give up.\n\n“Giving up” is a **graph break**. Dynamo stops tracing, compiles everything it has accumulated so far into a subgraph, emits code to call that compiled subgraph, and falls back to normal CPython execution for the unsupported instruction. Then it tries to resume tracing after the break.\n\n```\n1\n2\n3\n4\n5\n6\n7\npython\n@torch.compile\ndef f(x):\n    y = x * 2          # ← Dynamo traces this\n    z = y + 1           # ← Dynamo traces this\n    print(z.shape)      # ← graph break: side effect\n    w = z.relu()        # ← Dynamo starts a new subgraph\n    return w\n```\n\nThis function compiles into *two* subgraphs with a Python call to `print`\n\nbetween them. Each subgraph gets compiled and optimized independently. The `print`\n\nruns in eager Python. You get the optimization benefits for the two compiled regions, but you lose cross-region fusion and you pay two kernel-launch boundaries instead of one.\n\n### Guards\n\nBecause Python is dynamic, Dynamo cannot assume a compiled graph stays valid forever. A tensor’s shape might change on the next call. A module might get replaced. A global variable might flip.\n\nSo Dynamo generates a **guard function** alongside each compiled graph — a fast boolean check that verifies the assumptions the compilation made. Typical guards check tensor shapes, dtypes, device, and the identity of Python objects referenced during tracing. On each subsequent call, Dynamo runs the guard first. If it passes, the cached compiled code runs. If it fails, Dynamo recompiles.[2](#fn:2)\n\nGuard failures are a common performance problem. A model that sees different input shapes on every call recompiles on every call, and compilation is slow — often 10–60 seconds for a large model. This is the “cold start” problem, and it is the main reason `torch.compile`\n\nsometimes makes things slower in practice.\n\n### Why not TorchScript?\n\nPyTorch previously shipped TorchScript (`torch.jit.script`\n\n, `torch.jit.trace`\n\n) for ahead-of-time compilation. TorchScript required rewriting code into a restricted subset of Python. It could not handle data-dependent control flow, dynamic types, or most third-party libraries. Adoption was limited because the rewriting cost was high and the error messages were poor.\n\nDynamo takes the opposite approach: compile what you can, fall back for what you cannot. No rewriting needed. The tradeoff is that you get partial compilation instead of whole-program compilation, and you have to care about graph breaks. But the barrier to entry is a one-line decorator.\n\n## Stage 2: AOTAutograd — Tracing the Backward Pass\n\nIn standard PyTorch, the backward pass builds itself dynamically during the forward pass. Each forward op records a backward function onto the autograd tape. The tape is consumed once during `loss.backward()`\n\n.\n\nAOTAutograd changes this. 3 It traces both the forward and backward passes\n\n*ahead of time*, producing a joint graph. Then it partitions that graph into a separate forward graph and backward graph that can each be compiled by Inductor.\n\nThe process has three parts:\n\n### Decomposition\n\nHigh-level PyTorch ops get broken down into primitive ATen ops. A `nn.Linear`\n\nbecomes a `mm`\n\n(matrix multiply) plus an `add`\n\n(bias). A `nn.LayerNorm`\n\ndecomposes into `mean`\n\n, `sub`\n\n, `pow`\n\n, `sqrt`\n\n, `div`\n\n. This reduces the number of ops that Inductor needs to handle from thousands to a few hundred.\n\n### Functionalization\n\nPyTorch code is full of in-place mutations: `x.add_(1)`\n\n, `y[:, 0] = 3`\n\n, `buffer.copy_(new_data)`\n\n. A compiler cannot safely reorder or fuse ops that mutate shared state. Functionalization rewrites all mutations into pure, out-of-place equivalents.[4](#fn:4)\n\n```\n1\n2\n3\n4\n5\n6\n# Before functionalization:\nx.add_(1)       # mutates x in place\n\n# After functionalization:\nx_new = x + 1   # pure: creates a new tensor\n# all subsequent uses of x are replaced with x_new\n```\n\nThis is not optional. If the op schema lies about what it mutates (a missing `Tensor!`\n\nannotation), functionalization produces a wrong graph and the compiled model silently computes wrong results. I covered this failure mode in detail in the [vLLM custom ops post](/posts/vllm-custom-kernels-and-abi/).\n\n### Partitioning\n\nThe joint forward-backward graph gets split by a min-cut partitioner. It decides which forward activations to save for the backward pass (stored in “saved tensors”) and which to discard and recompute during backward. This is activation checkpointing built into the compiler, and it runs automatically. The partitioner’s objective is to minimize the memory high-water mark while keeping recomputation cost reasonable.[5](#fn:5)\n\n## Stage 3: TorchInductor — Fusion and Code Generation\n\nInductor takes the forward and backward FX graphs and turns them into executable code. This is where the actual optimization happens.\n\n### The fusion algorithm\n\nInductor’s primary job is reducing memory traffic. On a modern GPU, a single HBM read-write round trip takes orders of magnitude longer than an arithmetic operation. Fusing operations so that intermediate values stay in registers or shared memory instead of touching HBM is the single biggest lever Inductor has.\n\nInductor uses two fusion strategies:\n\n**Vertical fusion** chains dependent ops. If op B reads the output of op A and nothing else reads it, Inductor fuses A and B into a single kernel. The intermediate tensor never gets written to HBM. A chain like `matmul → add → relu → dropout`\n\nbecomes one kernel.\n\n**Horizontal fusion** groups independent ops that access the same memory. If two ops read from the same input tensor, Inductor can fuse them into one kernel that reads the input once instead of twice.\n\nThe fusion pass builds a dependency graph over the ops, groups fusible nodes, and emits one kernel per group. The output is Triton source code for GPU, or C++ source for CPU.\n\n### Triton vs. cuBLAS: who gets what\n\nNot every op goes through Triton. Inductor maintains a routing table:\n\n| Op pattern | Backend | Why |\n|---|---|---|\n| Pointwise ops (add, mul, relu, …) | Triton | Fusion is the win; Triton fuses these well |\n| Reductions (sum, mean, softmax) | Triton | Block-level reductions map naturally |\n| Dense GEMM (matmul, linear) | cuBLAS | NVIDIA hand-tunes cuBLAS per GPU; hard to beat |\n| Convolutions | cuDNN | Same story as cuBLAS |\n| Fused GEMM + epilogue (matmul + relu) | Triton template | When the epilogue fusion saves enough bandwidth |\n\nThe GEMM routing deserves a note. Inductor does not blindly call cuBLAS for every matrix multiply. It has a template system that benchmarks cuBLAS against Triton GEMM templates for each specific shape and picks the winner. For small or irregular shapes (common in LLM inference), Triton sometimes wins because cuBLAS’s tile sizes are tuned for large matrices. For large square GEMMs, cuBLAS almost always wins.\n\nFor more on how Triton generates the actual GPU binary, see the [Triton deep-dive](/posts/triton-compiler-deep-dive/).\n\n### Memory planning\n\nAfter fusion, Inductor analyzes tensor lifetimes and reuses buffers. If tensor A is dead before tensor B is born and they are the same size, B reuses A’s memory. This reduces peak memory and fragmentation, which matters for large models where every gigabyte of HBM counts.\n\n## Where It Breaks\n\n### Graph breaks\n\nGraph breaks are the most common reason `torch.compile`\n\nfails to help. Each break splits the graph into smaller subgraphs. Smaller subgraphs mean fewer fusion opportunities and more kernel launches. Common causes:\n\nDynamo cannot prove these do not affect tensor computation.`print`\n\n,`logging`\n\n, or any Python side effect.**Data-dependent control flow.**`if x.sum() > 0:`\n\nbranches on a runtime tensor value. Dynamo does not know which branch to trace.**C extensions.** Calling into a C library that Dynamo cannot see through.**Unsupported Python constructs.** Some`torch.autograd`\n\nfunctions, certain`__torch_function__`\n\noverrides, or exotic Python features.\n\nDebugging graph breaks requires setting `TORCH_LOGS='graph_breaks'`\n\nand reading the output, which tells you which bytecode instruction caused the break. This is not intuitive — the message refers to CPython bytecodes, not your source code. But it is the only diagnostic that works.\n\n### Dynamic shapes\n\nBy default, Dynamo specializes on input shapes. A model compiled with batch size 32 recompiles when it sees batch size 64. Setting `torch.compile(dynamic=True)`\n\ntells Dynamo to use symbolic shapes instead of concrete values. This avoids recompilation but limits some optimizations (the compiler cannot constant-fold shape-dependent expressions).\n\nEven with `dynamic=True`\n\n, highly variable shapes — say, a batch dimension that changes every call — can cause guard failures and thrash the compilation cache. The fix is `torch._dynamo.mark_dynamic()`\n\n, which explicitly tells Dynamo which dimensions vary. But this requires the user to know their data shapes ahead of time, which is not always possible.\n\n### Custom ops\n\nAny C++ or CUDA extension that Dynamo cannot trace is a graph break. If your model calls a custom CUDA kernel (common in LLM inference for PagedAttention, quantized GEMMs, or fused normalization), that kernel is a black box. Dynamo stops tracing at the call boundary.\n\nThe fix is `torch.library`\n\n, which lets you register a custom op with a schema, a meta function (for shape inference during tracing), and optional decomposition rules. 6 This makes the op visible to Dynamo and Inductor. I covered the registration mechanics in the\n\n[vLLM ABI post](/posts/vllm-custom-kernels-and-abi/)— the key requirement is that the schema must truthfully declare which tensors are mutated.\n\n### Compilation time\n\nCompilation is slow. A large LLM can take 30–120 seconds to compile on first run. This is split roughly evenly between Dynamo’s tracing, AOTAutograd’s graph manipulation, and Inductor’s code generation plus Triton’s own JIT compilation.\n\nFor production inference, compilation happens once and the result is cached. For training, the backward graph adds compilation time. For development iteration, the cold start is painful. Regional compilation — compiling only the hot inner module (e.g., the transformer block) rather than the entire model — is the practical workaround.\n\n## What torch.compile Sees and What It Does Not\n\nThe compilation pipeline sees a lot: op types, tensor shapes, dtypes, strides, device types, and the full dataflow graph including the backward pass. This is enough to fuse operators, eliminate dead code, reorder memory accesses, and pick the best kernel for each op.\n\nWhat it does not see:\n\n**Which physical device a tensor lives on.** A`cuda:0`\n\ntensor and a`cuda:1`\n\ntensor have the same dtype and shape. The compiler treats them identically. Whether a particular data movement between devices is correct, efficient, or even legal is not the compiler’s concern.**Memory space distinctions.** Host memory, device memory, pinned memory, and unified memory are all just tensors with a device tag. There is no type-level distinction between a pointer to HBM and a pointer to host DRAM. An accidental`.cpu()`\n\nin the wrong place is a runtime error or a silent performance cliff, not a compile error.**Placement decisions.** The compiler does not decide where an op should run. It compiles the graph as given. If a matmul should run on device 1 instead of device 0 for locality reasons, that is the user’s problem — or the framework’s (FSDP, pipeline parallelism, tensor parallelism all manage placement in Python).**Inter-device communication.** NCCL collectives (`all_reduce`\n\n,`all_gather`\n\n) interact with`torch.compile`\n\nbut are essentially opaque barriers. The compiler cannot fuse across a collective or reason about the communication cost.\n\nThis is not a criticism of `torch.compile`\n\n. The scope is deliberate. Compilation within a single device’s compute graph is a well-defined, solvable problem. Placement, routing, and memory-space management across a heterogeneous cluster are a different class of problem entirely — one where the constraints are topological, the costs are hardware-dependent, and the correctness properties cannot be inferred from tensor shapes alone.[7](#fn:7)\n\nToday, that outer layer lives in uncompiled Python: FSDP’s sharding logic, DeepSpeed’s pipeline scheduler, Megatron’s tensor-parallel annotations. These are written by hand, tested empirically, and verified by hoping the loss converges. No compiler checks them.\n\n## References\n\n*Disclaimer: This article was generated using the Gemini 3.1 Pro and Claude Opus 4.8 models.*\n\n**PEP 523 — Adding a frame evaluation API to CPython.** Python Enhancement Proposal enabling external tools to intercept frame evaluation. TorchDynamo uses this to capture FX graphs from running Python. ([Link](https://peps.python.org/pep-0523/))[↩︎](#fnref:1)**TorchDynamo: An Experiment in Dynamic Python Acceleration.** Ansel, J. et al. PyTorch team, 2022. Describes the guard mechanism and graph break strategy. ([Link](https://pytorch.org/docs/stable/torch.compiler_dynamo_overview.html))[↩︎](#fnref:2)**AOTAutograd: Ahead-of-Time Tracing for PyTorch.** Functorch / PyTorch team. Traces forward and backward passes ahead of time for compiler consumption. ([Link](https://pytorch.org/functorch/stable/notebooks/aot_autograd_optimizations.html))[↩︎](#fnref:3)**Functionalization in PyTorch.** Yang, E. (ezyang). Explains how in-place mutations are rewritten into pure functional form for compiler safety. ([Link](https://dev-discuss.pytorch.org/t/functionalization-in-pytorch-everything-you-wanted-to-know/965))[↩︎](#fnref:4)**Min-Cut Rematerialization Partitioning.** PyTorch Inductor’s activation checkpointing strategy that partitions the joint forward-backward graph to minimize memory usage. ([Link](https://pytorch.org/docs/stable/torch.compiler_aot_autograd.html))[↩︎](#fnref:5)**torch.library: Custom Operators for torch.compile.** PyTorch documentation on registering custom ops with schemas and meta functions for compiler visibility. ([Link](https://pytorch.org/docs/stable/library.html))[↩︎](#fnref:6)**PyTorch Distributed and torch.compile.** PyTorch team. Documents the interaction between FSDP, DDP, and the compilation stack. ([Link](https://pytorch.org/docs/stable/torch.compiler_faq.html))[↩︎](#fnref:7)\n\n[CC BY 4.0](https://creativecommons.org/licenses/by/4.0/)by the author.", "url": "https://wpnews.pro/news/how-torch-compile-actually-works", "canonical_source": "https://hiraditya.github.io/posts/how-torch-compile-works/", "published_at": "2026-07-21 19:00:00+00:00", "updated_at": "2026-07-22 14:00:22.993293+00:00", "lang": "en", "topics": ["machine-learning", "developer-tools"], "entities": ["PyTorch", "TorchDynamo", "AOTAutograd", "TorchInductor", "Triton", "cuBLAS", "cuDNN", "CPython"], "alternates": {"html": "https://wpnews.pro/news/how-torch-compile-actually-works", "markdown": "https://wpnews.pro/news/how-torch-compile-actually-works.md", "text": "https://wpnews.pro/news/how-torch-compile-actually-works.txt", "jsonld": "https://wpnews.pro/news/how-torch-compile-actually-works.jsonld"}}