Program capture, guards, and graph breaks
Explain graph reuse, guard failure, graph breaks, and dynamic-shape specialization
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Python execution model
- Dynamic shapes
- Pure versus side-effecting functions
Wrap a PyTorch transformer block with torch.compile and call it on a [8, 2048, 4096] tensor. The first call spends time capturing and compiling the operations; a later call with the same relevant properties reuses compiled code. Change the sequence length, execute a Python print, or branch on a tensor-derived value, and TORCH_LOGS can report new guards, a recompilation, or a graph break. Nsight Systems can then show extra launches or a host gap at the same request.
Program capture is the process that records tensor operations from an ordinary framework execution so a compiler can analyze them together. A guard is a runtime condition that must remain true for the recorded program to be valid. A graph break is the exact source location where capture stops and eager execution resumes. Compiler logs provide the observable evidence for all three terms.
Graph capture records tensor operations and the program facts needed to replay or transform a region. Ordinary Python permits arbitrary side effects, dynamic control, object mutation, and data-dependent behavior that a tensor graph might not represent. The capture system must specify which facts become guards, which values remain symbolic, and which operations return execution to Python.
A compiler can produce several graph variants. Guards select variants by shape, stride, datatype, device, or Python value. When a guard fails, the system recompiles or uses a fallback. Uncontrolled shape variation can therefore cause repeated compilation even when each compiled kernel is fast.
Use a small function such as y = sin(x) + cos(x) as the first case. Eager execution dispatches separate tensor operations. A compiler can record one graph, reason about the shared input and elementwise iteration space, and generate a fused implementation. The example also shows why capture needs assumptions about x before the generated code can be reused.
What capture records
Graph capture records tensor operations, their data dependencies, and enough metadata to reproduce the region. Capture does not ordinarily record every possible Python behavior. Arbitrary file input/output, object mutation, reflection, and data-dependent Python control can have effects that a tensor graph cannot safely reorder or reproduce.
The graph inputs and outputs define a region boundary. Values used only inside the region can be candidates for fusion or buffer reuse. Values that cross the boundary must be materialized in an agreed representation. Compiling a larger stable region can expose more optimization, but it also gives capture more Python behavior and shape variation to handle.
Guards make a captured graph conditional
Generated code is valid under assumptions. A guard can require a tensor datatype, device, rank, stride pattern, dimension value, or a Python constant observed during tracing. Before reusing a cached graph, the runtime checks these assumptions. If they hold, execution can skip tracing and compilation.
If a guard fails, the system can compile another variant or fall back. Repeated input variation can create repeated compilation and retain several generated modules in memory. Log the failing guard instead of treating all recompilation as one problem. A guard on an incidental Python value requires a different fix from a necessary specialization on matrix alignment.
Graph breaks divide the optimization region
A graph break occurs when capture reaches behavior that it cannot represent under the current mode. The system compiles the graph accumulated before the break, executes the unsupported portion outside that graph, and may resume capture afterward. The program can remain correct while losing fusion and adding dispatch at both boundaries.
Data-dependent Python branches, scalar extraction followed by Python logic, unsupported C calls, and side effects are common causes. Some behavior is fundamentally unsuitable for capture. Other behavior can be written with supported tensor control or moved outside the repeated compiled region. Use full-graph diagnostics or graph-break logs to identify the source line and reason before rewriting code.
Dynamic shapes and specialization
A symbolic dimension represents a family of runtime values and carries constraints such as bounds or divisibility. One symbolic graph can serve several sequence lengths, reducing recompilation. The generated schedule may be less specialized, and some operations still require guards or separate variants.
Static specialization can produce efficient fixed-size code for common shapes. Specialization becomes expensive when every request length creates a new variant. Bucketing or padding can reduce the number of shapes at the cost of extra work. Compare compile frequency, cache memory, padding work, and steady execution under the real length distribution before selecting a policy.
Effects, aliasing, and custom calls
The compiler may reorder or eliminate operations only when program effects allow it. Mutation and aliasing can create hidden dependencies: writing through one tensor view can change another view that shares storage. Randomness requires an explicit state model. A custom operator must declare mutation and output metadata so capture does not assume unsafe independence.
Successful capture does not guarantee strong optimization. An opaque custom call can remain as one graph node yet block fusion across its boundary. Conservative alias assumptions can prevent buffer reuse. Inspect the captured graph and generated launches, not only whether torch.compile returned without an exception.
Worked example: Finding a decode graph-break storm
A compiled decoder is fast after warm-up but periodically pauses and produces many compiled variants.
- Inspect capture logs and graph-break reasons. A Python branch depends on current sequence length and extracts a tensor scalar each step.
- Record which calls compile and which calls reuse a cached graph. Associate each pause with a specific failed guard or break rather than with compilation in general.
- Rewrite the decision with supported symbolic or tensor operations, or move it outside the repeated region if its value changes only at known boundaries.
- Inspect guards and variant keys. Bucketing cache lengths or padding to stable blocks may reduce recompilation while preserving efficient specialized shapes.
- Replay the real length distribution and measure compile time, cache hit rate, launch gaps, and steady device performance separately.
- Compare a symbolic-shape version with a small set of padded buckets. The preferred policy is the one that minimizes total request cost under the workload, not the one with the fewest graph variants in isolation.
Conclusion: The bottleneck was not generated kernel quality but instability in the captured program contract. Stabilizing shapes and control restored reuse of optimized regions.
Common errors
- Counting graph breaks without weighting their frequency and location. One break inside every decode step matters more than several in startup.
- Forcing fully dynamic shapes. Generality can prevent specialization and create slower code even when it avoids recompilation.
- Assuming capture means optimization. The graph may remain fragmented by opaque effects or lower to weak schedules.
- Suppressing compiler errors before reading the break or guard reason. Silent fallback can hide the exact region that needs attention.
Reference summary
A dynamic framework must discover a graph from ordinary program execution. The capture system records operations and assumptions called guards: types, shapes, values, or control-flow facts that make the captured graph valid. When a guard fails, the system may compile another variant. Unsupported or data-dependent behavior can cause a graph break, dividing optimization regions and reintroducing eager overhead.
- Inspect graph breaks instead of silently suppressing compiler errors.
- Track recompilation count and cache behavior across the real shape distribution.
- Avoid unnecessary Python-side value specialization.
- Use custom operators as explicit semantic boundaries for opaque kernels.
- Compile at a stable, sufficiently large region rather than every tiny helper.
Knowledge check
How can dynamic sequence lengths create both performance loss and memory growth through recompilation?
Show the answer guide
- Graph capture records shape assumptions as guards. When a new sequence length violates a guard, the runtime can recompile a new graph or fall back to eager execution, adding CPU work and a latency pause.
- A stream of many distinct lengths can create a recompilation storm. Kernel launches become fragmented across variants, compile time accumulates, and cache lookup or eviction adds overhead.
- Each compiled variant can retain generated code, constants, autotuning results, CUDA graphs, and shape-specific memory pools or workspaces. The process memory footprint therefore grows even when only one request is active at a time.
- Dynamic-shape support or length bucketing can reduce variants, but broader kernels may run less efficiently or pad extra tokens. The correct policy is determined by compile events, cache memory, and steady-state latency for the actual length distribution.
Observe recompilation under a variable-length trace
Compile a PyTorch attention or MLP function and call it with sequence lengths [128, 129, 256, 257, 512, 513] repeated in a fixed trace. Compare fully specialized capture, length buckets rounded to powers of two, and a dynamic-shape configuration. Enable compiler logs and record process and GPU memory after each first-seen length.
Evidence to keep: Submit the program and exact shape trace, compiler logs that count guards, graph breaks, and compilations, a latency trace marking cold and warm calls, and a plot of compiled variants and memory over time. Conclude which mechanism causes each latency spike and each persistent memory increase.
Captured graphs enter intermediate representations that preserve different kinds of knowledge as the compiler moves from tensors toward target instructions.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.