PaperlayerView exact revision on arXiv

Source-linked AI summary

RLX: A Unified Multi-Backend Tensor Compiler and Distributed Runtime in Rust

Eugene Hauptmann, Nataliya Kosmyna

arXiv:2609.37916v1cs.DCcs.AI

TL;DR

Production ML stacks split graph compilation and kernel execution across layers and languages, complicating end-to-end reasoning about deployment and backend behavior. RLX unifies compiler and runtime roles around a shared IR and explicit dispatch, and RLX-Metal is fastest at every tested batch for all-MiniLM-L6-v2.

  • Problem

    Production ML stacks often split graph compilation and kernel execution across layers and languages, complicating end-to-end reasoning about backend behavior and deployment.

    Source: §I · p. 1

  • Method

    RLX combines compiler and runtime roles in one Rust system around a shared IR, with explicit per-operator dispatch that fails compilation when an operator remains unsupported.

    Source: §I · p. 1, §IV · p. 3

  • Results

    16.6 ms at batch 32 for RLX-Metal versus 26.7 ms for PyTorch-MPS; RLX-Metal is fastest at every batch on all-MiniLM-L6-v2.

    Source: B. Inference: RLX backends vs. other… · p. 5

  • Takeaways & Limitations

    RLX’s shared system lets a model authored once target Apple Silicon, NVIDIA, AMD, TPU, web, and MCU without a rewrite.

    Source: §IX · p. 6

Abstract

from arXiv · show

Production machine learning (ML) stacks often split graph compilation and kernel execution across different layers and languages, making backend behavior, deployment guarantees, and performance fallbacks hard to reason about end-to-end. RLX addresses this gap with a single Rust codebase that combines compiler and runtime roles around one primitive-level, three-level intermediate representation (IR), plus a transparent dispatch contract that resolves each operator to native, common-IR, or rewritten lowering and fails compilation when legalization is not possible. The same IR targets fourteen runtime devices (cpu, metal, mlx, ane, cuda, rocm, oneapi, tpu, hexagon, gpu, vulkan, opengl, directx, webgpu) and two specialty codegen paths (Cortex-M INT8 and FPGA), ingests safetensors, GGUF, ONNX, and rten formats, supports F16/BF16/F64/C64 and quantized INT4/INT8 flows with AMP/PTQ/QAT, and scales via tensor-/pipeline-parallel collectives over TCP and RDMA transports. Beyond neural workloads, RLX also extends to scientific/physics-style domains through sparse and dense linear algebra extensions (e.g., CSR LU/CG/matvec and LAPACK- backed factorizations) and 3D Gaussian splatting operators. We evaluate RLX against PyTorch, TensorFlow, JAX, candle, burn, tch, rten, MLX, CoreML, IREE, Glow, TensorRT, and tinygrad under identical input generation and p50 measurement methodology on one host. On all-MiniLM-L6-v2, RLX-Metal is fastest at every batch (e.g., 16.6 ms at batch 32 vs. PyTorch-MPS 26.7 ms). In the MNIST training table, RLX also has the top-throughput entry (graph-fused MLP: 946,487 img/s), above NumPy+BLAS (787,349 img/s), while retaining 100% top-1 parity on reference checks (e.g., Qwen3).

I. INTRODUCTION

ML stacks combine graph compilers and kernel runtimes through adapter layers and often Python, making small-batch latency and deployment verification difficult. RLX brings compiler and runtime roles into one Rust system with a unified IR and transparent dispatch.

I. INTRODUCTION

  • Python’s interpreter lock, reference counting, and dynamic dispatch add latency at small batches and complicate deployment to limited or nontraditional targets.These factors also make it difficult to verify end-to-end whether the code will run.Source: §I · p. 1
  • RLX combines a graph compiler and kernel runtime in one memory-safe Rust system, with one IR spanning server GPUs to INT8 microcontrollers.Its optimizer includes fusion, automatic differentiation, batching, precision policies, and post-training quantization insertion.Source: §I · p. 1
  • RLX contributes a unified IR, optimizer, and runtime alongside transparent dispatch, portability, distributed tensor- and pipeline-parallel execution, and runnable cross-framework benchmarks.The dispatch fails rather than silently degrading when an operator cannot be legalized.Source: §I · p. 1

II. BACKGROUND AND MOTIVATION

RLX is designed for heterogeneous targets, a single-language compiler/runtime stack, transparent operator handling, and extensibility. Rust supports keeping its system components together while still calling vendor libraries.

A. Requirements

  • RLX’s requirements include reaching heterogeneous targets without per-device rewrites, keeping the system in one language, reporting operator lowering paths, and extending primitives to new domains.It must fail when an operator cannot be legalized for the selected device.Source: A. Requirements · pp. 1–2

