rlx_runtime/dist/

mod.rs

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! 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`](crate::dist::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`](crate::dist::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.
//!
//! ```no_run
//! # use rlx_runtime::dist;
//! # use rlx_driver::ProcessGroup; use std::sync::Arc;
//! # fn demo(group: Arc<ProcessGroup>) -> Result<(), String> {
//! // 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)?;
//! # Ok(()) }
//! # fn load_weights(_uri: &str) -> Vec<f32> { vec![] }
//! ```

use crate::Device;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

mod diagnostics;
mod inference;
pub mod node;
mod training;

pub use diagnostics::*;
pub use inference::*;
pub use training::*;

/// Where one parameter's data comes from. The worker resolves the `uri`
/// itself, so large weights stay node-local.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct WeightRef {
    pub name: String,
    /// Opaque to rlx core — the caller's resolver interprets it (e.g.
    /// `gguf:///models/x.gguf#blk.0.ffn`, `rlxp:///models/x.rlxp#w`, `file:///…`, `hf://…`).
    pub uri: String,
    /// If set, load the **quantized bytes as-is** (U8 param) instead of
    /// dequantizing to f32 — the graph's `DequantMatMul` decodes them at
    /// matmul time, so the weight stays at native quant memory (e.g. Q4_K)
    /// with no dequant round-trip. The coordinator sets this when it built a
    /// `DequantMatMul` graph for the stage.
    #[serde(default)]
    pub packed: bool,
}

fn resolve_device(spec: &str) -> Device {
    if spec.eq_ignore_ascii_case("auto") || spec.is_empty() {
        return crate::fastest_device();
    }
    match crate::parse_device(spec) {
        Ok(d) if crate::is_available(d) => d,
        _ => crate::fastest_device(),
    }
}

// ── weight source URIs ─────────────────────────────────────────────────────
//
// A model-agnostic resolver for the common on-disk formats, so the generic
// worker can load its own shard from a path the coordinator names — weights
// stay node-local, only the URI crosses the wire.

fn split_frag(rest: &str) -> Result<(&str, &str), String> {
    rest.split_once('#')
        .ok_or_else(|| format!("weight URI needs a #<tensor> fragment: {rest}"))
}

/// Error for a scheme this build recognises but cannot read.
///
/// `rlxp://`, `mlx://`, `npz://` and `dduf://` come from the `model-io`
/// feature; a node build drops it (see rlx-runtime's Cargo.toml). Saying
/// "unsupported scheme" there would send the reader hunting for a typo in a
/// URI that is perfectly well formed, so name the feature instead.
#[cfg(not(feature = "model-io"))]
fn needs_model_io(scheme: &str) -> String {
    format!(
        "weight URI scheme `{scheme}` needs rlx-runtime's `model-io` feature,          which this build does not have. gguf://, safetensors:// and file://          work without it."
    )
}

