EVT: Accelerating Deep Learning Training with Epilogue Visitor Tree

Zhaodong Chen, Andrew Kerr, Richard Cai, Jack Kosaian, Haicheng Wu, Yufei Ding and 1 other

Deep-learning compilers make trained models run fast, mostly by fusing many small operations into one so intermediate results never leave the GPU's fast memory — but this works far better for inference than for training, where the math is messier, the loss function goes unoptimized, and the tangle of forward-and-backward connections defeats the tools that decide what to fuse; this work from UC Santa Barbara and NVIDIA builds EVT, which keeps NVIDIA's expert-tuned matrix-multiply core but wraps it in a flexible, compiler-assembled tail, adds graph rewrites for the loss, and solves the fusion-grouping problem exactly, speeding up real training by up to roughly three times.

The argument in one glance

Problem

Compilers speed up inference well, but training resists

Deep-learning compilers fuse many small operators into one so intermediate results stay in the GPU's fast memory. This works for inference, but training is harder: operator compilers are fast or flexible but not both, loss functions go unoptimized, and the extra forward-to-backward edges in the training graph defeat the tools that decide what to fuse.
What people did

Reuse existing operator compilers and partitioners

Apply loop-based compilers for flexibility or template-based compilers for speed, and partition the training graph with the usual heuristics to decide which operators to fuse together.
Why it fell short

Each choice sacrifices speed, flexibility, or feasibility

Loop-based compilers cannot match expert kernels for matrix multiplication; template-based ones cannot handle diverse training fusion patterns. Loss functions stay unoptimized. And heuristic partitioning of the training graph, with its forward-to-backward edges, produces infeasible cyclic groupings or arbitrarily separates forward from backward, losing cross-pass fusion.
This paper's idea

Expert core, flexible epilogue, and an exact partitioner (EVT)

EVT keeps the expert-tuned matrix-multiply core of a GPU kernel but makes its epilogue flexible via a visitor tree of individually-built operators the compiler assembles, giving state-of-the-art speed and diverse fusion at once. Graph-level passes fuse the loss with its backward and reshape operators to be fusible, and an integer-linear-programming partitioner finds a feasible, globally-optimal fusion grouping. On diverse training workloads, EVT delivers a 1.26 to 3.1 times speedup, and is upstreamed into CUTLASS and usable from torch.compile.

Why this problem is worth solving

The gap between how models are written and how fast the hardware could run them is exactly what compilers exist to close.

Deep-learning frameworks like PyTorch are built for flexible, easy programming, which means a model as written often uses the GPU far below its potential, and hand-optimizing it demands rare expertise. Deep-learning compilers automate this. A typical one optimizes a model in three stages: it first traces the model into a computation graph, where nodes are operators and edges are the flow of data; then a graph-level stage simplifies the graph, for example by folding constants or eliminating repeated sub-expressions; a partitioner divides the graph into groups; and finally an operator-level stage turns each group into an optimized, fused kernel. The workhorse optimization is kernel fusion: when one operator feeds another, fusing them into a single kernel lets the intermediate result stay in the GPU’s fast on-chip memory — registers or shared memory — instead of being written out to slow global memory and read back. Because moving data to and from global memory is a dominant cost, fusion can dramatically cut the number of memory accesses and speed a model up. The quality of a fused kernel depends on how well each operator is implemented and how flexibly the compiler can align their loop structures.

Inference compilers do this well; training breaks all three stages at once.

Compilers have delivered excellent speedups for inference, but training is substantially harder, and the paper identifies three limitations. First, no operator compiler is both fast and flexible. Loop-based compilers, like TVM, represent operators as nested loops and are flexible, but they struggle to match the performance of expert-tuned kernels for the heavy core operations — matrix multiplication and convolution — because they cannot fully exploit the specialized accelerator features and datapaths of modern GPUs. Template-metaprogramming compilers, which build on expert-developed template libraries like NVIDIA’s CUTLASS, hit top performance but lack the flexibility to support the diverse fusion patterns that training graphs demand. Second, loss functions and their backward passes are largely unoptimized, yet in settings like extreme classification they can become the bottleneck. Third, partitioning the graph — deciding which operators to group and fuse — is much harder for training than inference. Inference is a stream of operators that simple heuristics can split, but a training graph has additional edges that pass saved activations from the forward pass to the backward pass for computing gradients. Heuristically partitioning these can easily produce an infeasible result containing cycles, while arbitrarily separating the forward and backward graphs hides the chance to fuse operators across them. And existing partitioners are tied to a specific operator compiler, so they cannot easily take on new constraints.

What people did before, and where it broke

