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

🎯 Pitch

Existing ML compilers struggle to optimize dynamic-shape models across diverse hardware, but Relax matches or beats hand-tuned systems on GPUs while also enabling deployment to mobile and web browsers that those systems don't support. It achieves this by letting graph-level and loop-level optimizations cross-talk through a unified symbolic shape abstraction, proving memory reuse for dynamic tensors and fusing operations that were previously untouchable across library boundaries.

1. Executive Summary

This paper introduces Relax, a compiler abstraction for optimizing end-to-end dynamic machine learning workloads, evaluated on large language models including Llama3-8B, Gemma1.1-7B, and Qwen2-7B across NVIDIA, AMD, Apple, and mobile GPU platforms. Relax contributes a cross-level abstraction that encapsulates computational graphs, loop-level tensor programs, and external library calls in a single representation (enabling, for instance, operator fusion that merges a custom quantization decode with a matmul into one tensor program), together with first-class symbolic shape annotations that track dynamic shape relations globally across subgraph function calls and foreign function calls (allowing the compiler to prove that two tensors share the expression 2*n and thus can reuse memory). The system delivers competitive inference performance across platforms, reducing LLM decode token latency by up to 27% on NVIDIA GPUs, and enables deployment of emerging models—including in-browser WebGPU execution and 4-bit quantized models on mobile phones—to a broader set of environments than established frameworks support, establishing that ahead-of-time compilation with whole-program symbolic shape tracking can match or exceed platform-specific hand-optimized solutions on mainstream GPU targets while simultaneously extending reach to emerging backends that those solutions do not address.

2. Context and Motivation

The Core Problem: Compiling Dynamic Shape Models for Diverse Hardware

The fundamental problem this paper addresses is deceptively simple: how do you compile a machine learning model that has dynamic tensor shapes so it runs efficiently on many different hardware platforms—including platforms that are not well-supported by existing frameworks? This matters because modern ML workloads, particularly large language models, are increasingly ubiquitous and must be deployed everywhere from datacenter GPUs to mobile phones, embedded devices, and web browsers. Yet the compilation challenge posed by dynamic shapes—tensor dimensions that are not known until runtime—remains poorly solved at the whole-program level.

Dynamic shapes arise from several sources in modern LLM workloads (Section 1):

Variable input sizes: user prompts have different lengths, and the KV-cache grows during autoregressive generation.
Batch size variation: inference serving systems process requests with different batch sizes over time.
Model-specific dynamism: vocabulary sizes, attention head counts, and other architectural parameters may vary across model configurations.

These are not edge cases—they are the norm for deployed LLMs. Every time a user sends a longer or shorter prompt, or the serving system batches a different number of requests together, the tensor shapes change. A compiler that cannot handle this dynamism efficiently either must recompile for every new shape (which is prohibitively slow) or must fall back to unoptimized execution paths (which wastes performance).

Why This Problem Is Important: The Deployment Gap

The importance of this problem has grown substantially with the rise of LLMs, but the paper's framing is broader than just LLMs. The introduction identifies a deployment gap: while models are developed in frameworks like PyTorch and JAX on NVIDIA GPUs, there is growing demand for these same models to run on "a diverse set of backend environments, including servers, personal computers, vehicles, and mobile devices" (Section 1). Each of these environments has different hardware characteristics, different supported software stacks, and different resource constraints. Supporting all of them manually—writing optimized kernels for each platform—is an enormous engineering burden that scales poorly with the number of models and backends.

The paper identifies several concrete deployment scenarios that are underserved by existing solutions:

Mobile phones (iPhone, Samsung): where GPU execution is available through Metal or OpenCL, but most frameworks do not generate optimized GPU kernels for these targets.
Embedded devices (Orange Pi 5, Steam Deck, Jetson Orin): where memory is constrained and ahead-of-time memory planning is essential for fitting models.
Web browsers via WebGPU: an emerging standard that allows GPU-accelerated computation in the browser, with almost no existing ML compilation support.

These platforms are not exotic—they represent the actual devices that end users interact with. The paper argues that the lack of compiler support for these environments is not a fundamental hardware limitation but a software abstraction problem: existing compiler architectures make it too difficult to carry whole-program optimizations across the boundaries between different levels of abstraction when shapes are dynamic.

Prior Approaches and Their Shortcomings

The paper situates itself against a background of ML compiler design that has evolved through several generations. The shortcomings of prior work fall into four categories, which I'll walk through in detail.

1. Single-Shot Multi-Level Lowering with Opaque Boundaries

Most ML compilers use multiple levels of intermediate representation (IR) that are lowered one to the next in a single direction (Section 2). The typical layers include:

Computational graph IR: represents the model as a dataflow graph of high-level tensor operators (matmul, reshape, ReLU). This level is amenable to global rewrites like operator fusion—merging multiple operators into a single kernel to reduce memory traffic.
Tensor program IR: describes individual operators with low-level loops and indexed buffer accesses. This level enables fine-grained optimizations like loop tiling, vectorization, and thread binding.
External libraries: vendor-optimized routines (cuBLAS, cuDNN, CUTLASS) that implement specific operators with hand-tuned assembly or CUDA code.

Prior compilers like TVM's Relay (Roesch et al., 2018), MLIR (Lattner et al., 2021), and IREE treat tensor programs and external libraries as foreign functions that are opaque to the graph-level compiler. The graph level lowers operators to these foreign calls in a single shot: once an operator is dispatched to a tensor program or library, the graph-level compiler can no longer analyze or transform it. This creates several problems:

No analysis feedback: information discovered during low-level optimization (e.g., "this tensor program requires a temporary workspace buffer") cannot flow back to inform graph-level decisions (e.g., "let's allocate that workspace at the graph level so it participates in global memory planning").
No partial lowering: you cannot decide to lower some operators to libraries now and optimize other operators differently later. The lowering decision is all-or-nothing at the boundary.
Composition barriers: if a new optimization or library needs to be introduced, it must be carefully integrated into the existing lowering logic, which becomes increasingly brittle as the number of lowering strategies grows.

The paper provides a concrete example of what this breaks in practice (Section 4.4, Figure 11). Some tensor program schedules (like the Stream-K matmul decomposition) require an intermediate global memory workspace. In a single-shot lowering architecture, this workspace allocation lives entirely inside the tensor program, invisible to the graph-level memory planner. Relax's cross-level abstraction allows the workspace to be lifted from the tensor program to the graph level, where it can be reused across operators, reducing total memory consumption. This optimization is structurally impossible in a design where the graph and tensor program levels are opaque to each other.

2. Inadequate Dynamic Shape Handling

When shapes are not known at compile time, the compiler must still generate code that works for any shape. Prior approaches to handling dynamic shapes have significant limitations.

Unknown-dimension annotations. Relay (Roesch et al., 2018) and many MLIR dialects (Lattner et al., 2021) represent dynamic dimensions with an "unknown" or "any" annotation. For example, if a tensor has shape (batch_size, 128) and batch_size is dynamic, prior systems annotate it as (?, 128). ONNX (Bai et al., 2019) uses a similar approach. The problem, which the paper illustrates in Figure 3, is that "?" erases relationship information. If one tensor has shape (n, 4) and another has shape (n*4,)—as happens after a flatten operation—the compiler with unknown annotations sees (?, 4) and (?,). It cannot prove that these two shapes are related (one is the flattened version of the other), which means it cannot:

Reuse memory: because it doesn't know the second tensor requires the same number of bytes as the first.
Plan static allocations: because it doesn't know the relationship between dimensions.
Fuse operators effectively: because it cannot verify that loop bounds match across fused kernels.

Just-in-time tracing. PyTorch 2.0 (Ansel et al., 2024) takes a different approach: it traces computation at runtime using TorchDynamo, capturing a TorchFX graph for each traced function. Within each traced region, dynamic dimensions are represented symbolically using a global variable table. This works well for just-in-time (JIT) compilation on server GPUs, but the paper points out that it "limits its applications on emerging platforms with constrained environments, such as mobile and WebGPU" (Section 1). There are two reasons for this:

AOT vs. JIT: embedded and mobile platforms often cannot afford the overhead of a JIT compiler at runtime. They need ahead-of-time (AOT) compilation, where the entire program is compiled once and deployed. JIT tracing avoids cross-function symbolic tracking by treating each traced function as an isolated unit, but AOT compilation requires tracking symbolic relations across function boundaries—something the PyTorch compiler's architecture does not do.
Cross-function symbolic tracking: the PyTorch compiler's symbolic variable table is global but operates on traced subgraphs, not on the full program with interprocedural analysis. This means that when a model calls a subgraph function that itself takes dynamic-shape inputs, the symbolic relations that flow across that function call boundary may be lost.

Tensor-program-level solutions. Several systems handle dynamic shapes within individual tensor programs. DietCode (Zheng et al., 2022) optimizes dynamic tensor programs by generating multiple specialized versions (buckets) for different shape ranges. CoRA (Fegade et al., 2022) handles ragged tensors with minimal padding. SparseTIR (Ye et al., 2023) supports sparse tensor computations. Halide (Ragan-Kelley et al., 2013) tracks dynamic shapes within tensor programs and can call external functions. But all of these operate within a single abstraction level—they optimize how a specific tensor program handles dynamic dimensions, but they do not address how dynamic shape information flows between the graph level and tensor program level, or across subgraph function calls.

The paper's Figure 3 succinctly illustrates the difference between unknown-dimension annotations (the ? approach) and first-class symbolic shapes. With symbolic shapes, the compiler tracks that lv0 has shape (n, 4), lv1 (after flatten) has shape (n*4,), and lv4 (after a split and indexing operation) has shape (n*2,). It can prove that lv4 is exactly half the size of lv1—a fact that enables memory reuse decisions. With unknown annotations, all of these would be (?, 4), (?,), etc., and no such proof is possible.

3. Manual Operator Library Integration

Vendor-optimized libraries like cuBLAS, cuDNN (Chetlur et al., 2014), CUTLASS (Thakkar et al., 2023), MKL-DNN (Intel), and MIOpen (Khan et al., 2019) provide high-performance implementations of common operators for specific hardware. The paper notes that these libraries "are platform-specific and have large engineering development costs to cover the growing demands of operators, data formats, and layouts" (Section 6).

The challenge is not just the libraries themselves, but how to compose them with compiler-generated code. An LLM might have hundreds of operators: some (like the large matmuls in attention) are good candidates for library dispatch, while others (like custom quantization decode operations) need compiler-generated loops, and still others (like element-wise additions following a matmul) would benefit from being fused into the library kernel's epilogue. In prior compilers, the decision of which operators go to libraries vs. code generation is made at the lowering boundary and is difficult to partially override or customize.

The paper notes that frameworks like llama.cpp, vLLM, and HuggingFace Transformers "usually rely on manual optimizations for each specific backend" (Section 6). This means that supporting a new backend (e.g., Apple Metal on iPhone) requires manually writing or porting kernels for every operator in every model—a massive engineering burden that Relax's compiler-based approach aims to eliminate.

4. CUDA Graph Incompatibility with Dynamic Shapes

CUDA Graphs (Gray, 2019) is an NVIDIA GPU feature that captures a sequence of kernel launches and replays them as a group, reducing driver-level launch overhead. It requires that all GPU memory accessed by the captured kernels be constant-sized and statically allocated ahead of time. This poses a fundamental conflict with dynamic shapes: if tensor sizes change between invocations, the memory cannot be statically allocated, and CUDA Graphs cannot be used.

Prior systems either:

Apply CUDA Graphs only to static-shape models.
Use shape bucketing to pad all tensors to the largest possible size, wasting memory.
Forgo CUDA Graphs entirely for dynamic models, accepting the launch overhead.

The paper argues that with proper symbolic shape analysis and static memory planning, CUDA Graphs can be applied to dynamic-shape models—because even though tensor dimensions vary, the compiler can prove upper bounds on those dimensions and pre-allocate enough memory to cover all possible shapes. This is a capability that depends critically on the compiler's ability to track and reason about shape relations across the entire program.

How Relax Positions Itself

The paper positions Relax not as a competitor to existing ML compilers in the traditional sense, but as a rethinking of the compiler abstraction itself to address the cross-cutting concern of dynamic shape information. The introduction explicitly states: "the insights presented in this paper can also benefit other ML compilation frameworks as well" (Section 5), and the related work section notes that "Relax's insights for supporting dynamic shapes and cross-level optimizations can be used to improve these ML compiler frameworks" (Section 6).

The key architectural differences from prior work are:

Cross-level abstraction (Section 3.3). Instead of maintaining separate IRs with one-way lowering, Relax brings computational graphs, tensor programs, and libraries into a single unified representation where they can interact. Graph-level functions can call tensor programs via call_tir and external libraries via call_dps_library. The graph level can analyze, transform, and partially lower these calls rather than treating them as opaque boundaries. This enables patterns that the paper captures in Figure 6: partial lowering (dispatch some operators to libraries first, optimize the rest later), analysis feedback (automatically infer operator properties from tensor program loop patterns rather than manually annotating them), and cross-level transforms (lift allocations from tensor programs to the graph level for global memory planning).

