rlx_runtime/
weights.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//! Weight-loading abstraction.
6//!
7//! Native targets typically `mmap` a `.safetensors` file and read tensors
8//! via byte-offset slices into the mapping. WASM has no `mmap`; weights
9//! arrive as `Vec<u8>` from `fetch()` or `Response.arrayBuffer()`. Both
10//! paths produce the same shape: a name → byte-slice lookup.
11//!
12//! `WeightLoader` is the contract. Concrete implementations live in
13//! in-tree model loaders (mmap-based) and here (`BytesWeightLoader` — works on
14//! every target including WASM).
15
16/// A name-keyed view of weight tensor bytes.
17///
18/// Implementations promise that the returned slice stays valid for the
19/// lifetime of `&self`. On native, this is the mmap region; on WASM, it
20/// is the in-memory `Vec<u8>` owned by the loader.
21pub trait WeightLoader {
22    /// Return the raw bytes for the tensor named `name`, or `None` if not
23    /// present. Bytes are in the source file's storage order (typically
24    /// row-major, dtype-native).
25    fn tensor_bytes(&self, name: &str) -> Option<&[u8]>;
26
27    /// All tensor names (for iteration / discovery). Order is
28    /// implementation-defined but stable for a given loader instance.
29    fn names(&self) -> Vec<String>;
30}
31
32/// Owned in-memory weight loader. The simplest, most portable variant —
33/// works on every target including WASM.
34///
35/// Construct via `BytesWeightLoader::from_safetensors(bytes)` once
36/// model builders integrate. For now the bare struct lets external callers
37/// build their own name → bytes mapping.
38pub struct BytesWeightLoader {
39    /// `(name, start_offset, len)` triples into `data`.
40    entries: Vec<(String, usize, usize)>,
41    data: Vec<u8>,
42}
43
44impl BytesWeightLoader {
45    /// Build a loader from a list of `(name, bytes)` pairs. Each tensor
46    /// is appended into a single backing `Vec<u8>`; `tensor_bytes` returns
47    /// a borrow into that vec.
48    pub fn from_pairs(pairs: Vec<(String, Vec<u8>)>) -> Self {
49        let total: usize = pairs.iter().map(|(_, b)| b.len()).sum();
50        let mut data = Vec::with_capacity(total);
51        let mut entries = Vec::with_capacity(pairs.len());
52        for (name, bytes) in pairs {
53            let start = data.len();
54            let len = bytes.len();
55            data.extend_from_slice(&bytes);
56            entries.push((name, start, len));
57        }
58        Self { entries, data }
59    }
60}
61
62impl WeightLoader for BytesWeightLoader {
63    fn tensor_bytes(&self, name: &str) -> Option<&[u8]> {
64        self.entries
65            .iter()
66            .find(|(n, _, _)| n == name)
67            .map(|(_, off, len)| &self.data[*off..*off + *len])
68    }
69
70    fn names(&self) -> Vec<String> {
71        self.entries.iter().map(|(n, _, _)| n.clone()).collect()
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn round_trip() {
81        let loader = BytesWeightLoader::from_pairs(vec![
82            ("w".into(), vec![1, 2, 3, 4]),
83            ("b".into(), vec![5, 6]),
84        ]);
85        assert_eq!(loader.tensor_bytes("w"), Some(&[1u8, 2, 3, 4][..]));
86        assert_eq!(loader.tensor_bytes("b"), Some(&[5u8, 6][..]));
87        assert_eq!(loader.tensor_bytes("missing"), None);
88        assert_eq!(loader.names(), vec!["w".to_string(), "b".to_string()]);
89    }
90}
rlx_runtime/weights.rs

rlx_runtime/
weights.rs
