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

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.

Authors

Institutions

Publication Details

Published
2025-03-27
DOI
https://doi.org/10.1145/3676641.3716249
Citations
8
Primary Topic
Machine Learning and Algorithms
Type
article
Field-Weighted Citation Impact
19.85
Controls
|||
ALL TIME
JAN
FEB
MAR
APR
MAY
JUN
JUL
AUG
SEP
OCT
article

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

Prakalp Srivastava, Junru Shao, Tianqi Chen, Ruihang Lai et al.
8 citations
Machine Learning and Algorithms
19.85
article

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

Prakalp Srivastava, Junru Shao, Tianqi Chen, Ruihang Lai, Sunghyun Park, Bohan Hou, Jared Roesch, Wuwei Lin, Todd C. Mowry, Li-Jie Jin, Steven Lyubomirsky, Jiawei Liu, Hongyi Jin, Zihao Ye, Zhenxuan Jiang, Y. Wu, S.M. Feng, Yaxing Cai, Yuchen Jin
article en
8 citations

Abstract

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.

University of Illinois Urbana-Champaign (US), Shanghai Jiao Tong University (CN), University of Washington (US), Nvidia (United States) (US), OpenAI (United States) (US), Seattle University (US), Carnegie Mellon University (US), Netflix (United States) (US)
Openalex Percentile: Top 1%
Machine Learning and Algorithms
19.85
DOIOpen Access PDFOpenAlex
AI Navigator

Ask Laika to Summarize, Analyze, and Connect papers live on the map.

Summarize Papers & Methodologies

Extract key findings, datasets, and comparative methods across publications.

Benchmark Rankings & Visual Analytics

Rank top research institutions, authors, funders, topics, and journals by Field-Weighted Citation Impact (FWCI) and paper volume with instant charts.

Connect Distant Disciplines

Bridge topological clusters on the map to find hidden collaborative intersections.