Inside PyTorch(1) — From Python to GPU Kernels

🔍 Following a PyTorch Program from Python to GPU Execution

A Transformer model is usually written as ordinary Python code, but its execution eventually becomes a sequence of optimized CPU or GPU kernels. PyTorch 2.x connects these layers through torch.compile, TorchDynamo, FX graphs, AOTAutograd, operator decomposition, TorchInductor, and backend-specific code generation. This post provides a high-level map of that compilation pipeline before examining each component in detail.

---

① The Compilation Pipeline

A PyTorch model normally executes in eager mode. Each Python operation immediately dispatches one or more PyTorch operators, which eventually launch kernels on the target device. This execution model is flexible and easy to debug, but the runtime sees only a small part of the program at a time, limiting opportunities such as operator fusion, memory planning, and kernel specialization.

outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()

With torch.compile, PyTorch attempts to capture larger regions of the program and transform them into optimized executable code.

compiled_model = torch.compile(model)
outputs = compiled_model(inputs)

The overall flow can be summarized as follows.

Figure 1. From Python to GPU Kernels.

This diagram is useful as a conceptual map, but the implementation is not always a single rigid sequence. Inference may not require a backward graph, graph breaks may divide one Python function into multiple compiled regions, and different backends can replace TorchInductor. The default torch.compile path, however, generally follows this structure.

---

② TorchDynamo — Capturing Python Execution

TorchDynamo is the frontend of torch.compile. Rather than requiring the user to rewrite the model in a separate graph language, Dynamo observes ordinary Python execution and extracts PyTorch operations into graph regions.

Dynamo integrates with CPython frame evaluation and interprets Python bytecode before normal execution. When it encounters operations that can be represented safely, it records them into an FX graph. Python behavior that cannot be captured may cause a graph break.

flowchart LR
    SOURCE["Python Source"]
    BYTECODE["Python Bytecode"]
    DYNAMO["TorchDynamo"]
    CAPTURED["Captured Tensor Operations"]
    EAGER["Regular Python Execution"]

    SOURCE --> BYTECODE
    BYTECODE --> DYNAMO
    DYNAMO --> CAPTURED
    DYNAMO --> EAGER

Consider a simple function.

def transform(x, bias):
    y = x + bias
    y = torch.relu(y)
    return y * 2

Dynamo does not primarily compile the source text itself. It observes the bytecode-level execution and captures the tensor operations that occur while the function runs.

The captured region is guarded by assumptions about the execution environment. These assumptions may include tensor dtype, device, rank, shape properties, Python object identity, or module attributes. When a later invocation still satisfies the guards, PyTorch can reuse the compiled result. When a guard fails, PyTorch may compile another specialization.

flowchart TD
    CALL["Function Call"]
    GUARD{"Existing Guards Match?"}
    CACHE["Reuse Compiled Graph"]
    TRACE["Trace and Compile"]
    SAVE["Store New Specialization"]

    CALL --> GUARD
    GUARD -->|Yes| CACHE
    GUARD -->|No| TRACE
    TRACE --> SAVE

Dynamo therefore performs more than graph extraction. It also preserves Python semantics by tracking the assumptions under which the captured graph remains valid.

---

③ FX Graph — The Captured Program

The result of Dynamo tracing is represented using PyTorch FX. FX provides a graph-based intermediate representation composed of nodes such as placeholders, function calls, method calls, module calls, and outputs.

The previous Python function may conceptually become a graph similar to this.

placeholder  x
placeholder  bias
call_function  add
call_function  relu
call_function  mul
output

flowchart LR
    X["x"]
    B["bias"]
    ADD["torch.add"]
    RELU["torch.relu"]
    MUL["multiply by 2"]
    OUT["output"]

    X --> ADD
    B --> ADD
    ADD --> RELU
    RELU --> MUL
    MUL --> OUT

The FX graph is not yet a GPU kernel. It is an intermediate representation describing tensor operations and their data dependencies. Compiler components can analyze this representation without repeatedly interpreting the original Python source.

At this level, the system can reason about questions that are difficult to answer from isolated eager operations:

• Which operators depend on each other?
• Which intermediate tensors are consumed only once?
• Which operations can be fused?
• Which values must be materialized in memory?
• Which dimensions are static or symbolic?
• Which operations belong to the forward or backward computation?

