GraphMend: Fixing Graph Breaks in PyTorch 2

• GraphMend is a source-level compiler that mitigates PyTorch 2 FX graph breaks by rewriting dynamic control flow and Python I/O side effects.
• It leverages Jac-based AST transformations to convert branch-dependent operations into predicated execution with torch.where and defers side effects to an epilogue.
• Evaluations on multiple Hugging Face models show significant latency improvements (30%-75% cold-start) and reduced CUDA graph captures, enhancing overall performance.

GraphMend is a source-level compiler transformation system for PyTorch 2 that eliminates FX graph breaks before TorchDynamo sees the program. It was introduced as a high-level compiler built on the Jac compilation framework, with the specific aim of rewriting graph-breaking source constructs into forms that PyTorch’s just-in-time compilation pipeline can trace as larger, uninterrupted FX graphs. In its primary sense, the term denotes the compiler described in "GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2" (Kashmira et al., 17 Sep 2025). In adjacent literature, closely related “graph mending” ideas appear in other senses, including interoperability repair between RDF and property graphs and graph generation or prediction framed as conditional generation; those usages are conceptually related but distinct from the PyTorch 2 compiler system.

1. Problem setting: FX graph breaks in PyTorch 2

PyTorch 2 uses TorchDynamo to intercept Python execution and symbolically trace PyTorch operations into an FX graph, and TorchInductor as the default backend compiler to optimize and lower FX graphs to GPU code. A graph break occurs when TorchDynamo encounters unsupported or hard-to-trace Python code and must stop graph capture. The forward pass is then fragmented into multiple FX graphs, and the unsupported segment is executed in eager mode (Kashmira et al., 17 Sep 2025).

The paper identifies two major costs of graph breaks. First, every fallback to eager mode forces control to move from GPU execution back to the Python interpreter, creating CPU/GPU synchronization overhead, device-to-host transfers, and scheduling overhead. Second, each FX fragment is compiled separately, so cross-fragment fusion and global optimization opportunities are lost. The resulting performance penalty is therefore both runtime and compiler-structural: execution is interrupted, and the optimization horizon is shortened.

The immediate causes discussed in the paper fall into two broad categories. One is data-dependent operations and dynamic control flow, including if statements, loops, and branches that depend on tensor values, as well as data accesses such as .item() or .data_ptr(). The other is unsupported Python built-ins and I/O, including print, logging, and warnings. TorchDynamo traces only operations it can safely model; when it cannot prove a condition statically or reason about a side effect, it inserts a break.

A central diagnosis is that TorchDynamo operates at the Python bytecode level rather than the source level. The paper emphasizes that two syntactically different functions can compile to the same bytecode, so higher-level intent is unavailable once the program has been lowered. This is the motivation for moving repair earlier in the compilation pipeline.

2. Compiler architecture and analysis pipeline

GraphMend’s defining design choice is to repair graph breaks at source or AST level, where semantic and structural information is still present, instead of reacting only after TorchDynamo encounters a break. The system analyzes the program beforehand, detects code patterns likely to produce graph breaks, rewrites them into graph-compatible forms, and then feeds the transformed program into the standard PyTorch 2 pipeline (Kashmira et al., 17 Sep 2025).

The implementation is built on Jac, a compiler and runtime framework that accepts Python and Jac code and lowers it through a structured intermediate representation. Jac provides an AST, a symbol table, a CFG, and a unified IR called UniiR. The paper states that Jac merges structural and control-flow information into a single representation, which makes higher-level program analysis and rewriting easier than bytecode-only tracing.

GraphMend adds three passes to the Jac pipeline. The Dynamo Entry Point Analysis Pass finds functions or nn.Modules decorated or wrapped with torch.compile; these become candidate entry points for graph-break analysis. The Graph Break Type Analysis Pass traverses the CFG from each entry point and detects dynamic control-flow breaks and print or logger.* side-effect breaks. The AST Transformation Pass rewrites the source using one of two transformations, after which the symbol table and CFG are rebuilt, the program is emitted back as Python bytecode, and CPython executes it normally.

The paper also describes a fixable graph-break detection algorithm. Starting from each Dynamo entry point in UniiR, the compiler performs CFG traversal. When it encounters an IfStmt, it checks the condition; if the condition contains a torch attribute known to be dynamic, such as .sum, or if use-def and symbol-table analysis show that the condition depends on input tensors, the node is tagged as GraphBreak[DynCtrlFl]. If it encounters a call to print or logger.*, it is tagged as GraphBreak[logger/print]. These tags determine which transformation is applied.

