Crate rlx_distributed 

rlx-distributed — model-agnostic multi-node inference

Run a model across machines two ways, both model-agnostic (RLX provides the API; model crates provide only the thin model-specific seam):

• Layer-block pipeline (crate root) — a model implements a BlockRunner for a contiguous range of its transformer layers; the PipelineCoordinator relays hidden states rank→rank over an rlx-driver ProcessGroup (TCP / Thunderbolt / MLX). Reverse-rank like mlx-lm’s PipelineMixin. config parses hosts.json, launch spawns a local N-rank cluster (torchrun-style). This is the idiomatic path for layer-structured LLMs (rlx-qwen3 ships a BlockRunner).
• Graph-node pipeline (graph) — partition an arbitrary compiled rlx_ir::Graph into weight-sharded subgraph stages joined by named boundary tensors (graph::partition), executed in-process (graph::run_pipeline_local) or over TCP (graph::serve_stage + graph::Pipeline::run_tcp). Use it when a model isn’t a clean layer stack (e.g. the DeepSeek-V4 graph) or you want automatic weight sharding.

Both seams share source::ParamSource — the model-agnostic weight seam: RLX asks for params by name, a model crate returns dense f32 or packed bytes (wrapping its own mlx / GGUF / safetensors loader). No model dependency leaks into this crate.

Re-exports

pub use config::DistConfig;

pub use config::Hostfile;

pub use config::ParallelMode;

pub use config::TransportBackend;

pub use launch::LocalCluster;

pub use launch::WorkerArgs;

pub use launch::free_loopback_ports;

pub use launch::worker_args;

pub use partition::BlockRole;

pub use partition::block_role;

pub use partition::pipeline_layer_range;

pub use pipeline::BlockInput;

pub use pipeline::BlockOutput;

pub use pipeline::BlockRunner;

pub use pipeline::PipelineCoordinator;

pub use experts::ExpertProvider;

pub use experts::ExpertShards;

pub use experts::dispatch_experts;

pub use experts::serve_expert_worker;

pub use experts::shutdown_expert_workers;

pub use paging::BankIx;

pub use paging::BankLocation;

pub use paging::ExpertPager;

pub use paging::PagerStats;

pub use paging::GGUF_MOE_BANKS;

pub use paging::register_gguf_expert_banks;

pub use graph::NamedTensor;

pub use graph::Pipeline;

pub use graph::Stage;

pub use graph::StageRunner;

pub use source::Param;

pub use source::ParamSource;

pub use cluster::Assignment;

pub use cluster::Cluster;

pub use cluster::ClusterConfig;

pub use cluster::ClusterRun;

pub use cluster::KvPolicy;

pub use cluster::ModelCost;

pub use cluster::NodeCaps;

pub use cluster::NodeConfig;

pub use cluster::PlacementPolicy;

pub use cluster::probe_local;

Modules

cluster

Cluster orchestration — the high-level DX over the pipeline primitives. Turn a declarative ClusterConfig into a probed, HW-planned, monitored run:

config

Launch configuration: hostfile parsing + process-group construction.

experts

Model-agnostic MoE expert-parallel offload.

graph

Graph-level pipeline parallelism: partition a compiled rlx_ir::Graph into weight-sharded subgraph stages connected by named boundary tensors. Use this for models NOT expressed as clean layer blocks (arbitrary graphs), or when you want automatic weight sharding without writing a BlockRunner.

launch

Local multi-process launcher — run an N-rank cluster on one host over loopback TCP, with no hand-written hosts.json and no N terminals.

paging

Disk-backed MoE expert paging.

partition

Pipeline-parallel layer partitioning.

pipeline

Pipeline-parallel forward orchestration.

source

The model-agnostic weight seam. rlx-distributed never knows how a model’s weights are stored — a stage worker asks a ParamSource for each parameter by name, and the source (implemented by a model crate, e.g. wrapping an mlx / GGUF / safetensors loader) returns it as dense f32 or as packed/quantized bytes. This is what keeps the crate model-agnostic: RLX provides the API, rlx-models provides the weights.

Structs

NetTransport

Full-mesh TCP transport implementing both Transport (two-sided send/recv) and SymmetricTransport (one-sided put/get/barrier). Construct via TcpTransport::bind or ThunderboltTransport::bind.

ProcessGroup

A handle to the collective-communication world: a rank, a size, and the transport that connects them.

TcpTransport

Portable TCP transport. Works over any reachable IP — Ethernet or the macOS Thunderbolt Bridge link. This is the baseline that pipeline- and tensor-parallel inference run on before any fast path.

ThunderboltTransport

Thunderbolt transport. Same wire protocol as TcpTransport today (TCP over the Thunderbolt Bridge IP link), but a distinct type so the symmetric one-sided heap path is the intended interface and so a future zero-copy Thunderbolt DMA backend can replace the engine behind the unchanged Transport + SymmetricTransport traits.

Enums

ReduceKind

Element-wise reduction operator for collectives.

Traits

Transport

Two-sided, tagged, byte-oriented point-to-point transport between world_size ranks.