RLX 深度拆解：MIT 如何用纯 Rust 写出一个「编译器+运行时合一」的 ML 框架——从 JAX 形状的张量 IR 到 15 种硬件后端的全栈工程哲学

引言：ML 基础设施的「两张皮」困境

如果你在 2024 年之前做过机器学习部署，你一定经历过这种痛苦：

用 PyTorch 训练完模型，要部署到边缘设备？先用 TorchScript 或 ONNX 导出，然后发现精度对不上。想在 Apple Silicon 上跑？得换 MLX。想在 NVIDIA GPU 上跑？得配 CUDA + cuDNN。想在浏览器里跑？得转 WebAssembly + ONNX Runtime。每换一个硬件，就要换一套工具链。

这就是 ML 基础设施的「两张皮」问题——图编译器（XLA、TVM、MLIR）和内核运行时（cuDNN、MPS、MLX）各管各的。你得在两个世界之间手动做模型转换、精度适配、内存布局协调。

2026 年，MIT 的一个研究项目 RLX（版本 0.2.14）试图用一种全新的方式解决这个问题：用纯 Rust 写一个编译器+运行时合一的 ML 框架，一个统一的 IR 贯穿从图优化到内核调度的全流程，覆盖从微控制器（Cortex-M）到超算（Cerebras WSE）的 15 种硬件后端。

这不是又一个「.me too」的 ML 框架。RLX 的核心洞察是：当编译器和运行时共享同一个数据结构和同一个调度逻辑时，跨后端的优化可以做得比任何分离架构都深。

本文将从架构哲学、IR 设计、编译流水线、内核调度、训练支持、量化系统、分布式执行等多个维度，深度拆解这个来自 MIT 的开源项目。

一、为什么又造一个轮子？RLX 的问题意识

在理解 RLX 的设计之前，我们先看看它要解决的问题。

1.1 现有 ML 框架的架构分裂

RLX 的回答是：为什么不把编译器和运行时放在一起？ 当编译器知道运行时有哪些内核可用、这些内核的延迟和吞吐是多少，它就能做出更好的融合决策。当运行时知道编译器做了哪些优化，它就能在调度时避开重复计算。

1.2 为什么选 Rust？

这不是「Rust 最酷所以用 Rust」的选择。RLX 选 Rust 有几个实际考量：

1. 零成本抽象：IR 构建、图变换都是编译期完成的，运行时零开销
2. 内存安全：ML 编译器最容易出内存问题的地方是张量生命周期管理
3. no_std 支持：边缘部署（Cortex-M）需要 no_std，Rust 原生支持
4. 跨后端 FFI：Rust 的 FFI 可以零成本调用 cuBLAS、Metal、Vulkan 等 C/C++ 库
5. 编译速度：Rust 2024 edition 的编译速度已经足够支撑大型项目

二、核心架构：IR → 编译器 → 运行时的三层一体化

RLX 的架构可以用一句话概括：一个可序列化的张量 IR，一套从 HIR → MIR → LIR 的编译流水线，一个知道所有后端内核成本的运行时调度器。

2.1 张量 IR：RLX 的「通用语言」

// RLX IR 的核心数据结构（简化）
pub struct Graph {
    pub name: String,
    pub nodes: Vec<Node>,
    pub inputs: Vec<NodeId>,
    pub outputs: Vec<NodeId>,
}

pub enum Op {
    MatMul { /* 矩阵乘法参数 */ },
    Attention { mask_kind: MaskKind, /* 注意力参数 */ },
    Scan { /* 循环/递归 */ },
    DenseSolve { /* 线性方程组 */ },
    FakeQuantize { /* 量化模拟 */ },
    Custom { /* 自定义算子扩展点 */ },
    // ... 其他算子
}

这个 IR 设计有几个关键特点：

1. JAX 形状的自动微分

RLX 的自动微分系统模仿了 JAX 的 grad、jvp（Jacobian-vector product）、hvp（Hessian-vector product）和 vmap（向量化映射）。这意味着：

use rlx_autodiff::grad;

// 前向图
let forward = build_forward_graph();
// 自动构建反向图：∂output/∂input
let grad_graph = grad(&forward, &[input_id]);

2. 可序列化

IR 是可序列化的，这意味着你可以：

• 把编译好的模型保存到磁盘，下次直接加载
• 在不同机器之间传输编译好的计算图
• 做增量编译（只重新编译变化的部分）

3. MaskKind 枚举

pub enum MaskKind {
    None,
    Causal,        // GPT 系列的因果掩码
    Prefix,        // 前缀注意力
    Bidirectional, // 双向注意力（BERT 系列）
}

