metrale_storage/
expert.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Expert identity and on-disk record geometry for MoE expert
4//! streaming. A record holds one (moe_layer, expert)'s gate, up and down NVFP4
5//! projections behind a fixed header; each MoE layer's file holds
6//! `num_experts` records at one stride, which the whole model shares.
7//!
8//! Owner: metrale-storage experts.
9//! Invariants:
10//! - Inside a record the header comes first, then gate, up and down (packed,
11//!   then scale, for each), every sub-buffer at a multiple of `sub_align` and
12//!   none overlapping.
13//! - `record_stride` is `raw_bytes` rounded up to a multiple of `fs_block_size`.
14
15/// 2026-09-25: Dense record id across the model:
16/// `layer · num_experts + expert`.
17#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
18pub struct ExpertRecordId(pub u64);
19
20/// 2026-09-25: One expert's record. `layer` is the dense MoE-layer index
21/// (`0..num_moe_layers`), not the model layer;
22/// `ExpertIndex::moe_layer_to_model_layer` maps one to the other.
23#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
24pub struct ExpertKey {
25    pub layer: u32,
26    pub expert: u32,
27}
28
29impl ExpertKey {
30    pub fn new(layer: u32, expert: u32) -> Self {
31        Self { layer, expert }
32    }
33}
34
35/// 2026-09-25: The three projections of an expert. Their order is the order
36/// of the sub-buffers in a record and of the header's per-projection fields,
37/// so changing it changes the on-disk format that
38/// [`ExpertRecordHeader::VERSION`] identifies.
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub enum Proj {
41    Gate = 0,
42    Up = 1,
43    Down = 2,
44}
45
46impl Proj {
47    pub const ALL: [Proj; 3] = [Proj::Gate, Proj::Up, Proj::Down];
48}
49
50/// 2026-09-25: Byte sizes of one NVFP4 projection's two buffers: `packed`,
51/// two 4-bit values per byte (`n·k/2`), and `scale`, one byte per
52/// `group_size` values (`n·k/group_size`). The per-projection scalars
53/// `weight_scale_2` and `input_scale` are in the record header instead.
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55pub struct ProjBytes {
56    pub packed_bytes: u64,
57    pub scale_bytes: u64,
58}
59
60impl ProjBytes {
61    /// 2026-09-25: `n` output rows, `k` contraction dim, `group_size` values
62    /// per scale.
63    pub fn nvfp4(n: u64, k: u64, group_size: u64) -> Self {
64        Self {
65            packed_bytes: n * k / 2,
66            scale_bytes: n * k / group_size,
67        }
68    }
69}
70
71/// 2026-09-25: Where each sub-buffer of an expert record sits, relative to the
72/// record base; `pack_record` and `unpack_record` place and read bytes by it.
73/// Every offset is a multiple of `sub_align`:
74///
75/// ```text
76///   [ header (ExpertRecordHeader::BYTES) ]
77///   [ gate.packed ][ gate.scale ]
78///   [ up.packed   ][ up.scale   ]
79///   [ down.packed ][ down.scale ]
80///   [ pad to record_stride ]
81/// ```
82#[derive(Clone, Copy, Debug, PartialEq, Eq)]
83pub struct ExpertRecordSpec {
84    pub inter: u64,
85    pub hidden: u64,
86    pub group_size: u64,
87    pub sub_align: u64,
88    /// 2026-09-25: `(packed_off, scale_off)` for gate, up, down.
89    offsets: [(u64, u64); 3],
90    bytes: [ProjBytes; 3],
91    /// 2026-09-25: Header plus all sub-buffers with their padding, before
92    /// `ExpertLayout` rounds it up to `record_stride`.
93    raw_bytes: u64,
94}
95
96/// 2026-09-25: `off` rounded up to a multiple of `align`, a power of two.
97#[inline]
98fn align_up(off: u64, align: u64) -> u64 {
99    (off + align - 1) & !(align - 1)
100}
101
102impl ExpertRecordSpec {
103    /// 2026-09-25: The record layout for experts of intermediate size `inter`
104    /// and hidden size `hidden`. Panics unless `sub_align` is a power of two.
105    pub fn new(inter: u64, hidden: u64, group_size: u64, sub_align: u64) -> Self {
106        assert!(
107            sub_align.is_power_of_two(),
108            "sub_align must be a power of two"
109        );
110        // 2026-09-25: gate and up are N=inter, K=hidden; down is N=hidden,
111        // K=inter. Both byte counts are symmetric in (N, K), so the three sizes
112        // are equal.
113        let bytes = [
114            ProjBytes::nvfp4(inter, hidden, group_size),
115            ProjBytes::nvfp4(inter, hidden, group_size),
116            ProjBytes::nvfp4(hidden, inter, group_size),
117        ];
118        let mut cursor = align_up(ExpertRecordHeader::BYTES, sub_align);
119        let mut offsets = [(0u64, 0u64); 3];
120        for i in 0..3 {
121            let packed_off = cursor;
122            cursor = align_up(packed_off + bytes[i].packed_bytes, sub_align);
123            let scale_off = cursor;
124            cursor = align_up(scale_off + bytes[i].scale_bytes, sub_align);
125            offsets[i] = (packed_off, scale_off);
126        }
127        Self {
128            inter,
129            hidden,
130            group_size,
131            sub_align,
132            offsets,
133            bytes,
134            raw_bytes: cursor,
135        }
136    }
137
138    pub fn proj_bytes(&self, p: Proj) -> ProjBytes {
139        self.bytes[p as usize]
140    }
141
142    /// 2026-09-25: Offset of `p`'s packed sub-buffer in the record.
143    pub fn packed_off(&self, p: Proj) -> u64 {
144        self.offsets[p as usize].0
145    }
146
147    /// 2026-09-25: Offset of `p`'s scale sub-buffer in the record.
148    pub fn scale_off(&self, p: Proj) -> u64 {
149        self.offsets[p as usize].1
150    }
151
152    /// 2026-09-25: See the `raw_bytes` field.
153    pub fn raw_bytes(&self) -> u64 {
154        self.raw_bytes
155    }
156
157    /// 2026-09-25: Sum of the six sub-buffers, without header and padding.
158    pub fn payload_bytes(&self) -> u64 {
159        self.bytes
160            .iter()
161            .map(|b| b.packed_bytes + b.scale_bytes)
162            .sum()
163    }
164}
165
166/// 2026-09-25: The fixed, versioned header at the front of every expert
167/// record. It carries what the dims do not give, each projection's
168/// `weight_scale_2` and `input_scale`, plus layer, expert and shape so a
169/// mismatched file can be recognised.
170#[derive(Clone, Copy, Debug, PartialEq)]
171pub struct ExpertRecordHeader {
172    pub layer: u32,
173    pub expert: u32,
174    pub inter: u32,
175    pub hidden: u32,
176    pub group_size: u32,
177    /// 2026-09-25: `weight_scale_2` of gate, up, down.
178    pub scale2: [f32; 3],
179    /// 2026-09-25: `input_scale` of gate, up, down; `None` when the
180    /// projection has none. Presence is a bit in the header's
181    /// `input_scale_flags` word, so `PartialEq` needs no NaN sentinel.
182    pub input_scale: [Option<f32>; 3],
183}
184
185impl ExpertRecordHeader {
186    pub const MAGIC: u32 = 0x5850_5254;
187    pub const VERSION: u32 = 1;
188    /// 2026-09-25: Reserved header size; the fields use the first 56 bytes,
189    /// so the header can grow without moving the sub-buffers.
190    pub const BYTES: u64 = 256;
191
192    /// 2026-09-25: Serialize into the fixed 256-byte header block,
193    /// little-endian:
194    ///   u32 magic, u32 version, u32 layer, u32 expert,
195    ///   u32 inter, u32 hidden, u32 group_size, u32 input_scale_flags,
196    ///   `f32 scale2[3]`, `f32 input_scale[3]` (0.0 where absent), zero pad to 256.
197    /// `input_scale_flags` bit `i` set => projection `i` has an activation scale.
198    pub fn to_bytes(&self) -> [u8; Self::BYTES as usize] {
199        let mut out = [0u8; Self::BYTES as usize];
200        let mut w = |off: usize, v: u32| out[off..off + 4].copy_from_slice(&v.to_le_bytes());
201        w(0, Self::MAGIC);
202        w(4, Self::VERSION);
203        w(8, self.layer);
204        w(12, self.expert);
205        w(16, self.inter);
206        w(20, self.hidden);
207        w(24, self.group_size);
208        let mut flags = 0u32;
209        for (i, s) in self.input_scale.iter().enumerate() {
210            if s.is_some() {
211                flags |= 1 << i;
212            }
213        }
214        w(28, flags);
215        for (i, s) in self.scale2.iter().enumerate() {
216            out[32 + i * 4..36 + i * 4].copy_from_slice(&s.to_le_bytes());
217        }
218        for (i, s) in self.input_scale.iter().enumerate() {
219            let v = s.unwrap_or(0.0);
220            out[44 + i * 4..48 + i * 4].copy_from_slice(&v.to_le_bytes());
221        }
222        out
223    }
224
225    /// 2026-09-25: Parse a header block. `None` when it is shorter than
226    /// `BYTES` or the magic or version differs.
227    pub fn from_bytes(buf: &[u8]) -> Option<Self> {
228        if buf.len() < Self::BYTES as usize {
229            return None;
230        }
231        let r = |off: usize| -> u32 {
232            u32::from_le_bytes([buf[off], buf[off + 1], buf[off + 2], buf[off + 3]])
233        };
234        let rf = |off: usize| -> f32 {
235            f32::from_le_bytes([buf[off], buf[off + 1], buf[off + 2], buf[off + 3]])
236        };
237        if r(0) != Self::MAGIC || r(4) != Self::VERSION {
238            return None;
239        }
240        let flags = r(28);
241        let iscale = |i: usize, off: usize| -> Option<f32> {
242            if flags & (1 << i) != 0 {
243                Some(rf(off))
244            } else {
245                None
246            }
247        };
248        Some(Self {
249            layer: r(8),
250            expert: r(12),
251            inter: r(16),
252            hidden: r(20),
253            group_size: r(24),
254            scale2: [rf(32), rf(36), rf(40)],
255            input_scale: [iscale(0, 44), iscale(1, 48), iscale(2, 52)],
256        })
257    }
258}
259
260/// 2026-09-25: File geometry of an expert store: one file per MoE layer
261/// (`ExpertIndex::file_name`), each holding `num_experts` records back to back
262/// at `record_stride`.
263#[derive(Clone, Copy, Debug, PartialEq, Eq)]
264pub struct ExpertLayout {
265    pub num_layers: u32,
266    pub num_experts: u32,
267    /// 2026-09-25: Record stride on disk, a multiple of `fs_block_size`.
268    pub record_stride: u64,
269    pub fs_block_size: u64,
270}
271
272impl ExpertLayout {
273    /// 2026-09-25: `record_stride` is `spec.raw_bytes()` rounded up to a
274    /// multiple of `fs_block_size`.
275    pub fn from_spec(
276        num_layers: u32,
277        num_experts: u32,
278        spec: &ExpertRecordSpec,
279        fs_block_size: u64,
280    ) -> Self {
281        let record_stride = spec.raw_bytes().div_ceil(fs_block_size) * fs_block_size;
282        Self {
283            num_layers,
284            num_experts,
285            record_stride,
286            fs_block_size,
287        }
288    }
289
290    /// 2026-09-25: Bytes of one MoE layer's file.
291    pub fn bytes_per_layer(&self) -> u64 {
292        (self.num_experts as u64) * self.record_stride
293    }
294
295    /// 2026-09-25: Offset of `key`'s record in its layer file. Debug builds
296    /// assert that `expert` is in range.
297    pub fn file_offset(&self, key: ExpertKey) -> u64 {
298        debug_assert!(key.expert < self.num_experts);
299        (key.expert as u64) * self.record_stride
300    }
301
302    pub fn record_id(&self, key: ExpertKey) -> ExpertRecordId {
303        ExpertRecordId((key.layer as u64) * (self.num_experts as u64) + (key.expert as u64))
304    }
305
306    /// 2026-09-25: Bytes of one record on disk: `record_stride`.
307    pub fn record_bytes(&self) -> u64 {
308        self.record_stride
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315
316    const A3B_INTER: u64 = 512;
317    const A3B_HIDDEN: u64 = 2048;
318    const GS: u64 = 16;
319
320    #[test]
321    fn a3b_per_expert_payload_matches_formula() {
322        // 2026-09-25: Three projections at 1/2 + 1/16 bytes per weight.
323        let spec = ExpertRecordSpec::new(A3B_INTER, A3B_HIDDEN, GS, 256);
324        let expected = 3 * A3B_INTER * A3B_HIDDEN * 9 / 16;
325        assert_eq!(spec.payload_bytes(), expected);
326        assert_eq!(expected, 1_769_472);
327    }
328
329    #[test]
330    fn projection_bytes_split_8_to_1() {
331        // 2026-09-25: 1/2 byte per weight against 1/16: packed is 8× scale.
332        let pb = ProjBytes::nvfp4(A3B_INTER, A3B_HIDDEN, GS);
333        assert_eq!(pb.packed_bytes, A3B_INTER * A3B_HIDDEN / 2);
334        assert_eq!(pb.scale_bytes, A3B_INTER * A3B_HIDDEN / 16);
335        assert_eq!(pb.packed_bytes, 8 * pb.scale_bytes);
336    }
337
338    #[test]
339    fn sub_buffers_are_aligned_and_non_overlapping() {
340        let align = 256;
341        let spec = ExpertRecordSpec::new(A3B_INTER, A3B_HIDDEN, GS, align);
342        let mut prev_end = ExpertRecordHeader::BYTES;
343        for p in Proj::ALL {
344            let po = spec.packed_off(p);
345            let so = spec.scale_off(p);
346            let pb = spec.proj_bytes(p);
347            assert_eq!(po % align, 0, "packed off aligned");
348            assert_eq!(so % align, 0, "scale off aligned");
349            assert!(po >= prev_end, "packed does not overlap previous");
350            assert!(so >= po + pb.packed_bytes, "scale does not overlap packed");
351            prev_end = so + pb.scale_bytes;
352        }
353        assert!(spec.raw_bytes() >= prev_end);
354    }
355
356    #[test]
357    fn layout_offsets_are_record_strided() {
358        let spec = ExpertRecordSpec::new(A3B_INTER, A3B_HIDDEN, GS, 256);
359        let layout = ExpertLayout::from_spec(40, 256, &spec, 4096);
360        assert_eq!(layout.record_stride % 4096, 0, "O_DIRECT alignment");
361        assert!(layout.record_stride >= spec.raw_bytes());
362        assert_eq!(layout.file_offset(ExpertKey::new(3, 0)), 0);
363        assert_eq!(
364            layout.file_offset(ExpertKey::new(3, 5)),
365            5 * layout.record_stride
366        );
367        assert_eq!(layout.bytes_per_layer(), 256 * layout.record_stride);
368    }
369
370    #[test]
371    fn record_id_is_dense_layer_major() {
372        let spec = ExpertRecordSpec::new(A3B_INTER, A3B_HIDDEN, GS, 256);
373        let layout = ExpertLayout::from_spec(40, 256, &spec, 4096);
374        assert_eq!(layout.record_id(ExpertKey::new(0, 0)).0, 0);
375        assert_eq!(layout.record_id(ExpertKey::new(0, 255)).0, 255);
376        assert_eq!(layout.record_id(ExpertKey::new(1, 0)).0, 256);
377    }
378
379    #[test]
380    fn header_round_trips() {
381        let h = ExpertRecordHeader {
382            layer: 7,
383            expert: 42,
384            inter: A3B_INTER as u32,
385            hidden: A3B_HIDDEN as u32,
386            group_size: GS as u32,
387            scale2: [0.5, 0.25, 1.5],
388            input_scale: [Some(2.0), None, Some(3.0)],
389        };
390        let bytes = h.to_bytes();
391        assert_eq!(bytes.len(), ExpertRecordHeader::BYTES as usize);
392        let back = ExpertRecordHeader::from_bytes(&bytes).expect("valid header");
393        assert_eq!(back, h);
394        assert_eq!(back.scale2, [0.5, 0.25, 1.5]);
395        assert_eq!(back.input_scale, [Some(2.0), None, Some(3.0)]);
396    }
397
398    #[test]
399    fn header_rejects_bad_magic_and_version() {
400        let mut bytes = ExpertRecordHeader {
401            layer: 0,
402            expert: 0,
403            inter: 1,
404            hidden: 1,
405            group_size: GS as u32,
406            scale2: [1.0; 3],
407            input_scale: [Some(1.0); 3],
408        }
409        .to_bytes();
410        bytes[0] ^= 0xFF;
411        assert!(ExpertRecordHeader::from_bytes(&bytes).is_none());
412        assert!(ExpertRecordHeader::from_bytes(&bytes[..10]).is_none());
413    }
414
415    #[test]
416    fn a3b_record_stride_is_4k_aligned_and_reasonable() {
417        let spec = ExpertRecordSpec::new(A3B_INTER, A3B_HIDDEN, GS, 256);
418        let layout = ExpertLayout::from_spec(40, 256, &spec, 4096);
419        assert!(layout.record_stride - spec.raw_bytes() < 4096);
420    }
421}
metrale_storage/expert.rs

metrale_storage/
expert.rs
