RLX:一个基于Rust的统一多后端张量编译器与分布式运行时
RLX: A Unified Multi-Backend Tensor Compiler and Distributed Runtime in Rust
- Massachusetts Institute of Technology(麻省理工学院)
机构由 AI 辅助整理,请以论文原文为准。
AI总结:
RLX用单一Rust代码库统一编译与运行时,通过三级IR和透明分派支持14种设备及多种格式,在基准测试中实现最快推理与最高训练吞吐,并保持精度一致。
AI中文摘要:
生产级机器学习(ML)技术栈通常将图编译和内核执行分散在不同的层和语言中,这使得后端行为、部署保证和性能回退难以进行端到端的推理。RLX通过一个单一的Rust代码库解决了这一差距,该代码库将编译器和运行时角色结合在一个基于原语级、三级中间表示(IR)上,并配有一个透明的分派契约,该契约将每个算子解析为原生、通用IR或重写降级,并在无法合法化时使编译失败。相同的IR面向十四个运行时设备(cpu、metal、mlx、ane、cuda、rocm、oneapi、tpu、hexagon、gpu、vulkan、opengl、directx、webgpu)和两条专用代码生成路径(Cortex-M INT8和FPGA),支持safetensors、GGUF、ONNX和rten格式的输入,支持F16/BF16/F64/C64以及量化的INT4/INT8流程(含AMP/PTQ/QAT),并通过基于TCP和RDMA传输的张量/流水线并行集合进行扩展。除神经负载外,RLX还通过稀疏和稠密线性代数扩展(例如CSR LU/CG/matvec和基于LAPACK的分解)以及3D高斯泼溅算子扩展到科学/物理风格领域。我们在单台主机上,以相同的输入生成和p50测量方法,将RLX与PyTorch、TensorFlow、JAX、candle、burn、tch、rten、MLX、CoreML、IREE、Glow、TensorRT和tinygrad进行了对比评估。在all-MiniLM-L6-v2上,RLX-Metal在每个批次下都是最快的(例如,批次32时为16.6毫秒,而PyTorch-MPS为26.7毫秒)。在MNIST训练表中,RLX也拥有最高吞吐量的条目(图融合MLP:946,487图像/秒),高于NumPy+BLAS(787,349图像/秒),同时在参考检查(例如Qwen3)上保持100%的top-1一致性。
英文摘要:
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).