metrale_storage/
expert_tier.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `ExpertTier`: fetch one expert record into a slot and return
4//! the device addresses of its six sub-buffers, with its scalars. The tiers are
5//! `PosixTier` (read, then copy into a device buffer), `UmaArenaTier` (read
6//! straight into a pinned arena slot) and `RdmaTier` (expert_tier_rdma.rs, from
7//! a peer). It is not a `StorageBackend`, whose reads take KV group keys.
8//!
9//! Owner: metrale-storage experts.
10//! Invariants:
11//! - Every tier builds its `ExpertResidency` with `residency_from`, which
12//!   rejects a record whose header names another (layer, expert).
13
14use crate::tier::pio;
15use anyhow::{Context, Result, bail};
16use std::fs::{File, OpenOptions};
17use std::path::Path;
18
19use crate::cuda_min::{DeviceBuffer, copy_h_to_d_async, stream_sync};
20use crate::expert::{ExpertKey, ExpertLayout, ExpertRecordHeader, ExpertRecordSpec, Proj};
21use crate::expert_arena::ExpertArena;
22use crate::expert_pack::{ExpertFileReader, ExpertIndex};
23
24/// 2026-09-25: The six sub-buffer device addresses and the scalars of one
25/// fetched expert, each array ordered gate, up, down.
26#[derive(Clone, Copy, Debug, PartialEq)]
27pub struct ExpertResidency {
28    pub packed_addr: [u64; 3],
29    pub scale_addr: [u64; 3],
30    pub scale2: [f32; 3],
31    /// 2026-09-25: `None` where the record has no input scale.
32    pub input_scale: [Option<f32>; 3],
33}
34
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
36pub enum TierKind {
37    Posix,
38    Uma,
39    Rdma,
40}
41
42/// 2026-09-25: A destination slot: which slab, and which slot in it.
43#[derive(Clone, Copy, Debug, PartialEq, Eq)]
44pub struct ArenaSlot {
45    pub slab: u32,
46    pub slot: u32,
47}
48
49impl ArenaSlot {
50    pub fn new(slab: u32, slot: u32) -> Self {
51        Self { slab, slot }
52    }
53}
54
55/// 2026-09-25: Fetch an expert record into a slot and return its residency.
56pub trait ExpertTier: Send {
57    fn fetch(&mut self, key: ExpertKey, slot: ArenaSlot, stream: u64) -> Result<ExpertResidency>;
58    fn kind(&self) -> TierKind;
59    /// 2026-09-25: Whether the tier can still serve fetches; the default is
60    /// `true`.
61    fn healthy(&self) -> bool {
62        true
63    }
64}
65
66/// 2026-09-25: Build an `ExpertResidency` from a record's header and the
67/// record's device base address, using the spec's sub-buffer offsets.
68pub(crate) fn residency_from(
69    spec: &ExpertRecordSpec,
70    record: &[u8],
71    base_dev_va: u64,
72    key: ExpertKey,
73) -> Result<ExpertResidency> {
74    let hdr = ExpertRecordHeader::from_bytes(record)
75        .context("expert record header magic/version mismatch")?;
76    // 2026-09-25: A record read from the wrong place fails here rather than
77    // being served.
78    if hdr.layer != key.layer || hdr.expert != key.expert {
79        bail!(
80            "expert record identity mismatch: header ({},{}) != requested {:?}",
81            hdr.layer,
82            hdr.expert,
83            key
84        );
85    }
86    let mut packed_addr = [0u64; 3];
87    let mut scale_addr = [0u64; 3];
88    for p in Proj::ALL {
89        packed_addr[p as usize] = base_dev_va + spec.packed_off(p);
90        scale_addr[p as usize] = base_dev_va + spec.scale_off(p);
91    }
92    Ok(ExpertResidency {
93        packed_addr,
94        scale_addr,
95        scale2: hdr.scale2,
96        input_scale: hdr.input_scale,
97    })
98}
99
100/// 2026-09-25: Read the record into a host buffer (buffered, through
101/// `ExpertFileReader`), copy it into the slot's part of one device buffer, and
102/// sync the stream.
103pub struct PosixTier {
104    reader: ExpertFileReader,
105    spec: ExpertRecordSpec,
106    layout: ExpertLayout,
107    /// 2026-09-25: One device buffer of `num_slabs * slots_per_slab` records.
108    dev: DeviceBuffer,
109    slots_per_slab: u32,
110    num_slabs: u32,
111}
112
113impl PosixTier {
114    pub fn open(dir: &Path, num_slabs: u32, slots_per_slab: u32) -> Result<Self> {
115        let reader = ExpertFileReader::open(dir)?;
116        let index: &ExpertIndex = reader.index();
117        let spec = index.spec();
118        let layout = index.layout();
119        if num_slabs == 0 || slots_per_slab == 0 {
120            bail!("PosixTier: zero geometry ({num_slabs},{slots_per_slab})");
121        }
122        let stride = layout.record_stride as usize;
123        // 2026-09-25: Checked, so an overflowed size cannot allocate less than
124        // the slots `slot_dev_va` addresses.
125        let total = (num_slabs as usize)
126            .checked_mul(slots_per_slab as usize)
127            .and_then(|v| v.checked_mul(stride))
128            .context("PosixTier: arena size overflow")?;
129        let dev = DeviceBuffer::new(total)?;
130        Ok(Self {
131            reader,
132            spec,
133            layout,
134            dev,
135            slots_per_slab,
136            num_slabs,
137        })
138    }
139
140    fn slot_dev_va(&self, slot: ArenaSlot) -> Result<u64> {
141        if slot.slab >= self.num_slabs || slot.slot >= self.slots_per_slab {
142            bail!("PosixTier: slot {:?} out of range", slot);
143        }
144        let i = (slot.slab as u64) * (self.slots_per_slab as u64) + (slot.slot as u64);
145        Ok(self.dev.ptr + i * self.layout.record_stride)
146    }
147}
148
149impl ExpertTier for PosixTier {
150    fn fetch(&mut self, key: ExpertKey, slot: ArenaSlot, stream: u64) -> Result<ExpertResidency> {
151        let record = self.reader.read_record_raw(key)?;
152        let dev_va = self.slot_dev_va(slot)?;
153        copy_h_to_d_async(dev_va, record.as_ptr() as *const _, record.len(), stream)?;
154        // 2026-09-25: The copy reads `record`, which is freed when this returns.
155        stream_sync(stream)?;
156        residency_from(&self.spec, &record, dev_va, key)
157    }
158    fn kind(&self) -> TierKind {
159        TierKind::Posix
160    }
161}
162
163/// 2026-09-25: Read the record straight into a pinned arena slot (O_DIRECT on
164/// Linux); the returned addresses point into the arena, and nothing is copied.
165pub struct UmaArenaTier {
166    // 2026-09-25: One file per MoE layer, O_DIRECT on Linux.
167    files: Vec<File>,
168    spec: ExpertRecordSpec,
169    layout: ExpertLayout,
170    arena: ExpertArena,
171}
172
173impl UmaArenaTier {
174    pub fn open(dir: &Path, num_slabs: u32, slots_per_slab: u32) -> Result<Self> {
175        let reader = ExpertFileReader::open(dir)?;
176        let index: &ExpertIndex = reader.index();
177        let spec = index.spec();
178        let layout = index.layout();
179        // 2026-09-25: The reader's files are buffered; these are opened again
180        // with `set_direct_flag`.
181        let mut files = Vec::with_capacity(index.num_moe_layers as usize);
182        for l in 0..index.num_moe_layers {
183            let p = dir.join(index.file_name(l));
184            let mut opts = OpenOptions::new();
185            opts.read(true);
186            set_direct_flag(&mut opts);
187            let f = opts
188                .open(&p)
189                .with_context(|| format!("open {}", p.display()))?;
190            files.push(f);
191        }
192        let arena = ExpertArena::new(num_slabs, slots_per_slab, layout.record_stride as usize)?;
193        Ok(Self {
194            files,
195            spec,
196            layout,
197            arena,
198        })
199    }
200
201    pub fn arena(&self) -> &ExpertArena {
202        &self.arena
203    }
204}
205
206impl ExpertTier for UmaArenaTier {
207    fn fetch(&mut self, key: ExpertKey, slot: ArenaSlot, _stream: u64) -> Result<ExpertResidency> {
208        let stride = self.layout.record_stride as usize;
209        let host = self.arena.slot_host_ptr(slot.slab, slot.slot)?;
210        let file = self
211            .files
212            .get(key.layer as usize)
213            .with_context(|| format!("UmaArenaTier: no file for layer {}", key.layer))?;
214        let off = self.layout.file_offset(key);
215        // 2026-09-25: The stride is a multiple of 4096 (`ExpertArena::new`), as
216        // O_DIRECT needs. SAFETY: `host` points at a slot of `stride` bytes
217        // inside the pinned arena, and the slice covers exactly that slot.
218        let dst = unsafe { std::slice::from_raw_parts_mut(host, stride) };
219        pio::read_exact_at(file, dst, off)
220            .with_context(|| format!("UmaArenaTier read {key:?} at {off}"))?;
221        // 2026-09-25: SAFETY: the slot holds the `stride` bytes just read.
222        let record = unsafe { std::slice::from_raw_parts(host, stride) };
223        let dev_va = self.arena.slot_dev_va(slot.slab, slot.slot)?;
224        residency_from(&self.spec, record, dev_va, key)
225    }
226    fn kind(&self) -> TierKind {
227        TierKind::Uma
228    }
229}
230
231/// 2026-09-25: Open the tier named by `backend`:
232///   * `posix`, `uma`: read the store in `dir`;
233///   * `rdma`: `RdmaTier` over TCP to `$METRALE_EXPERT_PEER`;
234///   * `rdma-verbs`: `RdmaTier` with one-sided RDMA READ from
235///     `$METRALE_EXPERT_PEER`.
236pub fn open_tier(
237    backend: &str,
238    dir: &Path,
239    num_slabs: u32,
240    slots_per_slab: u32,
241) -> Result<Box<dyn ExpertTier>> {
242    // 2026-09-25: The RDMA tiers exist on unix only; elsewhere asking for one
243    // is an error.
244    let rdma = |use_verbs: bool| -> Result<Box<dyn ExpertTier>> {
245        let flag = if use_verbs { "rdma-verbs" } else { "rdma" };
246        #[cfg(unix)]
247        {
248            let addr = std::env::var("METRALE_EXPERT_PEER").map_err(|_| {
249                anyhow::anyhow!("--expert-backend {flag} needs $METRALE_EXPERT_PEER=host:port")
250            })?;
251            Ok(Box::new(crate::expert_tier_rdma::RdmaTier::connect(
252                &addr,
253                num_slabs,
254                slots_per_slab,
255                use_verbs,
256            )?))
257        }
258        #[cfg(not(unix))]
259        {
260            bail!("--expert-backend {flag} requires rdma-core, which is unix-only")
261        }
262    };
263    match backend {
264        "posix" => Ok(Box::new(PosixTier::open(dir, num_slabs, slots_per_slab)?)),
265        "uma" => Ok(Box::new(UmaArenaTier::open(
266            dir,
267            num_slabs,
268            slots_per_slab,
269        )?)),
270        "rdma" => rdma(false),
271        "rdma-verbs" => rdma(true),
272        other => bail!("unknown expert backend '{other}' (want posix|uma|rdma|rdma-verbs)"),
273    }
274}
275
276/// 2026-09-25: Copy `len` bytes from a device address into a new Vec, then
277/// sync the stream.
278pub fn read_device(dev_va: u64, len: usize, stream: u64) -> Result<Vec<u8>> {
279    use crate::cuda_min::copy_d_to_h_async;
280    let mut out = vec![0u8; len];
281    copy_d_to_h_async(out.as_mut_ptr() as *mut _, dev_va, len, stream)?;
282    stream_sync(stream)?;
283    Ok(out)
284}
285
286/// 2026-09-25: O_DIRECT on Linux; other targets read buffered (the `layout`
287/// module header gives the reason for Windows).
288#[cfg(target_os = "linux")]
289fn set_direct_flag(opts: &mut OpenOptions) {
290    use std::os::unix::fs::OpenOptionsExt;
291    opts.custom_flags(libc::O_DIRECT);
292}
293
294#[cfg(not(target_os = "linux"))]
295fn set_direct_flag(_opts: &mut OpenOptions) {}
metrale_storage/expert_tier.rs

metrale_storage/
expert_tier.rs