B. Motivation behind choosing Rust

  • Rust provides memory and thread safety without garbage collection and zero-cost FFI to vendor libraries, keeping the IR, optimizer, and backend kernels in one tree.RLX differentiates itself from Burn and Candle by coupling a primitive-level compiler IR to transparent dispatch and MCU/FPGA targets.Source: B. Motivation behind choosing Rust · p. 2

C. Related work

RLX draws on prior array-transformation and whole-program compiler ideas, then organizes authoring, optimization, runtime, and hardware lowering around one IR. Its three IR levels share an operator set while supporting different stages of the compile pipeline.

C. Related work

  • RLX draws its representation from JAX and XLA, its whole-program approach from TVM, MLIR, IREE, and Glow, and its ergonomics from PyTorch.Unlike MLIR, which is compiler infrastructure, RLX is described as a complete system with all three layers.Source: C. Related work · p. 2

III. SYSTEM DESIGN

  • RLX’s system design traces front ends and weight formats into one IR, which the optimizer and runtime lower, dispatch, and distribute across hardware backends.The cited architecture description identifies safetensors, GGUF, and ONNX as loaders.Source: §III · p. 2

A. Intermediate representation

  • The primitive-level IR has roughly 113+ operators, including control flow, masked attention, recurrent operators, Fourier transforms, and quantized matrix multiplication.Custom operators can be registered by downstream crates without modifying the core.Source: A. Intermediate representation · p. 2
  • HIR represents model-building blocks, MIR is the fused primitive DAG consumed by optimization, and buffer planning produces LIR for backend lowering.The lazy Tensor DSL traces directly to MIR, while graph builders can author at HIR.Source: A. Intermediate representation · p. 2

B. Optimizer

The optimizer transforms the shared IR through fusion, differentiation, legalization, memory planning, and precision handling. These passes feed device-specific schedules while retaining low-precision and quantization flows in the IR pipeline.

B. Optimizer

  • Fusion combines MIR regions and can unfuse them to expose primitives for differentiation or backends that lack a fused form.Examples include fused attention, fused QK-RoPE, and a strided Narrow×3 →Attention pattern.Source: B. Optimizer · p. 2
  • The optimizer exposes fusion, autodiff, and compilation through one facade, while autodiff supports reverse- and forward-mode derivatives, Hessian-vector products, higher-order derivatives, and batching.Training-specific backward operators decompose to primitives so derivatives can be stacked to a defined order.Source: B. Optimizer · p. 2
  • Compilation legalizes operations, plans memory, applies AMP policy, inserts PTQ, and produces a device-specific schedule.Reshape, same-dtype Cast, and axis-0 Narrow can be aliased as views rather than copied.Source: B. Optimizer · p. 2, Figure 1 · p. 3
  • F16 compute leaves a small dispatch-bound convolutional net unchanged but speeds a wide 16k-hidden MLP matmul by 16%, with accuracy unmoved; quantization supports PTQ, QAT, and weight-only INT8-to-INT4 schemes.The IR tracks types and AMP policies, including F32 matmul accumulation over F16/BF16 inputs; low-precision graphs are parity-checked against F32 references.Source: B. Optimizer · pp. 2–3

C. Runtime

RLX exposes a user-facing runtime that compiles graphs for a selected device and caches per-device schedules for repeated execution. Distributed execution adds pipeline and tensor parallelism over local or network transports, with single-box correctness checks and multi-machine throughput left for future work.

C. Runtime

  • Per-device schedules are cached after graph pruning, placement, and scheduling, so repeated execution avoids recompilation.GraphDevices, DeviceRouter, DevicePolicy, and FlexibleSession support cost-based and environment-driven device selection, including fallback chains across devices.Source: C. Runtime · p. 3

D. Distributed execution

  • Pipeline parallelism assigns each node a layer block and sends only hidden-state tensors between graph runs, allowing models too large for one device to fit across several.The host-side relaying requires no compiler change.Source: D. Distributed execution · p. 3
  • Tensor parallelism shards layer operations and inserts in-graph all-reduce; the same collective synchronizes gradients for data-parallel training.Collectives use LocalTransport or TCP NetTransport; on Apple Silicon, all-reduce can remain device-resident through MLX ring/JACCL transports over RDMA.Source: D. Distributed execution · p. 3
  • Two- and four-way tensor-sharded layers reproduce single-node results bit-for-bit, and pipeline relaying is numerically exact; multi-machine throughput remains future work.The paper tests single-box multi-process TP/PP, including GPU-resident all-reduce, and says cross-machine transport is not yet hardened.Source: D. Distributed execution · p. 3

IV. TRANSPARENT KERNEL DISPATCH

RLX makes operator dispatch explicit: each operator follows a native, common-IR, rewritten, or unsupported path, and compilation fails when rewriting cannot legalize it. Backend portability is governed by the same supported-operator claims used for device selection and compilation.