注意力机制的掩码类型被提升为 IR 的一等公民，而不是运行时的参数。这意味着编译器可以在编译期根据掩码类型选择最优的内核实现。

三、编译流水线：HIR → MIR → LIR 的五阶段优化

RLX 的编译器不是一个简单的「图到图」变换，而是一个完整的多阶段优化流水线。

3.1 五阶段流水线

源码/模型输入
    ↓
[1] HIR (High-level IR)  — 语义级表示，保留用户意图
    ↓
[2] 合法化 (Legalization) — 把高级算子拆解为后端可执行的基础算子
    ↓
[3] MIR (Mid-level IR)   — 融合候选识别、精度策略、vmap 展开
    ↓
[4] 融合 + 内存规划     — 算子融合、内存布局优化、buffer 重用
    ↓
[5] LIR (Low-level IR)   — 后端特定的内核调度、编译缓存
    ↓
目标后端执行

3.2 合法化阶段：从「高级语义」到「基础算子」

合法化（Legalization）是编译器把用户写的高级算子拆解成后端能执行的基础算子的过程。

例如，一个 Transformer 的注意力层在 HIR 中可能是一个 Op::Attention 节点，但在合法化后，它会被拆解为：

Op::MatMul(Q, K^T) → Op::Scale(1/√d) → Op::Softmax → Op::MatMul(Attention, V)

这个拆解不是简单的展开——编译器会根据目标后端的内核库，选择最优的拆解方式。比如在 NVIDIA GPU 上，QKV 矩阵乘法可能会被合并成一个 GEMM 内核，而在 CPU 上则会走矩阵分块路径。

3.3 融合阶段：把多个操作合并成一个内核

算子融合（Operator Fusion）是 ML 编译器最重要的优化之一。RLX 的融合策略包括：

// 融合候选示例
// 优化前：3 次内存读写
let a = g.matmul(x, w1, shape1);  // 读 x, w1 → 写 a
let b = g.relu(a, shape1);        // 读 a → 写 b
let y = g.matmul(b, w2, shape2);  // 读 b, w2 → 写 y

// 融合后：1 次内存读写
// GPU 内核：fused_matmul_relu_matmul
// 所有中间结果留在寄存器/共享内存中

RLX 的融合器不只是简单的相邻融合——它会分析整个计算图的数据依赖，识别可以并行执行的分支，然后把它们合并成一个大的 GPU kernel。

3.4 内存规划：让 buffer 重用最大化

// 内存规划器的工作
// 优化前：需要 3 个 buffer
// buffer_a: 1MB (matmul 输出)
// buffer_b: 1MB (relu 输出) 
// buffer_c: 1MB (第二个 matmul 输出)

// 内存规划后：只需要 1 个 buffer
// buffer_a 和 buffer_b 的生命周期不重叠
// 所以可以重用同一块内存
// 只需要为最终输出保留一个 buffer

四、后端调度：15 种硬件的「最短路径」选择

这是 RLX 最让人兴奋的部分。它支持的后端数量和覆盖范围令人咋舌：

4.1 后端全景

4.2 自动设备选择：fastest_device()

RLX 不只是支持多后端——它知道每个后端的成本模型：

use rlx_runtime::fastest_device;

// 自动选择当前机器上最快的设备
let device = fastest_device(); // Metal / CUDA / wgpu / ... else CPU

// 这不是简单的「有 GPU 就用 GPU」
// RLX 会根据具体算子的内核性能来选择
// 比如某些算子在 CPU 上可能比在 GPU 上更快
// （特别是小规模矩阵运算，GPU 的启动开销可能大于计算收益）

4.3 跨平台 GPU：wgpu 的妙用

RLX 的 gpu 后端基于 wgpu，这意味着一个模型可以：

1. 在 macOS 上通过 Metal 运行
2. 在 Linux 上通过 Vulkan 运行
3. 在 Windows 上通过 DX12 运行
4. 在浏览器中通过 WebGPU 运行

// 同一份代码，不同后端
// macOS
let mut session = Session::new(Device::Metal).compile(graph);

// Linux
let mut session = Session::new(Device::Vulkan).compile(graph);

// 浏览器
let mut session = Session::new(Device::WebGpu).compile(graph);

4.4 编译缓存：不重复编译

// RLX 的编译缓存策略
// 1. 首次编译：IR → 目标代码（可能需要几秒到几十秒）
// 2. 缓存到磁盘：保存编译结果
// 3. 后续加载：直接从磁盘加载，毫秒级启动
// 4. 增量编译：只重新编译变化的部分

五、三种 API 风格：从底层到高层的渐进抽象