This architecture makes GraphMend a preprocessing compiler layer rather than a replacement for PyTorch 2. TorchDynamo and TorchInductor remain the runtime compiler stack; GraphMend complements them by making source programs more traceable before dynamic JIT compilation begins.

3. Core transformations

GraphMend focuses on two common and fixable sources of graph breaks: dynamic control flow based on tensor values and Python I/O side effects. Both rewrites are source-to-source transformations intended to preserve semantics while allowing TorchDynamo to capture a single larger FX graph (Kashmira et al., 17 Sep 2025).

The first transformation is Predicated Dynamic Control Flow. Its target is a branch whose outcome depends on a runtime tensor value, such as:

if x.sum() > 10: z = x_1 + y_1 else: z = x_1 * y_1

GraphMend rewrites this control flow into predicated tensor computation using torch.where:

cond = x.sum > 10 z_add = x_1 + y_1 z_mul = x_1 * y_1 z = torch.where(cond, z_add, z_mul)

The paper describes the rewrite as evaluating the condition once, storing it as a predicate, computing both branch results, and selecting with torch.where. The purpose is to transform dynamic control flow into predicated execution that remains GPU-compatible and traceable by Dynamo. In effect, the break disappears because the branch is no longer represented as Python control flow.

The second transformation is Graph-Epilogue Deferred Side Effects. Its target is a side effect such as print("tensor:", x) inside a compiled function. Because printing is a Python side effect, it forces a break when it appears in the middle of a traced region. GraphMend stores the message in a temporary variable, moves the side effect to the epilogue of the function, and performs the print or logging only at the end, after graph-executable computation is complete. The paper notes an important constraint: if the return value depends on values that would otherwise be delayed, those computations must be moved before the epilogue and stored in a temporary to avoid creating a new break.

The paper provides a representative logging rewrite:

@torch.compile def fn(x): x = torch.relu(x) to_print = "tensor:", x y = torch.sin(x) print(to_print) return y

The computational path remains contiguous for tracing, while the side effect is preserved semantically but displaced outside the central traced region.

A key feature of both transformations is that they act before runtime tracing. This suggests that GraphMend is motivated less by post hoc recovery than by proactive normalization of source constructs into graph-friendly forms.

4. Scope, assumptions, and limitations

GraphMend is deliberately scoped. The paper states that it fixes graph breaks caused by dynamic control flow based on tensor values and Python I/O side effects, especially print, logging, and warnings (Kashmira et al., 17 Sep 2025). These are the two break classes explicitly targeted by its rewrite rules.

The system assumes that graph breaks can be identified from source structure together with control-flow and data-flow analysis, that the problematic code belongs to one of the common fixable categories, and that the transformed code preserves semantics. Because the approach is source-level and AST-based, it also depends on source availability and Jac-based compilation.

The paper is equally explicit about what GraphMend does not solve. It does not fix tensor.item(), dynamic-shape operators, or unsupported constructs beyond its two rewrite rules. It is not presented as a general runtime repair mechanism for arbitrary bytecode fragments. In the paper’s framing, many graph breaks are not fundamentally unavoidable, but the system does not claim that all graph breaks are fixable by source-to-source transformation.

A common misconception is that GraphMend replaces PyTorch 2’s compilation stack. The paper rejects that interpretation: GraphMend complements TorchDynamo and TorchInductor rather than replacing them. Another misconception is that it repairs graph breaks by operating at the same level as TorchDynamo. Its distinguishing claim is precisely that bytecode-level tracing loses source-level intent, so the repair must occur earlier.

5. Evaluation and observed effects

The evaluation is based on a benchmark constructed by randomly sampling 65 Hugging Face models, filtering for models that trigger FX graph breaks, and retaining 8 models for study: biogpt, blenderbot-400M-distill, flan-t5-large, longformer-base-4096, moe-minicpm-x4-base, Phi-4-mini-instruct, Qwen-Audio-Chat, and tiny-random-PegasusForCausalLM (Kashmira et al., 17 Sep 2025).

The original graph-break counts and causes are reported as follows:

Model	Original breaks	Cause
biogpt	2	logger calls
blenderbot-400M-distill	3	logger calls
flan-t5-large	3	logger calls
longformer-base-4096	5	logger calls, tensor.item()
moe-minicpm-x4-base	15	dynamic shape operator
Phi-4-mini-instruct	5	dynamic control flow
Qwen-Audio-Chat	2	dynamic control flow
tiny-random-PegasusForCausalLM	2	logger calls

Experiments were run on NVIDIA RTX 3090 and NVIDIA A40 GPUs with TorchInductor as backend. Profiling used PyTorch Profiler with CPU activity, CUDA activity, and operator-level events; for each model and configuration, there were 7 iterations total, comprising 1 cold-start iteration and 6 warm iterations, and traces were exported for inspection with Chrome trace viewer.