The graph is therefore the boundary between flexible Python execution and compiler-oriented optimization.

---

④ FakeTensor and Symbolic Shapes

Graph capture and compilation do not require PyTorch to execute every tensor operation with real data. Instead, PyTorch uses FakeTensor execution to propagate tensor metadata—such as shape, dtype, stride, layout, and device—without allocating memory or performing the full numerical computation. This allows the compiler to analyze tensor properties while constructing the computation graph.

For example, the compiler may need to know that a matrix multiplication produces a tensor with shape [batch, sequence, hidden], but it does not need to compute the actual tensor values during graph construction. Metadata alone is sufficient for graph analysis and subsequent compiler optimizations.

FakeTensor execution can propagate both concrete and symbolic dimensions. Rather than requiring every dimension to be a fixed integer, PyTorch can represent dynamic dimensions using symbolic values. This allows a single compiled graph to support multiple compatible input sizes, provided that the generated guards and compiled code remain valid.

Symbolic shapes improve generality, but they also increase compiler complexity. Operations may introduce constraints on symbolic dimensions, and inputs that violate those constraints can still trigger recompilation. In practice, these symbolic assumptions are enforced by runtime guards, allowing compiled graphs to be reused whenever possible while preserving correctness.

---

⑤ AOTAutograd and Operator Decomposition

TorchDynamo captures executable Python regions as an FX graph. The graph is a container whose nodes represent PyTorch operations, such as calls to ATen operators. For inference, compiling the forward graph may be sufficient. Training, however, also requires the gradient computation performed during the backward pass. In eager execution, Autograd records the operations performed during the forward pass and dynamically invokes the corresponding backward functions when backward() is called.

AOTAutograd makes a larger portion of this training computation visible to the compiler. It takes the graph captured from the forward program and produces separate compiler-visible forward and backward FX graphs. These graphs still contain PyTorch operators as their nodes.

Exposing both graphs allows the compiler to optimize training operations that would otherwise be generated dynamically by Autograd. The amount of computation captured can still be affected by graph breaks, hooks, custom autograd functions, mutations, aliasing, and other dynamic behavior.

PyTorch exposes a large operator surface. A single high-level operation may combine broadcasting, type promotion, reductions, pointwise operations, indexing, or other lower-level behavior. Requiring every compiler backend to implement every PyTorch operator independently would make backend integration unnecessarily complex. PyTorch therefore uses decomposition rules to rewrite supported operators into a smaller and more regular operator vocabulary.

The important point is that decomposition changes the operators inside an FX graph. It does not replace the FX graph with a different graph format.

After decomposition, the graph may contain:

• Core ATen operators, which form a smaller subset of ATen suitable for compiler backends
• Prims operators, which express computation using lower-level primitive operations
• a mixture of operators appropriate for the selected backend

Core ATen operators do not always have to be decomposed completely into Prims. Some higher-level operations retain useful semantic information that allows a backend to select specialized implementations.

For example, preserving a matrix multiplication as an operation such as aten.mm allows a backend to recognize it as GEMM and select an optimized library call or kernel. Decomposing every matrix multiplication into scalar multiplication and addition would discard that useful structure.

⑥ PrimTorch

PrimTorch is best understood as the operator normalization effort surrounding Core ATen, Prims, reference implementations, and decomposition rules.

It is not a standalone compiler executable through which every graph must pass as a separate mandatory stage.

flowchart TD
    HIGH["PyTorch Operators"]
    RULES["Decomposition Rules"]

    subgraph PRIMTORCH["PrimTorch Operator Vocabulary"]
        CORE["Core ATen"]
        PRIMS["Prims"]
    end

    HIGH --> RULES
    RULES --> CORE
    RULES --> PRIMS
    CORE -. "Optional further decomposition" .-> PRIMS

The relationship between the graph and its operators can therefore be summarized as follows:

FX Graph
└── FX Nodes
    ├── Core ATen operators
    ├── Prims operators
    └── other backend-supported operators

An FX graph is the graph representation, while ATen and Prims are operator vocabularies used by nodes inside that graph.

The compiler backend receives the resulting forward and backward FX graphs and lowers their operators into its own internal representation.

For the default torch.compile backend, this means that TorchInductor lowers the FX graphs into TorchInductor’s loop-level IR before generating target-specific code.