RLX 提供了三种不同层次的 API，满足不同用户的需求。

5.1 底层图 API：完全控制

use rlx::prelude::*;

// 最底层的 API：手动构建计算图
let mut g = Graph::new("mlp");
let x  = g.input("x",  Shape::new(&[1, 784], DType::F32));
let w1 = g.param("w1", Shape::new(&[784, 256], DType::F32));
let b1 = g.param("b1", Shape::new(&[256], DType::F32));
let w2 = g.param("w2", Shape::new(&[256, 10], DType::F32));
let b2 = g.param("b2", Shape::new(&[10], DType::F32));

let h = g.matmul(x, w1, Shape::new(&[1, 256], DType::F32));
let h = g.add(h, b1);
let h = g.activation(Activation::Gelu, h, Shape::new(&[1, 256], DType::F32));
let y = g.matmul(h, w2, Shape::new(&[1, 10], DType::F32));
let y = g.add(y, b2);
g.set_outputs(vec![y]);

// 编译
let mut session = Session::new(Device::Cpu).compile(g);
session.set_param("w1", &w1_weights);
session.set_param("w2", &w2_weights);

// 运行
let out = session.run(&[("x", &input_data)]);

5.2 Tensor DSL：NumPy 风格

use rlx_tensor::Tensor;

// NumPy 风格的 API，自动追踪计算图
let a = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], [2, 2]);
let b = Tensor::from_vec(vec![5.0, 6.0, 7.0, 8.0], [2, 2]);

// 惰性求值：这里不执行计算，只构建 IR
let c = (&a + &b).relu();
let d = c.matmul(&Tensor::ones([2, 2]));

// 执行：编译 + 调度 + 计算
let result = d.to_vec(); // 自动选择最快的后端

5.3 rlx! 宏：最小化样板代码

use rlx::rlx;

// 用 rlx! 宏声明式地定义计算图
let g = rlx! {
    graph "transformer_block";
    input x: [?, 512, 768];              // 动态 batch，序列长度 512，隐藏维度 768
    param w_q: [768, 768];  param b_q: [768];
    param w_k: [768, 768];  param b_k: [768];
    param w_v: [768, 768];  param b_v: [768];
    param w_o: [768, 768];  param b_o: [768];

    // @ 是矩阵乘法，遵循 NumPy 优先级
    let q = (x @ w_q + b_q);
    let k = (x @ w_k + b_k);
    let v = (x @ w_v + b_v);
    let attn = softmax(q @ k.transpose() / 8.6) @ v;  // 8.6 ≈ √768
    let out = attn @ w_o + b_o;
    out out;
};

rlx! 宏的 @ 运算符表示矩阵乘法，? 表示动态维度，形状推导是自动的。

六、训练支持：不只是推理

RLX 的一个独特之处在于它同时支持推理和训练。这在边缘计算场景中特别有价值——你可能需要在边缘设备上做 fine-tuning 或 in-context learning。

6.1 自动微分

use rlx::prelude::*;
use rlx_autodiff::grad_with_loss;
use rlx_optim::{Adam, Optimizer};

// 前向图：第一个输出是标量 loss
let forward = build_forward_graph();

// 反向图：自动构建梯度计算
let backward = grad_with_loss(&forward, &[param_id]);

// 编译反向图
let mut compiled = Session::new(Device::Cpu).compile(backward);

// 训练循环
let mut opt = Adam::new(1e-3);
let mut weights = init_weights();

for step in 0..1000 {
    let out = compiled.run(&[
        ("x", &input),
        ("target", &target),
        ("w", &weights),
        ("d_output", &[1.0]),  // ∂loss/∂loss = 1
    ]);
    
    let (loss, grads) = (out[0][0], &out[1]);
    opt.step("w", &[weights.len()], &mut weights, grads);
    opt.end_iteration();
    
    if step % 100 == 0 {
        println!("step {step}: loss {loss:.6}");
    }
}

6.2 INT8 QAT（量化感知训练）

RLX 内置了 INT8 量化感知训练支持：

// INT8 QAT 示例
// 在训练过程中模拟量化误差
// 让模型在量化后也能保持精度

let backward = grad_with_loss(&forward, &[param_id]);
let mut compiled = Session::new(Device::Cpu)
    .compile_with_options(backward, CompileOptions {
        quantization: Some(QuantConfig::QatInt8),
        ..Default::default()
    });

6.3 优化器生态

use rlx_optim::{Adam, Lion, Muon};

// Adam - 经典选择
let mut opt = Adam::new(1e-3);

// Lion - Google 提出的优化器，内存效率更高
let mut opt = Lion::new(1e-4);

