rlx_distributed/
source.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//! The **model-agnostic weight seam**. `rlx-distributed` never knows how a
6//! model's weights are stored — a stage worker asks a [`ParamSource`] for each
7//! parameter by name, and the source (implemented by a *model* crate, e.g.
8//! wrapping an mlx / GGUF / safetensors loader) returns it as dense f32 or as
9//! packed/quantized bytes. This is what keeps the crate model-agnostic: RLX
10//! provides the API, `rlx-models` provides the weights.
11//!
12//! Ergonomics: a plain `HashMap<String, Vec<f32>>` is a `ParamSource` (great for
13//! tests and small models), and so is any closure `FnMut(&str) -> Option<Param>`
14//! (great for adapting an existing loader in one line):
15//!
16//! ```ignore
17//! // adapt any loader without a newtype:
18//! let mut src = |name: &str| loader.take_typed(name).map(|(b, dt)| Param::typed(b, dt));
19//! serve_stage(addr, stage, &mut src, Device::Cpu, &opts, 1)?;
20//! ```
21
22use rlx_ir::DType;
23use std::collections::HashMap;
24
25/// One parameter value: dense f32, or packed bytes with a dtype (quantized /
26/// f16 / bf16 weights fed to the runtime without a host cast).
27#[derive(Clone, Debug)]
28pub enum Param {
29    F32(Vec<f32>),
30    Typed(Vec<u8>, DType),
31}
32
33impl Param {
34    pub fn f32(v: Vec<f32>) -> Self {
35        Param::F32(v)
36    }
37    pub fn typed(bytes: Vec<u8>, dtype: DType) -> Self {
38        Param::Typed(bytes, dtype)
39    }
40}
41
42/// Supplies a stage's parameters on demand, by name. `None` means "not mine"
43/// (the runtime leaves the param unset — useful when a source only owns a shard).
44pub trait ParamSource {
45    fn get(&mut self, name: &str) -> Option<Param>;
46}
47
48/// A `HashMap<String, Vec<f32>>` is a source (dense f32; clones on read).
49impl ParamSource for HashMap<String, Vec<f32>> {
50    fn get(&mut self, name: &str) -> Option<Param> {
51        HashMap::get(self, name).map(|v| Param::F32(v.clone()))
52    }
53}
54
55/// A `HashMap<String, Param>` is a source (mixed f32 / packed).
56impl ParamSource for HashMap<String, Param> {
57    fn get(&mut self, name: &str) -> Option<Param> {
58        HashMap::get(self, name).cloned()
59    }
60}
61
62/// Any `FnMut(&str) -> Option<Param>` is a source — adapt an existing loader
63/// inline without defining a type.
64impl<F> ParamSource for F
65where
66    F: FnMut(&str) -> Option<Param>,
67{
68    fn get(&mut self, name: &str) -> Option<Param> {
69        self(name)
70    }
71}
rlx_distributed/source.rs

rlx_distributed/
source.rs
