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

arXiv: https://arxiv.org/abs/2311.02103 · HTML: https://mlsanity.com/paper/2311.02103/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 45 references that are in the map.
- First posted on arXiv 1 Nov 2023.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Visual Instruction Tuning](https://mlsanity.com/paper/2304.08485/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [WizardLM: Empowering large pre-trained language models to follow complex instructions](https://mlsanity.com/paper/2304.12244/) (2023)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python](https://mlsanity.com/paper/2112.08429/) (2021)
- [TensorIR: An Abstraction for Automatic Tensorized Program Optimization](https://mlsanity.com/paper/2207.04296/) (2022)
- [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022)
- [Relay: A New IR for Machine Learning Frameworks](https://mlsanity.com/paper/1810.00952/) (2018)
- [Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU](https://mlsanity.com/paper/2301.03598/) (2023)
- [Nimble: Efficiently Compiling Dynamic Neural Networks for Model Inference](https://mlsanity.com/paper/2006.03031/) (2020)
- [The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding](https://mlsanity.com/paper/2110.10221/) (2021)
- [Cortex: A Compiler for Recursive Deep Learning Models](https://mlsanity.com/paper/2011.01383/) (2020)
- [MIOpen: An Open Source Library For Deep Learning Primitives](https://mlsanity.com/paper/1910.00078/) (2019)

## Built on by (most cited first, 12 of 12)

- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [Towards Efficient Generative Large Language Model Serving: A Survey from Algorithms to Systems](https://mlsanity.com/paper/2312.15234/) (2023)
- [Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References](https://mlsanity.com/paper/2510.14719/) (2025)
- [Event Tensor: A Unified Abstraction for Compiling Dynamic Megakernel](https://mlsanity.com/paper/2604.13327/) (2026)
- [Productively Deploying Emerging Models on Emerging Platforms: A Top-Down Approach for Testing and Debugging](https://mlsanity.com/paper/2404.09151/) (2024)
- [Pruning Large Language Models to Intra-module Low-rank Architecture with Transitional Activations](https://mlsanity.com/paper/2407.05690/) (2024)
- [Local deployment of large-scale music AI models on commodity hardware](https://mlsanity.com/paper/2411.09625/) (2024)
- [Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference](https://mlsanity.com/paper/2605.11581/) (2026)
- [A Survey on Design Methodologies for Accelerating Deep Learning on Heterogeneous Architectures](https://mlsanity.com/paper/2311.17815/) (2023)
- [GeoT: Tensor Centric Library for Graph Neural Network via Efficient Segment Reduction on GPU](https://mlsanity.com/paper/2404.03019/) (2024)
- [RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators](https://mlsanity.com/paper/2603.10026/) (2026)
- [DanLing NestedTensor: Composable Multi-Ragged Tensors for Deep Learning](https://mlsanity.com/paper/2609.30379/) (2026)

## Abstract (arXiv)

Dynamic shape computations have become critical in modern machine learning workloads, especially in emerging large language models. The success of these models has driven the demand for their universal deployment across a diverse set of backend environments. In this paper, we present Relax, a compiler abstraction for optimizing end-to-end dynamic machine learning workloads. Relax introduces a cross-level abstraction that encapsulates computational graphs, loop-level tensor programs, and external library calls in a single representation. Relax also introduces first-class symbolic shape annotations to track dynamic shape computations globally across the program, enabling dynamic shape-aware cross-level optimizations. We build an end-to-end compilation framework using the proposed approach to optimize dynamic shape models. Experimental results on LLMs show that Relax delivers performance competitive with state-of-the-art systems across various GPUs and enables deployment of emerging models to a broader set of emerging environments, including mobile phones, embedded devices, and web browsers.