spark_storage/
expert_tier.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// ExpertTier — the residency abstraction the expert streamer fetches through.
4//
5// Sits ABOVE the record layer (`ExpertFileReader` / `ExpertIndex`), not as a
6// `StorageBackend` impl: that trait is GroupKey/KV-tile shaped and synchronizes
7// the stream on return, the wrong shape for pull-on-demand MoE experts. A tier
8// lands one expert record into a slot and returns the six device addresses the
9// fused kernels read (already resident / prefill-transposed layout — nothing is
10// transformed at fetch time, invariant D).
11//
12// Three tiers, one interface (residency order device < UMA-over-NVMe < RDMA):
13//   * `PosixTier`    — deterministic bounce oracle (pread -> copy_h2d into a
14//                      device buffer). The bit-identical acceptance reference.
15//   * `UmaArenaTier` — the zero-copy path: O_DIRECT NVMe fill straight into the
16//                      pinned arena; the ptr table points at the pinned VA, no
17//                      HtoD copy.
18//   * RdmaTier       — Stage 4: one-sided RDMA_READ into the SAME pinned arena.
19//
20// All three feed the identical ptr-table patch, so swapping tiers cannot change
21// a single byte the GEMM reads — which the Tier-1 parity test proves.
22
23use anyhow::{Context, Result, bail};
24use atlas_tier::pio;
25use std::fs::{File, OpenOptions};
26use std::path::Path;
27
28use crate::cuda_min::{DeviceBuffer, copy_h_to_d_async, stream_sync};
29use crate::expert::{ExpertKey, ExpertLayout, ExpertRecordHeader, ExpertRecordSpec, Proj};
30use crate::expert_arena::ExpertArena;
31use crate::expert_pack::{ExpertFileReader, ExpertIndex};
32
33/// The six sub-buffer device addresses (+ scalars) of one resident expert —
34/// exactly what the ptr-table patcher writes into the shadow tables.
35#[derive(Clone, Copy, Debug, PartialEq)]
36pub struct ExpertResidency {
37    /// gate/up/down B_packed device VA.
38    pub packed_addr: [u64; 3],
39    /// gate/up/down B_scale device VA.
40    pub scale_addr: [u64; 3],
41    /// gate/up/down per-tensor weight_scale_2.
42    pub scale2: [f32; 3],
43    /// gate/up/down input_scale (None = weight-only W4A16).
44    pub input_scale: [Option<f32>; 3],
45}
46
47#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48pub enum TierKind {
49    Posix,
50    Uma,
51    Rdma,
52}
53
54/// A destination slot in the residency ring: which slab, which slot within it.
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56pub struct ArenaSlot {
57    pub slab: u32,
58    pub slot: u32,
59}
60
61impl ArenaSlot {
62    pub fn new(slab: u32, slot: u32) -> Self {
63        Self { slab, slot }
64    }
65}
66
67/// Fetch an expert record into a slot; return the addresses to patch.
68pub trait ExpertTier: Send {
69    fn fetch(&mut self, key: ExpertKey, slot: ArenaSlot, stream: u64) -> Result<ExpertResidency>;
70    fn kind(&self) -> TierKind;
71    /// Graceful-degradation probe (RDMA link health at Stage 4). Default: up.
72    fn healthy(&self) -> bool {
73        true
74    }
75}
76
77/// Turn a record buffer's header + a base device VA into an `ExpertResidency`
78/// using the record spec's sub-offsets. Shared by every tier so they cannot
79/// disagree on layout.
80pub(crate) fn residency_from(
81    spec: &ExpertRecordSpec,
82    record: &[u8],
83    base_dev_va: u64,
84    key: ExpertKey,
85) -> Result<ExpertResidency> {
86    let hdr = ExpertRecordHeader::from_bytes(record)
87        .context("expert record header magic/version mismatch")?;
88    // Invariant D at fetch time: the header carries identity precisely so a
89    // misplaced/corrupt record is caught here, not served silently.
90    if hdr.layer != key.layer || hdr.expert != key.expert {
91        bail!(
92            "expert record identity mismatch: header ({},{}) != requested {:?}",
93            hdr.layer,
94            hdr.expert,
95            key
96        );
97    }
98    let mut packed_addr = [0u64; 3];
99    let mut scale_addr = [0u64; 3];
100    for p in Proj::ALL {
101        packed_addr[p as usize] = base_dev_va + spec.packed_off(p);
102        scale_addr[p as usize] = base_dev_va + spec.scale_off(p);
103    }
104    Ok(ExpertResidency {
105        packed_addr,
106        scale_addr,
107        scale2: hdr.scale2,
108        input_scale: hdr.input_scale,
109    })
110}
111
112/// Deterministic bounce oracle: pread the record (no O_DIRECT) into a host
113/// buffer, `copy_h2d` into a per-slot device buffer, stream-synced. This is the
114/// reference every other tier must match byte-for-byte.
115pub struct PosixTier {
116    reader: ExpertFileReader,
117    spec: ExpertRecordSpec,
118    layout: ExpertLayout,
119    /// One contiguous device buffer holding num_slabs*slots_per_slab records.
120    dev: DeviceBuffer,
121    slots_per_slab: u32,
122    num_slabs: u32,
123}
124
125impl PosixTier {
126    pub fn open(dir: &Path, num_slabs: u32, slots_per_slab: u32) -> Result<Self> {
127        let reader = ExpertFileReader::open(dir)?;
128        let index: &ExpertIndex = reader.index();
129        let spec = index.spec();
130        let layout = index.layout();
131        if num_slabs == 0 || slots_per_slab == 0 {
132            bail!("PosixTier: zero geometry ({num_slabs},{slots_per_slab})");
133        }
134        let stride = layout.record_stride as usize;
135        // Checked like the sibling ExpertArena ctor — a wrapped product must
136        // never yield a small alloc that slot_dev_va then addresses past.
137        let total = (num_slabs as usize)
138            .checked_mul(slots_per_slab as usize)
139            .and_then(|v| v.checked_mul(stride))
140            .context("PosixTier: arena size overflow")?;
141        let dev = DeviceBuffer::new(total)?;
142        Ok(Self {
143            reader,
144            spec,
145            layout,
146            dev,
147            slots_per_slab,
148            num_slabs,
149        })
150    }
151
152    fn slot_dev_va(&self, slot: ArenaSlot) -> Result<u64> {
153        if slot.slab >= self.num_slabs || slot.slot >= self.slots_per_slab {
154            bail!("PosixTier: slot {:?} out of range", slot);
155        }
156        let i = (slot.slab as u64) * (self.slots_per_slab as u64) + (slot.slot as u64);
157        Ok(self.dev.ptr + i * self.layout.record_stride)
158    }
159}
160
161impl ExpertTier for PosixTier {
162    fn fetch(&mut self, key: ExpertKey, slot: ArenaSlot, stream: u64) -> Result<ExpertResidency> {
163        let record = self.reader.read_record_raw(key)?; // host bytes
164        let dev_va = self.slot_dev_va(slot)?;
165        copy_h_to_d_async(dev_va, record.as_ptr() as *const _, record.len(), stream)?;
166        stream_sync(stream)?; // single bounce would be overwritten otherwise
167        residency_from(&self.spec, &record, dev_va, key)
168    }
169    fn kind(&self) -> TierKind {
170        TierKind::Posix
171    }
172}
173
174/// The zero-copy path: O_DIRECT the record straight into the pinned arena slot;
175/// the returned addresses point INTO the arena (GPU-addressable at the same
176/// VA), no `copy_h2d`.
177pub struct UmaArenaTier {
178    files: Vec<File>, // one O_DIRECT fd per MoE layer
179    spec: ExpertRecordSpec,
180    layout: ExpertLayout,
181    arena: ExpertArena,
182}
183
184impl UmaArenaTier {
185    pub fn open(dir: &Path, num_slabs: u32, slots_per_slab: u32) -> Result<Self> {
186        let reader = ExpertFileReader::open(dir)?;
187        let index: &ExpertIndex = reader.index();
188        let spec = index.spec();
189        let layout = index.layout();
190        // Re-open each layer file with O_DIRECT for aligned zero-copy reads.
191        let mut files = Vec::with_capacity(index.num_moe_layers as usize);
192        for l in 0..index.num_moe_layers {
193            let p = dir.join(index.file_name(l));
194            let mut opts = OpenOptions::new();
195            opts.read(true);
196            set_direct_flag(&mut opts);
197            let f = opts
198                .open(&p)
199                .with_context(|| format!("open {}", p.display()))?;
200            files.push(f);
201        }
202        let arena = ExpertArena::new(num_slabs, slots_per_slab, layout.record_stride as usize)?;
203        Ok(Self {
204            files,
205            spec,
206            layout,
207            arena,
208        })
209    }
210
211    pub fn arena(&self) -> &ExpertArena {
212        &self.arena
213    }
214}
215
216impl ExpertTier for UmaArenaTier {
217    fn fetch(&mut self, key: ExpertKey, slot: ArenaSlot, _stream: u64) -> Result<ExpertResidency> {
218        let stride = self.layout.record_stride as usize;
219        let host = self.arena.slot_host_ptr(slot.slab, slot.slot)?;
220        let file = self
221            .files
222            .get(key.layer as usize)
223            .with_context(|| format!("UmaArenaTier: no file for layer {}", key.layer))?;
224        let off = self.layout.file_offset(key);
225        // Positional read of the whole (4 KiB-aligned) record straight into the
226        // pinned, GPU-addressable slot — this is the "delete the bounce copy"
227        // step. O_DIRECT on Linux makes it zero-copy; elsewhere it is buffered.
228        // SAFETY: `host` points at a slot of `stride` bytes inside the pinned
229        // arena, and the slice covers exactly that slot.
230        let dst = unsafe { std::slice::from_raw_parts_mut(host, stride) };
231        pio::read_exact_at(file, dst, off)
232            .with_context(|| format!("UmaArenaTier read {key:?} at {off}"))?;
233        // SAFETY: the slot holds `stride` valid bytes just read from disk.
234        let record = unsafe { std::slice::from_raw_parts(host, stride) };
235        let dev_va = self.arena.slot_dev_va(slot.slab, slot.slot)?;
236        residency_from(&self.spec, record, dev_va, key)
237    }
238    fn kind(&self) -> TierKind {
239        TierKind::Uma
240    }
241}
242
243/// Open the tier named by `backend` over a built store:
244///   * `posix` / `uma` — read `dir` locally (bounce oracle / zero-copy).
245///   * `rdma`          — connect to `$ATLAS_EXPERT_PEER` over TWO-SIDED TCP.
246///   * `rdma-verbs`    — connect to `$ATLAS_EXPERT_PEER` over ONE-SIDED RDMA READ
247///     (verbs); device/GID from `$ATLAS_EXPERT_RDMA_DEV`/`$ATLAS_EXPERT_RDMA_GID`.
248///
249/// Both peer backends serve the store's records over the RoCE fabric.
250pub fn open_tier(
251    backend: &str,
252    dir: &Path,
253    num_slabs: u32,
254    slots_per_slab: u32,
255) -> Result<Box<dyn ExpertTier>> {
256    // The RDMA expert tiers need rdma-core, so they exist on unix only. The
257    // local `posix`/`uma` tiers below are portable; asking for an RDMA backend
258    // elsewhere fails with a clear message rather than being silently absent
259    // from the match.
260    let rdma = |use_verbs: bool| -> Result<Box<dyn ExpertTier>> {
261        let flag = if use_verbs { "rdma-verbs" } else { "rdma" };
262        #[cfg(unix)]
263        {
264            let addr = std::env::var("ATLAS_EXPERT_PEER").map_err(|_| {
265                anyhow::anyhow!("--expert-backend {flag} needs $ATLAS_EXPERT_PEER=host:port")
266            })?;
267            Ok(Box::new(crate::expert_tier_rdma::RdmaTier::connect(
268                &addr,
269                num_slabs,
270                slots_per_slab,
271                use_verbs,
272            )?))
273        }
274        #[cfg(not(unix))]
275        {
276            bail!("--expert-backend {flag} requires rdma-core, which is unix-only")
277        }
278    };
279    match backend {
280        "posix" => Ok(Box::new(PosixTier::open(dir, num_slabs, slots_per_slab)?)),
281        "uma" => Ok(Box::new(UmaArenaTier::open(
282            dir,
283            num_slabs,
284            slots_per_slab,
285        )?)),
286        "rdma" => rdma(false),
287        "rdma-verbs" => rdma(true),
288        other => bail!("unknown expert backend '{other}' (want posix|uma|rdma|rdma-verbs)"),
289    }
290}
291
292/// Copy `len` bytes from a device VA to a fresh host `Vec` (test/verify helper).
293pub fn read_device(dev_va: u64, len: usize, stream: u64) -> Result<Vec<u8>> {
294    use crate::cuda_min::copy_d_to_h_async;
295    let mut out = vec![0u8; len];
296    copy_d_to_h_async(out.as_mut_ptr() as *mut _, dev_va, len, stream)?;
297    stream_sync(stream)?;
298    Ok(out)
299}
300
301/// O_DIRECT on Linux (the zero-copy arena read depends on it); no equivalent
302/// flag elsewhere — see the `layout` module header for why Windows stays
303/// buffered rather than using FILE_FLAG_NO_BUFFERING.
304#[cfg(target_os = "linux")]
305fn set_direct_flag(opts: &mut OpenOptions) {
306    use std::os::unix::fs::OpenOptionsExt;
307    opts.custom_flags(libc::O_DIRECT);
308}
309
310#[cfg(not(target_os = "linux"))]
311fn set_direct_flag(_opts: &mut OpenOptions) {}
spark_storage/expert_tier.rs

spark_storage/
expert_tier.rs
