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

📄 arXiv: 2609.37916v1 📥 PDF

作者: Eugene Hauptmann, Nataliya Kosmyna

分类: cs.DC, cs.AI

发布日期: 2026-09-29

备注: 6 pages, 4 figures, peer-reviewed and presented at 2026 IEEE High Performance Extreme Computing Conference (HPEC)


💡 一句话要点

提出RLX以统一多后端张量编译和分布式运行时

🎯 匹配领域: 支柱三:空间感知与语义 (Perception & Semantics) 支柱八:物理动画 (Physics-based Animation)

关键词: 张量编译 分布式运行时 机器学习 Rust 中间表示 性能优化 科学计算 深度学习

📋 核心要点

  1. 现有的机器学习框架在图编译和内核执行之间存在分隔,导致性能和部署的复杂性增加。
  2. RLX提出了一个统一的Rust代码库,结合了编译和运行时功能,使用三层中间表示来简化操作符的调度和执行。
  3. 在与多个主流框架的比较中,RLX在性能上表现出色,尤其是在特定任务上显著提高了处理速度和吞吐量。

📝 摘要(中文)

生产机器学习(ML)堆栈通常在不同层和语言之间分割图编译和内核执行,这使得后端行为、部署保证和性能回退难以进行端到端的推理。RLX通过一个单一的Rust代码库解决了这一问题,该库将编译器和运行时角色结合在一个原始级别、三层中间表示(IR)周围,并提供透明的调度合同,能够将每个操作符解析为本地、通用IR或重写的降级,并在无法合法化时失败编译。该IR支持十四种运行时设备,并能够处理多种格式,支持多种数据流和量化方法。RLX在多个基准测试中表现优异,展示了其在性能上的优势。

🔬 方法详解

问题定义:现有的机器学习框架在图编译和内核执行之间存在分隔,导致后端行为和性能回退难以进行有效推理,增加了开发和部署的复杂性。

核心思路:RLX通过一个统一的Rust代码库,将编译器和运行时角色结合在一起,使用三层中间表示(IR)来简化操作符的调度和执行,确保在无法合法化时编译失败。

技术框架:RLX的整体架构包括一个原始级别的三层中间表示,支持十四种运行时设备,并能够处理多种输入格式。其设计允许通过透明的调度合同来解析操作符,支持多种数据流和量化方法。

关键创新:RLX的主要创新在于其统一的代码库和中间表示,能够同时处理多种后端设备和数据格式,显著简化了机器学习模型的编译和执行过程。

关键设计:RLX支持多种数据流(如F16/BF16/F64/C64和量化INT4/INT8),并通过AMP/PTQ/QAT等方法实现高效的模型训练和推理。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

在与多个主流框架的比较中,RLX在所有MiniLM-L6-v2任务中表现最佳,批量32时的处理时间为16.6毫秒,相较于PyTorch-MPS的26.7毫秒有显著提升。此外,在MNIST训练中,RLX的吞吐量达到946,487图像/秒,超越了NumPy+BLAS的787,349图像/秒,且保持了100%的准确率。

🎯 应用场景

RLX的设计使其在机器学习和科学计算领域具有广泛的应用潜力。其统一的编译和运行时框架能够支持多种硬件平台,适用于深度学习、科学计算以及物理模拟等领域,具有显著的实际价值和未来影响。

📄 摘要(原文)

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).