Local AI

torch.compile Recompilation: The Hidden Cost of Shape Guards

Why changing batch sizes can trigger new torch.compile graphs, how to log guard failures, and how to evaluate dynamic-shape trade-offs without confusing compile time with inference speed.

Approximately 5 min read

It is easy to benchmark a compiled PyTorch function on one tensor shape, observe a steady-state speedup, and then discover that a production workload behaves differently. A change from a batch of four to a batch of eight, or a different sequence length, can cause a previously compiled graph’s assumptions to stop holding.

The fundamental mechanism is guard-based specialization. A torch.compile graph is valid for inputs satisfying the conditions under which it was traced. If a later call violates one of those conditions, PyTorch may trace and compile another version. The work is not the same as ordinary eager inference.

This article explains the mechanism using the official PyTorch recompilation guide and dynamic-shapes documentation. Code below is a diagnostic illustration, not an independently measured performance result.

Three time costs that an average tokens-per-second number can hide

When measuring compiled inference, distinguish three regimes.

Regime What happens What to report
First invocation Graph capture and compilation can dominate latency Cold start separately
Repeated compatible inputs Reuse compiled graph Steady-state latency
New shape or changed guard Existing graph may not apply; another compilation may occur Recompile count, stall and subsequent reuse

A reported speedup that excludes compilation may be legitimate for a fixed-shape service; it does not establish that the entire service is faster for variable inputs. Conversely, a one-request latency spike does not prove the compiled execution itself is slower.

A useful minimal experiment changes only the shape of an otherwise identical function:

import torch

@torch.compile
def transform(x):
    return x.sin() + x.cos()

for n in (16, 16, 32, 32, 64, 16):
    x = torch.ones(n)
    y = transform(x)
    print(n, y.shape)

This tiny example is not representative of a language-model serving system. It exists to make a guard transition visible without mixing in tokenizer cost, GPU scheduling, networking, or model loading.

To see the reason for a recompile, launch a fresh Python process with:

TORCH_LOGS="recompiles" python shape_probe.py

PyTorch’s compiler troubleshooting reference shows messages describing guard mismatches, such as an input dimension having a different size from the previous trace. Use the actual log and version rather than assuming that every changed shape causes a new compile. The compiler may successfully generalize a dimension and reuse the graph.

Why dynamic=None is not the same as dynamic=False

The torch.compile interface has a dynamic argument whose options are easy to misread:

The current recompilation guide recommends checking that dynamic is not explicitly disabled when recompiles are driven by input shapes. But the dynamic-shapes guide also cautions that forcing dynamic=True everywhere can be error-prone and can regress performance.

An important distinction: dynamic shape is not dynamic rank. A graph that supports two-dimensional tensors with varying lengths does not automatically support an arbitrary change in tensor rank or unrelated Python-side control flow.

The misleading ‘fix’: make every dimension dynamic

A broad dynamic-shape setting can reduce shape-triggered recompiles but does not make every workload faster. Symbolic dimensions and their associated constraints may introduce different code generation or guards. Some operators are not supported for the shapes and control flow you want.

Treat this as an optimization decision:

  1. Identify which dimension actually varies: batch, sequence, crop size or another extent.
  2. Confirm its changes are the guard failures you see.
  3. Prefer a targeted dynamic-source approach or automatic generalization where supported.
  4. Compare cold-start, warmed steady-state, and shape-change latency separately.
  5. Confirm numerical outputs match an eager baseline for a representative set of inputs.

PyTorch documents the TORCH_COMPILE_DYNAMIC_SOURCES control for selecting dynamic sources and torch.compiler.config.patch(dynamic_sources=...) for more targeted experiments. These are advanced options; verify support in your installed release and do not paste an example source identifier into production without matching your own trace.

How to run a defensible comparison

A useful experimental table records the following fields for each request shape:

Field Why it matters
PyTorch and accelerator versions The compiler and backend evolve
torch.compile arguments Changes how graphs are specialized
Input dimensions and dtype Determines applicable guards and kernels
Cold versus warm run Separates compilation from execution
Total wall time and peak memory Captures costs hidden by a kernel-only metric
TORCH_LOGS=recompiles output Documents whether compilation occurred
Eager result comparison Checks that speed did not replace correctness

Run a fixed-shape baseline and a realistic shape mix. A fixed-shape batch benchmark answers whether reuse is efficient; a variable-shape stream answers whether compilation overhead is tolerable for that service.

For asynchronous CUDA execution, use synchronization appropriately when timing GPU work. Do not compare an unsynchronized host timer against synchronized end-to-end measurements and call the difference a kernel speedup.

Graph breaks are a separate diagnostic axis

Not every slowdown is a recompile. A graph break occurs when the compiler cannot keep the region inside a single traced graph and instead executes pieces under another mode. A function can have graph breaks even with constant shapes, and shape guards can invalidate compiled regions without implying graph breaks.

PyTorch’s troubleshooting reference documents:

TORCH_LOGS="graph_breaks,recompiles" python shape_probe.py

Read each category of log independently. If eager mode is slow for exactly the same workload, solving recompile count may not address the underlying performance problem. If the model changes Python objects, execution paths, or global state between calls, guard failures may not be reducible to one input dimension.

Practical decision rule

Use torch.compile when the workload is stable enough for compiled graphs to be reused and when the measured gain exceeds cold-start and recompilation costs for the lifetime of the process. If batch shapes vary constantly, measure that variability rather than copying a fixed-size benchmark.

A responsible report should say, for example, steady-state execution improved under repeated compatible inputs; shape transitions triggered recompilation, if and only if the recorded run supports both clauses. This article establishes the diagnostic method, not an improvement percentage for any particular GPU or model.

Primary references

Sources and further reading

Continue reading

More in Local AI ↗