flowchart TD
    PYTHON["Python / PyTorch Program"]
    DYNAMO["TorchDynamo"]
    CAPTURED["Forward FX Graph<br/>PyTorch Operators"]

    AOT["AOTAutograd"]

    FW["Forward FX Graph"]
    BW["Backward FX Graph"]

    DECOMP["Operator Decomposition"]

    FW_ATEN["Forward FX Graph<br/>Core ATen and/or Prims"]
    BW_ATEN["Backward FX Graph<br/>Core ATen and/or Prims"]

    INDUCTOR["TorchInductor"]
    INDUCTOR_IR["TorchInductor IR"]
    CODEGEN["Triton, C++, or Library Calls"]

    PYTHON --> DYNAMO
    DYNAMO --> CAPTURED
    CAPTURED --> AOT

    AOT --> FW
    AOT --> BW

    FW --> DECOMP
    BW --> DECOMP

    DECOMP --> FW_ATEN
    DECOMP --> BW_ATEN

    FW_ATEN --> INDUCTOR
    BW_ATEN --> INDUCTOR

    INDUCTOR --> INDUCTOR_IR
    INDUCTOR_IR --> CODEGEN

The sequence is therefore not:

FX Graph → ATen Graph → Prims Graph

as three mandatory graph formats.

Instead, it is:

FX Graph containing PyTorch operators
    ↓ AOTAutograd
Forward and backward FX graphs
    ↓ decomposition
FX graphs containing Core ATen and/or Prims operators
    ↓ backend lowering
Backend-specific IR

AOTAutograd should not be confused with AOTInductor.

• AOTAutograd exposes forward and backward computation as graphs that compiler backends can optimize.
• AOTInductor compiles exported models ahead of deployment and produces deployable compiled artifacts.

Official references:

• AOT Autograd
• PyTorch Compiler IRs
• Writing Graph Transformations on ATen IR
• PyTorch 2.x and TorchInductor

---

⑦ TorchInductor — Lowering and Optimization

TorchInductor is the default backend used by torch.compile. It consumes FX graphs, lowers their operators into its internal representations, analyzes data dependencies, plans loop execution, performs fusion, and generates target-specific code.

flowchart LR
    FX["FX Graph"]
    LOWER["Operator Lowering"]
    IR["Inductor IR"]
    FUSION["Fusion and Scheduling"]
    CODEGEN["Backend Code Generation"]

    FX --> LOWER
    LOWER --> IR
    IR --> FUSION
    FUSION --> CODEGEN

A major optimization opportunity is kernel fusion. In eager mode, a chain of elementwise operations may launch separate kernels and write intermediate results to global memory after every step.

flowchart LR
    INPUT["Input Tensor"]

    subgraph EAGER["Eager Execution"]
        direction LR
        ADD["Add<br/>Kernel"]
        RELU["ReLU<br/>Kernel"]
        MUL["Multiply<br/>Kernel"]

        ADD -->|"Write / Read<br/>Intermediate"| RELU
        RELU -->|"Write / Read<br/>Intermediate"| MUL
    end

    subgraph COMPILED["Compiled Execution"]
        direction LR
        FUSED["Fused Kernel<br/>Add + ReLU + Multiply"]
    end

    OUTPUT["Output Tensor"]

    INPUT --> ADD
    INPUT --> FUSED
    MUL --> OUTPUT
    FUSED --> OUTPUT

When the operations are compatible, Inductor can generate a fused loop or GPU kernel that computes the same result without materializing every intermediate tensor.

This can reduce :

• Kernel launch overhead
• Intermediate memory allocation
• Global memory reads and writes
• Python dispatch overhead
• Synchronization between individual eager operations

Fusion is not always possible or profitable. Tensor layouts, mutation, aliasing, reductions, device boundaries, unsupported operations, and scheduling constraints can divide the graph into multiple kernels.

---

⑧ Triton and Target-specific Code Generation

After optimization, Inductor generates code for the target device. For GPU workloads, Triton is a major code-generation path used to produce specialized GPU kernels. CPU graphs commonly use generated C++ code.