/// 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)
///   `mlx://<path>#<tensor>`          — MLX dir / safetensors / npz (dequant if needed)
///   `npz://<path>#<tensor>`          — alias for MLX `.npz` / `.npy`
///   `dduf://<path>#<qualified>`      — DDUF ZIP tensor (`component/name`)
///   `file://<path>`                  — raw little-endian f32 array
///   `rlxp://<path>#<tensor>`         — f32 tensor from an `.rlxp` package
///
/// This is model-agnostic (formats, not models); callers can still pass their
/// own closure to [`recv_stage`]/[`serve_stage`] for other schemes.
pub fn resolve_weight_uri(uri: &str) -> Result<Vec<f32>, String> {
    if let Some(rest) = uri.strip_prefix("gguf://") {
        let (path, tensor) = split_frag(rest)?;
        let mut f = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
        let gguf = rlx_gguf::GgufFile::from_reader(&mut f).map_err(|e| e.to_string())?;
        let (data, _dims) = gguf.dequant_f32(tensor).map_err(|e| e.to_string())?;
        Ok(data)
    } else if let Some(_rest) = uri.strip_prefix("rlxp://") {
        #[cfg(feature = "model-io")]
        {
            let (path, tensor) = split_frag(_rest)?;
            let pack = rlx_pkg::Package::open(path).map_err(|e| e.to_string())?;
            pack.tensor_f32(tensor).map_err(|e| e.to_string())
        }
        #[cfg(not(feature = "model-io"))]
        Err(needs_model_io("rlxp"))
    } else if let Some(rest) = uri.strip_prefix("safetensors://") {
        let (path, tensor) = split_frag(rest)?;
        let bytes = std::fs::read(path).map_err(|e| format!("read {path}: {e}"))?;
        read_safetensors_f32(&bytes, tensor)
    } else if let Some(_rest) = uri
        .strip_prefix("mlx://")
        .or_else(|| uri.strip_prefix("npz://"))
    {
        #[cfg(feature = "model-io")]
        {
            let (path, tensor) = split_frag(_rest)?;
            let w = rlx_mlx_io::load_path(path).map_err(|e| e.to_string())?;
            let map = w.into_f32_map().map_err(|e| e.to_string())?;
            map.get(tensor)
                .cloned()
                .ok_or_else(|| format!("tensor not found in mlx weights: {tensor}"))
        }
        #[cfg(not(feature = "model-io"))]
        Err(needs_model_io("mlx/npz"))
    } else if let Some(_rest) = uri.strip_prefix("dduf://") {
        #[cfg(feature = "model-io")]
        {
            let (path, tensor) = split_frag(_rest)?;
            let f = rlx_dduf::DdufFile::open(path).map_err(|e| e.to_string())?;
            f.tensor_f32(tensor)
                .map(|s| s.to_vec())
                .map_err(|e| e.to_string())
        }
        #[cfg(not(feature = "model-io"))]
        Err(needs_model_io("dduf"))
    } else if let Some(path) = uri.strip_prefix("file://") {
        let b = std::fs::read(path).map_err(|e| format!("read {path}: {e}"))?;
        Ok(b.chunks_exact(4)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect())
    } else {
        Err(format!("unsupported weight URI scheme: {uri}"))
    }
}

/// 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; `rlxp://<path>#<tensor>` returns the package shard range;
/// `file://<path>` returns the raw file. Feed to `set_param_typed(_, _,
/// U8)` behind a `DequantMatMul` graph.
pub fn resolve_weight_bytes(uri: &str) -> Result<Vec<u8>, String> {
    if let Some(rest) = uri.strip_prefix("gguf://") {
        let (path, tensor) = split_frag(rest)?;
        let mut f = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
        let gguf = rlx_gguf::GgufFile::from_reader(&mut f).map_err(|e| e.to_string())?;
        let t = gguf
            .get(tensor)
            .ok_or_else(|| format!("tensor not found: {tensor}"))?;
        Ok(gguf.tensor_bytes(t).map_err(|e| e.to_string())?.to_vec())
    } else if let Some(_rest) = uri.strip_prefix("rlxp://") {
        #[cfg(feature = "model-io")]
        {
            let (path, tensor) = split_frag(_rest)?;
            let pack = rlx_pkg::Package::open(path).map_err(|e| e.to_string())?;
            pack.tensor_bytes(tensor).map_err(|e| e.to_string())
        }
        #[cfg(not(feature = "model-io"))]
        Err(needs_model_io("rlxp"))
    } else if let Some(path) = uri.strip_prefix("file://") {
        std::fs::read(path).map_err(|e| format!("read {path}: {e}"))
    } else {
        Err(format!("packed load unsupported for scheme: {uri}"))
    }
}

/// Parse-once, serve-many weight cache. A GGUF / `.rlxp` 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.
#[derive(Default)]
pub struct WeightCache {
    gguf: HashMap<String, rlx_gguf::GgufFile>,
    #[cfg(feature = "model-io")]
    rlxp: HashMap<String, rlx_pkg::Package>,
}

impl WeightCache {
    pub fn new() -> Self {
        Self::default()
    }

    fn gguf_file(&mut self, path: &str) -> Result<&rlx_gguf::GgufFile, String> {
        if !self.gguf.contains_key(path) {
            let mut f = std::fs::File::open(path).map_err(|e| format!("open {path}: {e}"))?;
            let g = rlx_gguf::GgufFile::from_reader(&mut f).map_err(|e| e.to_string())?;
            self.gguf.insert(path.to_string(), g);
        }
        Ok(&self.gguf[path])
    }