Use a loop-based operator compiler for flexibility. Represent operators as nested loops that the compiler can fuse in many patterns. Where it broke. They cannot match expert-tuned kernels for matrix multiplication and convolution on modern GPUs, so core operations run slowly.
Use a template-based operator compiler for speed. Build on expert template libraries like CUTLASS to reach peak kernel performance. Where it broke. Their epilogues are rigid, so they cannot accommodate the diverse fusion patterns that training graphs require.
Partition the training graph with heuristics. Split the graph into groups for fusion using simple rules, as inference compilers do. Where it broke. The forward-to-backward edges make heuristic groupings infeasible (they form cycles) or force an arbitrary forward/backward split that loses cross-pass fusion, and the partitioners cannot take on new constraints.
The gap. No compiler combined state-of-the-art kernel performance with flexible training fusion, optimized the loss, and found feasible optimal partitions of the training graph. Providing all three is what EVT does.

How EVT works

From fast-or-flexible training compilation to both, with exact partitioning

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Keep the expert-tuned core, make the flexible part compiler-assembled. EVT’s central abstraction resolves the fast-or-flexible dilemma by splitting a GPU kernel into two parts: the heavy matrix-multiply mainloop, kept exactly as the expert-optimized library implementation so all its dedicated GPU optimizations carry over, and the epilogue — the post-processing applied to the result — made flexible through a visitor design in which each epilogue operator is developed and optimized individually and then assembled by the compiler into a tree for whatever fusion pattern a partition needs. Pairing an expert hand-tuned fast path with a compiler layer that composes and fuses around it, so peak performance and flexibility coexist, is the same design the encrypted-computation acceleration work covered elsewhere on this site follows, combining hand-optimized low-level GPU kernels with a compiler that fuses and composes them to fit each workload.

Reshape the graph so more of it can stay on-chip. Around the kernel abstraction, EVT applies graph-level passes that expose fusion the abstraction alone could not reach: it fuses each loss function with its backward pass by recognizing the loss value is not needed downstream, breaks complex operators like batch normalization into fusible pieces, and rewrites reductions into equivalent fusible forms — all so that more intermediate results can be kept in fast memory instead of spilling to global memory. Restructuring a tensor computation so that data stays local to the compute and is not repeatedly moved through slow memory is the same locality-driven transformation the tensor-dataflow compiler work covered elsewhere on this site performs, reorganizing the computation to keep intermediates close so the memory system stops being the bottleneck.

Turn partitioning into an exact optimization problem. For deciding which operators to group and fuse — hard because the training graph’s forward-to-backward edges make naive groupings cyclic and infeasible — EVT models partitioning as an integer linear program, so a solver finds a feasible, globally-optimal grouping and can absorb extra constraints, including the abstraction’s own fusion rules, with graph-splitting techniques to keep the NP-hard problem tractable at scale. Encoding a hard structural problem as an integer linear program and letting a solver return the provably optimal solution, rather than settling for a heuristic, is the same strategy the traversal-synthesis work covered elsewhere on this site takes, casting its problem as an integer program so a solver can find the best arrangement instead of a hand-guessed one.

Following a training step through EVT

Compiling a model's training step into fast fused kernels

A trace of EVT compiling a training workload, following the paper's system design.

Step 1 of 5

Build the joint forward-backward graph

EVT takes a PyTorch model and turns it into a joint forward-backward graph, whose inputs are the training data and model parameters and whose outputs are the gradients.

Both passes are in one graph so fusion can span them.

Why it works

Fusion’s value is avoided memory traffic. Because fusing operators keeps their intermediate results in fast on-chip memory rather than writing them to slow global memory → the kernel does far fewer memory accesses → so a fused kernel runs much faster than the separate operators.
Expert cores and compiler flexibility need not be exclusive. Because the heavy matrix-multiply mainloop is expert-tuned while the epilogue is assembled by the compiler → the kernel keeps state-of-the-art performance and gains flexible fusion → so training’s diverse patterns are handled without sacrificing speed.
Loss functions hide fusable work. Because the loss value does not participate in the backward pass → the loss and its backward can be fused once that is recognized → so a bottleneck prior compilers left alone becomes an optimization.
Training graphs demand feasible partitions. Because the forward-to-backward edges let naive partitioning create cycles or miss cross-pass fusion → modeling partitioning as an integer linear program finds a feasible, globally-optimal grouping → so partitions are both valid and maximally fused.
The optimum can be reached at scale. Because acyclic partitioning is NP-hard and networks have millions of variables → splitting the graph into smaller problems keeps the solver tractable → so the global optimum is found in minutes rather than abandoned for a heuristic.

What it costs (trade-offs and scope)

