spark_storage/
expert_arena.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// The UMA zero-copy expert arena — the one genuinely-new primitive of the
4// streaming-experts feature.
5//
6// A pinned LPDDR buffer (`cuMemAllocHost`) organized as a ring of `num_slabs`
7// slabs, each holding `slots_per_slab` fixed-stride expert records. On GB10 the
8// pinned allocation is GPU-addressable at the *same* virtual address (Gate
9// 0(b)), so a record read straight into a slot — by O_DIRECT NVMe today, by
10// one-sided RDMA_READ later — is immediately consumable by the fused MoE
11// kernels with no `cuMemcpyHtoD` bounce: the expert pointer table is simply
12// patched to point at the slot's device VA.
13//
14// This module owns the memory + geometry only. Who *fills* a slot (NVMe vs
15// RDMA) is an `ExpertTier` concern; the arena is registrable as an RDMA MR
16// (Stage 4) exactly because it is ordinary page-locked host memory.
17
18use anyhow::{Context, Result, bail};
19use std::ffi::c_void;
20
21use crate::cuda_min::PinnedBuffer;
22
23/// A ring of pinned per-layer slabs. Slot geometry is one expert record
24/// (`record_stride`, a 4 KiB multiple so O_DIRECT / RDMA landings are aligned).
25pub struct ExpertArena {
26    pinned: PinnedBuffer,
27    /// Device VA of the arena base. On GB10 this equals the host VA.
28    dev_base: u64,
29    num_slabs: u32,
30    slots_per_slab: u32,
31    record_stride: usize,
32}
33
34impl ExpertArena {
35    /// Allocate `num_slabs * slots_per_slab * record_stride` bytes of pinned
36    /// LPDDR and confirm the GB10 same-VA property (fail loudly otherwise — the
37    /// zero-copy patch would be silently wrong on a host without it).
38    pub fn new(num_slabs: u32, slots_per_slab: u32, record_stride: usize) -> Result<Self> {
39        if num_slabs == 0 || slots_per_slab == 0 || record_stride == 0 {
40            bail!("ExpertArena: zero geometry ({num_slabs},{slots_per_slab},{record_stride})");
41        }
42        if !record_stride.is_multiple_of(4096) {
43            bail!("ExpertArena: record_stride {record_stride} must be a 4 KiB multiple (O_DIRECT)");
44        }
45        let total = (num_slabs as usize)
46            .checked_mul(slots_per_slab as usize)
47            .and_then(|v| v.checked_mul(record_stride))
48            .context("ExpertArena: size overflow")?;
49        let pinned = PinnedBuffer::new(total)?;
50        let dev_base = pinned.device_ptr()?;
51        let host_base = pinned.ptr as u64;
52        if dev_base != host_base {
53            // Not fatal to correctness on a mapped device, but it means the
54            // ptr-table patch must use dev_base (not the host VA) and the
55            // zero-copy assumption behind our bandwidth model is weaker. On
56            // GB10 they are equal; assert so a non-UMA host is caught early.
57            bail!(
58                "ExpertArena: pinned host VA {host_base:#x} != device VA {dev_base:#x} \
59                 — host is not unified-addressing (UMA zero-copy unavailable)"
60            );
61        }
62        Ok(Self {
63            pinned,
64            dev_base,
65            num_slabs,
66            slots_per_slab,
67            record_stride,
68        })
69    }
70
71    pub fn num_slabs(&self) -> u32 {
72        self.num_slabs
73    }
74    pub fn slots_per_slab(&self) -> u32 {
75        self.slots_per_slab
76    }
77    pub fn record_stride(&self) -> usize {
78        self.record_stride
79    }
80
81    fn linear_slot(&self, slab: u32, slot: u32) -> Result<usize> {
82        if slab >= self.num_slabs || slot >= self.slots_per_slab {
83            bail!(
84                "ExpertArena: slot ({slab},{slot}) out of range ({},{})",
85                self.num_slabs,
86                self.slots_per_slab
87            );
88        }
89        Ok((slab as usize) * (self.slots_per_slab as usize) + (slot as usize))
90    }
91
92    /// Host pointer of a slot — the O_DIRECT / RDMA landing target.
93    pub fn slot_host_ptr(&self, slab: u32, slot: u32) -> Result<*mut u8> {
94        let i = self.linear_slot(slab, slot)?;
95        // SAFETY: i < num_slabs*slots_per_slab, so the offset is within the
96        // single pinned allocation.
97        Ok(unsafe { (self.pinned.ptr as *mut u8).add(i * self.record_stride) })
98    }
99
100    /// Device VA of a slot — what the expert pointer table is patched to.
101    pub fn slot_dev_va(&self, slab: u32, slot: u32) -> Result<u64> {
102        let i = self.linear_slot(slab, slot)?;
103        Ok(self.dev_base + (i as u64) * (self.record_stride as u64))
104    }
105
106    /// Raw pinned base as a `*mut c_void` (for future `ibv_reg_mr`).
107    pub fn base_ptr(&self) -> *mut c_void {
108        self.pinned.ptr
109    }
110    pub fn total_bytes(&self) -> usize {
111        self.pinned.bytes
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    // Pure-geometry checks that don't need a GPU are impossible here (the ctor
120    // allocates pinned CUDA memory), so these are GPU-gated.
121    #[test]
122    #[ignore = "requires GPU"]
123    fn arena_slots_are_strided_and_same_va() {
124        let _ctx = crate::cuda_min::CudaCtx::new(0).unwrap();
125        let stride = 8192; // 2x4KiB
126        let arena = ExpertArena::new(2, 3, stride).unwrap();
127        assert_eq!(arena.total_bytes(), 2 * 3 * stride);
128        // Slot device VAs are contiguous and stride-spaced.
129        let a = arena.slot_dev_va(0, 0).unwrap();
130        let b = arena.slot_dev_va(0, 1).unwrap();
131        let c = arena.slot_dev_va(1, 0).unwrap();
132        assert_eq!(b - a, stride as u64);
133        assert_eq!(c - a, 3 * stride as u64);
134        // Same-VA property already asserted in the ctor; re-confirm host==dev.
135        assert_eq!(
136            arena.slot_host_ptr(1, 2).unwrap() as u64,
137            arena.slot_dev_va(1, 2).unwrap()
138        );
139        // Out-of-range is an error, not a panic.
140        assert!(arena.slot_dev_va(2, 0).is_err());
141    }
142}
spark_storage/expert_arena.rs

spark_storage/
expert_arena.rs
