spark_runtime/weights/gguf/
expert_lru.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! The device-visible expert cache: fixed-size slots in one page-locked host
5//! arena, addressed by the GPU through its `cuMemHostGetDevicePointer` alias
6//! (on GB10 the two addresses name the same unified memory), filled by `pread`
7//! and evicted least-recently-used.
8//!
9//! A slot holds one routed expert in [`SlotLayout`](super::expert_stream::SlotLayout)
10//! order (`gate | up | down`, 12.22 MiB on the Q2_K checkpoint). A token
11//! touches 6 x 40 = 240 of them; [`ExpertLru::begin_token`] opens an epoch and
12//! every slot fetched inside it is pinned against eviction until the next
13//! epoch, so one token's working set can never evict itself. Misses inside
14//! [`ExpertLru::fetch_many`] are read on a scoped thread pool (positional reads
15//! share no file offset), which is what turns 240 x 12 MiB of NVMe traffic into
16//! one parallel gather instead of a serial walk.
17//!
18//! The arena is NOT owned here: the caller allocates it (see [`PinnedArena`])
19//! and keeps it alive for the cache's lifetime, and the unit tests back it with
20//! a plain `Vec<u8>`. The cache never reads or writes the arena itself except
21//! through [`ExpertSource::read_expert`](super::expert_stream::ExpertSource).
22
23use std::collections::HashMap;
24use std::sync::Arc;
25
26use anyhow::{Result, bail, ensure};
27
28use super::expert_prefetch::{ReaderPool, Ticket};
29use super::expert_stream::{ExpertSource, SlotLayout};
30use crate::gpu::{DevicePtr, GpuBackend};
31
32/// One page-locked host arena with its device alias.
33pub struct PinnedArena {
34    host: *mut u8,
35    dev: DevicePtr,
36    bytes: usize,
37}
38
39// SAFETY: a page-locked region with no thread affinity; the owner serialises access.
40unsafe impl Send for PinnedArena {}
41unsafe impl Sync for PinnedArena {}
42
43impl PinnedArena {
44    /// `alloc_host_pinned` zero-fills, so a large arena costs one memset at
45    /// load. The device alias is what the kernels are handed.
46    pub fn alloc(gpu: &dyn GpuBackend, bytes: usize) -> Result<Self> {
47        let host = gpu.alloc_host_pinned(bytes)?;
48        let dev = match gpu.host_ptr_to_device(host) {
49            Ok(d) => d,
50            Err(e) => {
51                let _ = gpu.free_host_pinned(host, bytes);
52                return Err(e);
53            }
54        };
55        Ok(PinnedArena { host, dev, bytes })
56    }
57
58    pub fn host(&self) -> *mut u8 {
59        self.host
60    }
61
62    pub fn dev(&self) -> DevicePtr {
63        self.dev
64    }
65
66    pub fn bytes(&self) -> usize {
67        self.bytes
68    }
69
70    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
71        gpu.free_host_pinned(self.host, self.bytes)
72    }
73}
74
75/// Where one cached expert's three projections are, as the GPU sees them.
76#[derive(Clone, Copy, Debug, PartialEq, Eq)]
77pub struct ExpertSlot {
78    pub gate: DevicePtr,
79    pub up: DevicePtr,
80    pub down: DevicePtr,
81    /// Host address of the slot (for oracles and CPU-side reads).
82    pub host: *const u8,
83}
84
85#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
86pub struct LruStats {
87    pub hits: u64,
88    pub misses: u64,
89    pub evictions: u64,
90    pub bytes_read: u64,
91    /// Hits on a slot whose read was still in flight (a prediction that
92    /// arrived in time to be waited on rather than missed).
93    pub waited: u64,
94    /// Experts read on a prediction, and of those, never asked for.
95    pub prefetched: u64,
96    pub prefetch_unused: u64,
97}
98
99pub(super) const NONE: u32 = u32::MAX;
100
101pub(super) struct SlotMeta {
102    pub(super) key: Option<(u32, u32)>,
103    pub(super) prev: u32,
104    pub(super) next: u32,
105    /// Epoch of the last fetch; equal to the current epoch = pinned.
106    pub(super) epoch: u64,
107    /// The read filling this slot, while it is in flight (pool only).
108    pub(super) ticket: Option<Arc<Ticket>>,
109    /// Filled by a prediction and not yet asked for by a real fetch.
110    pub(super) speculative: bool,
111}
112
113/// A raw slot address that may cross into a scoped reader thread: every miss
114/// in one `fetch_many` owns a distinct slot, so the regions never overlap.
115#[derive(Clone, Copy)]
116pub(super) struct SlotPtr(pub(super) *mut u8);
117
118impl SlotPtr {
119    /// The address, read through `self` so a closure captures the `Send`
120    /// wrapper and not the raw pointer.
121    pub(super) fn get(self) -> *mut u8 {
122        self.0
123    }
124}
125// SAFETY: see above; the pointer is only dereferenced as a disjoint `&mut [u8]`
126// by the one thread whose chunk names it (or, split into byte ranges, by
127// several threads on disjoint ranges of it).
128unsafe impl Send for SlotPtr {}
129unsafe impl Sync for SlotPtr {}
130
131/// One byte range of a missed expert's slot, as handed to a reader thread by
132/// `fetch_many`: `(slot index, (layer, expert) key, slot host pointer, byte
133/// offset of the range within the slot, byte length of the range)`.
134type MissRange = (u32, (u32, u32), SlotPtr, usize, usize);
135
136pub struct ExpertLru {
137    host: *mut u8,
138    pub(super) dev: u64,
139    layout: SlotLayout,
140    pub(super) n_slots: usize,
141    pub(super) meta: Vec<SlotMeta>,
142    pub(super) map: HashMap<(u32, u32), u32>,
143    /// Most recently used.
144    pub(super) head: u32,
145    /// Least recently used.
146    pub(super) tail: u32,
147    pub(super) epoch: u64,
148    pub(super) stats: LruStats,
149    /// The persistent reader pool and the source it reads from (see
150    /// `set_pool` in `expert_lru_pool.rs`); without it, misses are read on
151    /// scoped threads inside `fetch_many`.
152    pub(super) pool: Option<(Arc<dyn ExpertSource + Send + Sync>, ReaderPool)>,
153    /// Slots with a ticket that may still be in flight.
154    pub(super) in_flight: Vec<u32>,
155    /// The victim is drawn from the oldest this-many percent (0 = strict LRU).
156    pub(super) evict_pct: usize,
157    pub(super) rng: u64,
158    /// `ATLAS_DS41_ROUTE_TRACE`: one line per `fetch_many` (see `set_trace`).
159    pub(super) trace: Option<std::io::BufWriter<std::fs::File>>,
160    pub(super) t0: std::time::Instant,
161    /// Set when the slots are device memory fed through a page-locked ring
162    /// (`expert_arena.rs`); `fetch_many_on` is then the only fetch.
163    pub(super) staging: Option<super::expert_arena::Staging>,
164    /// `(layer, expert, slot | -1)` since the last `drain_slot_changes`: what
165    /// a device-side slot table has to learn (assignments and evictions).
166    pub(super) slot_changes: Vec<(u32, u32, i32)>,
167}
168
169// SAFETY: the raw arena pointers are addresses into memory the caller owns
170// and keeps alive; the cache holds no thread-affine state.
171unsafe impl Send for ExpertLru {}
172
173impl ExpertLru {
174    /// `bytes / layout.bytes` slots over the arena at `host` (device alias
175    /// `dev`). Every slot starts empty and at the LRU end.
176    pub fn new(host: *mut u8, dev: DevicePtr, bytes: usize, layout: SlotLayout) -> Result<Self> {
177        ensure!(layout.bytes > 0, "expert slot layout has zero bytes");
178        let n_slots = bytes / layout.bytes;
179        ensure!(
180            n_slots >= 1,
181            "arena of {bytes} bytes holds no {}-byte expert slot",
182            layout.bytes
183        );
184        ensure!(n_slots < NONE as usize, "too many slots");
185        let mut meta = Vec::with_capacity(n_slots);
186        for i in 0..n_slots as u32 {
187            meta.push(SlotMeta {
188                key: None,
189                prev: if i == 0 { NONE } else { i - 1 },
190                next: if i as usize + 1 == n_slots {
191                    NONE
192                } else {
193                    i + 1
194                },
195                epoch: 0,
196                ticket: None,
197                speculative: false,
198            });
199        }
200        Ok(ExpertLru {
201            host,
202            dev: dev.0,
203            layout,
204            n_slots,
205            meta,
206            map: HashMap::new(),
207            head: 0,
208            tail: n_slots as u32 - 1,
209            epoch: 1,
210            stats: LruStats::default(),
211            pool: None,
212            in_flight: Vec::new(),
213            evict_pct: std::env::var("ATLAS_DS41_EVICT_RANDOM_PCT")
214                .ok()
215                .and_then(|v| v.parse().ok())
216                .filter(|p| *p <= 100)
217                .unwrap_or(8),
218            rng: 0x9E37_79B9_7F4A_7C15,
219            trace: None,
220            t0: std::time::Instant::now(),
221            staging: None,
222            slot_changes: Vec::new(),
223        })
224    }
225
226    pub fn n_slots(&self) -> usize {
227        self.n_slots
228    }
229
230    pub fn layout(&self) -> SlotLayout {
231        self.layout
232    }
233
234    /// Experts currently cached.
235    pub fn resident(&self) -> usize {
236        self.map.len()
237    }
238
239    pub fn contains(&self, layer: u32, expert: u32) -> bool {
240        self.map.contains_key(&(layer, expert))
241    }
242
243    pub fn stats(&self) -> LruStats {
244        self.stats
245    }
246
247    pub fn reset_stats(&mut self) {
248        self.stats = LruStats::default();
249    }
250
251    /// Open a new epoch: slots fetched from now on are pinned until the next
252    /// `begin_token`; everything fetched earlier becomes evictable again.
253    pub fn begin_token(&mut self) {
254        self.epoch += 1;
255        if self.pool.is_some() {
256            self.reap();
257        }
258    }
259
260    pub fn slot(&self, i: u32) -> ExpertSlot {
261        let base = i as usize * self.layout.bytes;
262        let l = self.layout;
263        ExpertSlot {
264            gate: DevicePtr(self.dev + (base + l.gate_off) as u64),
265            up: DevicePtr(self.dev + (base + l.up_off) as u64),
266            down: DevicePtr(self.dev + (base + l.down_off) as u64),
267            // SAFETY: `base + bytes <= n_slots * bytes <= arena bytes`; a
268            // device arena has no host alias (null stays null).
269            host: if self.host.is_null() {
270                std::ptr::null()
271            } else {
272                unsafe { self.host.add(base) as *const u8 }
273            },
274        }
275    }
276
277    pub(super) fn slot_ptr(&self, i: u32) -> SlotPtr {
278        // SAFETY: as in `slot`.
279        SlotPtr(unsafe { self.host.add(i as usize * self.layout.bytes) })
280    }
281
282    pub(super) fn unlink(&mut self, i: u32) {
283        let (p, n) = (self.meta[i as usize].prev, self.meta[i as usize].next);
284        if p == NONE {
285            self.head = n;
286        } else {
287            self.meta[p as usize].next = n;
288        }
289        if n == NONE {
290            self.tail = p;
291        } else {
292            self.meta[n as usize].prev = p;
293        }
294    }
295
296    fn push_front(&mut self, i: u32) {
297        let m = &mut self.meta[i as usize];
298        m.prev = NONE;
299        m.next = self.head;
300        if self.head != NONE {
301            self.meta[self.head as usize].prev = i;
302        }
303        self.head = i;
304        if self.tail == NONE {
305            self.tail = i;
306        }
307    }
308
309    pub(super) fn touch(&mut self, i: u32) {
310        self.unlink(i);
311        self.push_front(i);
312        self.meta[i as usize].epoch = self.epoch;
313    }
314
315    /// Assign a slot for `key` and map it (the bytes are not read yet).
316    pub(super) fn assign(&mut self, key: (u32, u32)) -> Result<u32> {
317        let i = self.take_victim()?;
318        self.meta[i as usize].key = Some(key);
319        self.map.insert(key, i);
320        self.slot_changes.push((key.0, key.1, i as i32));
321        self.touch(i);
322        Ok(i)
323    }
324
325    pub(super) fn unmap(&mut self, i: u32) {
326        if let Some(k) = self.meta[i as usize].key.take() {
327            self.map.remove(&k);
328            self.slot_changes.push((k.0, k.1, -1));
329        }
330        self.meta[i as usize].epoch = 0;
331        self.meta[i as usize].ticket = None;
332        self.meta[i as usize].speculative = false;
333        self.unlink(i);
334        // back to the LRU end so it is the next victim
335        let m = &mut self.meta[i as usize];
336        m.next = NONE;
337        m.prev = self.tail;
338        if self.tail != NONE {
339            self.meta[self.tail as usize].next = i;
340        }
341        self.tail = i;
342        if self.head == NONE {
343            self.head = i;
344        }
345    }
346
347    /// One expert. Returns the slot and whether it was already resident.
348    pub fn fetch<S: ExpertSource + ?Sized>(
349        &mut self,
350        src: &S,
351        layer: u32,
352        expert: u32,
353    ) -> Result<(ExpertSlot, bool)> {
354        ensure!(
355            self.staging.is_none(),
356            "device expert cache: use fetch_many_on"
357        );
358        let key = (layer, expert);
359        if let Some(&i) = self.map.get(&key) {
360            self.touch(i);
361            self.stats.hits += 1;
362            return Ok((self.slot(i), true));
363        }
364        let i = self.assign(key)?;
365        let p = self.slot_ptr(i);
366        // SAFETY: slot `i` is exactly `layout.bytes` long and mapped to no other key.
367        let dst = unsafe { std::slice::from_raw_parts_mut(p.0, self.layout.bytes) };
368        if let Err(e) = src.read_expert(layer, expert, dst) {
369            self.unmap(i);
370            return Err(e);
371        }
372        self.stats.misses += 1;
373        self.stats.bytes_read += self.layout.bytes as u64;
374        Ok((self.slot(i), false))
375    }
376
377    /// One token's gather: every key resolved to a slot, misses read on up to
378    /// `threads` scoped threads. Duplicate keys are one read. Slots are pinned
379    /// for the current epoch, so call [`Self::begin_token`] first.
380    pub fn fetch_many<S: ExpertSource + ?Sized>(
381        &mut self,
382        src: &S,
383        keys: &[(u32, u32)],
384        threads: usize,
385    ) -> Result<Vec<ExpertSlot>> {
386        ensure!(
387            self.staging.is_none(),
388            "device expert cache: use fetch_many_on"
389        );
390        if self.pool.is_some() {
391            return self.fetch_many_prefetching(keys, &[]);
392        }
393        let mut out = Vec::with_capacity(keys.len());
394        let mut misses: Vec<(u32, (u32, u32))> = Vec::new();
395        let evictions0 = self.stats.evictions;
396        for &key in keys {
397            if let Some(&i) = self.map.get(&key) {
398                self.touch(i);
399                self.stats.hits += 1;
400                out.push(i);
401                continue;
402            }
403            let i = self.assign(key)?;
404            misses.push((i, key));
405            out.push(i);
406        }
407        let bytes = self.layout.bytes;
408        let mut failures: Vec<(u32, anyhow::Error)> = Vec::new();
409        if threads <= 1 {
410            for &(i, (l, e)) in &misses {
411                let p = self.slot_ptr(i);
412                // SAFETY: as in `fetch`.
413                let dst = unsafe { std::slice::from_raw_parts_mut(p.0, bytes) };
414                if let Err(err) = src.read_expert(l, e, dst) {
415                    failures.push((i, err));
416                }
417            }
418        } else if !misses.is_empty() {
419            // A layer's decode step misses one or two experts: one thread
420            // per expert leaves the other readers idle, so every miss is
421            // cut into `parts` byte ranges and the ranges are spread over
422            // the threads (>= 1 MiB a range, up to `threads` per expert).
423            let parts = (threads / misses.len())
424                .clamp(1, threads)
425                .min((bytes / (1 << 20)).max(1));
426            let per_part = bytes.div_ceil(parts);
427            let work: Vec<MissRange> = misses
428                .iter()
429                .flat_map(|&(i, k)| {
430                    let p = self.slot_ptr(i);
431                    (0..parts).map(move |j| {
432                        let off = j * per_part;
433                        (i, k, p, off, per_part.min(bytes - off))
434                    })
435                })
436                .filter(|w| w.4 > 0)
437                .collect();
438            let per = work.len().div_ceil(threads.min(work.len()));
439            let results: Vec<Vec<(u32, anyhow::Error)>> = std::thread::scope(|s| {
440                let handles: Vec<_> = work
441                    .chunks(per)
442                    .map(|chunk| {
443                        s.spawn(move || {
444                            let mut errs = Vec::new();
445                            for (i, (l, e), p, off, len) in chunk {
446                                // SAFETY: each miss owns a distinct slot and
447                                // the ranges of one slot are disjoint.
448                                let dst =
449                                    unsafe { std::slice::from_raw_parts_mut(p.0.add(*off), *len) };
450                                if let Err(err) = src.read_expert_range(*l, *e, *off, dst) {
451                                    errs.push((*i, err));
452                                }
453                            }
454                            errs
455                        })
456                    })
457                    .collect();
458                handles
459                    .into_iter()
460                    .map(|h| h.join().expect("expert reader thread panicked"))
461                    .collect()
462            });
463            // one failure per slot: a slot whose ranges failed twice is
464            // unmapped once
465            let mut seen = std::collections::HashSet::new();
466            failures = results
467                .into_iter()
468                .flatten()
469                .filter(|(i, _)| seen.insert(*i))
470                .collect();
471        }
472        let n_ok = misses.len() - failures.len();
473        self.stats.misses += n_ok as u64;
474        self.stats.bytes_read += (n_ok * bytes) as u64;
475        if self.trace.is_some() {
476            let ev = self.stats.evictions - evictions0;
477            self.trace_line(keys, misses.len(), ev, 0);
478        }
479        if let Some((_, first)) = failures.first() {
480            let msg = format!(
481                "{} of {} expert reads failed; first: {first:#}",
482                failures.len(),
483                misses.len()
484            );
485            for (i, _) in &failures {
486                self.unmap(*i);
487            }
488            bail!(msg);
489        }
490        Ok(out.into_iter().map(|i| self.slot(i)).collect())
491    }
492}
spark_runtime/weights/gguf/expert_lru.rs

spark_runtime/weights/gguf/
expert_lru.rs