The epilogue abstraction still constrains what can be fused. The visitor-tree epilogue supports only its set of epilogue operators, and fusing reductions with their consumers in the epilogue is difficult and sometimes not beneficial. The graph-level passes exist partly to work around these limits, but the constraints are real.
Optimality has a cost. Acyclic partitioning is NP-hard, so even with the graph-splitting simplifications, solving the integer program takes real time — described as minutes for networks with millions of variables — which is compile-time cost rather than run-time.
It leans on expert kernel libraries. EVT’s peak performance comes from reusing expert-optimized mainloops from CUTLASS and similar libraries. Its strength is tied to the availability and quality of those library implementations for the target hardware.
It targets GPU training with matrix-multiply-heavy partitions. The design centers on the mainloop-epilogue structure of heavy operations like matrix multiplication on modern GPUs; workloads or hardware that do not fit that structure benefit less.
The speedups are workload- and baseline-specific. The 1.26 to 3.1 times range comes from specific training workloads across domains measured against specific baselines; other models and hardware would shift the numbers.
It is a compiler pass to adopt. Realizing the benefit means running a model through EVT (usable via torch.compile); it is a real toolchain component, though its integration into CUTLASS and torch.compile lowers the barrier.

The evaluation

Three ways to compile training kernels

Property	Loop-based compiler	Template-based compiler	EVT
State-of-the-art GPU kernel performance
Flexible fusion patterns for training
Optimizes loss and its backward	No	No
Feasible, optimal partitioning		Heuristic
Outcome	Flexible but slow kernels	Fast but rigid kernels	Fast and flexible, 1.26 to 3.1 times speedup

Three ways to compile fused kernels for deep-learning training, as framed by the paper. Loop-based compilers (like TVM) represent operators as nested loops and are flexible, but cannot match expert-tuned kernels for the heavy matrix multiplications at a model's core. Template-based compilers (like AITemplate, using CUTLASS) reach peak kernel performance but have rigid epilogues that cannot handle training's diverse fusion patterns. EVT gets both: its Epilogue Visitor Tree keeps the expert-optimized matrix-multiply mainloop and assembles a flexible epilogue from individually-built operators. It adds graph-level passes (loss elimination, decomposition, reduction elimination) that optimize the loss and expose more fusion, and an integer-linear-programming partitioner that finds a feasible, globally-optimal grouping of the joint forward-backward graph, scaling to millions of variables in minutes via graph-splitting. Across diverse training workloads, EVT achieves a 1.26 to 3.1 times speedup. It is open-source, upstreamed into NVIDIA's CUTLASS library, and usable as a torch.compile backend with a single line of code.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

EVT is a strong, practical compiler contribution that attacks training optimization on all three fronts where it had been stuck. Its central idea elegantly dissolves the fast-versus-flexible dilemma: keep the expert-optimized matrix-multiply mainloop that gives state-of-the-art GPU performance, and make the epilogue — the flexible tail where training’s diverse fusion patterns live — a compiler-assembled tree of individually-built operators, so speed and flexibility no longer trade off. Around that abstraction, the graph-level passes are well-chosen: loss elimination optimizes a genuinely neglected bottleneck by noticing the loss value does not feed the backward pass, and decomposition and reduction elimination reshape awkward operators into fusible ones. And the integer-linear-programming partitioner is the right tool for the training graph, whose forward-to-backward edges make heuristic grouping infeasible; modeling it exactly finds feasible, globally-optimal, cross-pass fusions, with graph-splitting to keep the NP-hard problem solvable in minutes. The results — a 1.26 to 3.1 times speedup across diverse training workloads — are backed by real adoption: the work is open-source, upstreamed into NVIDIA’s CUTLASS, and usable directly from torch.compile. Read the scope precisely. The epilogue abstraction still limits which patterns are fusible, which is why the graph passes are needed; the exact partitioning costs compile-time minutes; peak performance leans on the quality of expert libraries like CUTLASS; and the approach centers on matrix-multiply-heavy GPU training, with workload- and baseline-specific numbers. The durable idea is that the fast-or-flexible choice in code generation is a false one — a system can retain an expert-tuned core and wrap it in compiler-generated flexibility — and that the messy structural decisions around it, like how to partition a training graph, are often best handed to an exact solver rather than a heuristic.

Background

Why co-designing the whole compiler stack around the hardware keeps winning

