AutoFuser

Many operations in a neural network are limited by memory bandwidth rather than arithmetic: each one reads its inputs from GPU memory and writes its result back, only for the next operation to read it again. Fusing a chain of operations into one kernel keeps the intermediate values on chip. AutoFuser does this automatically for a PyTorch model, in the same spirit as torch.compile: it finds fusable chains in the model’s graph, generates a Triton kernel for each, tunes it for the actual input shapes, and rewrites the model to call the fused kernels.

I proposed the project and its pipeline for the Advanced Deep Learning Systems course at Imperial, and built it with Ching Bon Tang, Bruno Duaso and Oliver Hannibal. We used Imperial’s MASE framework for the graph representation and evaluated on scientific machine learning models such as neural operators and transformer-based PDE solvers.

Exposing the whole graph

Models are built from nested modules, and a fusion opportunity often spans a module boundary. I wrote the pass that recursively traces each submodule and splices it into the parent graph, so the fusion analysis sees individual operations instead of opaque module calls. Submodules that cannot be traced (FFT-based layers, for example) are left intact. I also patched einops’ rearrange, repeat and reduce so that models using them could be traced at all.

Generating kernels

Before the pipeline could generate anything, I wrote fused Triton kernels by hand for the operations we cared about (an MLP, layer normalisation, a convolution and an upsampling block) to give the generated code a target. I then wrote the code generator for chains that contain reductions and matrix multiplications: it lowers a chain of linear layers, layer norms, softmax and MLP blocks into a single row-tiled Triton kernel, with its own block sizes for each operation. On the transformer model, the automatically fused kernels matched the performance of my hand-written ones. I also wrote the dispatch layer and most of the per-operation lowering handlers that the code generators share.

Tiling and autotuning

Each kind of chain gets a tiling strategy that defines its tunable block sizes, its launch grid and which configurations are valid: elementwise, row-wise reductions, layer norm, matrix multiply, linear plus activation, layer norm plus linear, and full MLP blocks. The autotuner takes the generated kernel source as a string, loads it into a temporary module, generates candidate configurations from the strategy, and uses Triton’s own benchmarking to pick the fastest. Results are cached on disk, so a model only pays for tuning once.

Putting it together

I wrote the graph rewrite that replaces each fused chain with a single call, checking that fusion groups do not overlap, handling kernels with several outputs, and removing the nodes that are no longer used. The orchestrator ties the stages together behind one call: the model is compiled lazily on its first forward pass, so tuning happens with real input shapes, and it reports which candidates were fused and which were skipped. Finally, I wrote the benchmark harness that times each model in eager mode, with AutoFuser and with torch.compile, and checks the fused model’s outputs against the original.

Bruno built the loading of the scientific models into MASE and many of the operation handlers, Ching Bon the original fusion specification and a separate 3D-tiling code generator, and Oliver the MXINT8 weight quantisation and parts of the elementwise code generation.

Results

We benchmarked eight models from the Neural-Solver-Library on its Airfoil dataset, timing each in eager mode, with AutoFuser and with torch.compile (100 iterations after warm-up), and checking the fused model’s test loss against the original. torch.compile is a production compiler, so we treated it as an upper bound rather than a target. Fusion sped up almost every model. The largest gain was on the Galerkin Transformer: 33.9 ms to 9.1 ms, a 3.7× speedup (torch.compile reached 5.7×). On F-FNO, torch.compile failed to trace the model, while AutoFuser’s inlining recovered a usable graph and gave a 1.14× speedup. With MXINT8 quantisation on a transformer, run on an RTX 3070, AutoFuser reached 1.71× against torch.compile’s 1.61×. The one slowdown was FNO, where kernel launch overhead outweighed the benefit of fusion, as it did for torch.compile. The change in test loss stayed small throughout, of the same order as torch.compile’s.

The main limitation is coverage: AutoFuser only fuses operations it has handlers and tiling strategies for, so a model built from new kinds of operations gets little benefit. A single n-dimensional backend in place of the separate 1D, 2D and 3D generators would be the natural next step.