Relax: Composable Abstractions for End-to-End Dynamic Machine Learning

Compilers for machine learning were built when tensor sizes were known at compile time, and language models broke that — prompt length, cache length and batch size are all unknown until the moment you run — so existing compilers mark such dimensions "unknown" and lose the fact that one tensor is exactly four times another; Relax (CMU and collaborators) keeps them as algebraic expressions instead, and puts the computation graph, the loop-level code and the vendor library calls into a single representation so optimisations can cross between them.

Why this problem is worth solving

Everything about these models became variable at once.

A machine learning compiler is really several compilers stacked.

What people did before, and where it broke

• Annotate unknown dimensions as any. The approach in several established compilers and in the common model interchange format. Where it broke. It records that a dimension is unknown and nothing else, so the relation between a tensor of length n and one of length four n disappears — and that relation is precisely what would prove the second buffer can reuse the first.
• Trace the program just in time and specialise to the shapes you see. What the major framework’s compiler does: run the model, capture what happened, compile that. Where it broke. It sidesteps tracking shapes across function boundaries, which is elegant, but it needs the tracing machinery present at run time — so it does not reach constrained environments like phones or a browser, and it still requires a fixed-size cache, which in practice means rewriting the model.
• Handle dynamic shapes inside a single tensor program. Several systems reason carefully about dynamic loop bounds within one kernel. Where it broke. The reasoning stops at the function boundary, so nothing propagates to the graph level where memory planning and fusion decisions are made.
• Let tensor programs call out to libraries. One long-established system tracks dynamic shapes in loop-level code and provides a way to invoke external routines from there. Where it broke. The direction is wrong for this purpose: the graph level still cannot see into or reason about what the library call does.
• Lower through the levels in one direction, one shot. The standard architecture, and it keeps each level simple. Where it broke. A fact discovered while optimising loop-level code — that a kernel needs a scratch buffer, say — cannot travel back up to inform graph-level memory planning, and nothing can be partially lowered, examined, and then lowered further.
• The gap. Every approach handles dynamism within one level and one function. What breaks is everything that needs to reason across those boundaries, and shape information is exactly what gets destroyed at each crossing.

How a compiler keeps hold of a number it does not know

The idea

Recording that a number is unknown is not the same as recording nothing about it. The placeholder approach is entirely reasonable in isolation: a dimension whose value arrives at run time cannot be a constant, so mark it and move on. What it discards is that dynamic dimensions in a real model are almost never independent. A tensor of shape n by 4, flattened, has 4n elements; reshaped again it may be n by 2 by 2; the batch dimension threads through fifty operations unchanged. Every one of those relations is a fact the compiler could have used — to prove two buffers hold the same number of elements and can share storage, to know a fused region’s output size, to specialise every loop bound except the one that genuinely varies. Writing the dimension as a variable and letting shapes be algebraic expressions keeps all of it, and the effect is direct: static memory planning becomes possible under dynamism, cutting activation memory by 22 percent during prefill and 40 percent during decoding. Finding that a convenient abstraction is too blunt to express what the optimiser needed is the same complaint the tensor program work covered elsewhere on this site makes about declaratively reshaping loops, which is easy to write and cannot express the techniques that make hand-written kernels fast.

Lowering in one direction means a fact discovered late cannot be used early. The conventional stack is three representations and two one-way trips between them, which keeps each level tractable and forecloses a whole category of optimisation. The paper gives three concrete shapes it enables by refusing that separation. Partial lowering means a pass can send some operations to a vendor library and leave the rest for a later pass to handle differently, rather than committing to one dispatch strategy for everything. Analysis feedback means the compiler can inspect the loop structure of a kernel and derive that it is element-wise, or broadcasting, or a reduction, rather than requiring a human to annotate that property on every operator — which is how a custom quantisation routine somebody wrote gets fused correctly without anyone teaching the graph level about it. And cross-level transformation means that when analysis of a kernel concludes it needs a scratch buffer, that allocation can be lifted up to the graph level and included in the global memory plan instead of being hidden inside the kernel. Supplying the missing layer that lets general programs meet specialised implementations, rather than treating each combination as its own point solution, is the contribution the sparse tensor abstraction work covered elsewhere on this site makes for accelerators that had each been built around a single expression.

Compiling ahead of time is what reaches the platforms nobody else does. The alternative approach — trace the program as it runs and compile what you saw — genuinely simplifies the shape problem, because within a single traced region there are no function boundaries to lose information across. It also requires the tracing and compilation machinery to be present wherever the model runs, which rules out a phone, an embedded board, or a browser. Committing to compiling everything in advance is what forces the hard version of the problem, since ahead-of-time compilation needs whole-program optimisation across functions and therefore needs the shape relations to survive every boundary. The payoff is the deployment table: the same model, compiled once for arbitrary batch sizes and sequence lengths, running on an iPhone, an Android phone, a small single-board computer, a games handheld, an embedded developer kit, and inside a web browser. On Android it reaches 55 percent higher throughput than the leading hand-written inference engine, for the unglamorous reason that the hand-written engine has no Android graphics kernels and falls back to the processor while a compiler simply generates them. Getting a model onto a constrained device by not requiring the whole conventional stack to come with it is the same move the compact graphics runtime work covered elsewhere on this site makes when it replaces 48 megabytes of vendor runtime with 50 kilobytes.

Following one fused operator through a program with an unknown length

Shape expressions reuse the same integer expression system the loop-level representation uses, so the compiler can prove equalities between them with machinery that already exists.

Why it works

