Neural Network Compilation, Made Simple — How a PyTorch Model Becomes Hardware Language

This post is a concept explainer that branched off from Porting an Unsupported Model to a Korean NPU. Where that post is about “how we actually broke through,” this one unpacks from scratch the “what does it even mean to compile a model” that underlies it. It’s written to read smoothly for anyone with a bit of deep learning background.

Starting from the Familiar Word “Compile”

In programming, compilation is the process of translating ahead of time the human-readable source code (C, Rust, etc.) into machine code that the CPU executes directly. You do the translation up front, and in return, execution is fast.

Neural network compilation is exactly the same idea. You take a model written in PyTorch and translate it ahead of time into a form that a specific accelerator (GPU, NPU, etc.) can execute directly. The one thing that differs is the starting point of the translation.

• Input to a regular compiler = text source code
• Input to a neural network compiler = computation graph

A computation graph is a directed graph of chained tensor operations, like “matrix-multiply this tensor → then attention → then normalization → …”. Think of it as the model’s forward() unfolded into a picture.

That said, “a pipeline flowing in a straight line” is a story about the simple models in textbooks; a real model is a graph where multiple modules branch and merge. For example, an image editing model (of the Qwen-Image-Edit family) sends the input image, edit prompt, and timestep through their own encoders, merges them into a single backbone, and loops through denoising with feedback.

flowchart TD
    IMG["Input image"] --> VAE["VAE encoder"]
    IMG --> VIT["Vision encoder<br/>Qwen2.5-VL"]
    TXT["Edit prompt"] --> TENC["Text encoder<br/>LLM"]
    TS["timestep t"] --> TEMB["Time embedding"]

    VAE --> LAT["Image latent tokens"]
    VIT --> COND["Semantic condition tokens"]
    TENC --> TTOK["Text tokens"]

    LAT --> BB
    COND --> BB
    TTOK --> BB
    TEMB --> BB

    subgraph BB["MMDiT backbone · multimodal joint attention × N blocks"]
        direction TB
        JA["Joint self-attention<br/>image ↔ text"] --> FF["MLP"]
        FF -. residual .-> JA
    end

    BB --> EPS["Predicted noise ε"]
    EPS --> LOOP{"Denoising loop"}
    LOOP -->|"next step t−1"| BB
    LOOP -->|"done"| VDEC["VAE decoder"]
    VDEC --> OUT["Edited image"]

The same input (the image) splits into two branches feeding different encoders, merges at the backbone together with the text and timestep, and inside the blocks there are residual connections while outside there’s a denoising feedback loop. The graph a compiler actually faces is this kind of tangle of branching, merging, and repetition.

Here’s one important fact. The front end of a regular compiler — lexical analysis (splitting source text into tokens) and parsing (turning tokens into a syntax tree) — does not exist in neural network compilation. That’s because the input is already given as a structured graph (equivalent to the AST/IR in a regular compiler). So neural network compilation skips the entire front end of the compiler pipeline and effectively starts from the IR optimization and code generation stages.

So Why Compile at All? — The Cost of Eager Execution

PyTorch’s default execution mode is eager execution. Each time a line of Python code runs, it immediately invokes the corresponding operator kernel on the spot. It’s closer to an interpreter that translates one sentence at a time as you converse — flexible, and easy to debug by printing intermediate values along the way.

The problem is cost. For every single operation, Python interpretation → kernel dispatch → execution repeats, and this overhead adds up a little differently each time (i.e., latency spikes, jitter). General-purpose GPUs have kernels that are so fast and flexible that this approach runs well enough, but specialized accelerators like NPUs deliver their real performance only when they know in advance what to run and in what order.

Compilation is the act of turning this interpreter into a translated book. If you receive the whole graph before execution and freeze it into a single deterministic execution unit, there’s nothing to decide on the fly for each request, so execution doesn’t spike.

Seeing it is faster than talking about it. Below, run the same model many times in both eager and compiled modes and compare how the latency distributions differ.

Eager dispatches ops one at a time, so its average is higher and its distribution is spread out widely. Compilation executes a single graph deterministically, so it converges to nearly a single point. This is where the value of compilation shows up in real services that need low latency and low jitter.

What the Compiler Does in Between

When you hand over a single model, the compiler goes through roughly five stages.

flowchart TD
    A["① Graph capture<br/>trace / export"] --> B["② Operator fusion & optimization"]
    B --> C["③ Lowering to hardware kernels<br/>lowering"]
    C --> D["④ Fix static shapes"]
    D --> E["⑤ Memory planning"]
    E --> F(["Execution binary"])

① Graph Capture — Turning Python into a Graph

First, you extract a graph of what operations the model actually performs. There are two approaches.

• trace (torch.jit.trace): Actually run an example input through once and record the operations it passed through. Simple, but for things like if branches, only “the path taken that time” is captured.
• export (torch.export): Statically analyze the code to extract a graph that includes control flow. More robust.

② Operator Fusion — Combining Multiple Operations into One