// Muon - 新兴优化器，针对大模型训练优化
let mut opt = Muon::new(1e-2);

七、量化系统：从训练到部署的精度阶梯

量化是边缘部署的关键技术。RLX 的量化系统覆盖了从训练到部署的完整链路。

7.1 量化格式支持

7.2 GGUF 支持

// 加载 GGUF 格式的模型
use rlx_gguf::GgufModel;

let model = GgufModel::load("model.gguf")?;

// 自动反量化
let tensor = model.tensor("layer.0.weight")?;
// tensor 的 dtype 会自动匹配 GGUF 中的量化格式

7.3 模型导入

// 从 PyTorch 导入
use rlx_onnx::import_pytorch;

let graph = import_pytorch("model.pt")?;

// 从 ONNX 导入
use rlx_onnx::import_onnx;

let graph = import_onnx("model.onnx")?;

八、分布式执行：跨机器的梯度同步

RLX 支持数据并行训练，而且设计了一个巧妙的 NAT 穿透方案。

8.1 Iroh 拓扑：不需要配 IP 地址

# 在两台机器上运行，不需要手动配置网络
# 机器 1（Mac，Metal GPU）
RANK=0 WORLD=2 TOPOLOGY=iroh RLX_IROH_SEED=cafe1234 \
  cargo run -p rlx-vision-bench --features iroh

# 机器 2（Linux，CUDA GPU），即使在防火墙后面
RANK=1 WORLD=2 TOPOLOGY=iroh RLX_IROH_SEED=cafe1234 \
  cargo run -p rlx-vision-bench --features iroh

Iroh 使用公钥基础设施（PKARR）做 NAT 穿透，两台机器只要共享一个 seed 就能自动发现对方，不需要配置 ip:port。

8.2 图内集合操作

use rlx::distributed;

// 在计算图内部执行 all_reduce
// 梯度同步是计算图的一部分，不是外部操作
let grads_reduced = distributed::all_reduce(grads, ReduceOp::Sum);

8.3 MoE 专家并行

// MoE（Mixture of Experts）的专家并行
// 不同专家分布在不同机器上
let moe_output = distributed::moe_dispatch(
    input,
    expert_ids,
    expert_weights,
    num_experts_per_node,
);

九、边缘部署：从浏览器到微控制器

RLX 的「edge-first」理念意味着它对边缘场景有特别的优化。

9.1 微控制器支持（Cortex-M）

// rlx-cortexm：no_std 的 ARM 内核
// 直接在 Cortex-M 系列 MCU 上运行推理
// INT8 量化，最小化内存占用

use rlx_cortexm::CortexMBackend;

let backend = CortexMBackend::new(Device::CortexM7);
let mut session = Session::new(backend).compile(graph);
let output = session.run(&[("input", &sensor_data)]);

9.2 FPGA 综合

// rlx-fpga：IR → SystemVerilog → FPGA bitstream
// 把 ML 模型直接综合成硬件电路

use rlx_fpga::export_fpga;

// 编译后的模型可以直接导出为 Verilog
let verilog = export_fpga(&compiled_graph, FpgaTarget::XilinxVcu118)?;
std::fs::write("model.v", verilog)?;

9.3 WebAssembly / 浏览器

// 通过 wgpu 后端，RLX 可以在浏览器中运行
// 编译为 WebAssembly 后，利用 WebGPU API 加速

let mut session = Session::new(Device::WebGpu).compile(graph);

十、rlx-models：开箱即用的模型生态

RLX 的模型定义不在核心仓库中，而是在独立的 rlx-models 仓库里。这保持了核心的轻量和模型无关性。

# 模型仓库结构
rlx-models/
├── asr/          # 自动语音识别
├── tts/          # 文本到语音
├── llm/          # 大语言模型
├── vlm/          # 视觉语言模型
├── vision/       # 计算机视觉
└── audio/        # 音频编解码

每个模型都是独立的 crate，可以按需引入：

[dependencies]
rlx-models = { git = "https://github.com/MIT-RLX/rlx-models", features = ["whisper"] }

十一、扩展机制：LayerStage 和 Flow

RLX 的可扩展性设计值得深入分析。

11.1 LayerStage：不改核心代码添加新架构

use rlx_extend::prelude::*;

// 定义一个新的 Transformer block
struct GatedMlp {
    gate: String,
    up: String,
    down: String,
}

impl LayerStage for GatedMlp {
    fn name(&self) -> &str { "gated_mlp" }
    