It helps to see EVT within a broad theme in machine-learning systems: the largest speedups come not from one clever kernel but from co-designing the whole compilation stack — how the graph is rewritten, how it is partitioned, and how each kernel is generated — around what the hardware does efficiently, above all keeping data movement low. EVT’s three parts reinforce each other: the graph passes expose fusion, the partitioner groups operators so that fusion is feasible and maximal, and the operator compiler realizes each group as a fast fused kernel that keeps intermediates on-chip. The device-placement and partitioning work covered elsewhere on this site pursues the same instinct at a different scale, deciding how to split a training computation into pieces that map well onto the hardware, so that the arrangement itself — not just the individual operations — is what unlocks efficiency. The common thread is that a deep network’s training is a large graph with enormous freedom in how it is grouped, transformed, and mapped down to kernels, and treating those decisions jointly, with the hardware’s real costs in view, extracts far more than optimizing any one stage alone. As models keep growing and training costs mount, compilers that co-design graph transformation, partitioning, and code generation — and that reuse expert-tuned cores while adding automated flexibility around them — are likely to be central to making training affordable.

If you remember only three things

1. Deep-learning compilers speed up models mainly through kernel fusion: merging producer-consumer operators into one kernel so intermediate results stay in the GPU's fast on-chip memory instead of round-tripping through slow global memory. This works well for inference but training resists, for three reasons: no operator compiler is both fast and flexible (loop-based ones like TVM cannot match expert kernels for matrix multiplication, while template-based ones on CUTLASS reach top speed but have rigid epilogues); loss functions and their backward passes go unoptimized yet can dominate runtime; and the training graph's extra edges (saving activations from forward to backward) make heuristic partitioning either infeasible (cyclic) or force an arbitrary forward/backward split that loses cross-pass fusion.
2. EVT addresses all three. Its core Epilogue Visitor Tree abstraction keeps the expert-optimized matrix-multiply mainloop of a GPU kernel (so all its CUDA optimizations carry over and performance stays state-of-the-art) while making the epilogue flexible through a visitor design: each epilogue operator is built and optimized individually, then the compiler assembles them into a tree for whatever fusion pattern a partition needs — resolving the fast-or-flexible dilemma. Graph-level passes expose more fusion: loss elimination fuses each loss with its backward pass (the loss value is not needed downstream), decomposition breaks operators like batch normalization into fusible pieces, and reduction elimination rewrites reductions into fusible forms. An integer-linear-programming partitioner models the joint forward-backward graph and finds a feasible, globally-optimal fusion grouping, with graph-splitting to solve millions of variables in minutes.
3. Across diverse training workloads, EVT achieves a 1.26 to 3.1 times speedup, and it is real, adopted infrastructure: open-source, upstreamed into NVIDIA's CUTLASS library, and usable as a torch.compile backend with one line of code. Scope: the epilogue abstraction still limits which patterns are fusible (so the graph passes are needed), the exact partitioning costs compile-time minutes, peak performance depends on expert kernel libraries like CUTLASS, and the design centers on matrix-multiply-heavy GPU training with workload-specific numbers. The durable lesson is that the fast-versus-flexible choice in code generation is a false dilemma — keep an expert-tuned core and wrap it in compiler-generated flexibility — and that hard structural decisions like partitioning a training graph are often best handed to an exact solver rather than a heuristic.

Check yourself

Check your understanding

1. What is kernel fusion, and why does it speed up deep-learning models?
2. Why is compiling training harder than compiling inference?
3. What is the key idea of the Epilogue Visitor Tree abstraction?
4. Why does EVT use an integer linear program to partition the training graph?

Sources

The paper is “EVT: Accelerating Deep Learning Training with Epilogue Visitor Tree” by Zhaodong Chen, Andrew Kerr, Richard Cai, Jack Kosaian, Haicheng Wu, Yufei Ding, and Yuan Xie, published at ASPLOS 2024, DOI 10.1145/3620666.3651369. The authors are at UC Santa Barbara and NVIDIA.

This explainer is based on the full text of the paper, obtained from an archived copy of the publisher’s open-access (CC-BY) version of record. The problem framing (deep-learning compilers and their three-stage optimization of graph-level optimization, partitioning, and operator-level optimization; kernel fusion keeping intermediate results in fast on-chip memory; and the three limitations for training — operator compilers being fast or flexible but not both, unoptimized loss functions, and infeasible or shallow partitioning of the training graph due to its forward-to-backward edges), the EVT design (the Epilogue Visitor Tree abstraction pairing an expert-optimized mainloop with a compiler-assembled visitor-tree epilogue; the graph-level passes of loss elimination, decomposition, and reduction elimination; the integer-linear-programming partitioner with graph-splitting simplifications for the NP-hard acyclic partitioning problem; and the EVT operator compiler generating fused kernels using CUTLASS and other libraries), and the results (a 1.26 to 3.1 times speedup across diverse training workloads, open-source availability, upstreaming into NVIDIA’s CUTLASS, and usability as a torch.compile backend) are all drawn directly from the paper. Background explanations of computation graphs, kernel fusion, matrix multiplication on GPUs, the forward and backward passes, and integer linear programming are standard and provided to make the paper’s contribution understandable.