    #[cfg(feature = "model-io")]
    fn rlxp_pack(&mut self, path: &str) -> Result<&rlx_pkg::Package, String> {
        if !self.rlxp.contains_key(path) {
            let p = rlx_pkg::Package::open(path).map_err(|e| e.to_string())?;
            self.rlxp.insert(path.to_string(), p);
        }
        Ok(&self.rlxp[path])
    }

    /// Dequantized f32 for `uri`, reusing the parsed file for `gguf://` / `rlxp://`.
    pub fn f32(&mut self, uri: &str) -> Result<Vec<f32>, String> {
        if let Some(rest) = uri.strip_prefix("gguf://") {
            let (path, tensor) = split_frag(rest)?;
            let (data, _dims) = self
                .gguf_file(path)?
                .dequant_f32(tensor)
                .map_err(|e| e.to_string())?;
            Ok(data)
        } else if let Some(_rest) = uri.strip_prefix("rlxp://") {
            #[cfg(feature = "model-io")]
            {
                let (path, tensor) = split_frag(_rest)?;
                self.rlxp_pack(path)?
                    .tensor_f32(tensor)
                    .map_err(|e| e.to_string())
            }
            #[cfg(not(feature = "model-io"))]
            Err(needs_model_io("rlxp"))
        } else if uri.starts_with("safetensors://")
            || uri.starts_with("file://")
            || uri.starts_with("mlx://")
            || uri.starts_with("npz://")
            || uri.starts_with("dduf://")
        {
            resolve_weight_uri(uri) // header parse is cheap / file is the tensor
        } else {
            Err(format!("unsupported weight URI scheme: {uri}"))
        }
    }

    /// Packed (quantized) bytes for `uri`, reusing the parsed file for `gguf://` / `rlxp://`.
    pub fn bytes(&mut self, uri: &str) -> Result<Vec<u8>, String> {
        if let Some(rest) = uri.strip_prefix("gguf://") {
            let (path, tensor) = split_frag(rest)?;
            let g = self.gguf_file(path)?;
            let t = g
                .get(tensor)
                .ok_or_else(|| format!("tensor not found: {tensor}"))?;
            Ok(g.tensor_bytes(t).map_err(|e| e.to_string())?.to_vec())
        } else if let Some(_rest) = uri.strip_prefix("rlxp://") {
            #[cfg(feature = "model-io")]
            {
                let (path, tensor) = split_frag(_rest)?;
                self.rlxp_pack(path)?
                    .tensor_bytes(tensor)
                    .map_err(|e| e.to_string())
            }
            #[cfg(not(feature = "model-io"))]
            Err(needs_model_io("rlxp"))
        } else {
            resolve_weight_bytes(uri)
        }
    }
}

/// 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.
pub fn uri_resolver(uri: &str) -> Vec<f32> {
    match resolve_weight_uri(uri) {
        Ok(v) => v,
        Err(e) => {
            eprintln!("rlx dist: weight resolve failed [{uri}]: {e}");
            Vec::new()
        }
    }
}

// bf16 is the top 16 bits of an f32.
fn bf16_to_f32(b: u16) -> f32 {
    f32::from_bits((b as u32) << 16)
}

// IEEE-754 half → f32 (handles subnormal / inf / nan).
fn f16_to_f32(h: u16) -> f32 {
    let sign = (h >> 15) & 1;
    let exp = (h >> 10) & 0x1f;
    let mant = h & 0x3ff;
    let v = match exp {
        0 => (mant as f32) * 2f32.powi(-24), // 0 or subnormal
        0x1f if mant == 0 => f32::INFINITY,
        0x1f => f32::NAN,
        _ => (1.0 + mant as f32 / 1024.0) * 2f32.powi(exp as i32 - 15),
    };
    if sign == 1 { -v } else { v }
}

