RMLX

Rust ML runtime for Apple Silicon — zero-copy Metal GPU pipeline with RDMA distributed inference

E2E 32-layer: Prefill 1.1–1.6× · Decode ~2.1× vs MLX · 24.05T FP16 GEMM · TP=2 5.0× · EP 34×

Install · Quickstart · Performance · Features · Architecture · Roadmap · Docs

🇰🇷 한국어 문서: docs/README_ko.md

🧠 What is RMLX?

RMLX reimplements Apple's MLX Metal GPU pipeline entirely in Rust, built on objc2-metal / objc2 / block2 / objc2-foundation. The framework is organized into 7 crates spanning GPU compute, memory allocation, neural network layers, and RDMA-based distributed inference.

End-to-end on a 32-layer Qwen 7B-style model (M3 Ultra, f16, random weights): prefill 1.1–1.6× faster, decode ~2.1× faster (~44 tok/s vs 21 tok/s) than MLX compiled. Kernel-level: FP16 GEMM 24.05 TFLOPS (MLX parity), QMM Q4 17.43 TFLOPS (+28% vs MLX). Distributed: TP=2 5.0× faster decode than MLX JACCL, EP 34× vs MLX end-to-end. Phase 9 adds 7 model architectures and generate() API; Phase 10 delivers 20 kernel/infra optimizations.

Single Rust binary. Zero-copy unified memory. No Python runtime, no framework overhead.

📦 Install

Prerequisites: macOS 14+ on Apple Silicon (M1 or later). See full prerequisites.

git clone https://github.com/0xDaizz/RMLX.git
cd RMLX
cargo build --workspace

🚀 Quickstart

Build & Test

cargo build --workspace           # Build all 7 crates
cargo test  --workspace           # Run 1,298 tests

Benchmark

cargo bench -p rmlx-nn --bench pipeline_bench

Distributed RDMA (2-node)

# One-time install
cargo install --path crates/rmlx-cli

# Auto-detect TB5 topology, assign IPs, configure interfaces
rmlx config --hosts node1,node2 --auto-setup --output rmlx-hosts.json --verbose

# Launch distributed job
rmlx launch --backend rdma --hostfile rmlx-hosts.json -- ibv_devices

--auto-setup discovers Thunderbolt connections via system_profiler, assigns point-to-point IPs, and configures RDMA interfaces automatically.

📊 Performance

⚡ E2E Model Performance (PRIMARY) — 32-layer Qwen 7B-style, M3 Ultra, f16

All E2E numbers measured on 32-layer model with random weights (Qwen 7B-style config), M3 Ultra, f16.

Prefill (32 layers):

RMLX Prefill TFLOPS:

Decode (32 layers, M=1, kv=128 tokens):

Note: RMLX decode is estimated from 60-layer pipeline_bench (704 μs/layer × 32 = 22,528 μs). A direct 32-layer E2E decode bench should be added for exact comparison.

⚡ Kernel-Level Benchmarks

All kernel-level numbers measured on Apple Silicon (M3 Ultra), single transformer layer, Qwen 3.5 MoE A22B config, f16 unless noted.

Decode — 110 ms → 703 μs (156× internal speedup)

BW efficiency 73.6% at 60L pipeline floor.

Prefill GEMM — 24.05 TFLOPS

⚡ Distributed (2-node TB5 RDMA)

TP Decode — 5.0× vs MLX at TP=2

Split-CB TP batches an entire layer into 2 command buffers with inter-CB allreduce, eliminating per-op GPU sync overhead.

TP Prefill — Split-CB

Expert Parallelism — 30–178× vs MLX

Buffer-pooled grouped_forward: 32 allocations → 4 bulk allocations (14 ms → 359 μs, 39×). commandBufferWithUnretainedReferences removes CB retain/release overhead.

EP-2 End-to-End (4 experts/rank + RDMA)

EP-2 grouped seq=2: was 212.0 μs baseline, now 200.8 μs (-5.3%). EP-2 grouped seq=32: was 578.8 μs, now 549.0 μs (-5.1%).

RDMA vs JACCL Transport

✨ Features

RMLX vs MLX vs CUDA

🔧 Key Capabilities

🏗️ Architecture

graph TD
    CLI[🖥️ rmlx-cli] --> NN[🧠 rmlx-nn]
    NN --> CORE[⚙️ rmlx-core]
    CORE --> METAL[🔩 rmlx-metal]
    CORE --> ALLOC[📦 rmlx-alloc]
    DIST[🌐 rmlx-distributed] --> CORE
    DIST --> RDMA[🔗 rmlx-rdma]
    ALLOC --> RDMA
    METAL -.-> ALLOC

🗺️ Roadmap

See full roadmap and benchmark report for details.

📚 Docs

🙏 Acknowledgments

• MLX by Apple — the Metal GPU compute framework that RMLX reimplements in Rust
• mlx-lm by Apple — LLM inference patterns and Metal kernel references
• vllm-mlx — distributed inference architecture and RDMA transport patterns

📄 License

MIT — see LICENSE.