    fn emit_layer(&self, ctx: &mut FlowCtx<'_>, x: FlowValue) 
        -> anyhow::Result<(FlowValue, StageArtifacts)> 
    {
        let g = ctx.linear(&x, &self.gate, false)?;
        let g = ctx.silu(&g);  // SiLU 激活
        let u = ctx.linear(&x, &self.up, false)?;
        let h = ctx.mul(&g, &u);  // 门控机制
        let out = ctx.linear(&h, &self.down, false)?;
        Ok((out.clone(), StageArtifacts::hidden_only(out.shape().clone())))
    }
}

// 把自定义 block 放入模型 flow
let flow = ModelFlow::new(vec![
    Box::new(GatedMlp { gate: "gate".into(), up: "up".into(), down: "down".into() }),
    // ... 其他层
]);

这个设计的精妙之处在于：新架构通过 Flow 的 builder 方法构建基础算子，所以它自动对所有后端可见，不需要为每个后端单独实现。

11.2 CustomOp 扩展点

// 如果需要完全自定义的算子
let custom_op = Op::CustomFn {
    name: "my_custom_attention".into(),
    inputs: vec![q_id, k_id, v_id],
    outputs: vec![out_shape],
    func: Box::new(|inputs| {
        // 自定义实现
        // 可以调用任意 CUDA/Metal/Vulkan 内核
    }),
};

十二、性能优化实战：RLX 的优化策略

12.1 混合精度

// RLX 支持 AMP（自动混合精度）
let options = CompileOptions {
    amp: Some(AmpConfig {
        fp16_ops: vec![OpKind::MatMul, OpKind::Attention],
        fp32_ops: vec![OpKind::Softmax, OpKind::LayerNorm],
    }),
    ..Default::default()
};
let session = Session::new(Device::Cuda).compile_with_options(graph, options);

12.2 编译缓存

// 编译结果自动缓存到磁盘
// 第一次编译可能需要 10 秒
let session = Session::new(Device::Cuda).compile(graph);

// 第二次加载只需 50 毫秒
let session = Session::load_from_cache("model_hash_abc123")?;

12.3 调度报告

// 开启调度报告，查看每个算子用了哪个内核
std::env::set_var("RLX_DISPATCH_REPORT", "1");
let session = Session::new(device).compile(graph);

// 输出示例：
// MatMul_0 → native (cuBLAS gemm, 0.23ms)
// Relu_1 → common_ir (fused, 0.01ms)
// MatMul_2 → native (cuBLAS gemm, 0.25ms)

十三、与主流框架的对比

十四、局限性与未来方向

14.1 当前局限

1. 预训练不支持：RLX 明确表示「不是从零开始训练前沿模型」的工具。它专注于推理和 fine-tuning
2. 生态成熟度：相比 PyTorch 的庞大生态，RLX 的模型库和社区还很小
3. 文档和示例：核心文档完善，但模型级别的文档还在完善中
4. GPU 驱动兼容性：某些后端需要特定版本的驱动（如 CUDA 需要 580.65.06+）

14.2 未来方向

1. 更多模型支持：rlx-models 仓库正在快速增加模型
2. 编译器优化：更激进的算子融合和内存优化
3. FPGA 后端成熟：当前 FPGA 支持还处于早期阶段
4. 移动端部署：iOS/Android 的原生部署路径

总结：ML 基础设施的「Rust 革命」

RLX 不是另一个 ML 框架——它是一个ML 基础设施的架构实验。

它的核心贡献不在于某个具体的性能数字，而在于证明了一个架构命题：当编译器和运行时共享同一个数据结构和调度逻辑时，跨后端的优化可以做到什么程度。

对于不同角色的开发者，RLX 意味着：

• 边缘 AI 工程师：终于有了一个从 MCU 到浏览器的统一推理框架
• ML 编译器研究者：一个干净的、Rust 实现的编译器+运行时参考架构
• Rust 系统程序员：一个展示 Rust 在高性能计算领域能力的范例
• Apple Silicon 开发者：一个比 MLX 更通用的推理框架

RLX 还很年轻（0.2.x），但它代表的方向是对的：ML 基础设施正在从「一个框架管一个硬件」走向「一个框架管所有硬件」。 Rust 的零成本抽象和内存安全特性，让这种「统一」不再是性能上的妥协。

如果你正在做边缘 AI 部署，或者对 ML 编译器感兴趣，RLX 值得关注。它可能还不能替代 PyTorch 在训练领域的地位，但在推理和边缘部署这个赛道上，它已经展示了令人印象深刻的技术深度。

项目地址：https://github.com/MIT-RLX/rlx
模型仓库：https://github.com/MIT-RLX/rlx-models
文档：https://github.com/MIT-RLX/rlx#documentation
协议：MIT OR Apache-2.0
状态：v0.2.14，活跃开发中