metrale_storage/
expert_pack.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The expert-store format. `pack_record` and `unpack_record` build
4//! and parse one record in memory, with no I/O. `ExpertIndex` is the store's
5//! `manifest.json`, from which a reader rebuilds the record spec and layout.
6//! The layer files are written and read in `expert_pack_fs.rs`.
7//!
8//! Owner: metrale-storage experts.
9//! Invariants:
10//! - `ExpertIndex::load` and `ExpertFileReader::open` reject a manifest whose
11//!   `version` is not `ExpertRecordHeader::VERSION`.
12
13use anyhow::{Context, Result, bail};
14use serde::{Deserialize, Serialize};
15
16use crate::expert::ExpertKey;
17use crate::expert::{ExpertLayout, ExpertRecordHeader, ExpertRecordSpec, Proj};
18
19/// 2026-09-25: One projection's packed and scale bytes, as `pack_record`
20/// writes them.
21#[derive(Clone, Copy, Debug)]
22pub struct ProjData<'a> {
23    pub packed: &'a [u8],
24    pub scale: &'a [u8],
25}
26
27/// 2026-09-25: One projection's sub-buffers inside a parsed record.
28#[derive(Clone, Copy, Debug)]
29pub struct ProjView<'a> {
30    pub packed: &'a [u8],
31    pub scale: &'a [u8],
32}
33
34/// 2026-09-25: Build one `stride`-byte record: the header at offset 0, each
35/// projection's bytes at the spec's offsets, zeros elsewhere. Fails when
36/// `stride` is below `spec.raw_bytes()` or a projection's lengths differ from
37/// the spec.
38pub fn pack_record(
39    spec: &ExpertRecordSpec,
40    stride: u64,
41    header: &ExpertRecordHeader,
42    projs: &[ProjData; 3],
43) -> Result<Vec<u8>> {
44    let stride = stride as usize;
45    if (spec.raw_bytes() as usize) > stride {
46        bail!(
47            "record stride {} smaller than raw record bytes {}",
48            stride,
49            spec.raw_bytes()
50        );
51    }
52    let mut buf = vec![0u8; stride];
53    let hdr = header.to_bytes();
54    buf[..hdr.len()].copy_from_slice(&hdr);
55
56    for p in Proj::ALL {
57        let pb = spec.proj_bytes(p);
58        let d = &projs[p as usize];
59        if d.packed.len() as u64 != pb.packed_bytes {
60            bail!(
61                "{:?} packed len {} != expected {}",
62                p,
63                d.packed.len(),
64                pb.packed_bytes
65            );
66        }
67        if d.scale.len() as u64 != pb.scale_bytes {
68            bail!(
69                "{:?} scale len {} != expected {}",
70                p,
71                d.scale.len(),
72                pb.scale_bytes
73            );
74        }
75        let po = spec.packed_off(p) as usize;
76        let so = spec.scale_off(p) as usize;
77        buf[po..po + d.packed.len()].copy_from_slice(d.packed);
78        buf[so..so + d.scale.len()].copy_from_slice(d.scale);
79    }
80    Ok(buf)
81}
82
83/// 2026-09-25: Parse a record of at least `spec.raw_bytes()` bytes into its
84/// header and projection views. Fails on a short buffer or a header whose
85/// magic or version differs.
86pub fn unpack_record<'a>(
87    spec: &ExpertRecordSpec,
88    buf: &'a [u8],
89) -> Result<(ExpertRecordHeader, [ProjView<'a>; 3])> {
90    if (buf.len() as u64) < spec.raw_bytes() {
91        bail!(
92            "record buffer {} smaller than raw record bytes {}",
93            buf.len(),
94            spec.raw_bytes()
95        );
96    }
97    let header = ExpertRecordHeader::from_bytes(buf)
98        .context("record header magic/version mismatch (wrong file or format version?)")?;
99    let mut views = [ProjView {
100        packed: &[],
101        scale: &[],
102    }; 3];
103    for p in Proj::ALL {
104        let pb = spec.proj_bytes(p);
105        let po = spec.packed_off(p) as usize;
106        let so = spec.scale_off(p) as usize;
107        views[p as usize] = ProjView {
108            packed: &buf[po..po + pb.packed_bytes as usize],
109            scale: &buf[so..so + pb.scale_bytes as usize],
110        };
111    }
112    Ok((header, views))
113}
114
115/// 2026-09-25: The store manifest, `manifest.json` beside the per-layer
116/// files: the geometry a reader needs to rebuild the spec and layout and open
117/// the files.
118#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
119pub struct ExpertIndex {
120    /// 2026-09-25: Must equal [`ExpertRecordHeader::VERSION`] for a reader to
121    /// accept the store.
122    pub version: u32,
123    pub num_moe_layers: u32,
124    pub num_experts: u32,
125    pub inter: u64,
126    pub hidden: u64,
127    pub group_size: u64,
128    pub sub_align: u64,
129    pub fs_block_size: u64,
130    pub record_stride: u64,
131    pub record_raw_bytes: u64,
132    /// 2026-09-25: Set to `FILE_TEMPLATE` by `new`; `file_name` does not read
133    /// it.
134    pub file_template: String,
135    /// 2026-09-25: Dense MoE-layer index → model layer index; `new` sets
136    /// `num_moe_layers` to its length.
137    pub moe_layer_to_model_layer: Vec<u32>,
138}
139
140impl ExpertIndex {
141    pub const FILE_TEMPLATE: &'static str = "experts_{:05}.xpr";
142    pub const MANIFEST_NAME: &'static str = "manifest.json";
143
144    pub fn new(
145        inter: u64,
146        hidden: u64,
147        group_size: u64,
148        sub_align: u64,
149        fs_block_size: u64,
150        moe_layer_to_model_layer: Vec<u32>,
151        num_experts: u32,
152    ) -> Self {
153        let spec = ExpertRecordSpec::new(inter, hidden, group_size, sub_align);
154        let layout = ExpertLayout::from_spec(
155            moe_layer_to_model_layer.len() as u32,
156            num_experts,
157            &spec,
158            fs_block_size,
159        );
160        Self {
161            version: ExpertRecordHeader::VERSION,
162            num_moe_layers: moe_layer_to_model_layer.len() as u32,
163            num_experts,
164            inter,
165            hidden,
166            group_size,
167            sub_align,
168            fs_block_size,
169            record_stride: layout.record_stride,
170            record_raw_bytes: spec.raw_bytes(),
171            file_template: Self::FILE_TEMPLATE.to_string(),
172            moe_layer_to_model_layer,
173        }
174    }
175
176    /// 2026-09-25: Read and version-check `manifest.json` from `dir`, without
177    /// opening the layer files.
178    #[cfg(unix)]
179    pub fn load(dir: &std::path::Path) -> Result<Self> {
180        let p = dir.join(Self::MANIFEST_NAME);
181        let json = std::fs::read_to_string(&p).with_context(|| format!("read {}", p.display()))?;
182        let index: ExpertIndex =
183            serde_json::from_str(&json).with_context(|| format!("parse {}", p.display()))?;
184        if index.version != ExpertRecordHeader::VERSION {
185            bail!(
186                "manifest version {} != supported {}",
187                index.version,
188                ExpertRecordHeader::VERSION
189            );
190        }
191        Ok(index)
192    }
193
194    pub fn spec(&self) -> ExpertRecordSpec {
195        ExpertRecordSpec::new(self.inter, self.hidden, self.group_size, self.sub_align)
196    }
197
198    pub fn layout(&self) -> ExpertLayout {
199        ExpertLayout::from_spec(
200            self.num_moe_layers,
201            self.num_experts,
202            &self.spec(),
203            self.fs_block_size,
204        )
205    }
206
207    /// 2026-09-25: `experts_{moe_layer:05}.xpr`, whatever `file_template` holds.
208    pub fn file_name(&self, moe_layer: u32) -> String {
209        format!("experts_{moe_layer:05}.xpr")
210    }
211
212    /// 2026-09-25: Bytes of all layer files together.
213    pub fn total_bytes(&self) -> u64 {
214        (self.num_moe_layers as u64) * self.layout().bytes_per_layer()
215    }
216}
217
218pub use fs_impl::{ExpertFileReader, ExpertFileWriter};
219
220#[path = "expert_pack_fs.rs"]
221mod fs_impl;
222
223#[cfg(test)]
224#[path = "expert_pack_tests.rs"]
225mod tests;
metrale_storage/expert_pack.rs

metrale_storage/
expert_pack.rs
