rlx_distributed/
lib.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! # rlx-distributed — model-agnostic multi-node inference
6//!
7//! Run a model across machines two ways, both **model-agnostic** (RLX provides
8//! the API; model crates provide only the thin model-specific seam):
9//!
10//! - **Layer-block pipeline** (crate root) — a model implements a
11//!   [`BlockRunner`] for a contiguous range of its transformer layers; the
12//!   [`PipelineCoordinator`] relays hidden states rank→rank over an
13//!   `rlx-driver` [`ProcessGroup`] (TCP / Thunderbolt / MLX). Reverse-rank like
14//!   mlx-lm's `PipelineMixin`. [`config`] parses `hosts.json`, [`launch`] spawns
15//!   a local N-rank cluster (torchrun-style). This is the idiomatic path for
16//!   layer-structured LLMs (rlx-qwen3 ships a `BlockRunner`).
17//! - **Graph-node pipeline** ([`graph`]) — partition an arbitrary compiled
18//!   [`rlx_ir::Graph`] into weight-sharded subgraph stages joined by named
19//!   boundary tensors ([`graph::partition`](fn@graph::partition)), executed in-process
20//!   ([`graph::run_pipeline_local`]) or over TCP ([`graph::serve_stage`] +
21//!   [`graph::Pipeline::run_tcp`]). Use it when a model isn't a clean layer
22//!   stack (e.g. the DeepSeek-V4 graph) or you want automatic weight sharding.
23//!
24//! Both seams share [`source::ParamSource`] — the model-agnostic weight seam:
25//! RLX asks for params by name, a model crate returns dense f32 or packed bytes
26//! (wrapping its own mlx / GGUF / safetensors loader). No model dependency leaks
27//! into this crate.
28
29pub mod cluster;
30pub mod config;
31pub mod experts;
32pub mod graph;
33pub mod launch;
34pub mod paging;
35pub mod partition;
36pub mod pipeline;
37pub mod source;
38
39// ── Layer-block seam (public API used by rlx-qwen3 / rlx-protocol) ──
40pub use config::{DistConfig, Hostfile, ParallelMode, TransportBackend};
41pub use launch::{LocalCluster, WorkerArgs, free_loopback_ports, worker_args};
42pub use partition::{BlockRole, block_role, pipeline_layer_range};
43pub use pipeline::{BlockInput, BlockOutput, BlockRunner, PipelineCoordinator};
44
45// ── MoE expert-parallel offload seam (model-agnostic; DeepSeek / Llama-4 / Kimi) ──
46pub use experts::{
47    ExpertProvider, ExpertShards, dispatch_experts, serve_expert_worker, shutdown_expert_workers,
48};
49pub use paging::{BankIx, BankLocation, ExpertPager, PagerStats};
50#[cfg(feature = "gguf")]
51pub use paging::{GGUF_MOE_BANKS, register_gguf_expert_banks};
52
53// Re-export the transport primitives so model crates depend only on
54// `rlx-distributed`, not `rlx-driver` directly.
55pub use rlx_driver::{
56    NetTransport, ProcessGroup, ReduceKind, TcpTransport, ThunderboltTransport, Transport,
57};
58
59// ── Graph-node seam + shared weight source ──
60pub use graph::{NamedTensor, Pipeline, Stage, StageRunner};
61pub use source::{Param, ParamSource};
62
63// ── Cluster orchestration (config-driven, HW-aware placement + monitoring) ──
64pub use cluster::{
65    Assignment, Cluster, ClusterConfig, ClusterRun, KvPolicy, ModelCost, NodeCaps, NodeConfig,
66    PlacementPolicy, probe_local,
67};
rlx_distributed/lib.rs

rlx_distributed/
lib.rs