IV. TRANSPARENT KERNEL DISPATCH

  • Unsupported operators that remain illegal after structural rewriting cause compilation to fail with a diagnostic report rather than silently falling back to CPU.The other paths are native backend kernels, registered logical kernels lowered to portable primitive MIR, and structural unfuse/lower rewrites that preserve semantics.Source: §IV · p. 3
  • PreferNative selects native kernels when declared in supported_ops and otherwise uses common-IR; the claim set also gates device selection so routed graphs are guaranteed to compile.Dispatch can be overridden globally or per compile, and a report can show each operator's selected path.Source: §IV · pp. 3–4

V. BACKEND PORTABILITY

  • RLX targets fourteen runtime device variants and two specialty codegen paths: no_std ARMv7E-M INT8 for Cortex-M and IR-to-Verilog-to-bitstream for FPGA.The specialty targets consume a narrower operator set and are outside the runtime registry.Source: §V · p. 4
  • Among eight production backends with published coverage, CPU claims 104+ of 113+ operator kinds, while GPU backends claim 57–84 and the INT8-first TPU path claims 50.Backends expose an open trait for registering additional runtime targets, and the legalizer checks the same claim set before compilation.Source: §V · p. 4

VI. EXTENSIBILITY CASE STUDIES

RLX builds advanced capabilities from primitives, including higher-order differentiation and quantized inference, while specialty codegen paths extend the IR to microcontrollers and FPGA. Its import pipeline lowers four model formats into the same IR and applies a common processing pipeline.

VI. EXTENSIBILITY CASE STUDIES

  • RLX provides first- through third-order scalar derivatives, with third-order parity between CPU and CUDA, Metal, MLX, and wgpu for six listed activations and functions.The passage also describes quantized inference through DequantMatMul for llama.cpp schemes and a fused INT4 GEMV on Metal; ScaledMatMul supports block-scaled minifloats and is described as bit-exact against an F32 oracle.Source: §VI · p. 4
  • Cortex-M codegen emits no_std ARMv7E-M INT8 kernels, while FPGA codegen lowers the IR to Verilog and a bitstream using integer-only Q0.31 requantization.The Cortex-M path is validated against reference implementations; the same IR also drives server-GPU and custom-logic paths.Source: §VI · p. 4

VII. IMPLEMENTATION AND FORMATS

  • Safetensors, GGUF, ONNX, and rten imports lower to one IR, after which checkpoints use the same fusion, quantization, memory-planning, and dispatch pipeline.The standalone GGUF crate covers llama.cpp schemes and can write and convert GGUF; the ONNX importer has a conformance suite.Source: §VII · p. 4

VIII. EVALUATION

RLX is evaluated on inference latency, differentiation breadth, model coverage, and correctness, with benchmarks run on one Apple Silicon host. On all-MiniLM-L6-v2, RLX-Metal leads the reported backends at every batch, including against PyTorch-MPS.

VIII. EVALUATION

  • The evaluation covers inference latency, differentiation breadth, model coverage, and end-to-end correctness, and discusses how the architecture supports training as well as inference.Operator coverage is quantified in Section V.Source: §VIII · p. 4

A. Setup

  • Benchmark numbers use the latest RLX on one Apple Silicon host; latency is the median of repeated encoder-forward trials, with identical synthetic seq-128 inputs across embedding runners.The differentiation benchmark uses CNTVCT_EL0 and refuses to run under thermal throttling.Source: A. Setup · p. 4

B. Inference: RLX backends vs. other frameworks

  • The inference comparison uses all-MiniLM-L6-v2 with seq-128 synthetic inputs across batch sizes 1–32, measuring p50 encoder-forward latency on one Apple Silicon host.The same compiled RLX graph runs unchanged across the RLX backends; the figure compares five RLX backends with PyTorch and ONNX Runtime on CPU and Apple GPU/ANE.Source: B. Inference: RLX backends vs. other… · p. 4, Figure 2 · p. 5
  • 16.6 ms at batch 32: RLX-Metal is fastest at every batch, ahead of PyTorch-MPS at 26.7 ms on all-MiniLM-L6-v2.At batch 32, RLX-Metal is also about 3.3× faster than RLX-CPU at 55.2 ms; RLX-CPU matches or beats PyTorch-CPU and leads ONNX Runtime from batch 4.Source: Figure 2 · p. 5, B. Inference: RLX backends vs. other… · p. 5

C. Training: what the architecture enables

RLX trains through transformations of the same IR, with autodiff, fusion, and device routing shared by forward and backward passes. Its reported results include top training throughput, parity across training paths, and inference on RLX backends, with larger convolutional backward graphs bounded by Core ML lowering.