flowchart LR
    INDUCTOR["TorchInductor"]
    TARGET{"Target Device"}
    TRITON["Triton Kernel"]
    CPP["Generated C++"]
    GPU["GPU Execution"]
    CPU["CPU Execution"]

    INDUCTOR --> TARGET
    TARGET -->|GPU| TRITON
    TARGET -->|CPU| CPP
    TRITON --> GPU
    CPP --> CPU

Triton provides a Python-based language and compiler for expressing blocked parallel computations. TorchInductor can generate Triton programs specialized for tensor shapes, layouts, dtypes, and hardware properties.

The Triton compiler then lowers the program through lower-level compiler representations until executable GPU code is produced. Triton is therefore not the component that captures Python or creates the Autograd graph. It operates near the lower end of the pipeline, where compiler-visible tensor computations are converted into executable GPU kernels.

---

⑨ Runtime Execution and Caching

Compilation does not occur on every operation. During the first invocation of a compiled region, PyTorch traces the program, creates graphs, runs compiler passes, generates code, and compiles the backend kernel. This makes the first execution slower than steady-state execution.

flowchart TD
    FIRST["First Invocation"]
    TRACE["Graph Capture"]
    COMPILE["Backend Compilation"]
    CACHE["Compiled Artifact Cache"]
    EXECUTE["Execute Compiled Code"]
    NEXT["Later Invocation"]

    FIRST --> TRACE
    TRACE --> COMPILE
    COMPILE --> CACHE
    CACHE --> EXECUTE
    NEXT --> CACHE

Later calls can reuse the cached result while the associated guards remain valid. A guard failure, unsupported input variation, or graph change may cause recompilation.

This creates two performance phases:

• Compilation phase: tracing, lowering, code generation, compilation, and autotuning
• Steady-state phase: repeated execution of cached compiled kernels

Benchmarks must separate these phases. Measuring only the first iteration can make compiled execution appear slower even when its steady-state throughput is higher.

---

⑩ Graph Breaks

Dynamo cannot always represent an entire Python function as one FX graph. Unsupported Python behavior, data-dependent control flow, calls into unsupported libraries, explicit graph-breaking APIs, or values that must be materialized in Python can interrupt graph capture.

flowchart LR
    PY1["Python Region"]
    G1["Compiled Graph 1"]
    EAGER["Eager Python Region"]
    G2["Compiled Graph 2"]
    PY2["Remaining Program"]

    PY1 --> G1
    G1 --> EAGER
    EAGER --> G2
    G2 --> PY2

Execution remains correct because PyTorch can return to eager mode and resume graph capture later. However, excessive graph breaks reduce optimization scope and introduce transitions between compiled and eager execution.

A graph break does not necessarily mean that torch.compile has completely failed. It means that one continuous graph was divided into smaller regions.

flowchart LR
    PROGRAM["Python Function"]
    BREAK{"Graph Break?"}
    SINGLE["One Large Compiled Region"]
    MULTIPLE["Multiple Compiled Regions"]
    EAGER["Eager Segments Between Graphs"]

    PROGRAM --> BREAK
    BREAK -->|No| SINGLE
    BREAK -->|Yes| MULTIPLE
    MULTIPLE --> EAGER

Graph-break analysis is therefore one of the first steps when compiled performance does not match expectations.

---

⑪ What Happens to a Transformer Model?

A Transformer contains matrix multiplications, normalization, attention, activation functions, indexing, communication operations, and loss computation. Under eager execution, these operations are dispatched individually or through pre-existing fused kernels.

With torch.compile, compatible regions can be captured and compiled together.

flowchart TD
    MODEL["Transformer Python Code"]
    DYNAMO["TorchDynamo Capture"]
    FX["FX Graph Regions"]
    AOT["Forward and Backward Graphs"]
    INDUCTOR["Inductor Optimization"]
    KERNELS["Generated and Library Kernels"]
    GPU["GPU Execution"]

    MODEL --> DYNAMO
    DYNAMO --> FX
    FX --> AOT
    AOT --> INDUCTOR
    INDUCTOR --> KERNELS
    KERNELS --> GPU

Not every Transformer operation necessarily becomes a newly generated Triton kernel. The final execution may contain a combination of:

• Inductor-generated Triton kernels
• Existing ATen CUDA kernels
• Vendor library kernels such as cuBLAS or cuDNN
• Specialized attention kernels
• NCCL communication kernels
• Eager operations outside captured regions
• Host-side Python and runtime work