• Dynamic dimensions in real models are related, not independent. The batch dimension threads unchanged through dozens of operations and a flattened tensor has exactly the product of its input dimensions → so writing each dimension as a variable and each shape as an expression preserves relations a placeholder erases → letting the compiler prove two buffers hold the same number of elements and plan their storage statically.
• Static memory planning is the payoff that matters most. Without shape relations, memory must be allocated whenever an input size changes, which under varying lengths and batch sizes means allocating repeatedly and unpredictably → whereas relations let the compiler plan once against an upper bound and reuse the same memory across all sizes → cutting activation memory by 22 percent in one phase and 40 percent in the other, and making pre-recorded execution graphs possible at all.
• Forward deduction is cheap enough to run after every pass. Shape information is only useful if it survives transformation, and it is destroyed by any pass that does not maintain it → so each operator carries a rule producing output annotations from input ones, applied forward through the program in time linear in the operation count → which means every pass can rely on shapes being current rather than recomputing or losing them.
• Some properties are easier to read off the loops than to write down. Fusion needs to know whether an operator is element-wise or a reduction, and conventionally a human annotates that for every operator the compiler knows about → but the read and write index patterns of the loop-level code determine it → so a custom kernel nobody annotated is classified automatically and fuses correctly.
• Facts discovered low down are sometimes needed high up. Deciding a kernel requires a scratch buffer happens during loop-level analysis, long after graph-level memory planning would conventionally have run → but with all levels in one representation the allocation can be lifted to the graph → where it joins the global plan instead of being an opaque allocation inside a kernel.

What it costs (trade-offs and scope)

• The performance claim is competitive, not faster. Against hand-optimised engines and vendor-tuned serving systems the result is parity, and the paper says so. What is won is breadth — one compilation path reaching platforms the others do not — rather than peak speed on the platforms they do.
• Several comparison cells are empty, which is part of the argument and complicates the numbers. Two baselines have no Apple graphics support; the hand-written engine is strong on Apple and weak on NVIDIA. So consistently competitive is partly a claim about coverage, since no single baseline appears everywhere.
• The largest reported win is against a system not using the accelerator. The 55 percent advantage on an Android phone arises because the hand-written engine has no Android graphics kernels and runs on the processor. That is a genuine argument for compilation, and it is not a comparison of two graphics implementations.
• Symbolic tracking has a manual escape hatch. Data-dependent operators produce shapes nothing can deduce, so the annotation degrades to a dimension count and a construct is provided for a developer or a pass to assert a symbolic shape. The assertion is checked at run time, and it is still an assertion someone has to make.
• The runtime checks are asserted to be cheap. They are described as lightweight and not impacting performance, with no measurement, and they run at every function boundary.
• Memory planning is against an upper bound. The plan is computed for a maximum sequence length and batch size rather than for unbounded dynamism, so the savings depend on choosing that bound, and behaviour beyond it is not discussed.
• The pass pipeline is fixed and not run to a fixed point, which the paper states. Ordering is decided by hand, so the phase-ordering question that cross-level optimisation opens up is left where it was.
• The execution-graph optimisation contributes 1 to 2 percent, and it only becomes possible because all memory is planned in advance, so it is more a consequence of the memory work than an independent gain.
• The ablation is one model on one graphics processor. Effects of fusion, library dispatch and graph capture are measured on a single eight-billion-parameter model on a single consumer card.
• The mobile results use smaller and quantised models. Four-bit quantisation throughout, and a seven-billion-parameter model rather than eight on one phone to fit its memory, so deployment everywhere carries conditions.
• Operator coverage is not quantified. The long tail of operators is a large part of why compilers historically lose to hand-written kernels, and there is no accounting of how many needed individual attention.
• Everything is implemented within one framework, with the claim that the insights transfer to others stated rather than demonstrated.

The evaluation

Check yourself

Sources

The paper is “Relax: Composable Abstractions for End-to-End Dynamic Machine Learning” by Ruihang Lai, Junru Shao, Siyuan Feng, Steven Lyubomirsky, Bohan Hou, Wuwei Lin, Zihao Ye, Hongyi Jin, Yuchen Jin, Jiawei Liu, Lesheng Jin, Yaxing Cai, Ziheng Jiang, Yong Wu, Sunghyun Park, Prakalp Srivastava, Jared Roesch, Todd C. Mowry and Tianqi Chen, published at ASPLOS 2025, pages 998-1013, DOI 10.1145/3676641.3716249. The lead authors are at Carnegie Mellon University, with collaborators across several institutions and companies.

This explainer is written from the full arXiv text, arXiv:2311.02103. Everything above comes from the paper’s own sections: the survey of how existing compilers, interchange formats and framework compilers handle dynamic shapes, and why handling them within one level or one traced function leaves cross-function and cross-level optimisation unaddressed; the language constructs, covering the annotation system with its object, shape, tensor, tuple and callable forms, the side-effect-free dataflow blocks, and function calls within and across levels; the first-class symbolic shape design, contrasted directly with unknown-dimension annotations, with the fallback to coarse annotations for data-dependent operators and the runtime-checked assertion construct; the cross-level abstraction with its two foreign call primitives in destination-passing style and the reasoning about reconciling pure graph semantics with mutating low-level interfaces; the three optimisation patterns cross-level abstraction enables; the shape deduction algorithm with its forward direction, its treatment of function boundaries and first-class functions, and its support for arithmetic expressions in parameter annotations including the worked fusion example that must pass a symbolic variable explicitly; the loop-pattern classification algorithm and the fusion algorithm built on it; dynamic shape-aware memory planning, cross-level workspace lifting, execution graph offloading, and tensor operator optimisation by partial lowering; the fixed-order pipeline; and the evaluation covering three language models on three graphics vendors against four baselines, the ablation of fusion, library dispatch and graph capture, the memory planning measurements, deployment across six constrained platforms, the Android comparison, and results on a speech model and a vision-language model.