/// Minimal safetensors reader: `[u64 header_len][JSON header][data]`; returns
/// the named tensor as f32 (F32 / F16 / BF16). No external crate needed.
fn read_safetensors_f32(buf: &[u8], name: &str) -> Result<Vec<f32>, String> {
    if buf.len() < 8 {
        return Err("safetensors file too small".into());
    }
    let hlen = u64::from_le_bytes(buf[..8].try_into().unwrap()) as usize;
    let header = buf
        .get(8..8 + hlen)
        .ok_or("safetensors: truncated header")?;
    let data = &buf[8 + hlen..];
    let v: serde_json::Value = serde_json::from_slice(header).map_err(|e| e.to_string())?;
    let t = v
        .get(name)
        .ok_or_else(|| format!("tensor not found: {name}"))?;
    let dtype = t.get("dtype").and_then(|d| d.as_str()).ok_or("no dtype")?;
    let off = t
        .get("data_offsets")
        .and_then(|o| o.as_array())
        .ok_or("no data_offsets")?;
    let start = off[0].as_u64().ok_or("bad data_offsets")? as usize;
    let end = off[1].as_u64().ok_or("bad data_offsets")? as usize;
    let raw = data
        .get(start..end)
        .ok_or("safetensors: data_offsets out of range")?;
    Ok(match dtype {
        "F32" => raw
            .chunks_exact(4)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect(),
        "F16" => raw
            .chunks_exact(2)
            .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]])))
            .collect(),
        "BF16" => raw
            .chunks_exact(2)
            .map(|c| bf16_to_f32(u16::from_le_bytes([c[0], c[1]])))
            .collect(),
        other => return Err(format!("unsupported safetensors dtype: {other}")),
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn resolve_safetensors_f32() {
        let vals = [1.0f32, 2.0, -3.5, 4.25];
        let data: Vec<u8> = vals.iter().flat_map(|v| v.to_le_bytes()).collect();
        let header = format!(
            r#"{{"w":{{"dtype":"F32","shape":[4],"data_offsets":[0,{}]}}}}"#,
            data.len()
        );
        let mut buf = (header.len() as u64).to_le_bytes().to_vec();
        buf.extend_from_slice(header.as_bytes());
        buf.extend_from_slice(&data);
        let path =
            std::env::temp_dir().join(format!("rlx_dist_{}.safetensors", std::process::id()));
        std::fs::write(&path, &buf).unwrap();
        let got = resolve_weight_uri(&format!("safetensors://{}#w", path.display())).unwrap();
        std::fs::remove_file(&path).ok();
        assert_eq!(got, vals);
    }

    #[test]
    fn resolve_gguf_f32() {
        use rlx_gguf::{GgmlType, GgufWriter};
        let vals = [1.0f32, 2.0, -3.5, 4.25];
        let bytes: Vec<u8> = vals.iter().flat_map(|v| v.to_le_bytes()).collect();
        let mut w = GgufWriter::new();
        w.add_tensor_bytes("w", vec![4], GgmlType::F32, bytes)
            .unwrap();
        let path = std::env::temp_dir().join(format!("rlx_dist_{}.gguf", std::process::id()));
        w.write_to_path(&path).unwrap();
        let got = resolve_weight_uri(&format!("gguf://{}#w", path.display())).unwrap();
        std::fs::remove_file(&path).ok();
        assert_eq!(got, vals);
    }

    #[test]
    fn weight_cache_serves_many_tensors_from_one_file() {
        use rlx_gguf::{GgmlType, GgufWriter};
        let a = [1.0f32, 2.0, 3.0, 4.0];
        let b = [10.0f32, 20.0];
        let mut w = GgufWriter::new();
        w.add_tensor_bytes(
            "a",
            vec![4],
            GgmlType::F32,
            a.iter().flat_map(|v| v.to_le_bytes()).collect(),
        )
        .unwrap();
        w.add_tensor_bytes(
            "b",
            vec![2],
            GgmlType::F32,
            b.iter().flat_map(|v| v.to_le_bytes()).collect(),
        )
        .unwrap();
        let path = std::env::temp_dir().join(format!("rlx_cache_{}.gguf", std::process::id()));
        w.write_to_path(&path).unwrap();
        let p = path.display();

        // One cache, two tensors from the same parsed file.
        let mut cache = WeightCache::new();
        assert_eq!(cache.f32(&format!("gguf://{p}#a")).unwrap(), a);
        assert_eq!(cache.f32(&format!("gguf://{p}#b")).unwrap(), b);
        // packed bytes of the F32 tensor are exactly its little-endian f32 bytes.
        let raw = cache.bytes(&format!("gguf://{p}#a")).unwrap();
        let raw_f32: Vec<f32> = raw
            .chunks_exact(4)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect();
        assert_eq!(raw_f32, a);
        std::fs::remove_file(&path).ok();
    }

    #[test]
    #[cfg(feature = "model-io")]
    fn resolve_rlxp_f32() {
        use rlx_ir::op::BinaryOp;
        use rlx_ir::{DType, Graph, Shape};
        use rlx_pkg::{BakeWeight, WriteOptions, package_from_bake};

        let s = Shape::new(&[4], DType::F32);
        let mut g = Graph::new("rlxp_uri");
        let x = g.input("x", s.clone());
        let w = g.param("w", s.clone());
        let y = g.binary(BinaryOp::Mul, x, w, s);
        g.set_outputs(vec![y]);
        let vals = [1.0f32, 2.0, -3.5, 4.25];
        let weights = [BakeWeight {
            name: "w".into(),
            shape: vec![4],
            encoding: "f32".into(),
            data: vals.iter().flat_map(|v| v.to_le_bytes()).collect(),
        }];
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("m.rlxp");
        package_from_bake(
            &path,
            &g,
            &weights,
            WriteOptions {
                container: rlx_pkg::ContainerKind::Flat,
                name: "rlxp_uri".into(),
                ..WriteOptions::default()
            },
        )
        .unwrap();
        let got = resolve_weight_uri(&format!("rlxp://{}#w", path.display())).unwrap();
        assert_eq!(got, vals);
        let mut cache = WeightCache::new();
        assert_eq!(
            cache.f32(&format!("rlxp://{}#w", path.display())).unwrap(),
            vals
        );
    }

    #[test]
    fn dequant_matmul_gguf_end_to_end() {
        // Quantize a weight to Q8_0, feed the PACKED bytes as a U8 param behind
        // a DequantMatMul graph, and check the output matches dequant-then-
        // matmul — the quantized-on-device path, weight at native quant memory.
        use crate::Session;
        use rlx_ir::quant::QuantScheme;
        use rlx_ir::{DType, Graph, Shape};

        let (m, k, n) = (2usize, 32usize, 4usize); // K multiple of 32 for Q8_0
        let w_floats: Vec<f32> = (0..n * k).map(|i| ((i % 17) as f32 - 8.0) * 0.1).collect();
        let packed = rlx_gguf::quantize(&w_floats, rlx_gguf::GgmlType::Q8_0).unwrap();
        let x: Vec<f32> = (0..m * k).map(|i| ((i % 13) as f32) * 0.05).collect();

        let mut g = Graph::new("dq");
        let xin = g.input("x", Shape::new(&[m, k], DType::F32));
        let wp = g.param("W", Shape::new(&[packed.len()], DType::U8));
        let out = g.dequant_matmul_packed(
            xin,
            wp,
            QuantScheme::GgufQ8_0,
            Shape::new(&[m, n], DType::F32),
        );
        g.set_outputs(vec![out]);

        let mut c = Session::new(Device::Cpu).compile(g);
        c.set_param_typed("W", &packed, DType::U8);
        let got = c.run(&[("x", x.as_slice())]).into_iter().next().unwrap();

        // Reference: dequant the SAME bytes, then BT-matmul (W is [n, k]).
        let w = rlx_gguf::dequant_q8_0(&packed, n * k).unwrap();
        let mut expect = vec![0f32; m * n];
        for i in 0..m {
            for j in 0..n {
                let mut acc = 0f32;
                for c2 in 0..k {
                    acc += x[i * k + c2] * w[j * k + c2];
                }
                expect[i * n + j] = acc;
            }
        }
        let err = got
            .iter()
            .zip(&expect)
            .map(|(a, b)| (a - b).abs())
            .fold(0f32, f32::max);
        assert!(err < 1e-3, "max_err={err} got={got:?} expect={expect:?}");
    }
}