The reported fix rates were 100% fixed for biogpt, blenderbot-400M-distill, flan-t5-large, Phi-4-mini-instruct, Qwen-Audio-Chat, and tiny-random-PegasusForCausalLM; 40% fixed for longformer-base-4096, reducing breaks from 5 to 2; and 0% fixed for moe-minicpm-x4-base, where dynamic-shape operator breaks remained. The abstract summarizes the aggregate outcome by stating that GraphMend removes all fixable graph breaks due to dynamic control flow and Python I/O functions, drives the break count to 0 in 6 models, and reduces it from 5 to 2 in another model.

Latency improvements are reported separately for cold-start and steady-state execution. Cold-start latency is lower by about 30% to 75% on both RTX 3090 and A40, which the paper attributes to the fact that graph breaks cause multiple separate CUDA Graph recordings and caches, whereas fixing breaks reduces the number of captures and idle GPU time. Steady-state latency improves by about 2.5% to 25%; the smaller gain is attributed to graph-capture cost being amortized after warm-up even though graph breaks still impose CPU/GPU switching and scheduling overhead.

Throughput gains are reported as about 5% to 8% across models on both GPUs. The paper notes model-specific variation: Qwen-Audio-Chat shows about 7.5%–8%, while Phi-4-mini-instruct shows about 5%–6%, because the Qwen break occurs in the middle of the forward pass and is therefore more disruptive than the Phi-4-mini-instruct break nearer the beginning. For tiny-random-PegasusForCausalLM, latency improves by more than 25% but throughput gains are smaller, because throughput depends not only on GPU compute but also on prefill, decode, CPU-side sampling and tokenization, and data transfer.

The paper also includes a case study of Qwen-Audio-Chat on A40. In steady state, the original model executes as three CUDA graphs, with idle GPU time between graphs while the CPU performs eager fallback and D2H memcpy; after GraphMend rewrites the dynamic branch into torch.where, execution becomes one continuous CUDA graph, CPU fallback disappears, and GPU activity is uninterrupted. In cold start, the original model incurs additional CUDA Graph recording and caching events at each break; after GraphMend, a single capture suffices. A kernel-level analysis on Phi-4-mini-instruct reports 404 kernels in the original model and 393 kernels after repair, an 11 kernel reduction, which the paper interprets as evidence of more fusion, fewer launches, and better scheduling or reordering.

6. Relation to broader “graph mending” concepts

Although GraphMend most directly refers to the PyTorch 2 compiler system, the supplied literature places it within a broader conceptual family of graph-mending tasks. One neighboring line of work concerns interoperability repair between graph data models. "Mapping RDF Graphs to Property Graphs" defines G2GML (Graph to Graph Mapping Language) and a converter called G2G Mapper for transforming RDF graphs into property graphs that can be loaded into engines such as Neo4j, Oracle Labs PGX, and Amazon Neptune. That paper does not use the term GraphMend by name, but it is described as conceptually close because it “mends” an ecosystem gap by converting RDF triples into a property-graph representation suitable for traversal and analytics (Matsumoto et al., 2018).

A second neighboring line appears in graph diffusion. "Toward a Unified Geometry Understanding: Riemannian Diffusion Framework for Graph Generation and Prediction" introduces GeoMancer, which is presented as a Riemannian reformulation of the latent graph diffusion idea behind GraphMend-style graph generation and prediction. In that usage, “graph mending” refers not to compiler repair but to reconstructing or predicting missing graph attributes or labels by treating prediction as conditional generation, with node-, edge-, and graph-level signals decoupled onto task-specific manifolds rather than embedded in a single Euclidean latent space (Gao et al., 6 Oct 2025).

These neighboring usages are important because they show that “graph mending” is not confined to one substrate. In the compiler setting, the object being repaired is an execution graph fragmented by unsupported source constructs. In RDF-to-property-graph conversion, the repair target is interoperability between graph models. In GeoMancer, the target is the conditional reconstruction or prediction of graph properties under a geometry-aware latent representation. A plausible implication is that the term functions as a family resemblance label for techniques that restore continuity, compatibility, or inferential completeness in graph-structured systems, even when the underlying technical objects are different.

Within that broader family, the defining feature of GraphMend proper remains narrower and more concrete: it is a Jac-based source-to-source compiler pass system for eliminating specific PyTorch 2 FX graph breaks by rewriting dynamic control flow into torch.where-based predication and deferring Python side effects to a function epilogue.