flowchart LR
    FX["FX Graph"]

    INDUCTOR["TorchInductor"]

    TRITON["Generated Triton Kernels"]
    LIB["CUDA Library Kernels"]
    EAGER["Fallback"]
    NCCL["NCCL"]

    TIMELINE["GPU Runtime"]

    FX --> INDUCTOR

    INDUCTOR --> TRITON
    INDUCTOR --> LIB
    INDUCTOR --> EAGER
    INDUCTOR --> NCCL

    TRITON --> TIMELINE
    LIB --> TIMELINE
    NCCL --> TIMELINE
    EAGER --> TIMELINE

This is why torch.compile should not be understood as simply translating an entire model into one GPU kernel. It restructures captured regions and generates an optimized execution plan containing multiple types of kernels and runtime operations.

---

⑫ Compilation Does Not Eliminate Every Synchronization

Compiler optimization can reduce Python overhead, kernel launches, and intermediate memory traffic, but it does not automatically remove every host synchronization.

A CUDA tensor may still be materialized on the host when Python requires its value. Common examples include scalar extraction, printing a tensor, converting a CUDA tensor into a Python boolean, or branching on device-resident data outside a captured tensor graph.

if cuda_tensor > 0:
    run_operation()

When Python must evaluate the condition, the host may need to wait until the GPU value becomes available.

sequenceDiagram
    participant CPU as Python Host
    participant GPU as GPU
    CPU->>GPU: Launch tensor operation
    CPU->>GPU: Request scalar value
    GPU-->>CPU: Device-to-Host result
    Note over CPU,GPU: Host waits for GPU completion
    CPU->>CPU: Evaluate Python condition

Whether a compiler can transform such code depends on how the operation is captured and whether its control flow can be represented safely. Data-dependent Python behavior may instead produce a graph break or remain outside the compiled graph.

This connects directly to the host synchronization previously observed in Transformers Sequence Parallel loss aggregation.

total_loss = sum(
    losses_per_rank[rank] * good_tokens_per_rank[rank]
    for rank in range(sp_world_size)
    if good_tokens_per_rank[rank] > 0
)

The condition evaluates CUDA tensor values from Python. Nsight Systems exposed the resulting synchronization as repeated small Device-to-Host copies and host-side waiting.

The optimized implementation replaced Python branching with tensor operations so that the filtering and aggregation remained on the device.

---

⑬ Observing the Pipeline

PyTorch provides logging and debugging options for inspecting graph capture and compiler output.

A minimal example is:

import torch

class Model(torch.nn.Module):
    def forward(self, x, bias):
        return torch.relu(x + bias) * 2

model = Model().cuda()
compiled_model = torch.compile(model)

x = torch.randn(1024, 1024, device="cuda")
bias = torch.randn(1024, 1024, device="cuda")

output = compiled_model(x, bias)

Dynamo and graph-break logs can be enabled using TORCH_LOGS.

TORCH_LOGS="dynamo,graph_breaks,recompiles" python example.py

Generated Inductor artifacts can be inspected using the compile debug option.

TORCH_COMPILE_DEBUG=1 python example.py

The exact debug output depends on the PyTorch version and configuration, but it may include FX graphs, Inductor intermediate representations, generated Triton code, generated C++ code, and compilation metadata.

The compiler view and the profiler view answer different questions.

flowchart LR
    SOURCE["Python Source"]
    LOGS["Dynamo and Inductor Logs"]
    GRAPH["Captured and Generated Code"]
    NSIGHT["Nsight Systems"]
    RUNTIME["Actual CPU and GPU Timeline"]

    SOURCE --> LOGS
    LOGS --> GRAPH
    GRAPH --> NSIGHT
    NSIGHT --> RUNTIME

Compiler logs explain how the program was transformed. Nsight Systems explains when CPU work, memory copies, CUDA kernels, synchronization, and communication actually occurred.

---

⑭ The Complete Mental Model

The entire path can now be summarized in one diagram.

