rlx-distributed

Multi-node distributed inference (pipeline / tensor parallel) for RLX models. It layers on rlx-driver's transports (TcpTransport, ThunderboltTransport, and rlx-mlx's MlxTransport) and ProcessGroup to run one model split across several hosts, coordinating per-rank layer blocks with a model-agnostic relay.

Placing a model on the cluster

Machines can be listed, discovered, or both — hand-written [[node]] entries are kept verbatim and keep their place in the pipeline; discovered hosts are appended and deduped by addr. See cluster.example.toml.

# use rlx_distributed::cluster::*;
# fn tensor_index() -> Vec<TensorEntry> { unimplemented!() }
let mut cx = Cluster::from_path("cluster.toml")?;
// discover -> probe -> derive cost from the checkpoint -> plan
cx.autoplan(&tensor_index(), &CostOptions::default().with_experts(8, 288), worker_bin)?;
print!("{}", cx.plan_table());
# Ok::<(), anyhow::Error>(())

node                 device    layers      resident      disk   stage
10.0.0.5:9100        cuda        0..14     8.9/34.0G    28.0G   0.042s paged
node-b:9100          metal      14..30     9.9/58.0G    32.0G   0.031s paged
node-c:9100          cpu        30..45    10.1/60.0G    30.0G   0.088s paged
critical path: node-c:9100 at 0.088s/token

policy = "auto" decides capacity as the better of experts resident in RAM/VRAM and dense weights resident, expert banks on disk, then shares layers out by predicted per-layer time — compute plus the disk time to stream in the experts each token touches. For a fine-grained MoE that is the difference between a plan that fits and one that says the model is too big: GLM-5.3-Flash needs ~10 GB resident per node, not the 93 GB a RAM-only view reports.

Objectives, KV and precision

[placement.objective] says what to make small, separately from how to split: latency uses the fewest nodes that hold the model (each stage boundary is a hop on the critical path), throughput evens out stage times, memory spreads thin. context reserves KV-cache bytes before weights, so a plan that fits at seq = 8 cannot OOM at 128 k.

The cache width is inferred from the checkpoint's tensor shapes — latent (MLA) layers cache only the compressed latent, ordinary attention caches K and V, and linear-attention layers have a fixed state that does not grow with context. On a hybrid stack like GLM-5.3-Flash (34 recurrent + 11 latent) that is the difference between pricing the cache at 3 GB and at 12 GB. The inference is anchored on hidden_size, so it reads GGUF and HF dim orders alike.

It is kept per block, not just averaged, because on a hybrid stack the average is not what any node pays: a stage holding GLM-5.3-Flash layers 0..3 caches nothing, while one holding 3..7 caches a full layer's worth. Each assignment reports its exact est_kv_bytes, and capacity — decided before the range is — budgets for the worst contiguous window of that length, so the plan fits wherever the range lands. CostOptions::kv_dtype scales the whole thing for an f16 or int8 cache.

When it cannot be read, planning stops rather than reserving nothing — a silent zero looks exactly like "no cache" and fails much later, far from its cause. The error names the three ways out: declare kv_bytes_per_layer_token, set ModelCost::kv from the model crate, or on_unknown_kv = "assume" for a pessimistic MHA-width bound. assume needs the model's real width and errors without one: estimating width from byte counts is not conservative, because the "≥1 byte per parameter" premise inverts under quantization and the bound comes out at half the real cache. precision_ladder = ["f16", "mxfp4"] steps down until the model fits and reports which width it settled on — including any fNeXmY minifloat, whose width is parsed from the name rather than looked up in a table.

device and precision are validated when the config loads, so device = "cuda+cpu" or a mistyped precision fails naming the file rather than surfacing as a puzzling compile flag on the third node. The TOML is unchanged: both are still written as plain strings.

Standby nodes and retirement

A node is active, standby or draining. A standby is planned with the same layer range as the node it backs, so it holds the same weights and stays warm:

# use rlx_distributed::cluster::*;
# fn cost() -> ModelCost { unimplemented!() }
# let mut cx: Cluster = unimplemented!();
cx.retire("10.0.0.5:9100", cost())?;   // promote its standby, or re-plan
# Ok::<(), anyhow::Error>(())

Promotion is a pointer swap rather than a cold load — for a 90 GB stage that is seconds instead of minutes. With no standby, retire still works by re-planning across the remaining nodes; it just reloads weights, which is why a cluster you intend to drain should carry spares.

ModelCost::from_tensor_index derives the per-layer numbers from a checkpoint's tensor listing (a GGUF header is enough — no weights are read), keeping routed expert bytes separate from resident ones. With the gguf feature, ModelCost::from_gguf takes the shard paths directly.

Modules

• config — hosts.json parsing + process-group construction ([DistConfig::connect]).
• partition — pipeline-parallel layer assignment ([pipeline_layer_range], [block_role]).
• pipeline — the model-agnostic relay ([PipelineCoordinator]) driving per-rank [BlockRunner]s.
• launch — local multi-process cluster helpers ([LocalCluster], [worker_args]).

A model family plugs in by implementing [BlockRunner] for its layer block.

Public API

use rlx_distributed::{DistConfig, ParallelMode, PipelineCoordinator, pipeline_layer_range};

let cfg = DistConfig::from_hostfile("hosts.json", /*rank*/ 0)?;
let group = cfg.connect()?;                     // ProcessGroup over the chosen transport
let (start, end) = pipeline_layer_range(rank, world, num_layers);
// build a BlockRunner for [start,end) and hand it to PipelineCoordinator
# anyhow::Ok(())

Quick start

cargo run -p rlx-distributed --example transport_bench

How it fits

Built on rlx-driver (ProcessGroup, transports). Model crates such as rlx-qwen3 provide the per-rank block runners.

License

MIT OR Apache-2.0.