Module dist 

Ship-graph distributed execution (build the worker once, run any model). Ship-graph distributed execution — build the worker once, run any model.

A worker is a generic executor: it compiles and runs a graph the coordinator ships at runtime, so no model is baked into the worker binary and there is no per-node / per-model recompile. Build one worker binary per architecture (Node::from_env().connect() + [serve_stage]), deploy it to every node, and drive any model from a coordinator.

The coordinator — which owns the model’s graph builders (e.g. a crate in rlx-models) — partitions the model into stages and ships each worker a [StageSpec]: the serialized subgraph, its I/O node names, where to fetch each weight, and a device directive. Weights are resolved locally by a caller-provided closure (GGUF / safetensors / HF — rlx core stays model-agnostic), so only specs (KB) and activations cross the wire — never the weights.

// Generic worker binary — the SAME build runs any model:
let mut stage = dist::recv_stage(&group, |uri| load_weights(uri))?;
let x = dist::recv_activation(&group, group.rank() - 1)?;
let y = stage.run(&x);
dist::send_activation(&group, /*next*/ 0, &y)?;

Structs

StageSpec

A self-describing unit of work for a worker: everything needed to run one stage of a model, with nothing implied by convention.

WeightCache

Parse-once, serve-many weight cache. A GGUF file is parsed a single time and then serves any number of its tensors (f32 or packed) — the right shape for a worker whose stage pulls many tensors (q/k/v/o/ffn) from one model file. Created per stage inside recv_stage; also usable standalone.

WeightRef

Where one parameter’s data comes from. The worker resolves the uri itself, so large weights stay node-local.

WorkerStage

A compiled stage sitting on this worker, ready to run activations.

Functions

recv_activation

Receive an activation tensor from from.

recv_stage

Worker: receive a StageSpec from the coordinator (rank 0), resolve its weights via resolve (uri -> f32), pick the directed device, and compile. This is the generic worker’s setup — the same code runs any model’s stage.

resolve_weight_bytes

Load the quantized bytes of a weight as-is (no dequant), for the on-device-quant path. gguf://<path>#<tensor> returns the packed block bytes; file://<path> returns the raw file. Feed to set_param_typed(_, _, U8) behind a DequantMatMul graph.

resolve_weight_uri

Resolve a weight tensor to f32 from its source URI. Supported schemes: gguf://<path>#<tensor> — dequantize a GGUF tensor (any K-quant) safetensors://<path>#<tensor> — read a safetensors tensor (F32/F16/BF16) file://<path> — raw little-endian f32 array

send_activation

Send an activation tensor to to (a worker, or the next pipeline stage).

serve_stage

Worker convenience: set up the stage, then run one activation cycle — receive from the previous rank (or an external feed for rank 1), run, and forward to the next rank (or back to the coordinator if last). Returns the device it ran on. For a serving loop, call recv_stage once and drive WorkerStage::run + recv_activation/send_activation yourself.

serve_stage_uri

Same as serve_stage but with no custom resolver — a worker binary that serves GGUF / safetensors / raw-f32 models out of the box (recv_stage’s built-in WeightCache handles those schemes; anything else warns).

ship_stage

Coordinator: ship a stage spec to worker rank.

uri_resolver

Built-in resolver for serve_stage_uri: resolve_weight_uri, logging and returning empty on failure so a missing shard surfaces loudly rather than panicking mid-collective.