flowchart TD
    PY["Python Model"]
    DYNAMO["TorchDynamo<br/>Bytecode Analysis and Graph Capture"]
    FX["FX Graph<br/>Tensor Program Representation"]
    META["FakeTensor and Symbolic Shapes<br/>Metadata Propagation"]
    AOT["AOTAutograd<br/>Forward and Backward Graphs"]
    DECOMP["ATen and Primitive Decomposition<br/>Operator Simplification"]
    INDUCTOR["TorchInductor<br/>Lowering, Fusion, Scheduling"]
    BACKEND{"Code-generation Target"}
    TRITON["Triton GPU Kernels"]
    CPP["Generated C++"]
    LIBRARY["Existing CUDA Libraries"]
    DEVICE["GPU or CPU Execution"]
    PROFILE["Nsight Systems and Runtime Profiling"]
    OPT["Source or Compiler Optimization"]

    PY --> DYNAMO
    DYNAMO --> FX
    FX --> META
    META --> AOT
    AOT --> DECOMP
    DECOMP --> INDUCTOR
    INDUCTOR --> BACKEND
    BACKEND -->|GPU| TRITON
    BACKEND -->|CPU| CPP
    INDUCTOR --> LIBRARY
    TRITON --> DEVICE
    CPP --> DEVICE
    LIBRARY --> DEVICE
    DEVICE --> PROFILE
    PROFILE --> OPT
    OPT --> PY

The key responsibilities are:

Component	Primary responsibility
torch.compile	User-facing entry point for compilation
TorchDynamo	Captures Python tensor execution into FX graphs
FX	Represents captured operations and data dependencies
FakeTensor	Propagates tensor metadata without full computation
Symbolic Shapes	Represents dimensions that may vary at runtime
AOTAutograd	Produces compiler-visible forward and backward graphs
Decomposition / PrimTorch	Reduces broad operators into smaller operator sets
TorchInductor	Lowers, optimizes, fuses, schedules, and generates code
Triton	Compiles generated GPU programs into executable kernels
CUDA Runtime and Driver	Launches and manages GPU execution
Nsight Systems	Reveals the resulting CPU and GPU runtime timeline

---

⑮ Conclusion

Modern PyTorch execution involves much more than launching CUDA kernels from Python. Between the original model and the final runtime execution lies a compiler stack responsible for graph capture, intermediate representations, automatic differentiation, operator decomposition, optimization, and backend-specific code generation.

Understanding this pipeline provides a mental model for reasoning about performance. Instead of viewing GPU execution as a black box, it becomes possible to identify where Python execution ends, where graph capture begins, how kernels are generated, and why certain optimizations—or performance bottlenecks—occur.

In the following posts, each stage of the pipeline will be explored in greater detail:

• TorchDynamo — Bytecode analysis, graph capture, guards, and graph breaks
• FX Graph — Intermediate representation, FakeTensor, and symbolic shapes
• AOTAutograd — Forward/backward graph extraction and operator decomposition
• TorchInductor — Lowering, scheduling, kernel fusion, and optimization
• Triton — GPU kernel generation and interaction with CUDA
• Profiler — Connecting compiler internals with Nsight Systems timelines and real optimization cases

Ultimately, the goal is not simply to understand each compiler component in isolation, but to connect the entire journey—from Python source code to GPU execution—and to use that understanding to analyze and optimize real-world AI workloads.

---

References

• PyTorch Compiler Documentation
• torch.compile API
• TorchDynamo Overview
• Dynamo Core Concepts
• Torch Compiler Troubleshooting
• AOTInductor
• Hugging Face Transformers Issue #47068
• Hugging Face Transformers Pull Request #47073

---

TL;DR

• torch.compile captures compatible Python tensor operations and sends them through the PyTorch compiler stack.
• TorchDynamo extracts tensor operations from Python bytecode into FX graphs.
• FakeTensor and symbolic shapes provide tensor metadata for graph analysis and specialization.
• AOTAutograd exposes forward and backward computation as compiler-visible graphs.
• Operator decomposition reduces complex PyTorch operations into smaller ATen or primitive operations.
• TorchInductor performs lowering, fusion, scheduling, and target-specific code generation.
• GPU workloads commonly use generated Triton kernels together with existing CUDA library kernels.
• Graph breaks divide Python execution into multiple compiled and eager regions.
• Compilation can reduce dispatch and memory overhead, but it does not automatically eliminate every host synchronization.
• Compiler logs explain graph transformation, while Nsight Systems reveals the actual CPU and GPU runtime behavior.