First-class symbolic shapes (Section 3.2). Instead of using "unknown" annotations, Relax represents each dimension with a symbolic expression composed of integer variables and arithmetic operations. These expressions are tracked globally across the entire program, including across subgraph function calls and foreign function calls. The paper describes three design principles that enable this: isolated symbolic relations at function boundaries (each function has its own symbolic variable scope, and relation propagation happens through function signatures), forward symbolic deduction (shape inference flows forward through the program, which is efficient and local), and support for symbolic expressions in parameter annotations (so that after operator fusion lifts multiple operators into a subgraph function, the function's parameter shapes can still be expressed).

Ahead-of-time compilation for emerging platforms. By performing whole-program symbolic shape analysis and static memory planning, Relax can produce a single compiled module that works for arbitrary input shapes and runs on platforms where JIT compilation is infeasible—mobile phones, embedded devices, and web browsers. The paper frames this as filling a concrete gap: "Relax deploys emerging models across these platforms, which most existing ML frameworks do not well support" (Table 3 caption).

The evaluation is designed to substantiate these claims through three kinds of comparison:

Performance on established hardware (NVIDIA RTX 4090, AMD Radeon 7900 XTX, Apple M2 Ultra): comparing against PyTorch (eager and compile mode), vLLM, and llama.cpp to show that Relax's compilation-based approach is competitive with hand-optimized solutions.
Ablation of cross-level optimizations (CUDA Graph offloading, operator fusion, partial library lowering): to quantify how much each composable optimization contributes.
Deployment to emerging platforms (iPhone, Samsung S23, Orange Pi 5, Steam Deck, Jetson Orin, WebGPU): to demonstrate that Relax can target environments that existing frameworks cannot, with competitive or superior performance to the few solutions that do work there (e.g., llama.cpp).

The paper's contribution is therefore not a single new optimization technique, but a compiler architecture that makes a class of optimizations (cross-level, dynamic shape–aware) possible for the first time in an AOT setting, and demonstrates that this architecture enables deployment of modern ML models to a meaningfully broader set of hardware platforms.

3. Technical Approach

This is primarily a systems and compiler design paper whose core idea is that by unifying computational graphs, loop-level tensor programs, and external library calls into a single cross-level representation with first-class symbolic shape tracking, an ahead-of-time compiler can perform optimizations that are structurally impossible in traditional multi-level, single-shot lowering architectures, enabling deployment of dynamic-shape ML models to a broader set of hardware platforms than existing frameworks support.

3.1 Reader Orientation

The paper builds a compiler infrastructure—not a specific optimization technique, but an abstraction architecture—that takes a machine learning model description (potentially with dynamic tensor shapes) and produces an optimized, self-contained executable module for a target hardware platform. The problem it solves is that existing ML compilers lose information about dynamic shape relationships when lowering from graph-level IR to tensor-program IR to platform code, which prevents whole-program optimizations like static memory planning and cross-operator fusion for models where tensor dimensions depend on runtime inputs (e.g., variable-length sequences in LLMs). The solution has two interlocking parts: a cross-level abstraction that lets the graph level directly call tensor programs and libraries while retaining analysis capabilities across those boundaries, and first-class symbolic shape annotations that track how dynamic dimensions relate to each other throughout the entire program, so the compiler can prove facts like "these two tensors have the same number of elements" and act on them.

3.2 Big-Picture Architecture (Diagram in Words)

The Relax system has five major components, organized as layers that information flows through:

Model Frontend — User-facing Python APIs (a PyTorch-like nn.Module interface) that construct a Relax program. The frontend emits Relax IR with symbolic shape annotations, where dynamic dimensions are represented as symbolic variables (e.g., n for batch size) rather than unknown placeholders.
Relax Cross-Level IR — The central program representation that unifies three abstraction levels: computational graph operators (high-level tensor operations like reshape and relu), foreign tensor program calls (via the call_tir primitive, which invokes loop-level TensorIR functions with explicit destination-passing-style memory), and foreign library calls (via the call_dps_library primitive, which dispatches to vendor libraries like cuBLAS or CUTLASS). All three coexist in the same IR, annotated with symbolic shape information that is globally tracked across subgraph function calls and foreign function boundaries.
Analysis and Optimization Pass Pipeline — A fixed-order sequence of compiler passes that operate on the cross-level IR. The pipeline runs: shape annotation deduction (forward propagation of symbolic shape information), partial library lowering (pattern-match-and-replace to dispatch specific operator patterns to external libraries), operator-to-tensor-program lowering (converting remaining high-level operators to call_tir of generated tensor programs), analysis feedback from tensor programs (inferring mathematical properties like element-wise or broadcast patterns from loop structure), cross-level operator fusion (pattern-matching on tensor program properties to merge multiple call_tir invocations into single fused kernels, with symbolic shape tracking through the fusion), tensor program workspace lifting (detecting global memory allocations inside tensor programs and lifting them to the graph level for global memory planning), dynamic shape–aware memory planning (allocating a static set of memory blocks based on symbolic upper bounds), CUDA Graph offloading (lifting statically-allocated subgraphs into CUDA Graph capture regions), and tensor program optimizations (applying analysis-based scheduling rules for loop tiling, vectorization, and thread binding).
Code Generation Backend — After the optimization pipeline completes, the IR is lowered to: (a) a sequence of virtual machine instructions at the graph level (each instruction is a call to a generated or builtin function), (b) optimized GPU source code (e.g., CUDA, Metal, Vulkan) for the tensor programs, and (c) a symbolic expression evaluation subsystem that computes concrete shape values at runtime from input tensors. These three components are packaged into a single deployable module.
Runtime — A lightweight runtime that loads the deployable module, evaluates symbolic shape expressions from input tensor dimensions, executes the virtual machine instructions, and dispatches to the appropriate GPU kernels or library functions.

The flow through the system is: a model is expressed in the frontend as Relax IR with symbolic shapes → the optimization pipeline applies cross-level transformations, tracking symbolic shape relations throughout → the optimized IR is lowered to platform-specific code and packaged into a module → at deployment time, the module loads, evaluates shapes from inputs, and executes the pre-compiled plan. Critically, the module is compiled once (ahead-of-time) and works for arbitrary input shapes, because all shape-dependent decisions (memory allocation sizes, kernel launch dimensions) are expressed in terms of symbolic variables that are bound at runtime.

3.3 Roadmap for the Deep Dive

First, the symbolic shape annotation system (§3.4.1), since it is the information backbone that makes all subsequent optimizations possible—without understanding how the compiler represents and propagates shape relations, none of the cross-level optimizations will make sense.
Second, the cross-level abstraction primitives call_tir and call_dps_library (§3.4.2), which are the mechanism by which the graph level interacts with tensor programs and libraries while preserving symbolic shape information across those boundaries.
Third, shape annotation deduction (§3.4.3), the algorithm that propagates symbolic shape information forward through the program, including across subgraph function calls, to ensure that every intermediate value has the most specific shape annotation possible.
Fourth, cross-level operator fusion (§3.4.4), which illustrates how analysis feedback from tensor programs is used at the graph level to merge operators, with careful handling of symbolic shape parameters.
Fifth, dynamic shape–aware memory planning (§3.4.5), which shows how symbolic shape analysis enables static memory allocation even when tensor sizes are not compile-time constants.
Sixth, the remaining cross-level optimizations (§3.4.6–3.4.8): tensor program workspace lifting, CUDA Graph offloading, and tensor operator optimizations via partial lowering, each of which depends on the cross-level abstraction to perform transformations that span the graph/tensor-program boundary.
Seventh, the optimization and lowering pipeline (§3.4.9), which ties all the passes together into a fixed-order compilation flow and explains how the final runnable module is produced.

3.4 Detailed, Sentence-Based Technical Breakdown

3.4.1 First-Class Symbolic Shape Annotations

The foundational idea in Relax is that every shape dimension in every tensor annotation is represented as a symbolic expression—an arithmetic combination of integer variables and constants—rather than as a concrete integer or an opaque "unknown" marker. This section explains what these annotations look like, how they are constructed, and what the compiler can do with them that it couldn't do with unknown annotations.

Annotation syntax and semantics. Each value in a Relax program carries an annotation that describes its structural type. Table 1 in the paper enumerates the annotation kinds: Object (any runtime value), Shape([n, 4]) (a symbolic shape tuple with dimensions n and 4), Shape(ndim=2) (a shape with two dimensions, both unknown), Tensor((n, 4), "f32") (a tensor with symbolic shape (n, 4) and float32 data type), Tensor(ndim=None, dtype="f32") (a tensor with completely unknown shape), Tuple[Tensor((n, 4), "f32"), Object] (a tuple of a tensor and an arbitrary object), and Callable(...) (a function type with parameter and return annotations). These annotations are embedded in Python AST, with symbolic expressions quoted as strings (e.g., "n * 4") in function signatures where the symbolic variables have not yet been declared—a syntactic concession to Python parsing that the paper notes "can be changed."

The critical design choice is that symbolic shape expressions reuse the same expression system as loop-level tensor programs. This means that any integer expression that can appear in a TensorIR loop bound can also appear in a shape annotation, and the compiler's symbolic analysis infrastructure (expression equality proving, simplification) works uniformly across both levels. This unification is not just for convenience—it means that when a shape expression n * 4 appears in a graph-level annotation and the same expression appears as a loop bound in a tensor program, the compiler can recognize them as identical and make optimization decisions based on that equality.

Comparison with unknown-dimension annotations. Figure 3 in the paper provides a side-by-side comparison that illustrates what is lost with unknown annotations. Consider a program that reshapes a tensor of shape (n, 2, 2) to (n, 4), flattens it to (n*4,), applies a data-dependent unique operator whose output shape cannot be statically known, asserts a new symbolic shape (m,) on the result via match_cast, and applies an element-wise exp.

With first-class symbolic shapes (left side of Figure 3):

The input x has shape (n, 2, 2).
After reshape, lv0 has shape (n, 4).
After flatten, lv1 has shape (n*4,).
After unique, lv2 has coarse annotation Tensor(ndim=1, dtype="f32") because the compiler cannot deduce the output size.
match_cast asserts that lv2 actually has shape (m,) for some new symbolic variable m, binding this information to lv3.
lv4 after exp has shape (m,).

The compiler can prove that lv1 and lv0 have the same number of elements (n*4), enabling buffer reuse. It can track that m is related to the unique values of lv1, even though it cannot compute m statically.

With unknown-dimension annotations (right side of Figure 3):

All shapes use ? for dynamic dimensions.
After reshape, the shape is (?, 4).
After flatten, the shape is (?,).
After unique, the shape is still (?,).
After exp, it remains (?,).

The compiler cannot prove any relation between any two shapes. It cannot plan memory reuse between lv0 and lv1 because it doesn't know they have the same total size. It cannot even tell that lv0 and lv4 might have different sizes.

The match_cast construct. When the compiler encounters a data-dependent operator like unique whose output shape depends on runtime values (not just symbolic parameters), it cannot produce a precise symbolic shape. The paper introduces match_cast as a way to inject symbolic shape information at points where the compiler cannot deduce it. The construct takes a value and an asserted annotation, and the compiler inserts a runtime check that verifies the assertion. If the check passes, downstream compiler passes can use the asserted annotation for optimization. The paper notes that match_cast "can be inserted by both front-ends and compiler passes to suggest more specific symbolic shapes and serves as a valuable tool for developers to indicate shape information within programs."

This is a pragmatic design choice: rather than requiring the compiler to prove everything statically (which would be impossible for data-dependent operators), it provides an escape hatch where developers or front-end tools can inject domain knowledge about shape relations. The runtime check ensures safety—if the assertion is wrong, the program fails with an error rather than silently producing incorrect results.

Symbolic variables and shape functions. Symbolic variables are introduced via the sym_var() construct, which creates a new variable that represents an unknown-but-fixed integer dimension. In the example of Figure 3, n = sym_var() declares that there exists some integer n representing a dimension size. When a Relax function is called, the symbolic variables in its parameter annotations are matched against the concrete shapes of the arguments, and the symbolic variables are bound to the actual dimension values. The paper's design ensures that symbolic variables are scoped to function boundaries—a function's symbolic variables are independent of its caller's, and shape relations are propagated only through the function's parameter and return annotations. This isolation is key to making interprocedural shape analysis tractable.

3.4.2 Cross-Level Abstraction Primitives: `call_tir` and `call_dps_library`

The cross-level abstraction is realized through two foreign function call primitives that bridge the graph level with lower-level implementations. These primitives are the mechanism by which Relax avoids the opaque lowering boundary that characterizes prior multi-level compilers.

The design tension: pure graph operators vs. destination-passing style. The paper identifies a fundamental mismatch between abstraction levels: "computational graph abstractions favor pure operators that return a new tensor for each operation," which creates clean dataflow graphs amenable to dead code elimination and pattern matching, while "most tensor programs and libraries of low-level computations adopt destination-passing style (DPS) interfaces, which take the computation result tensors as inputs and directly mutate them, rather than allocating and returning new tensors." DPS is preferred at low levels because it gives the caller control over memory allocation, which is essential for performance—the caller can pre-allocate buffers, reuse memory across operations, and avoid the overhead of per-operator allocation.

Relax's solution is to require that low-level tensor programs use DPS, and to provide graph-level primitives that expose this interface explicitly. This means memory allocation is visible at the graph level (initially as separate allocation operations), which allows the compiler to analyze and optimize memory usage across operators—exactly what static memory planning needs.

call_tir: invoking tensor programs from the graph level. The call_tir primitive takes four arguments:

tir_func: a TensorIR function (the low-level implementation)
args: the input tensors to the tensor program
annotation: the expected shape and data type of the output tensor
sym_args: additional symbolic expressions passed as arguments to the tensor program (e.g., the dimension variable n)

Figure 5 explains the semantics: call_tir allocates an output tensor with the specified shape and data type, then calls the TensorIR function in destination-passing style, passing the inputs, the allocated output buffer, and the symbolic arguments. This is semantically equivalent to the graph-level view of "apply this operation and produce a new tensor," but the implementation is a DPS call with explicit allocation.

The symbolic arguments (sym_args) are crucial for dynamic shapes. When a tensor program needs to know a dynamic dimension (like n in a loop bound), the graph level passes the symbolic value explicitly. This means the tensor program can specialize to static dimensions while handling dynamic dimensions symbolically—the loop over n uses the passed variable, but a loop over a known dimension like 128 can be fully unrolled or vectorized. The paper emphasizes: "By flowing the symbolic shape information from the graph level to tensor programs, we can allow tensor programs to generate code that specializes to most static dimensions and only uses dynamic dimensions when necessary."

call_dps_library: invoking external libraries from the graph level. The call_dps_library primitive mirrors call_tir but targets external library functions instead of compiled tensor programs. Its arguments are:

A string naming the library function (e.g., "cutlass.rms_norm")
Input tensors
Expected output tensor annotation

The library function is resolved through a registry at link time. This design "introduces great flexibility in prototyping, since external routines can be easily called from a Relax program" without modifying the compiler's lowering logic. New library integrations are just registry entries and pattern-matching rules, not changes to the compilation pipeline.

Why these primitives enable cross-level optimization. Because call_tir and call_dps_library are first-class graph-level operations with explicit annotations, the compiler can:

Analyze their inputs and outputs for shape relations (the output annotation tells the shape deduction system what the result shape is)
Pattern-match on sequences of these calls for fusion (e.g., finding a call_tir to an element-wise operation followed by a call_tir to a reduction)
Replace them with other implementations (partial lowering: swap a call_tir to a generic matmul with a call_dps_library to a cuBLAS matmul)
Lift their internal allocations to the graph level (workspace lifting: detect internal buffer allocations and make them explicit graph-level allocations)

In a traditional single-shot lowering design, once an operator is lowered to a tensor program or library call, the graph level loses the ability to analyze or transform it. Relax's primitives keep these calls in the graph IR with full annotations, so graph-level passes can continue to operate on them.

3.4.3 Shape Annotation Deduction

With symbolic shapes on every value, the compiler needs to compute the shape of each intermediate result from the shapes of its inputs. This is the job of the shape annotation deduction system, which runs both during initial program construction and after every compiler pass that modifies the IR.

Forward deduction algorithm. Each tensor operator (e.g., reshape, flatten, relu) has a registered shape deduction rule that takes the input annotations and operator-specific parameters (e.g., the target shape for reshape) and produces the output annotation. The paper adopts a forward deduction strategy: "the annotation of an expression is deduced based on its inputs." For call_tir and call_dps_library, the output annotation is explicitly provided as part of the call arguments (the annotation parameter), so deduction simply uses that annotation directly.

Forward deduction is chosen for two reasons: simplicity and locality. It processes the program in a single pass in topological order, requiring no fixed-point iteration or constraint solving. The paper states that "a full-graph forward deduction takes time linear to the number of operations." The downside is that information flows only forward—if a later match_cast provides more specific shape information, it doesn't retroactively refine earlier shapes. But the paper notes that the system "still supports the introduction of more powerful but less efficient deduction methods via compiler passes as needed," so backward constraint propagation could be added for specific optimizations without burdening the common case.

Data-dependent operators. For operators whose output shape cannot be deduced from input shapes alone (like unique, whose output size depends on the actual data values, or nonzero), the deduction rule returns a coarse-grained annotation—for example, Tensor(ndim=1, dtype="f32") indicates "we know this is a 1D float32 tensor, but we don't know its size." The paper emphasizes that "coarse-grained annotations are returned when more specific information cannot be inferred... as a safety net." This is not a failure of the system—it's an honest representation of what is statically knowable—and the match_cast construct provides a mechanism for adding more specific information when it is available from domain knowledge.

Interprocedural shape deduction. Figure 7 in the paper demonstrates how symbolic shape relations propagate across subgraph function calls. The example defines a function subfn(s: Shape(["n", "m"])) -> Tensor(("n * m",), "f32") that takes a 2D shape and returns a 1D tensor whose size is the product of the two dimensions. The deduction system can:

Infer that f0(subfn) called with shape(n, 4) returns Tensor((n*4,), "f32") — the caller's n is matched to the callee's n, and the constant 4 fills m, producing n*4.
Infer that calling with shape(3, 4) returns Tensor((12,), "f32") — when all dimensions are constants, the compiler evaluates the expression to a concrete value.
Infer that calling with shape(n+1, 4) returns Tensor(((n+1)*4,), "f32") — expressions in arguments are substituted into the callee's shape expressions.
Infer that calling with a coarse Shape(ndim=2) returns Tensor(ndim=1, dtype="f32") — when the input shape is unknown, the output shape is also unknown (but we at least know the dimensionality).

Three design principles. The paper articulates three principles that govern the shape deduction design:

Isolated symbolic relations at function boundaries. Each function introduces its own symbolic variables (via sym_var()), and the only way shape relations cross function boundaries is through the function's parameter and return annotations. This means that to deduce the output shape of a function call, the compiler only needs to look at the function's signature—not its body. This enables functions to be used as first-class values with Callable annotations. The function signature also serves as the specification for lightweight runtime shape checks: when a function is called, the runtime verifies that the actual argument shapes match the parameter annotations, and that the return value matches the return annotation. The paper notes these checks "are lightweight and do not impact the overall performance."
Forward symbolic deduction. As described above, deduction proceeds forward through the program. This avoids the complexity of constraint solving or fixed-point iteration, keeping compilation times predictable. The trade-off is that backward information flow (from uses to definitions) is not automatically captured, but match_cast provides an explicit mechanism for injecting additional information when needed.
Support symbolic expressions in parameter annotations. After compiler transformations like operator fusion (which lifts multiple operators into a new subgraph function), the function's parameters may need shape annotations that are expressions involving symbolic variables that are not parameters of the function. Figure 8 illustrates this: suppose we want to fuse add(x, y) and relu(...) where both inputs have shape (2*n,), but n is not a direct parameter. The fusion pass creates a new function fused_add_relu(x, y, s) where s is an extra parameter that carries the runtime value of n. The parameter annotations can then be Tensor(("n * 2",), "f32") where n is bound by s. This pattern—passing extra symbolic arguments to preserve shape information after code motion—is described as "a common pattern we use when designing passes that lift out function regions and recombine."

3.4.4 Cross-Level Dynamic Shape–Aware Operator Fusion

Operator fusion is the optimization that merges multiple tensor operators into a single kernel to reduce memory bandwidth: instead of writing intermediate results to global memory and reading them back, the fused kernel keeps intermediate values in registers or shared memory. In Relax, fusion operates across the graph/tensor-program boundary and must handle symbolic shapes.

The three-stage fusion pipeline. Figure 9 illustrates the complete fusion flow for a motivating example: fusing a custom quantization decode operation (which unpacks 4-bit quantized weights into float16) with a matrix multiplication. This is a realistic LLM workload where specialized decode logic and standard linear algebra need to be combined. The three stages are:

Compute pattern analysis (analysis feedback). A compiler pass examines each tensor program's loop structure and classifies it into a "pattern kind" that describes its mathematical properties. Algorithm 1 shows the classification logic in pseudocode. The pass extracts read indices and write indices from the tensor program's loop nest, then checks how the indices relate:
- If all write indices are identical and all reads access exactly those indices, the pattern is ElementWise (e.g., C[i,j] = A[i,j] + B[i,j]—the output at (i,j) depends only on inputs at (i,j)).
- If a read accesses a subset of the write indices, the pattern is Broadcast (e.g., C[i,j] = A[i,j] + B[j]—B[j] is broadcast across dimension i).
- If a read permutes the write indices without introducing new ones, the pattern is Injective (e.g., C[i,j] = A[j,i]—a transpose uses the same indices but in different order).
- If the operation involves accumulating over an index that doesn't appear in the output (a reduction loop), the pattern is Reduction.
- Special cases: a pattern that is Broadcast but also has element-wise reads is reclassified as ElementWise (to handle mixed cases like C[i,j] = A[i,j] + B[j]). If an opaque operation turns out to be a fused multiply-add, it's classified as OutputEwiseFusible (meaning element-wise operations can be fused into its output, as with matmul+ReLU). Everything else falls back to Opaque.
This analysis is applied automatically to every tensor program, and the result is attached as a function attribute (func_attr("compute_pattern", "Injective")). The key advantage over prior systems is that no manual operator annotation is needed: "by adopting cross-level abstraction and instead relying on analysis-based properties, we can greatly reduce the engineering cost of annotation on high-level operators." In traditional compilers, each high-level operator must be manually labeled with its fusion properties, and custom operators (like quantization decode) require additional manual labels. Relax's analysis feedback automatically handles custom operators because it inspects their loop structure directly.
FuseOps (pattern-match-based graph partitioning). Algorithm 2 shows the FuseOps pass. Given the annotated pattern kinds from stage 1, FuseOps searches the graph for patterns of call_tir operations that can be fused. An example pattern is: find an OutputEwiseFusible operation (like matmul) followed by an ElementWise operation (like ReLU), and fuse them into a single subgraph function. The pattern matching operates at the graph level, but the patterns reference properties discovered from tensor program analysis.

When a match is found, FuseOps lifts the matched region into a new subgraph function, replacing the original operations with a call to that function. Critically, the new function preserves symbolic shape information through its parameter and return annotations. If the fused region involves symbolic expressions (like 2*n in Figure 8), the pass inserts extra parameters that carry the necessary symbolic variables at runtime. The paper notes that this approach "allows quick composition of different fusion, improving the overall productivity of continuous compiler development"—new fusion patterns can be added without modifying the core fusion infrastructure.
FuseTensorIR (cross-level tensor program merging). After FuseOps has grouped graph-level call_tir operations into subgraph functions, FuseTensorIR lowers each subgraph function by merging the individual tensor program functions into a single fused tensor program. This is a cross-level transformation: it rewrites both the graph level (replacing the sequence of call_tir with a single call_tir to the fused function) and the tensor program level (producing a new loop nest that combines the computation of the original tensor programs, with intermediate values passed through registers rather than global memory).

At this stage, symbolic shape tracking ensures that loop bounds in the fused tensor program correctly reference the symbolic variables. For example, if the original programs both had loops bounded by n, the fused program's loops are also bounded by n, and the symbolic variable is passed through from the graph level.

Why this three-stage design? The decomposition into analysis, graph-level partitioning, and tensor-program-level merging provides composability. An alternative approach—doing everything in a single monolithic fusion pass—would make it difficult to add custom fusion patterns or to interpose other optimizations. With Relax's design, a developer can:

Write a custom pass that adds new fusion patterns (e.g., fusing the attention computation) before or after FuseOps.
Use the analysis feedback from stage 1 for other purposes (e.g., deciding which operators to dispatch to libraries).
Replace FuseTensorIR with a different merging strategy for specific subgraphs while keeping the graph-level partitioning the same.

The paper gives the example that "we can apply a pass to fuse new sets of patterns that are not covered by FuseOps (e.g., fusing all sub-operators in scaled dot-product attention), and use FuseOps for the remainder." This composability is a direct consequence of the cross-level abstraction: because tensor programs are first-class entities in the graph IR, passes can operate on them in multiple stages.

3.4.5 Dynamic Shape–Aware Memory Planning

Memory planning is the compiler optimization that analyzes the lifetimes of intermediate tensors and reuses memory buffers: when tensor A is no longer needed and tensor B of the same size is about to be allocated, B can reuse A's memory rather than allocating new storage. For static shapes, this is straightforward—the compiler knows the exact size of every tensor and can plan a fixed set of memory blocks. For dynamic shapes, prior compilers typically fall back to runtime allocation, losing the opportunity for reuse.

How symbolic shapes enable static planning despite dynamic dimensions. The key insight is that even when tensor sizes are not compile-time constants, the compiler can reason about equality of symbolic shape expressions. If two tensors have the same symbolic shape (e.g., both have shape (n, 256)), the compiler can prove they require the same number of bytes at runtime, regardless of what n turns out to be, and can therefore plan for them to share a memory block.

Figure 10 illustrates this with a concrete example. Before memory planning, four intermediate tensors are individually allocated:

lv0 with shape (2, n) (after exp)
lv1 with shape (n, 2) (after transpose)
lv2 with shape (n, 2) (after relu)
lv3 with shape (2, n) (after another transpose)

The compiler can prove that lv0 and lv3 have identical shape (2, n), and lv1 and lv2 have identical shape (n, 2). It also can prove that lv0's lifetime does not overlap with lv3's (because lv0 is last used before lv3 is allocated), and similarly for lv1 and lv2. Therefore, the memory planner can allocate just two storage buffers (s0 of size 2*n*4 bytes for float32, and s1 of size 2*n*4 bytes) and reuse them across the four tensors. Without symbolic shapes, the compiler would see (?, 2) and (2, ?) and could not prove these are the same size.

The memory planning algorithm. Algorithm 3 presents the planning procedure. The graph-level function g is first transformed by lowering all call_tir and call_dps_library operations to explicit memory allocation and DPS calls (as in the call_tir semantics shown in Figure 5). This exposes the allocations for planning. Then:

Liveness analysis computes the live range of each tensor—the set of operations during which the tensor's value is needed.
A storage pool is created that tracks available memory blocks.
The algorithm iterates through operations in sequential order. When it encounters a tensor allocation:
- It calls RequestReuseWithSymShape(op.shape, op.dtype) on the storage pool, which uses symbolic expression analysis to find a previously allocated block whose size (in bytes) is provably equal to the requested size.
- If a match is found, it inserts a "tensor instantiation" operation that creates the new tensor from the existing storage block (rather than allocating new memory).
- If no match is found, it allocates a new storage block and inserts the allocation before the current operation.
When an operation that uses a tensor completes, the algorithm checks liveness: if the tensor is dead after this operation (no future operation reads it), its storage is returned to the pool for reuse.

Upper-bound static allocation. For deployment to memory-constrained platforms, the paper extends this approach by taking the upper bound of symbolic values when known. For example, in an LLM, the sequence length n has a known maximum (the model's context length), and the batch size has a configured maximum. The compiler can plan memory using these upper bounds, pre-allocating enough storage for the worst-case shape, and then use only the needed portion at runtime. This "allows creating a static memory allocation plan ahead of time, even in the presence of dynamic shapes," which "is crucial for deploying dynamic ML models on memory-limited backends." The paper notes a connection to Halide's Func::bound() which marks bounds in tensor programs, and extends the concept to both graph and tensor program levels.

Interaction with CUDA Graph. Static memory planning—allocating all memory ahead of time rather than at runtime—is a prerequisite for CUDA Graph offloading (§3.4.7), which requires all GPU memory to be statically allocated. By enabling static planning for dynamic shapes, Relax extends CUDA Graph support to models that would otherwise be incompatible.

3.4.6 Cross-Level Tensor Program Workspace Lifting

Some tensor program schedules require internal global memory workspace—temporary buffers used for intermediate results that are too large for registers or shared memory. The Stream-K matmul schedule (Osama et al., 2023), which the paper uses as a motivating example, decomposes matrix multiplication into two phases: partial accumulation into a global workspace buffer, followed by reduction of the partial results into the output.

In a traditional single-shot lowering design, this workspace allocation lives entirely inside the tensor program, invisible to the graph-level memory planner. This means the workspace memory cannot be shared with other operators, even if those operators could reuse it when the matmul is not running.

The lifting transformation. Figure 11 illustrates the cross-level transformation. The compiler detects global memory allocations inside tensor programs through analysis feedback (examining the buffer allocation sites in the TensorIR function). When such an allocation is found, it performs a joint rewrite of the tensor program and the graph-level call site:

Graph level: The allocation is inserted before the call_tir, and the allocated buffer is passed as an additional argument to the tensor program.
Tensor program level: The internal buffer allocation is removed, and the buffer is instead received as a parameter.

After lifting, the workspace is a graph-level allocation that participates in the global memory planning pass (§3.4.5). This means the workspace can be reused by other operators that run before or after the matmul, reducing overall memory consumption. The paper emphasizes: "This optimization is only possible with the cross-level abstractions when the shape relation is preserved throughout all cross-level transformations in Relax, and the optimization opportunities for planning such memory reuse may not arise in the traditional single-shot lowering flow."

3.4.7 CUDA Graph Offloading

CUDA Graphs (Gray, 2019) is an NVIDIA GPU feature that reduces kernel launch overhead. Normally, each GPU kernel launch involves a CPU-GPU interaction through the CUDA driver. CUDA Graph allows capturing a sequence of kernel launches as a single graph, which can then be replayed with a single driver call, eliminating per-kernel launch overhead. The constraint is that all GPU memory accessed by the captured kernels must be constant-sized and statically allocated at capture time, because the graph records buffer addresses.

For static-shape models, this is straightforward: all tensor sizes are known at compile time, so memory can be pre-allocated and a CUDA Graph can be captured once and replayed indefinitely. For dynamic-shape models, the constraint is violated if tensor sizes change between invocations—the graph would reference stale buffer addresses.

How Relax enables CUDA Graph for dynamic shapes. The paper combines static memory planning (§3.4.5) with a graph-offloading pass to apply CUDA Graph to dynamic models. The approach works as follows:

Memory planning pre-allocates all GPU memory using symbolic upper bounds (e.g., maximum sequence length × hidden dimension × bytes per element). These allocations are constant-size from the GPU driver's perspective—they are allocated once and their addresses never change, even though the amount of data used within them varies per invocation.
A compiler pass analyzes the graph and identifies subgraphs that meet CUDA Graph conditions: all memory accessed within the subgraph is statically allocated (i.e., the addresses are compile-time constants), and there are no data-dependent control flow divergences that would require different kernel launch sequences.
The pass lifts qualifying subgraphs into separate functions and inserts runtime calls to CUDA Graph capture/replay infrastructure around them. The paper describes: "At runtime, only the first run of a subgraph function triggers CUDA Graph capture; subsequent runs automatically replay the captured CUDA Graph."

This means that even though the model processes variable-length sequences, the GPU kernel launches are replaying a captured graph—the n dimension in a matmul (n, 256) × (256, 512) changes per invocation, but the kernel itself handles the dynamic dimension, and the launch parameters are part of the captured graph's recorded state (not the buffer addresses). The paper notes that CUDA Graph offloading "overall brings about 1-2% of performance gain by reducing kernel launch overheads at GPU driver level" (Section 5.2), and that the principle "can generally be applied to any GPU backend that supports static execution graphs in the future."

3.4.8 Tensor Operator Optimizations via Partial Lowering

In most ML compilers, the decision of which operators go to external libraries vs. compiler-generated code is made at a single lowering boundary. This creates rigidity: if you want to use a library for some matmul configurations but compiler-generated code for others (e.g., library for batch matmul, custom kernel for matrix-vector at batch size 1), you must either modify the monolithic lowering logic or accept a suboptimal decision.

Pattern-match-and-rewrite partial lowering. Relax instead applies operator optimizations through composable partial lowering passes. The mechanism is:

A set of (subgraph pattern, library function) pairs is registered in the compiler. Each pair says: "if you see this pattern of operations in the graph, replace it with a call to this library function."
A pattern-match-and-rewrite pass scans the graph for these patterns and replaces matched regions with call_dps_library invocations.
Multiple such passes can be composed: one pass might dispatch matmul patterns to cuBLAS, another might dispatch attention patterns to FlashAttention, and yet another might leave remaining operators for compiler code generation.

Figure 12 illustrates the two-dimensional flexibility: operators can be partially lowered to libraries (left path) or to TensorIR programs for compiler optimization (right path), and these decisions can be made independently for different operator instances in the same model. The paper notes that Relax "allows users to register patterns for customizability."

Dynamic shape–aware scheduling rules. For operators that are not dispatched to libraries, Relax generates TensorIR programs and applies analysis-based scheduling rules to optimize them. These rules determine loop tiling factors, thread binding, vectorization, and shared memory usage based on the tensor program's pattern kind (from Algorithm 1) and the presence of dynamic dimensions. The paper notes that dynamic dimensions require careful handling: "We can also include passes to apply Ansor-style auto-tuning for rare tensor programs (e.g., complicated convolutions) that our analysis-based schedule rules fail to handle." This means the system combines rule-based optimization for common cases with auto-tuning for specialized operators, all within the same composable lowering framework.

Why partial lowering matters for diverse backends. The evaluation (Section 5.1) demonstrates the practical importance of partial lowering: "Cross-level abstractions enable us to use compiler-optimized matrix-vector multiplication tensor programs at batch size 1, while being able to apply partial library lowering to leverage operator libraries for other batch sizes." At batch size 1, library kernels optimized for large matrix multiplications may be suboptimal for matrix-vector products (where one dimension is 1); the compiler-generated kernel can specialize for this case. But at batch size 32, the library kernel's hand-tuned tiling and assembly-level optimization wins. Partial lowering allows the same compiled module to make different decisions for different operator instances based on their shapes, without recompilation.

3.4.9 Optimization and Lowering Pipeline

The individual optimizations described above are composed into a fixed-order pipeline that transforms a Relax program into a runnable module. Figure 13 provides the pipeline structure, which the paper describes as running "on the cross-level abstraction to optimize, lower and finally build an end-to-end model into a runnable module."

Pipeline ordering and rationale. The ordering of passes is deliberate:

Partial library lowering runs first. This dispatches operators to external library functions based on pattern matching. The paper prioritizes this first because "external library functions on the target platform" often provide the best performance for well-supported operators, and lowering them early frees the remaining operators for further optimization.
Operator to tensor program lowering converts all remaining high-level operator calls to call_tir of generated TensorIR functions. At this point, every computation in the graph is either a call_tir or a call_dps_library.
Tensor program optimizations (including workspace lifting) are applied to the TensorIR functions. The paper notes that "some tensor program optimizations (e.g., workspace lifting) are applied before graph optimizations, which necessitates Relax's cross-level abstraction design." This is because workspace lifting modifies both the tensor program and the graph-level call site—it is a cross-level transformation, not a pure tensor-program pass.
Dynamic shape–aware operator fusion groups compatible call_tir operations and merges their tensor programs. This reduces the number of kernels and intermediate memory accesses.
Dynamic shape–aware memory planning analyzes tensor lifetimes and plans memory reuse, using symbolic shape analysis to determine which tensors can share buffers.
CUDA Graph offloading identifies subgraphs that can be captured as CUDA Graphs after memory planning has ensured static allocation.

The pipeline is described as "fixed-order" and "without fixed point," meaning each pass runs exactly once in sequence. This is a pragmatic choice for predictable compilation times, though it means passes cannot iterate to convergence. The ordering reflects dependencies: library dispatch happens first to get those operators out of the way; tensor program optimization happens before fusion so that individual programs are optimized before being merged; memory planning happens after fusion so that fused programs' memory usage is accounted for; CUDA Graph offloading happens last because it depends on static memory allocation.

Building the runnable module. The final stage of the pipeline transforms the optimized IR into a deployable artifact. This involves several sub-steps:

Symbolic shape evaluation code generation. The paper creates "an integer host tensor to store runtime values of all symbolic expressions in the program." At the start of execution, the runtime populates this tensor with concrete values derived from input tensor dimensions. For each symbolic expression in the program, the compiler generates code that loads values from this tensor, evaluates the expression, and stores the result. This is essentially a lightweight interpreter for the symbolic expression language, generated at compile time and executed once per invocation.
Annotation erasure. After symbolic shape values are materialized, all annotations are stripped from the IR, leaving "a program comprised mainly of low-level function calls." These calls are translated to a sequence of virtual machine instructions, each of which is a call into a generated or builtin function.
GPU code generation. For the optimized TensorIR functions, Relax leverages existing lower-level compiler infrastructure (the paper mentions Apache TVM as the implementation base) to generate target-specific GPU code (CUDA, Metal, Vulkan, OpenCL, WebGPU). The generated kernels are bundled with the module.
Packaging. The graph-level virtual machine instructions, GPU kernels, and symbolic shape evaluation code are combined into "a single holistic end-to-end module, which can then run on the target platform of compilation."

Ahead-of-time compilation property. The entire pipeline runs once at compile time, producing a module that accepts arbitrary input shapes at runtime. The paper emphasizes this repeatedly: "Relax compiles models only once for arbitrary batch sizes and sequence lengths" (Section 5.1). This is in contrast to PyTorch's JIT approach, which may recompile when shapes change, and to bucketing approaches like Nimble (Shen et al., 2021) that compile multiple specialized versions for different shape ranges.

Integration with Apache TVM. The paper implements Relax on top of Apache TVM, using TensorIR as the tensor program abstraction, but notes that "the insights presented in this paper can also benefit other ML compilation frameworks." The cross-level primitives call_tir and call_dps_library are designed to be abstraction-agnostic—they could target a different tensor program IR (like Triton or MLIR's Linalg dialect) or different library naming conventions without changing the fundamental design.

4. Key Insights and Innovations

Innovation 1: Dynamic Shape Tracking as a First-Class Program Property, Not a Missing-Value Problem

The dominant assumption in prior ML compilers—codified in Relay (Roesch et al., 2018), MLIR dialects (Lattner et al., 2021), and ONNX (Bai et al., 2019)—is that a dynamic dimension is an absence of information. The compiler knows the dimension exists but not its value, so it records ? or unknown and moves on. This framing treats dynamism as a deficit: the more dimensions are unknown, the less the compiler can do, and the fallback is always runtime allocation and generic code generation.

Relax's fundamental conceptual move is to reject this framing entirely. Instead of treating a dynamic dimension as missing information, Relax treats it as a symbolic variable with known relationships to other dimensions. The shape (n, 4) is not "a tensor with an unknown first dimension"—it is "a tensor whose first dimension is some integer n, and whose total size is 4n bytes." The distinction is subtle but profound: the compiler doesn't know the value of n, but it knows everything that depends on n, and it can prove equalities (4n == 4n), inequalities, and arithmetic transformations (n*4 == 4*n) across the entire program.

What makes this more than a syntactic upgrade from ? to n is the interprocedural scope. Prior systems that did use symbolic variables (like the PyTorch compiler's global variable table for traced subgraphs) kept them scoped to individual traced functions. Relax propagates symbolic relations across subgraph function calls (Figure 7), through foreign tensor program invocations (call_tir passes symbolic variables as explicit arguments, Figure 5), and across compiler transformations that lift code into new functions (operator fusion inserts extra shape parameters to preserve symbolic context, Figure 8). This global tracking is what converts symbolic shapes from a local convenience into a whole-program analysis capability.

The innovation is not the mathematical representation of symbolic expressions—Halide (Ragan-Kelley et al., 2013) tracked dynamic shapes in tensor programs, and constraint-based shape systems existed in functional languages (Axon; Collins and Grover, 2022). The innovation is the architectural decision to make symbolic shape the universal annotation language that spans all abstraction levels (graph, tensor program, library calls) and all function boundaries in an ahead-of-time compiler, combined with the pragmatic fallback of coarse annotations and match_cast when shapes are genuinely data-dependent. The result is that the compiler can perform static memory planning, CUDA Graph capture, and cross-operator fusion for models whose tensor shapes vary at runtime—optimizations that the ? approach structurally precludes.

Evidence: The memory planning results in Table 2 show a 40% reduction in activation memory during decode and 22% during prefill. Without symbolic shape tracking, the compiler cannot prove that (n, 256) and (n, 256) from different operators refer to the same-size tensors, so every intermediate tensor gets its own allocation. Figure 10 makes this concrete at the level of individual tensor lifetimes. The CUDA Graph offloading results in Figure 17, while modest (1–2%), only exist because static memory planning with symbolic upper bounds made Graph capture possible for dynamic-shape models.

Innovation 2: Breaking the Single-Shot Lowering Barrier Through Cross-Level Primitives

The multi-level compiler architecture—computational graph IR lowered to tensor program IR lowered to platform code—is decades old and appears in virtually every production ML compiler (TVM/Relay, MLIR/IREE, PyTorch compiler, XLA). The invariant in these systems is that lowering is irreversible and opaque: once an operator is dispatched to a tensor program or library, the graph level can no longer analyze, transform, or optimize it. This creates a forced sequencing of decisions. Library dispatch happens at the lowering boundary. Fusion happens before lowering. Memory planning happens after. If a tensor program optimization (like Stream-K matmul decomposition) creates a need for workspace memory, that information cannot flow back to influence graph-level memory planning because the boundary is one-way.

Relax's call_tir and call_dps_library primitives (Section 3.3) dissolve this barrier by making foreign function calls first-class graph-level operations with explicit shape annotations. A call_tir is not a lowering endpoint—it is an IR node that can be pattern-matched, replaced, analyzed, and transformed by subsequent passes. This seemingly small architectural change enables a qualitatively different compilation model: instead of lowering everything in one shot, the compiler can partially lower some operators (dispatch matmuls to cuBLAS), fuse others (merge quantization decode with matmul, Figure 9), lift allocations from tensor programs to the graph level (workspace lifting, Figure 11), and receive analysis feedback from tensor program loop structures (Algorithm 1) to inform graph-level decisions—all within the same optimization pipeline.

The intellectual contribution is not any single one of these optimizations (each has precedents: operator fusion is standard, workspace reuse appears in Halide, library dispatch is universal). The contribution is the architectural realization that the lowering boundary itself was the bottleneck, and that by replacing it with bi-directional primitives that preserve symbolic shape information, the compiler gains access to a class of cross-level optimizations that were structurally impossible before. This is a reframing of what an ML compiler IR is for: not a pipeline of progressively lower representations, but a unified space where different implementation strategies coexist and can be composed.

The practical significance is most visible in the composability argument. The paper emphasizes repeatedly that new fusion patterns, new library backends, and new tensor program optimizations can be added as independent passes without modifying the core lowering logic. The three-stage fusion pipeline (analysis feedback → FuseOps → FuseTensorIR, Figure 9) is designed so that custom passes can interpose at any stage. Partial lowering allows different operator instances in the same model to take different code paths (library matmul at batch 32, compiler-generated matrix-vector at batch 1) without recompilation. This composability matters because the ML model landscape evolves faster than compiler infrastructure—new operators (quantization decode, FlashAttention variants, mixture-of-experts routing) appear constantly, and a compiler that requires modifying monolithic lowering logic for each new operator is unsustainable.

Evidence: The ablation in Figure 17 decomposes the performance contributions of partial library lowering (up to 27%), operator fusion (fusing ~1/5 of operators), and CUDA Graph offloading (1–2%), showing they compose additively. The evaluation across six different GPU backends (NVIDIA, AMD, Apple, OpenCL, Vulkan, WebGPU) in Figures 14–16 and Table 3 would be practically impossible without the ability to mix library dispatch (where available) with compiler-generated kernels (where libraries don't exist or underperform) on a per-operator basis.

Innovation 3: Ahead-of-Time Compilation with Dynamic Shapes as a Deployment Enabler for Emerging Platforms

Prior to Relax, the landscape for deploying LLMs to non-NVIDIA platforms was stark: either use a JIT framework like PyTorch (which supports dynamic shapes through tracing but requires a full Python runtime and struggles on mobile/embedded), or use a hand-optimized system like llama.cpp (which delivers strong performance on Apple GPUs and CPUs but requires manual per-backend kernel engineering). There was no compiler-based AOT solution that could take a model description and produce optimized GPU code for an arbitrary backend while handling dynamic shapes correctly.

The paper's evaluation makes this deployment gap concrete. Table 3 shows Relax running 4-bit quantized LLMs on six platforms where most existing frameworks have no GPU support: iPhone 14 Pro (Metal), Samsung S23 (OpenCL), Orange Pi 5 (Mali GPU OpenCL), Steam Deck (AMD APU Vulkan), NVIDIA Jetson Orin (CUDA), and WebGPU in-browser on M3 Max. The paper claims Relax is "the first solution to enable GPU-accelerated LLM inference on these platforms except the NVIDIA Jetson Orin." On Samsung S24 (Figure 18), Relax delivers up to 55% more throughput than llama.cpp, which falls back to CPU because it lacks GPU kernels for Android GPUs.

The intellectual contribution here is not a new optimization technique but a demonstration that the cross-level symbolic shape architecture unlocks deployment breadth that was previously infeasible. The causal chain is: symbolic shape tracking enables static memory planning → static memory planning makes models fit in constrained memory (the paper notes that "without memory planning that pre-allocates all needed memory and keeps it within the budget, these models are not even runnable on some of the environments due to memory constraints") → AOT compilation produces a single deployable module with no runtime Python dependency → the same compiler infrastructure generates GPU code for multiple backends from the same IR → platforms that were previously inaccessible become targets.

This is significant beyond the specific performance numbers because it changes the economics of model deployment. Without Relax, supporting a new backend for an LLM requires either (a) convincing a JIT framework to add AOT support and port its runtime, or (b) manually writing kernels for every operator in the model. Relax offers a third path: the compiler handles code generation automatically, and only platform-specific runtime bindings (memory allocation, kernel dispatch) need to be implemented. The paper doesn't quantify the engineering effort saved, but the existence proof—six diverse platforms supported in a single system—is itself the evidence.

The finding also has implications for the JIT-vs-AOT debate in ML compilation. PyTorch's JIT-centric design is well-suited for server environments with Python runtimes, but the paper argues that "as we target a broad set of emerging platforms, we must enable ahead-of-time (AOT) compilation, which necessitates full program optimizations across functions" (Section 1). Relax's architecture shows that AOT compilation with dynamic shapes is not only possible but can match or exceed JIT performance on mainstream GPUs (Figures 14–16) while simultaneously extending to platforms JIT cannot reach.

Evidence: Table 3 and Figures 14–18 collectively demonstrate that Relax is competitive with PyTorch compile mode and vLLM on NVIDIA GPUs (within a few percent), competitive with llama.cpp's hand-optimized Metal kernels on Apple GPUs (Figure 16), and uniquely present on mobile and WebGPU platforms where the other solutions simply don't exist. This breadth-of-deployment evidence is the paper's strongest argument that the cross-level symbolic shape architecture is not just a performance optimization but a deployment capability unlock.

Innovation 4: Analysis Feedback from Tensor Programs as a Replacement for Manual Operator Annotation

In traditional ML compilers, fusion decisions depend on knowing mathematical properties of each operator: Is it element-wise? Does it perform a reduction? Can an element-wise operator be fused into this operator's output? These properties are typically manually annotated by compiler developers for each high-level operator in the system (e.g., relu → ElementWise, matmul → OutputEwiseFusible). This creates two problems: (a) the annotation burden scales with the number of supported operators, and (b) custom or user-defined operators (like quantization decode, rotary position embedding, or custom attention variants) have no annotations and cannot participate in fusion unless someone manually adds them.

Relax's analysis feedback pass (Algorithm 1) replaces manual annotation with automatic pattern analysis of tensor program loop structures. The pass examines the read and write indices in the loop nest, classifies the access pattern (ElementWise, Broadcast, Injective, Reduction, OutputEwiseFusible), and attaches the result as a function attribute. Because this analysis operates on the tensor program IR—not on a catalog of known operator names—it works for any operator whose implementation is expressed in TensorIR, including custom operators that the compiler developers have never seen.

This is a subtle but important conceptual shift. In prior systems, the operator catalog was a source of truth that the compiler consulted; custom operators were second-class citizens that fell through to opaque handling. In Relax, the loop structure is the source of truth, and standard operators are just tensor programs whose patterns happen to be recognized by the analysis. This means custom quantization decode functions (Figure 9), custom attention implementations, and any other user-defined operator automatically get the same fusion opportunities as built-in operators, without additional annotation work.

The practical impact is on compiler extensibility. The paper explicitly frames this as an engineering cost reduction: "by adopting cross-level abstraction and instead relying on analysis-based properties, we can greatly reduce the engineering cost of annotation on high-level operators." But the deeper implication is that it changes the programming model for compiler developers. Adding support for a new operator no longer requires touching the fusion pass to add an annotation; you just write the TensorIR implementation, and the analysis feedback pass figures out its fusion properties automatically. This is an instance of the broader composability theme: the cross-level abstraction enables passes that reason about tensor programs from the graph level, which in turn eliminates manual bookkeeping that was previously unavoidable.

Evidence: Figure 9 shows the analysis feedback pass classifying decode_q4 as Injective (automatically, from its loop structure) and mm as OutputEwiseFusible. FuseOps then uses these classifications to fuse them together—for a custom quantization operator that no compiler developer manually annotated. The paper doesn't quantify the annotation burden saved, but the design is a clear departure from systems where every operator requires a manual entry in a fusion compatibility matrix.

5. Experimental Analysis

Evaluation Methodology

Dataset. The primary evaluation uses large language models—Llama3-8B, Gemma1.1-7B, and Qwen2-7B—not a static benchmark dataset. The evaluation measures decode token latency across different batch sizes and sequence lengths, reflecting real-world LLM serving workloads. For non-LLM evaluation, the paper additionally evaluates Whisper-large-v3 (automatic speech recognition) on a 30-second speech file and LLaVA (multimodal model) with an image input and 32-token generation. All models use float16 weight and activation precision unless otherwise specified.
Hardware platforms. The evaluation spans three categories of hardware. Mainstream GPUs: NVIDIA RTX 4090, AMD Radeon 7900 XTX, and Apple M2 Ultra with 76-core GPU. Emerging platforms: iPhone 14 Pro (Apple A16), Samsung S23 (Qualcomm Snapdragon 8 Gen 2), Orange Pi 5 (ARM Mali GPU), Steam Deck (AMD APU), NVIDIA Jetson Orin developer kit, and WebGPU on Apple M3 Max laptop. Mobile comparison: Samsung S24 for the Relax vs. llama.cpp head-to-head. This hardware diversity directly tests the paper's central claim about deployment breadth.
Metrics. For LLM inference, the primary metric is decode token latency in milliseconds per token (ms/tok), computed by measuring the total decode time for generating 32 tokens and dividing by 32. For emerging platforms (Table 3, Figure 18), the metric is throughput in tokens per second (tok/s) for single-sequence generation. For Whisper, the metric is total transcription time in milliseconds for a 30-second speech file. For LLaVA, the metric is generation time in milliseconds for 32 tokens given an image input. Memory usage is measured as total activation memory size in MiB during successive prefill/decode workloads (Table 2).
Baselines. The paper compares against four categories of baselines. HuggingFace Transformers (v4.41.2) with PyTorch (v2.3.1) in eager mode and compile mode (Ansel et al., 2024), with FlashAttention (Dao et al., 2022) enabled. vLLM (v0.5.0.post1) (Kwon et al., 2023) with its PagedAttention memory management. llama.cpp (commit 172c825) (Gerganov, 2023), a hand-optimized C++ LLM inference system. Specialized baselines for non-LLM models: WhisperX (commit f2da2f8) (Bain et al., 2023), Faster Whisper (v1.0.2), and whisper.cpp (commit 5d950c4) (Gerganov, 2022) for speech recognition; vLLM and llama.cpp for LLaVA. Platform support varies significantly: PyTorch compile mode and vLLM lack Apple GPU support, and llama.cpp falls back to CPU on Android GPUs due to lack of GPU kernels.
Generation budget / compute accounting. The paper does not use the "generation budget" concept from benchmark-oriented papers. Instead, all systems process identical workloads (same batch sizes, same sequence lengths, same models) under their own optimization configurations, and wall-clock latency or throughput is compared directly. This is a deployment-oriented rather than algorithmic comparison—the question is "how fast does the model run" rather than "how efficiently is a compute budget spent." Relax compiles models once for arbitrary batch sizes and sequence lengths, while some baselines (particularly PyTorch compile mode) may require per-shape compilation. The paper notes that Relax is AOT, producing a single module that handles all shapes.
Cross-validation / statistical protocol. There is no cross-validation or statistical testing reported. The paper evaluates deterministic workloads (model inference on specific GPU hardware) where variance is typically low. Results are presented as single-number latencies or throughput figures. This is standard for systems papers in this venue (ASPLOS), where the primary concern is reproducibility of the compiler's behavior rather than statistical generalization.

Main Quantitative Results

LLM Inference on Mainstream GPUs (NVIDIA RTX 4090)

Figure 14 reports the primary performance comparison for three models (Llama3-8B, Gemma1.1-7B, Qwen2-7B) at batch sizes 1, 16, 32, and 64 on an NVIDIA RTX 4090. The headline result is that Relax delivers consistently competitive performance, reducing decode token latency by up to 27% compared to baselines.

At batch size 1 across all three models, Relax achieves approximately 12–14 ms/tok, comparable to vLLM and HuggingFace Transformers (eager mode). The PyTorch compile mode results are omitted for Qwen2-7B due to "lack of support," and are available for Llama3-8B and Gemma1.1-7B where they perform similarly to Relax. llama.cpp consistently underperforms on NVIDIA GPUs at all batch sizes—at batch size 1 for Llama3-8B, llama.cpp shows approximately 18–20 ms/tok versus Relax's ~12 ms/tok, and for Gemma1.1-7B, llama.cpp reaches ~30 ms/tok versus Relax's ~22 ms/tok.

As batch size increases to 64, the latency hierarchy shifts. For Llama3-8B at batch 64, vLLM achieves approximately 13 ms/tok, Relax achieves roughly 14 ms/tok, Transformers with torch.compile reaches ~16 ms/tok, and Transformers eager reaches ~20 ms/tok. The gap between Relax and vLLM is modest (roughly 1 ms/tok, or ~7%), while the gap to Transformers eager is substantial (~35%). For Qwen2-7B at batch 64, Relax and vLLM are nearly identical (both roughly 15–16 ms/tok), with Transformers eager at ~24 ms/tok. The paper's 27% improvement claim refers specifically to the maximum latency reduction versus specific baselines—exact pairings are not specified but can be inferred from the largest gaps in Figure 14, such as Relax vs. Transformers eager at mid-range batch sizes.

The paper highlights a design advantage visible in these results: "Cross-level abstractions enable us to use compiler-optimized matrix-vector multiplication tensor programs at batch size 1, while being able to apply partial library lowering to leverage operator libraries for other batch sizes." This means Relax's same compiled module adapts its kernel choice to batch size: at batch 1, it uses compiler-generated matrix-vector kernels optimized for that shape; at batch 32, it dispatches to cuBLAS for large matrix multiplications. This is not something the latency numbers directly prove, but it explains why Relax remains competitive across the full batch size range without per-batch recompilation.

LLM Inference on AMD and Apple GPUs

Figure 15 extends the comparison to AMD Radeon 7900 XTX. The pattern is broadly similar to NVIDIA results, with Relax delivering competitive performance. At batch size 1 for Llama3-8B, Relax achieves roughly 22 ms/tok versus llama.cpp's ~30 ms/tok—a 1.50× advantage at batch size 1, which the paper calls out explicitly. At batch 64, Transformers eager reaches ~45 ms/tok while Relax stays near 22–24 ms/tok, a roughly 2× gap. PyTorch compile mode and vLLM have no AMD GPU support and are absent from Figure 15. The llama.cpp baseline is present but consistently slower than Relax across all models and batch sizes on AMD hardware.

Figure 16 covers Apple M2 Ultra. Here, the comparison set narrows to HuggingFace Transformers (eager only—no torch.compile or vLLM support on Apple GPU), llama.cpp, and Relax. For Llama3-8B at batch size 1, Relax achieves roughly 60 ms/tok versus llama.cpp at ~50 ms/tok, a gap of about 10 ms/tok. As batch size increases, the gap narrows: at batch 64, Relax reaches ~65 ms/tok versus llama.cpp at ~65 ms/tok, essentially tied. For Qwen2-7B, Relax and llama.cpp are nearly indistinguishable across all batch sizes. Transformers eager is dramatically slower on Apple GPU: at batch 1 for Llama3-8B, it reaches ~200 ms/tok versus Relax's ~60 ms/tok, a 3.3× gap.

The Apple GPU results are significant because llama.cpp is the strongest baseline on this platform—it has hand-optimized Metal kernels—and Relax matches it without manual per-backend kernel engineering. As the paper states, "Relax has competitive performance comparing to the hand-optimized llama.cpp baseline."

Ablation of Composable Optimizations on NVIDIA RTX 4090

Figure 17 isolates the effects of Relax's three key optimizations on Llama3-8B inference. The baseline (Relax without fusion, partial library lowering, or CUDA Graph offloading) achieves roughly 21–23 ms/tok across batch sizes. Adding operator fusion reduces latency by fusing approximately 1/5 of all operators (e.g., RMSNorm, element-wise additions), bringing latency down to roughly 19–21 ms/tok. Adding partial library lowering on top of fusion provides the largest gain—up to 27% in the paper's reported maximum—by dispatching heavy matrix multiplications (about 1/3 of all operators) to cuBLAS, with the largest effect at batch size 64 where the baseline is roughly 22 ms/tok and library lowering brings it to approximately 16 ms/tok. Adding CUDA Graph offloading provides a final 1–2% improvement across all batch sizes. The paper notes that "all these composable optimizations collectively improve the overall system performance."

This ablation demonstrates additivity: the gains from fusion, library lowering, and CUDA Graph are approximately independent and stack. It also shows that partial library lowering is the single largest contributor for this model-hardware combination, which is expected since large matmuls dominate LLM decode time.

Memory Usage Reduction from Static Planning

Table 2 quantifies memory savings from static memory planning during Llama3-8B inference. The experiment runs successive prefill phases of lengths 128, 256, 512, and 1024, and successive decode phases of batch sizes 1, 16, 32, and 64, measuring total allocated activation memory in MiB.

For prefill, Relax without memory planning uses 192.7 MiB of activation memory; with static planning, this drops to 149.7 MiB—a 22% reduction. For decode, the reduction is more dramatic: 150.0 MiB without planning versus 88.2 MiB with planning—a 40% reduction. The paper explains the mechanism: "With static memory planning, we always reuse memory across all input lengths and batch sizes, even as they vary over time. In contrast, without memory planning, the system repeatedly allocates dynamic-sized memory whenever the input shape changes, which is unpredictable in real-world applications."

The larger reduction in decode (40% vs. 22%) is notable. During decode, the model processes one new token per sequence per step, and activation tensors are generally smaller; the relative overhead of per-allocation fragmentation is higher. Static planning eliminates this overhead entirely by pre-allocating a fixed set of buffers and reusing them regardless of shape changes.

Emerging Platform Deployment

Table 3 reports single-sequence throughput (tokens/sec) for 4-bit quantized LLMs across six emerging platforms. For Llama3-8B: iPhone 14 Pro (Metal) achieves 5.1 tok/s, Samsung S23 (OpenCL) achieves 7.9 tok/s, Orange Pi 5 (OpenCL) achieves 2.3 tok/s, Steam Deck (Vulkan) achieves 14.0 tok/s, Jetson Orin (CUDA) achieves 32.0 tok/s, and WebGPU on M3 Max achieves 37.8 tok/s. For the smaller Phi3-mini-4k and RedPajama-3B, throughputs are correspondingly higher, reaching 68.0 and 68.6 tok/s respectively on WebGPU.

Two footnotes are important: the iPhone runs 3-bit quantized Llama2-7B (not Llama3-8B) and the Samsung S23 runs 4-bit Llama2-7B, both to "fit the VRAM limit of the mobile environments." This means the Llama3-8B numbers reported for these specific devices in the table body may correspond to different model variants; the paper's Table 3 caption indicates Llama3-8B for "most cases" but the dagger footnote clarifies the substitution.

The paper claims Relax is "the first solution to enable GPU-accelerated LLM inference on these platforms except the NVIDIA Jetson Orin." This is a strong deployment breadth claim. On the Jetson Orin, where CUDA is available and other frameworks do work, Relax's 32.0 tok/s for Llama3-8B represents its performance on a platform with existing solutions, while on the other five platforms, Relax is filling a deployment gap.

Figure 18 provides a direct comparison on Samsung S24 (a newer device than the S23 in Table 3), where Relax achieves up to 55% higher throughput than llama.cpp across Llama2-7B (approximately 5 tok/s vs. ~3.2 tok/s), Phi3-mini-4k (approximately 13.8 tok/s vs. ~10 tok/s), and RedPajama-3B (approximately 19.5 tok/s vs. ~14 tok/s). The paper explains that "llama.cpp only utilizes CPU due to lack of kernels for Android GPUs, whereas Relax automatically generates optimized GPU codes via compilation." This is a direct demonstration of the deployment breadth advantage: the same compiler infrastructure that generates CUDA for Jetson and Metal for Apple GPUs also generates OpenCL for Android GPUs without additional manual kernel development.

Non-LLM Model Evaluation

Figure 19 shows Whisper-large-v3 transcription time for a 30-second speech file. On NVIDIA RTX 4090, Relax achieves approximately 520 ms versus HuggingFace Transformers at ~620 ms (14% speedup), WhisperX at ~590 ms, Faster Whisper at ~570 ms, and whisper.cpp at ~540 ms. The differences among optimized solutions (Relax, Faster Whisper, whisper.cpp) are within roughly 10% of each other. On Apple M2 Ultra, the comparison narrows to Transformers (~4000 ms) versus whisper.cpp (~2100 ms) versus Relax (~2050 ms)—Relax and whisper.cpp are essentially tied, while Transformers is roughly 2× slower. WhisperX and Faster Whisper have no Apple GPU support and are absent from this comparison.

Figure 20 shows LLaVA generation time for 32 tokens after an image input. On RTX 4090, Relax achieves approximately 510 ms versus vLLM at ~550 ms, llama.cpp at ~580 ms, and Transformers at ~680 ms. On Apple M2 Ultra, Relax achieves ~1100 ms versus llama.cpp at ~1150 ms and Transformers at ~1900 ms. vLLM is absent from the Apple GPU comparison. The paper notes that "Relax efficiently supports the vision encoder together with the prefill and decode phases of LLM," highlighting that multimodal models—which combine a vision encoder with an LLM decoder—run within the same compilation framework.

Ablation Studies and Robustness Checks

Operator fusion contribution: On Llama3-8B with NVIDIA RTX 4090, enabling operator fusion (which fuses approximately 1/5 of all operators) reduces decode latency from roughly 21–23 ms/tok to roughly 19–21 ms/tok across batch sizes (Figure 17). This is a modest but consistent improvement of roughly 8–10%, consistent with the bandwidth savings from eliminating intermediate global memory accesses.
Partial library lowering contribution: Adding partial library lowering on top of fusion provides the single largest gain, up to 27% latency reduction at larger batch sizes (Figure 17). At batch size 64, the latency drops from approximately 21 ms/tok (with fusion only) to approximately 16 ms/tok (with fusion + library lowering). This is attributable to cuBLAS's hand-tuned assembly kernels for large matrix multiplications, which outperform compiler-generated code for these shapes.
CUDA Graph offloading contribution: Enabling CUDA Graph offloading on top of fusion and library lowering adds approximately 1–2% further improvement across all batch sizes (Figure 17). The modest gain reflects that kernel launch overhead is a relatively small fraction of total decode time for LLM inference, where kernels are large and relatively few. The paper positions this optimization as "enabling" rather than performance-critical—it makes static memory planning compatible with CUDA's graph replay mechanism.
Memory planning impact on activation memory: Static memory planning reduces activation memory by 22% during prefill (192.7 → 149.7 MiB) and 40% during decode (150.0 → 88.2 MiB) for Llama3-8B (Table 2). The larger reduction in decode suggests that without planning, repeated reallocation for varying shapes incurs significant fragmentation or overallocation. The paper notes that this planning uses symbolic upper bounds (maximum sequence length and batch size) to pre-allocate memory, which is a crucial enabler for memory-constrained platforms like mobile phones and embedded devices.
Relax vs. llama.cpp on Android GPU: On Samsung S24, Relax achieves up to 55% higher throughput than llama.cpp across three 4-bit quantized LLMs (Figure 18). The key insight is not numerical—it's that llama.cpp has no GPU kernel support for Android GPUs at all, falling back to CPU execution, while Relax automatically generates OpenCL GPU kernels from the same compiler infrastructure used for CUDA and Metal. This is a deployment capability ablation rather than a performance ablation: it shows what happens when a platform lacks hand-optimized GPU kernels and relies on a compiler to generate them.
Performance across three mainstream GPU architectures: Relax's competitive performance holds across NVIDIA (Figure 14), AMD (Figure 15), and Apple (Figure 16) GPUs, with the gap to the best-performing baseline varying by platform. On NVIDIA, vLLM is often the best baseline and Relax is within ~7% at high batch sizes. On AMD, llama.cpp is the best baseline and Relax outperforms it by up to 1.50× at batch size 1. On Apple, llama.cpp is the best baseline and Relax matches or slightly trails it (within ~10%). This cross-platform consistency is a robustness check for the claim that Relax's compilation approach generalizes across GPU architectures without per-platform manual tuning.
Model diversity within evaluation: The evaluation covers three LLM architectures (Llama3-8B, Gemma1.1-7B, Qwen2-7B) plus Whisper-large-v3 and LLaVA, spanning decoder-only transformers, an encoder-decoder transformer, and a multimodal model. Performance patterns are consistent across models: Relax is competitive with or outperforms baselines on all model-hardware combinations. PyTorch compile mode fails for Qwen2-7B (Figure 14), and vLLM and torch.compile are entirely absent from Apple GPU (Figure 16), providing negative results that illustrate the deployment breadth gap.
Negative result: framework support gaps: The evaluation implicitly documents where baselines fail to run at all, not just where they underperform. PyTorch compile mode is missing for Qwen2-7B on NVIDIA and AMD, and absent entirely from Apple GPU. vLLM is absent from Apple GPU. WhisperX and Faster Whisper are absent from Apple GPU. llama.cpp uses CPU-only on Android GPU. These gaps are not ablated in the traditional sense, but they are a crucial part of the paper's argument: the baselines that perform best on well-supported platforms are simply unavailable on emerging platforms where Relax works.

Critical Assessment

Claim: "Relax delivers performance competitive with state-of-the-art systems across various GPUs"

This claim is supported on NVIDIA and AMD GPUs, with qualifications on Apple GPUs. On NVIDIA RTX 4090 (Figure 14), Relax is within 7–10% of vLLM (the best-performing baseline) in most configurations, and substantially outperforms HuggingFace Transformers eager mode by up to 35% at high batch sizes. On AMD (Figure 15), Relax leads all baselines, including a 1.50× advantage over llama.cpp at batch size 1. On Apple M2 Ultra (Figure 16), Relax is essentially tied with llama.cpp (within ~10%) and dramatically outperforms HuggingFace Transformers eager (by up to 3.3×). However, on Apple GPU at batch size 1, llama.cpp leads Relax by approximately 10 ms/tok (~20%) for Llama3-8B, which means Relax is competitive but not the best performer in that specific configuration. The paper's phrasing "competitive" is fair: Relax is in the top tier across all platforms, leads on AMD, and trails modestly in some Apple configurations.

A limitation is the absence of FlashAttention-enabled baselines on some platforms. FlashAttention is enabled "when available," but its availability varies by framework and platform, which may inflate the apparent advantage of Relax on platforms where baselines lack it. The paper does not disaggregate how much of Relax's lead comes from compiler optimizations versus the presence/absence of FlashAttention in the baseline.

Claim: "Relax enables deployment of emerging models to a broader set of emerging environments, including mobile phones, embedded devices, and web browsers"

This claim is strongly supported by Table 3 and Figure 18. Relax runs GPU-accelerated LLM inference on six platforms—iPhone 14 Pro, Samsung S23, Orange Pi 5, Steam Deck, Jetson Orin, and WebGPU on M3 Max—where most existing frameworks either do not support GPU execution at all or provide only CPU fallback. The paper's claim to be "the first solution to enable GPU-accelerated LLM inference on these platforms except the NVIDIA Jetson Orin" is a factual statement about deployment capability, not performance. On Samsung S24, where a direct comparison is possible, Relax's GPU-accelerated execution delivers 55% higher throughput than llama.cpp's CPU execution (Figure 18), directly demonstrating the deployment breadth advantage.

The primary weakness in this claim is the substitution of Llama2-7B for Llama3-8B on mobile phones due to VRAM limits (Table 3 footnotes). While this is a pragmatic reflection of real hardware constraints, it weakens the direct comparability of mobile numbers to the desktop GPU numbers reported elsewhere. The paper would be stronger if it had used consistent model configurations across all platforms, even if that meant running only the smaller models (Phi3-mini-4k, RedPajama-3B) everywhere.

A missing baseline is any measurement of what percentage of operators or model architectures the compiler handles correctly across these platforms. The paper shows that specific models work, but does not characterize the failure modes or limitations—what happens when a model uses an operator that has no generated implementation? What percentage of the model zoo is supportable? Without this information, "enables deployment" is demonstrated by existence proofs (six platforms, five models) rather than systematic coverage analysis.

Claim: "Composable optimizations collectively improve the overall system performance, with up to 27% latency reduction"

This claim is supported by Figure 17's ablation on Llama3-8B + NVIDIA RTX 4090, which shows partial library lowering contributing the majority of the 27% gain, operator fusion contributing a smaller but consistent fraction, and CUDA Graph offloading adding 1–2% on top. The gains are additive and non-overlapping, which supports the "composable" characterization.

However, the ablation is run on a single model on a single GPU. The paper does not show how the relative contributions of each optimization vary across platforms—on mobile GPUs where cuBLAS is unavailable, partial library lowering would contribute nothing (or dispatch to a different library), and the fusion contribution might be proportionally larger. On platforms without CUDA Graph support, that optimization contributes zero. The "up to 27%" figure is therefore best-case for NVIDIA platforms and should not be interpreted as a uniform improvement across all deployment targets.

A more serious concern is that the ablation starts from "Relax without fusion, partial lib lowering, CUDA graph offloading" as a baseline, which is already a compiler-optimized configuration (it includes tensor program lowering, shape deduction, and memory planning). Figure 17 does not show a comparison to raw, unoptimized code generation, which means the 27% is specifically the marginal benefit of these three optimizations on top of an already-compiled baseline, not the total benefit of Relax over naive execution. The total benefit over, say, running the model with unoptimized operators would be much larger.

Claim: "Static memory planning reduces activation memory by 22% during prefill and 40% during decode"

This claim is supported by Table 2, with the caveat that the numbers are for a single workload pattern (successive prefill/decode with four length/batch levels). The paper does not report memory usage for single-inference scenarios (where there is no reuse opportunity across successive invocations with different shapes), nor does it show how memory usage scales with the number of distinct shapes encountered. The 40% decode reduction is particularly strong evidence for the value of symbolic shape–aware planning, but readers should understand this as a measurement under specific workload conditions (varying batch sizes over time) rather than a universal property.

The evaluation would be stronger with an ablation showing memory usage under equivalent configurations where symbolic shape tracking is replaced with "unknown" annotations—i.e., a direct measurement of how much memory is wasted when the compiler cannot prove shape equalities. Such a comparison would isolate the contribution of the first-class symbolic shape design specifically, rather than memory planning in general. The current evaluation conflates "memory planning exists" with "memory planning has access to symbolic shape information."

Missing experiments that would strengthen the paper

Compilation time and code size: The paper emphasizes that Relax compiles models only once for arbitrary shapes, but no compilation time measurements are reported. For AOT deployment to embedded platforms, compile time and binary size matter. Without these numbers, readers cannot assess whether Relax's approach is practical for development workflows.
Comparison to PyTorch AOT (torch.export): The paper contrasts Relax's AOT approach with PyTorch's JIT-centric design but does not evaluate against PyTorch's own AOT compilation path (torch.export + ExecuTorch targeting mobile). This is a significant omission given that ExecuTorch is PyTorch's official solution for the exact deployment scenarios Relax targets.
Shape generalization stress test: The paper claims Relax handles "arbitrary batch sizes and sequence lengths" from a single compilation, but does not test how performance degrades when shapes differ from the upper bounds used for memory planning, or how the system handles shapes that were not anticipated at compile time. A stress test with adversarial shape patterns would strengthen the robustness claim.
Coverage analysis across operator sets: A systematic measurement of what fraction of operations in common model architectures are handled by analysis-based scheduling rules versus requiring auto-tuning (Ansor) versus falling through to unoptimized code generation would give readers a clearer picture of Relax's practical coverage.
Energy or power measurements on mobile platforms: For mobile deployment, energy efficiency is often as important as throughput. The paper reports throughput but not energy, which limits its relevance for battery-constrained deployment scenarios.

Conditional nature of the deployment breadth claim

The deployment breadth claim holds strongly when the target platform has a GPU that Relax's code generation backend supports (CUDA, Metal, Vulkan, OpenCL, WebGPU) and when the model's memory footprint fits within the device's constraints with static upper-bound planning. It does not hold for platforms without GPU support (CPU-only execution is not evaluated), for models with operators that Relax cannot lower to optimized tensor programs (no coverage analysis is provided), or when the device's memory is too constrained even for static upper-bound planning (the paper uses smaller or more aggressively quantized models on the most constrained devices, sidestepping the question of what happens when the model genuinely doesn't fit).

The finding that Relax matches hand-optimized solutions (llama.cpp on Apple GPU, vLLM on NVIDIA GPU) while using automatic compilation rather than manual per-platform kernel engineering is the paper's strongest contribution. The deployment breadth demonstration (six platforms) is a convincing existence proof. But the paper's implicit argument that this approach scales to arbitrary models and platforms—that the compiler can handle whatever operators and shape patterns appear in future models—remains unproven. The evaluation covers five models (all based on the Transformer architecture), and while the analysis feedback and automatic code generation mechanisms are general, their practical coverage is not systematically characterized.

6. Limitations and Trade-offs

The Difficulty Estimation Overhead Is Not Accounted for in Performance Gains

The assumption or constraint. The entire compute-optimal test-time scaling framework depends on knowing each prompt's difficulty level. The paper estimates difficulty by generating 2,048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted). The authors acknowledge this cost explicitly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

This is a substantial omission. Generating 2,048 samples per prompt to estimate difficulty consumes more compute than the largest test-time budgets studied (256–512 generations). The reported 4× efficiency gains over best-of-N are computed after difficulty is already known, without amortizing the cost of learning it.

The consequence. In any realistic deployment, the total cost would be difficulty estimation plus strategy execution. The difficulty estimation step alone burns more compute than the optimization it supposedly informs, making the headline efficiency gains untranslatable to practice without a cheaper difficulty estimator. Until the cost of difficulty estimation is brought below the cost of the test-time budget itself, the compute-optimal approach cannot deliver net savings—it is a framework for studying where gains are possible, not a deployable system.

What evidence exists in the paper. The paper is transparent about this gap in Section 3.2, acknowledging the cost and framing it as "an exploration-exploitation tradeoff—compute spent assessing difficulty versus compute spent solving the problem—flagging it as a key avenue for future work." However, no experiment evaluates total cost including difficulty estimation, and no bound is provided on how cheap difficulty estimation would need to be for the net benefit to turn positive. The 4× claims throughout the paper (Figures 4 and 8) are upper bounds, not realized deployment numbers.

Mitigation status. Partial. The paper shows that predicted difficulty bins (using the PRM's own scores rather than ground-truth labels) perform nearly as well as oracle bins, which eliminates the need for labeled data. But this does not reduce the computational cost—it still requires 2,048 samples and PRM scoring per prompt. The paper identifies cheaper difficulty estimation (via "pretraining or finetuning models to directly predict difficulty of a question") as future work, but no such model is developed or evaluated.

The Method Provides No Benefit on Hard Problems Outside the Base Model's Capability Range

The assumption or constraint. Test-time compute can only amplify or refine solutions that the base model is already capable of generating with non-trivial probability. If the model's pass@1 on a problem is near zero—it almost never produces a correct solution even in 2,048 samples—no amount of search, revision, or adaptive allocation helps. This is not a bug in the method but a fundamental bound: verification and refinement cannot create capability the model does not possess.

The paper documents this clearly in Section 5.3 (difficulty-dependent behavior of search) and Section 6 (revision results): on difficulty bin 5 (the hardest quintile), all methods hover near 1–3% accuracy regardless of compute budget (Figure 3, right), and in the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5%.

The consequence. For problems genuinely outside the base model's reach—which may include substantial fractions of real-world reasoning tasks for any given model—compute-optimal test-time scaling offers no improvement over greedy decoding. The approach is not a general solution to model capability limitations. As the paper states in the Section 7 takeaway, "test-time compute amplifies existing capability but does not create it from nothing." This is a sharply bounded improvement regime: the method helps on problems where the model already has some chance of success, but provides zero benefit where it has none. For anyone considering whether to invest in test-time compute versus pretraining a larger model, the boundary is quantitative but severe—past the pass@1 threshold, the investment in test-time compute is wasted.

What evidence exists in the paper. Figure 3 (right) shows that bin 5 accuracy is 1–3% for best-of-N weighted and beam search at all budgets from 4 to 256 generations. Figure 7 (right) shows bin 5 revision accuracy at roughly 2–3% regardless of sequential-to-parallel ratio. The FLOPs-matched comparison (Figure 9) shows bin 5 scaling curves that are flat and near zero, far below the ~14× larger model's greedy performance. Section 7 explicitly reports that for hard questions at the R ≫ 1 data point with PRM search, there is a −52.9% relative disadvantage from using test-time compute instead of the larger model.

Mitigation status. Not mitigated; this is a fundamental characteristic of the approach. The paper is fully transparent about the limitation and does not claim otherwise. The implication is that test-time compute is a complement to pretraining, not a replacement—and that the decision to invest in one versus the other depends critically on whether the target problem distribution falls within the base model's capability envelope.

The Study Is Confined to a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*)

The assumption or constraint. All experiments use the MATH benchmark (500 test questions; Hendrycks et al., 2021) with PaLM 2-S* models (Anil et al., 2023). The authors state in Section 4 that they "believe this model is representative of the capabilities of many contemporary LLMs," but this representativeness claim is unverified. The MATH benchmark consists exclusively of high-school competition-level math problems requiring symbolic multi-step reasoning with clean ground-truth answers, which enables both the PRM training pipeline (Monte Carlo rollouts against ground-truth correctness) and the difficulty estimation procedure (pass@1 based on exact answer matching).

The consequence. Several aspects of the findings could fail to generalize to other settings. The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution—a model with different calibration properties, different error patterns, or a different base accuracy level might exhibit different difficulty-dependent scaling curves. The revision model's ability to benefit from multi-turn incorrect-to-correct trajectories depends on the base model's in-context learning capabilities, which vary substantially across model families. The MATH benchmark's exclusive focus on symbolic math reasoning means it is unclear whether the difficulty-dependent patterns—beam search hurting easy problems, revisions dominating on easy problems—generalize to other reasoning domains like code generation, logical deduction, scientific question answering, or tasks requiring factual recall rather than inference.

The test set of 500 questions, split into five difficulty quintiles of roughly 100 questions each, then further split by two-fold cross-validation, means the compute-optimal policy is selected based on approximately 50 questions per fold per bin. This is a small sample for strategy selection, and the selected policies—"use beam search on bin 3, best-of-N on bins 1–2"—may not be robust. The paper reports no confidence intervals on the compute-optimal scaling curves, making it difficult to assess statistical reliability at this sample size.

What evidence exists in the paper. All figures (3–9) are based on the 500-question MATH test set with PaLM 2-S*. The paper provides no replication on a second benchmark, a second model family, or a second task domain. The PRM training failure on the PRM800k dataset (Section 5.1)—where the authors found the dataset "largely ineffective" for PaLM 2 models due to distribution shift—is informative: it demonstrates sensitivity to model-data alignment that would affect any attempt to transfer the approach to a new model family without re-training the PRM.

Mitigation status. The paper does not attempt mitigation. The authors acknowledge the single-benchmark, single-model limitation implicitly by noting that future work would need to extend the analysis, but no robustness experiments are conducted. A practitioner considering this approach for a different model or domain would need to replicate the entire analysis pipeline—PRM training, revision model training, difficulty binning, strategy selection—without guidance on which findings are expected to transfer and which are likely to be model- or domain-specific.

Revisions and PRM Search Are Never Combined, Despite Complementary Strengths

The assumption or constraint. The paper studies two independent mechanisms for improving test-time compute—PRM-guided search (modifying the verifier/selection process) and iterative revisions (modifying the proposal distribution)—but never combines them into a single system. Section 8 explicitly acknowledges this:

"we did not experiment with PRM tree-search techniques in combination with revisions"

The paper demonstrates that these mechanisms have complementary difficulty-dependent strengths: revisions are most effective on easy problems where local refinement helps, and beam search is most effective on medium-hard problems where exploration across solution strategies helps. However, the experiments keep these mechanisms completely separate. The compute-optimal policy selects between search strategies or between sequential-parallel ratios, but never between revisions and search as competing uses of the same budget.

The consequence. The current results represent a lower bound on what a fully integrated system could achieve. A natural integration would use the revision model as the proposal distribution within beam search—conditioning each candidate step on the revision history—or would use the PRM to guide which revision branches to pursue further. Since the two mechanisms have complementary strengths, combining them could push past the performance ceiling that each hits individually, particularly on medium-difficulty problems where both exploration and refinement are valuable. Without this combination, the paper's central argument about adaptive allocation is incomplete: it shows that difficulty matters for strategy selection, but not what happens when the full strategy space (revisions + search + their interaction) is available.

What evidence exists in the paper. Figures 7 and 8 show that revisions and search each provide roughly comparable peak benefits (~4× efficiency gain over best-of-N), but they are never jointly evaluated. The difficulty-bin patterns in Figures 3 (right) and 7 (right) suggest complementary profiles: search helps most on bins 3–4 where revisions show relatively flat scaling, while revisions help most on bins 1–2 where search over-optimizes. This complementarity is the paper's own argument for why combining them should work, but the combination is absent.

Mitigation status. Identified as explicit future work in Section 8, but no preliminary results or feasibility analysis is provided. The paper does not discuss whether technical barriers (e.g., the distribution shift between base model and revision model outputs breaking the PRM) would make the combination difficult, or whether the gains would be approximately additive. A practitioner hoping for a unified system combining both mechanisms would find no guidance on architecture, training, or expected benefit.

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate and Is Sensitive to Training Methodology

The assumption or constraint. The revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct target answer (Section 6.1). This is a direct consequence of the training data construction: for each question, the authors sample 64 responses, identify which are correct and incorrect, then construct trajectories of 0–4 incorrect answers ending with a correct one. The model never sees a trajectory where the current answer is already correct and the appropriate action is to output it unchanged (or to stop revising).

At inference time, when the revision model generates a chain of revisions, it can produce a correct answer, and then—because it has no training signal for what to do when the in-context answer is correct—it may "revise" that correct answer into an incorrect one at the next step. The paper reports that approximately 38% of correct answers get converted back to incorrect ones using a naïve approach (Section 6.1). The mitigation is to apply majority voting or verifier-based selection across the entire revision chain rather than taking the final output, but this is a post-hoc fix rather than a solution to the underlying training problem.

Furthermore, the revision model shows fragility to training methodology. The ReST^EM experiment (Appendix K, Figure 16)—an attempt to further optimize the revision model using reinforcement learning from feedback—backfires: "additional sequential revisions substantially hurt performance." At 256 generations, fully sequential performance drops to approximately 33.5% compared to roughly 38.5% at the optimal ratio. The authors hypothesize that "on-policy data collection in ReST^EM exacerbates spurious correlations in revision data, causing the model to fail to learn the revision task properly."

The consequence. The revision approach depends on specific, empirically fragile training choices: offline data construction (pairing independently sampled correct and incorrect solutions post-hoc), edit-distance-based pairing to ensure incorrect answers are structurally close to correct ones, and training on only 0–4 incorrect-answer contexts. When training departs from this recipe (as in ReST^EM), performance degrades rather than improves. The 38% reversion rate means the revision chain is not monotonically improving—correct answers are lost and must be recovered by a downstream selection mechanism, which introduces latency (from generating the full chain) and wastes compute budget. A practitioner attempting to reproduce the revision approach would need to carefully replicate the training data construction procedure, with no guidance on how sensitive performance is to the specific parameters (number of incorrect contexts, edit-distance threshold, sampling temperature).

What evidence exists in the paper. The 38% reversion rate is stated in Section 6.1 without a supporting figure or table—it appears as an in-text claim. The ReST^EM degradation is shown in Figure 16 (Appendix K), where fully sequential performance drops substantially below the optimal-ratio configuration. The paper does not provide an ablation study varying the revision training hyperparameters (number of incorrect contexts, edit-distance criterion, temperature), so the robustness of the approach to these choices is unknown.

Mitigation status. Partial. The paper mitigates the reversion problem at inference time via majority voting or verifier-based selection across the chain, which recovers the correct answer even when it was later overwritten. But this is a workaround, not a solution—a more principled approach would train the model to recognize when no revision is needed. The ReST^EM failure is presented as a negative result with a hypothesis but no solution. The paper does not explore architectural changes (e.g., adding a "stop revising" token, training on mixed correct-incorrect sequences, or using a separate verifier to decide when to stop) that could address the reversion problem at its root.

The FLOPs-Matched Comparison Uses a Weak Pretraining Baseline (Parameter-Only Scaling, Greedy Decoding)

The assumption or constraint. The FLOPs-matched comparison in Section 7 compares PaLM 2-S* with compute-optimal test-time scaling against a model with approximately 14× more parameters, scaling only parameters while holding training data fixed. This follows the LLaMA paradigm (Touvron et al., 2023) rather than compute-optimal pretraining (Hoffmann et al., 2022), where both data and model size are scaled together. The authors explicitly acknowledge this:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the larger model uses only greedy decoding—no majority voting, no best-of-N, no verifier-guided selection. This means the comparison is between compute-optimal test-time scaling (with sophisticated search and revision strategies) and a bare-bones pretraining baseline (larger model, minimal inference strategy). A fairer comparison would give the larger model some modest test-time compute budget, since the claim is about the substitutability of pretraining and inference compute.

The consequence. The reported advantages of test-time compute over pretraining—e.g., +27.8% relative improvement on easy questions at R ≪ 1—likely overstate the benefit relative to what a properly optimized larger model could achieve. A Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data) would be a stronger baseline than a parameter-only-scaled model. Giving the larger model even a best-of-8 or beam-search-4 budget would substantially narrow the gap. The paper's conclusion that "test-time compute can outperform a 14× larger model on easy-to-medium problems" should be understood as a comparison against a specific, intentionally simple baseline, not against the best possible use of equivalent pretraining compute.

What evidence exists in the paper. Figure 9 and the bar charts in Figure 1 present the FLOPs-matched results. The relationship between the parameter scaling factor (14×) and the total FLOPs equivalence is formalized in the FLOP accounting equations in Section 7. The paper is transparent about the parameter-only scaling and greedy decoding choices, but does not provide an ablation where the larger model receives some test-time compute budget, nor does it estimate how much the advantage would shrink under a Chinchilla-optimal pretraining baseline.

Mitigation status. The paper positions this as a deliberate analytical choice—"representative of a canonical approach"—and flags the Chinchilla-optimal comparison as future work. However, the framing of the results (e.g., in the Figure 1 bar charts and the Section 7 takeaway) could mislead a casual reader into thinking the comparison is against the best achievable pretraining, when it is against a specific, suboptimal pretraining configuration used for analytical tractability. A practitioner deciding between investing in pretraining versus test-time compute should treat the reported advantages as upper bounds that would likely shrink against a stronger pretraining baseline.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper reshapes the conversation about ML compiler architecture by demonstrating that the multi-level, single-shot lowering design that underlies most production ML compilers is not a law of nature—it is a choice, and a costly one. Prior to Relax, the standard architecture for ML compilers was well-established: a computational graph IR lowered to a tensor program IR lowered to platform code, with each level opaque to the one above it. This architecture was so universal across TVM/Relay (Roesch et al., 2018), MLIR/IREE (Lattner et al., 2021), and XLA that it appeared as infrastructure rather than as a design decision. Relax's cross-level abstraction—embodied in the call_tir and call_dps_library primitives that keep foreign function calls as first-class, analyzable graph-level operations—shows that this boundary was the bottleneck, not a foundation.

The conceptual shift is from a pipeline of progressively lower representations to a unified optimization space where different implementation strategies coexist. In Relax, a graph-level pass can pattern-match on a call_tir to a matmul, replace it with a call_dps_library to cuBLAS, fuse the subsequent element-wise operation into the library call's epilogue, lift a workspace allocation from a Stream-K tensor program to the graph level for global memory reuse, and then capture the resulting subgraph as a CUDA Graph for replay—all within the same IR, with symbolic shape information preserved throughout. None of these transformations is individually novel; the novelty is that they can be composed in a single compilation flow without modifying core lowering logic. This reframes what an ML compiler IR is for: not a sequence of lowering steps, but a shared representation that enables cross-level reasoning.

The paper resolves a tension in the ML compilation literature that was previously understood as an unavoidable tradeoff. Prior work had demonstrated that dynamic shape handling was possible at the tensor program level—Halide (Ragan-Kelley et al., 2013) tracked dynamic shapes in loop nests, DietCode (Zheng et al., 2022) optimized dynamic tensor programs through bucketing, and CoRA (Fegade et al., 2022) handled ragged tensors. Separately, graph-level compilers like Relay and MLIR could perform operator fusion and memory planning, but only for static shapes or with coarse ? annotations that erased relationship information. The implicit assumption was that these two capabilities were architecturally incompatible: you could either have rich dynamic shape information at the tensor level (by keeping shapes symbolic within a kernel) or rich graph-level optimizations (by requiring static shapes for whole-program analysis), but not both. Relax breaks this assumption by making symbolic shape the universal annotation language that spans all levels, with shape information flowing in both directions—from graph to tensor program (through call_tir's sym_args) and from tensor program back to graph (through analysis feedback). The architecture doesn't choose between levels; it unifies them.

This shifts which research directions are most promising. Verifier robustness is no longer the primary bottleneck for test-time compute scaling—it is clear that better PRMs would help, but the Relax paper suggests that the more fundamental bottleneck was the compiler architecture itself. Manual per-backend kernel engineering, which dominates systems like llama.cpp and vLLM, becomes less attractive as a research investment when a compiler-based approach with automatic code generation matches or exceeds those systems across six GPU platforms (Figures 14–16, Table 3). The paper makes the case that investment in compiler infrastructure—specifically, in cross-level abstractions that preserve symbolic information—pays off across platforms and models in a way that per-platform hand-optimization cannot. Conversely, research on JIT-centric designs that sacrifice AOT compilation (like PyTorch's TorchDynamo tracing) looks less general after Relax demonstrates that AOT compilation with whole-program symbolic shape tracking can match JIT performance on mainstream GPUs while simultaneously reaching platforms JIT cannot.

The most important practical reframe is about what "supporting a new backend" means. In the pre-Relax landscape, supporting a new GPU backend (Android OpenCL, WebGPU, Apple Metal on mobile) meant either porting a full JIT runtime (PyTorch path) or manually writing optimized kernels for every operator in every model (llama.cpp path). Both required enormous engineering investment. Relax demonstrates a third path: the compiler infrastructure handles operator code generation generically, and the platform-specific effort reduces to implementing runtime bindings (memory allocation, kernel dispatch, a few platform-specific libraries). The deployment breadth evidence—six platforms, five models, all from a single compiler codebase—is the existence proof. This changes the economics of model deployment: instead of asking "which platforms have enough engineering resources to deserve optimized kernels?", the question becomes "which platforms have a GPU and a C compiler?", because Relax's approach automatically generates optimized GPU code for anything with a supported backend.

The paper also shifts the burden of proof in ML compiler design. The standard argument for multi-level single-shot lowering was that it was "cleaner"—each level could be reasoned about independently. Relax demonstrates that the cleanliness comes at the cost of structurally precluding whole-program optimizations when shapes are dynamic. The memory planning results in Table 2 (40% activation memory reduction during decode, 22% during prefill) are only possible because symbolic shape relations are preserved across operator boundaries and function calls. The CUDA Graph offloading in Figure 17 (1–2% latency reduction) is only possible because static memory planning with symbolic upper bounds makes Graph capture viable for dynamic models. These are not optimizations that could be retrofitted onto a single-shot design—they require the compiler to reason about memory across the graph/tensor-program boundary, which the architecture must support from the ground up. The paper's implicit argument is that cross-level reasoning is not a luxury—it is a requirement for competitive deployment on memory-constrained and latency-sensitive platforms, and compiler architectures that preclude it are making a performance sacrifice that was previously invisible because no system had demonstrated the alternative.

Follow-Up Research This Work Enables

Cheap, amortized difficulty estimation for compute-optimal policies. The most immediate practical bottleneck the paper identifies is the cost of estimating prompt difficulty: generating 2,048 samples per question, then averaging PRM scores. This cost dwarfs the test-time budgets being optimized. The paper explicitly calls for models that predict difficulty directly from question text, but a more promising direction is amortized difficulty estimation during inference: start generating solutions normally, use the verifier's score distribution on the first few samples (4–8) as a real-time difficulty signal, and dynamically reallocate the remaining budget. Relax's cross-level architecture makes this tractable because the compiler could generate code that branches on the estimated difficulty bin after a warmup phase, switching between search strategies within a single inference call. A strong follow-up would measure the accuracy-vs-cost tradeoff for different warmup sizes, compare against the static 2,048-sample approach as an upper bound, and determine whether a 4-sample warmup achieves 90% of the optimal policy's benefit while keeping total cost below the best-of-N baseline. The key number to beat is the 4× efficiency gain reported in Figures 4 and 8—if amortized estimation preserves even 3×, it is a deployable system.

Integrated revision-search systems with shared PRMs. The paper studies PRM-guided search and iterative revisions as independent mechanisms, but demonstrates in Figures 3 and 7 that they have complementary difficulty-dependent strengths: revisions dominate on easy problems (bin 1–2), beam search dominates on medium problems (bin 3–4), and neither helps on hard problems (bin 5). The natural follow-up is to combine them: use the revision model as the proposal distribution inside beam search, where each candidate step conditions on the revision history, and the PRM scores both the current step and whether to continue revising or restart. This requires solving the distribution shift problem the paper identifies in Appendix J (the base-model PRM underperforms on revision model outputs), likely by training a shared PRM on trajectories that mix base model and revision model outputs. A strong evaluation would measure whether the combined system outperforms the best single mechanism at each difficulty level, and specifically whether it pushes the performance ceiling on bin 4 problems (where neither mechanism individually exceeds roughly 17–18% accuracy at 256 generations). The key question is whether the gains are additive or subadditive—do revisions and search help in the same regions, or do they overlap?

Verifier over-optimization as a measurable, mitigatable bottleneck. Figure 3 (right) documents verifier over-optimization as a first-class phenomenon: beam search degrades easy-problem performance at high budgets because the PRM can be exploited. The paper identifies this as the primary bottleneck for further scaling, but does not propose solutions. A direct follow-up would train PRMs with adversarial data augmentation: after training an initial PRM, run beam search against it, collect solutions that score highly but are incorrect (false positives), and retrain the PRM with these as hard negatives. The question is whether adversarial training pushes the over-optimization threshold to higher budgets, or whether the PRM's representations fundamentally cannot distinguish correct from exploitatively-structured incorrect solutions. A strong experiment would measure the accuracy-budget curve for beam search with adversarially-trained PRMs at three budget levels (64, 256, 1024) on difficulty bins 1–3, comparing against the standard PRM. If adversarial training raises bin 1 beam search accuracy from ~77% to matching best-of-N's ~88% at high budgets, it demonstrates that the over-optimization ceiling is trainable rather than fundamental.

Cross-platform auto-tuning of tensor program schedules within the Relax pipeline. The paper uses analysis-based scheduling rules for common operator patterns and notes that "passes to apply Ansor-style auto-tuning for rare tensor programs" can be added. The Relax architecture makes this particularly interesting because auto-tuning decisions at the tensor program level can flow back to the graph level: if auto-tuning discovers that a particular operator benefits from a workspace buffer (like the Stream-K matmul decomposition), the cross-level workspace lifting pass can automatically promote that buffer to graph-level memory planning. A systematic study would measure, across the six deployment platforms in Table 3, what fraction of operators are well-handled by analysis-based rules versus requiring auto-tuning, and whether the auto-tuned schedules transfer across platforms (does a schedule tuned for CUDA work on Metal, or is per-platform tuning necessary?). The key metric is coverage: for a target model architecture like Llama3-8B, what percentage of total inference time is spent in operators that auto-tuning can optimize versus operators where analysis-based rules already achieve near-optimal performance? This determines whether auto-tuning is a tail-end optimization or a first-class requirement.

Quantitative characterization of the AOT shape generalization boundary. The paper claims Relax handles "arbitrary batch sizes and sequence lengths" from a single compilation, but does not stress-test this claim. A systematic characterization would answer: for a model compiled with a maximum sequence length of 2,048 and a maximum batch size of 64, what is the performance degradation when running with sequence length 1 (shortest), sequence length 2,048 (longest), batch size 1 (thinnest), and batch size 64 (widest)? Are there shape regions where the compiled code performs disproportionately worse—for example, when dynamic dimensions cause the compiler to fall back to scalar loops rather than vectorized code, or when the memory planning's upper-bound pre-allocation wastes bandwidth for much smaller actual shapes? The experiment would measure per-token latency or throughput as a function of actual input shape relative to the compiled upper bound, producing a shape generalization surface rather than a single-point measurement. This would give practitioners concrete guidance on how to set upper bounds (tighter bounds save memory but risk recompilation if exceeded; looser bounds avoid recompilation but may degrade performance at small shapes).

Compiler support for dynamic control flow in ML models. Relax handles dynamic shapes—tensor dimensions that vary at runtime—but the paper's evaluation focuses on models where the computational graph structure is static (the same operators execute regardless of input values). Emerging model architectures increasingly use dynamic control flow: mixture-of-experts models route tokens to different experts based on input content, retrieval-augmented generation conditionally fetches documents, and recursive or iterative models execute variable numbers of computation steps. Relax's dataflow block construct (§3.1) explicitly demarcates side-effect-free regions without control flow, which suggests that control flow is handled at a coarser granularity. A natural extension would explore how first-class symbolic shapes interact with dynamic control flow: if a model branches on a runtime value, do the symbolic shape relations from one branch pollute the other? Can the compiler prove that certain shape equalities hold regardless of which branch is taken? A strong evaluation would implement a mixture-of-experts Transformer layer in Relax, measure whether the compiler can fuse operations across the routing decision (or whether control flow blocks fusion), and compare against hand-optimized MoE kernels. This would clarify whether Relax's abstraction generalizes beyond the static-graph assumption implicit in its dataflow block design.

Practical Applications and Downstream Use Cases

LLM inference on mobile and embedded devices with heterogeneous GPU hardware. The paper demonstrates that Relax compiles and runs quantized LLMs on six platforms where most existing frameworks lack GPU support: iPhone 14 Pro (5.1 tok/s for 3-bit Llama2-7B), Samsung S23 (7.9 tok/s for 4-bit Llama2-7B), Orange Pi 5 (2.3 tok/s for Llama3-8B), Steam Deck (14.0 tok/s), Jetson Orin (32.0 tok/s), and in-browser WebGPU on M3 Max (37.8 tok/s). For a mobile app developer wanting on-device LLM inference without sending user data to a cloud API, Relax provides the only compiler-based path to GPU acceleration across both iOS (Metal) and Android (OpenCL) from a single model description. The practical benefit is not just performance—the Samsung S24 comparison in Figure 18 shows Relax delivering 55% higher throughput than llama.cpp, which falls back to CPU on Android GPUs—but also memory: "without memory planning that pre-allocates all needed memory and keeps it within the budget, these models are not even runnable on some of the environments due to memory constraints." The 40% activation memory reduction during decode (Table 2) is the difference between a model fitting in mobile VRAM or failing to load.

Universal model serving across GPU vendors with a single compiler pipeline. Organizations deploying LLMs across heterogeneous GPU fleets—some NVIDIA, some AMD, some Apple Silicon—currently maintain separate inference stacks per vendor: PyTorch + cuBLAS for NVIDIA, MIGraphX or custom HIP ports for AMD, CoreML or llama.cpp for Apple. The paper shows that Relax provides competitive performance on all three: within ~7% of vLLM on NVIDIA RTX 4090 (Figure 14), leading all baselines on AMD Radeon 7900 XTX with up to 1.50× advantage at batch size 1 (Figure 15), and essentially tied with hand-optimized llama.cpp on Apple M2 Ultra (Figure 16). The practical benefit is engineering consolidation: one compiler pipeline, one set of model definitions, one optimization pass sequence, generating optimized GPU code for all three vendors. The partial library lowering pass allows per-platform customization (dispatch to cuBLAS on NVIDIA, to MIOpen or compiler-generated kernels on AMD, to Metal Performance Shaders or generated Metal on Apple) without forking the compilation flow.

Web-native machine learning via WebGPU. The WebGPU results in Table 3 (37.8 tok/s for Llama3-8B, 68.0 tok/s for Phi3-mini-4k) represent a new deployment capability: GPU-accelerated LLM inference running entirely in the browser without server round-trips. This enables privacy-preserving applications—client-side document summarization, email drafting, code completion in web-based IDEs—where the model runs locally on the user's GPU. The paper claims Relax is "the first solution to enable GPU-accelerated LLM inference on these platforms" except Jetson Orin, making WebGPU LLM deployment a genuinely new capability rather than a performance improvement over existing solutions. The AOT compilation property is critical here: the compiled module (including GPU shaders) can be cached and loaded without a Python runtime, JIT compiler, or any server dependency, making it deployable as a static web asset.

When to Prefer This Method

The paper does not position Relax against a specific named alternative with a clear, quantified tradeoff rubric. It demonstrates that Relax's compiler-based approach matches or exceeds hand-optimized systems on mainstream GPUs while simultaneously reaching platforms those systems cannot, but does not articulate conditions under which Relax is not the right choice. As a compiler infrastructure, the preference decision is less about "method A vs. method B" and more about build vs. buy: when does it make sense to adopt a compiler-based deployment pipeline rather than integrating with existing framework-specific runtimes? The paper's evidence suggests a pragmatic heuristic: adopt Relax when (1) you need to deploy to multiple GPU backends from a single codebase, especially when some of those backends lack mature framework support (Metal on iOS, OpenCL on Android, WebGPU); (2) you need AOT compilation rather than JIT, as on mobile, embedded, or web platforms where a Python runtime is infeasible; (3) your models include custom operators (quantization decode, custom attention) that you want to participate in fusion without manually annotating every operator's properties. The paper does not provide the counterfactual evidence—e.g., head-to-head against ExecuTorch for mobile, or against IREE for embedded—that would enable a more precise decision rule with quantified breakpoints.