C. Training: what the architecture enables

  • 946,487 img/s for the graph-fused MLP exceeds NumPy+BLAS at 787,349 img/s, the top entry in the full training table.The table's top entry is reported at higher accuracy.Source: C. Training: what the architecture… · p. 5
  • Rewriting HIR for autodiff lets reverse-, forward-, and higher-order gradients use the same fusion, legalization, and device-routing pipeline as the forward pass.The optimizer can fuse SGD or Adam into the schedule with weights resident via JIT.Source: C. Training: what the architecture… · p. 5
  • 67,476 img/s for RLX-Metal leads the plotted throughput order, followed by RLX-Vk MLP at 63,967 img/s and RLX g-fused at 56,507 img/s.The identical LeNet CNN runs at SGD, batch 128, for two epochs; all frameworks reach 95–98% accuracy.Source: Figure 3 · p. 5, C. Training: what the architecture… · p. 5
  • Trained models serve inference on every RLX backend at identical accuracy, reaching up to 161,660 img/s on the Neural Engine.The gradients match the CPU bit-for-bit; CPU and Metal are the most exercised paths.Source: C. Training: what the architecture… · p. 5
  • Some larger convolutional backward graphs remain bounded by Core ML lowering.CPU and Metal are the most exercised paths.Source: C. Training: what the architecture… · p. 5

D. Higher-order differentiation

RLX benchmarks third-order scalar differentiation and reports parity across several backends and test functions. Its supported-model examples also report reference parity for model families and a Cortex-M INT8 path, alongside a deployment-reach and capability comparison.

D. Higher-order differentiation

  • At batch 200, CPU differentiation time grows ∼O(N^2), from 0.5 µs at N=64 to 12.7 µs at N=4096; launch-bound Metal and MLX stay nearly flat and overtake only at much larger N.The benchmark uses a cubic-sum objective and reports medians over 50 runs on Apple Silicon.Source: D. Higher-order differentiation · p. 5
  • Third-order parity holds between the CPU reference and CUDA, Metal, MLX, and wgpu on x^3, ReLU, tanh, GELU, and SiLU.At the tested sizes, CPU is fastest; the same transforms differentiate to third order across the reported backends.Source: D. Higher-order differentiation · p. 5

E. Supported models

  • The companion repository provides roughly 100+ model-family crates spanning LLMs, state-space models, vision and diffusion, embeddings, ASR, and TTS, all at reference parity.Examples include Qwen3, Mamba, BioCLIP-2, and Whisper.Source: E. Supported models · p. 5
  • Figure 4 compares deployment reach across twelve platforms with four capability dimensions, placing RLX at the top-right and CoreML and MLX within the Apple family.The dimensions are native-Rust authoring and compilation, fused compiler IR, training-grade autodiff, and MCU/FPGA codegen.Source: Figure 4 · p. 6
  • Qwen3 reaches 100% top-1 parity with Hugging Face, BioCLIP-2 reaches 100% parity on CPU, Metal, MLX, and wgpu, and Cortex-M INT8 reaches 96.6% on nRF52840.The TPU path matches MiniLM-L6 through PJRT.Source: E. Supported models · p. 6

IX. COMPARISON WITH ML FRAMEWORKS, APPLE LIBRARIES, AND INFERENCE ENGINES

RLX is compared with frameworks, kernel libraries, inference engines, LLM stacks, and compiler stacks. The conclusion distinguishes it by combining training-grade differentiation, transparent dispatch, and edge codegen in native Rust, while reporting top throughput results and identifying future work.

IX. COMPARISON WITH ML FRAMEWORKS, APPLE LIBRARIES, AND INFERENCE ENGINES

  • The compared systems cover different parts of the design space: many lack whole-program compiler IR, while inference engines and most compiler stacks lack first-class differentiation.The passage also identifies Apple-only reach, heavy or single-vendor runtime binding, and non-Rust implementations among many systems.Source: §IX · p. 6
  • JAX and TensorFlow share RLX's transform-based, fused-IR view of differentiation, but neither is native Rust or reaches Neural Engine, microcontrollers, or FPGAs from one IR.The comparison does not claim that either has lower overhead or better performance than RLX.Source: §IX · p. 6
  • RLX combines native-Rust authoring and compilation, a primitive-level IR with higher-order differentiation and batching, transparent multi-backend dispatch, and edge codegen.MLX and CoreML are RLX backends, and the same model can target Apple Silicon, NVIDIA, AMD, TPU, web, and MCU without a rewrite.Source: §IX · p. 6

X. CONCLUSION AND FUTURE WORK

  • RLX reports the fastest inference from batch 4 up on a common embedding model and the top training-table entry at 946k img/s.The conclusion identifies these as reported benchmark results across its comparisons.Source: §X · p. 6
  • Future work includes GPU and TPU optimizations, better training, and better multi-machine transport.Source: §X · p. 6
Loading 2609.37916v1…