Adjacent operations are merged into one. For example, running matrix-multiply → add bias → activation function separately writes and reads the intermediate result to memory three times. Merging these into a single fused kernel reduces the memory round trips and makes it much faster.

flowchart LR
    subgraph before["Before fusion — 3 kernels, many memory round trips"]
        direction LR
        M1["matmul"] --> B1["+bias"] --> A1["gelu"]
    end
    subgraph after["After fusion — 1 kernel"]
        direction LR
        F1["fused_matmul_bias_gelu"]
    end
    before -.optimize.-> after

③ Lowering to Hardware Kernels

Each standard operation (matrix multiply, attention…) is translated into instructions and kernels specific to that chip. A GPU example gives you the feel of it. The default backend of PyTorch’s torch.compile (TorchInductor) lowers operations to OpenAI Triton. Triton is a language for describing GPU kernels with Python-like syntax; the compiler generates a kernel matching the fused operations as Triton code on the spot, then lowers that further into a GPU binary (PTX/cubin). Instead of calling a vendor’s prebuilt library kernel (cuBLAS, etc.), it generates a kernel tailored to the graph. NPU compilers follow the same idea — the only difference is that the code generation backend is specific to that chip rather than Triton.

④ Fix Static Shapes

You nail down the input sizes in advance (e.g., [1, 512]). Once the sizes are fixed, branches at runtime disappear and a deterministic path is created, and fusion and memory planning can be done more aggressively too. This is why “fixed 512 tokens” keeps showing up in the NPU-related posts. The cost is flexibility — inputs that go beyond the size decided at compile time can’t be fed in (unless you recompile or pad).

⑤ Memory Planning

The placement and reuse of the buffers each operation will use are finalized before execution. Since memory isn’t allocated and freed on the fly during execution, it’s that much faster and more consistent.

What Does the Compiled Artifact Look Like

The result is a binary the hardware consumes directly (for Rebellions, a .rbln). It’s serialized, so you can’t just open it and read it, but if you conceptually unfold the internal graph, it looks like this.

# compiled_model — internal graph (conceptual representation)
target: RBLN-CA22   dtype: fp16
in  %input_ids      : i64[1, 512]
in  %attention_mask : i64[1, 512]

  %emb  = embedding       %input_ids                    -> f16[1,512,768]
  %h0   = layernorm       %emb
  # × N encoder layers:
  %a    = sdpa_attention  %h{i}, mask=%attention_mask   @kernel=attn_fused
  %m    = geglu_mlp       %a                            @kernel=geglu_fused
  ...
  %pool = mean_pool       %enc, mask=%attention_mask    -> f16[1,768]
  %log  = linear          %pool, W=const[768,2]         -> f16[1,2]
out %log

The stages we saw earlier show up directly — operations are mapped to dedicated kernels in the @kernel=..._fused form, and every tensor shape is statically fixed at [1,512,...].

What Already Exists — ONNX and the Compilation Ecosystem

So far we’ve lumped “the compiler” into one thing, but in reality several standards and tools form layers. The one that frequently appears at the hub is ONNX (Open Neural Network Exchange).

ONNX is not a compiler but an open standard format for holding a computation graph. No matter where a model is built — PyTorch, TensorFlow, wherever — once you export it to ONNX, various runtimes and compilers that can read that graph execute it with their own backends. It’s a common language that decouples the framework a model was built in from where it gets deployed.

flowchart TD
    PT["PyTorch"] --> ONNX
    TF["TensorFlow"] --> ONNX
    JAX["JAX"] --> ONNX
    ONNX["ONNX · common graph standard"] --> ORT["ONNX Runtime"]
    ONNX --> TRT["TensorRT · NVIDIA"]
    ONNX --> OV["OpenVINO · Intel"]
    ONNX --> NPU["Vendor NPU compiler"]

Here are the tools you’ll often run into, organized by their nature.

The big picture is usually a flow of ① extracting a graph from a framework (capture/export) → ② passing through a common representation like ONNX, or going straight → ③ compiling for the target backend. Most NPU vendors also accept ONNX input. In this series, however, instead of going through the extra ONNX layer, we chose the path of handing the PyTorch-native graph (torch.export) directly to the vendor compiler — when dealing with a model whose architecture is subtly misaligned, cutting one step of format conversion made the debugging surface area smaller.

Trade-offs — What You Gain and What You Lose

In short, compilation is turning a “flexible interpreter” into a “deterministic execution binary.” You give up flexibility and gain low latency, low jitter, and efficiency.

So When Do You Compile

• Research and prototyping, when input sizes change every time → eager is enough. Flexibility is valuable.
• Real-service low-latency inference, edge and on-device, specialized accelerators (NPUs) → compilation wins. For NPUs in particular, compilation is effectively mandatory.

In other words, “compilation is always the answer” is not true. When flexibility matters, eager is the fit; when latency, stability, and efficiency matter, compilation is. How we navigated this trade-off while actually running a real model on a Korean NPU continues in the porting experiment writeup.