spark_runtime/weights/gguf/
expert_arena.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! The expert cache's slots in DEVICE memory, filled through a page-locked
5//! staging ring.
6//!
7//! The GB10's expert GEMV reads the same bytes at 215 GB/s from `cudaMalloc`
8//! memory and at 176 GB/s from the page-locked arena the GPU reads in place
9//! (`kquant_mmvq_q2_k_experts_w8` at the real shapes, 09-19: same L2 hit
10//! rate, same occupancy, same long-scoreboard stall share; only the path
11//! differs). A decode step reads 40 layers x 6 experts = 2.93 GB of them, so
12//! the address the bytes live at is worth 3.2 ms of a 44 ms step.
13//!
14//! The CPU cannot write `cudaMalloc` memory on GB10 (SIGSEGV), so a miss is
15//! read by the reader threads into a slot of a small page-locked ring and
16//! then copied device-to-device-speed by the copy engine (`cuMemcpyHtoDAsync`,
17//! 217 us for one 12.22 MiB expert) INTO its device slot on the compute
18//! stream, ahead of the token's launches: stream order is the whole
19//! synchronisation, no launch can read a slot before its copy has landed,
20//! and a victim's new copy is enqueued after every launch that read the old
21//! bytes. The ring is two halves used alternately; a half is rewritten only
22//! after the event recorded behind its last copies has completed, so the
23//! copy engine never reads a ring slot the readers are refilling. A prefill
24//! layer's hundreds of misses go through in half-ring batches with the
25//! previous half's copies in flight.
26//!
27//! [`ExpertArena`] is the owner's handle for either kind of arena;
28//! [`ExpertLru::fetch_many_on`] is the one fetch that works on every kind.
29
30use anyhow::{Context, Result, bail, ensure};
31
32use super::expert_direct::{StagedRead, direct_plan, ring_slot_bytes};
33use super::expert_lru::{ExpertLru, ExpertSlot, PinnedArena, SlotPtr};
34use super::expert_stream::{DirectSeg, ExpertSource, SlotLayout, pread_at_least};
35use crate::gpu::{DevicePtr, GpuBackend};
36
37/// One device arena plus its page-locked staging ring.
38pub struct DeviceArena {
39    dev: DevicePtr,
40    bytes: usize,
41    ring: *mut u8,
42    ring_bytes: usize,
43    events: [u64; 2],
44}
45
46// SAFETY: raw addresses of memory the owner keeps alive; access is serialised
47// by the `ExpertLru` that uses them.
48unsafe impl Send for DeviceArena {}
49unsafe impl Sync for DeviceArena {}
50
51/// The staging state an [`ExpertLru`] over a [`DeviceArena`] carries.
52pub struct Staging {
53    pub(super) ring: *mut u8,
54    /// Slots in the ring (even, >= 2); a half is `ring_slots / 2`.
55    pub(super) ring_slots: usize,
56    /// Ring bytes per slot (`expert_direct::ring_slot_bytes`: the slot image
57    /// plus the padding the direct windows need).
58    pub(super) slot_bytes: usize,
59    /// One per half, recorded on the compute stream behind the half's copies.
60    pub(super) events: [u64; 2],
61    /// The event was recorded since the half was last waited for.
62    pub(super) armed: [bool; 2],
63    /// The half the next batch of misses goes through.
64    pub(super) half: usize,
65}
66
67// SAFETY: as for `DeviceArena`.
68unsafe impl Send for Staging {}
69
70impl DeviceArena {
71    /// `bytes` of device memory (off the allocation ledger) and a ring of
72    /// `ring_slots` page-locked expert slots (rounded up to an even count).
73    pub fn alloc(
74        gpu: &dyn GpuBackend,
75        bytes: usize,
76        layout: SlotLayout,
77        ring_slots: usize,
78    ) -> Result<(Self, ExpertLru)> {
79        let ring_slots = ring_slots.max(2).next_multiple_of(2);
80        let dev = gpu.alloc_arena(bytes)?;
81        let slot_bytes = ring_slot_bytes(&layout);
82        let ring_bytes = ring_slots * slot_bytes;
83        let ring = match gpu.alloc_host_pinned(ring_bytes) {
84            Ok(p) => p,
85            Err(e) => {
86                let _ = gpu.free_arena(dev);
87                return Err(e);
88            }
89        };
90        let events = [gpu.create_event()?, gpu.create_event()?];
91        let mut lru = ExpertLru::new(std::ptr::null_mut(), dev, bytes, layout)?;
92        lru.staging = Some(Staging {
93            ring,
94            ring_slots,
95            slot_bytes,
96            events,
97            armed: [false; 2],
98            half: 0,
99        });
100        Ok((
101            DeviceArena {
102                dev,
103                bytes,
104                ring,
105                ring_bytes,
106                events,
107            },
108            lru,
109        ))
110    }
111
112    pub fn dev(&self) -> DevicePtr {
113        self.dev
114    }
115
116    pub fn bytes(&self) -> usize {
117        self.bytes
118    }
119
120    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
121        for e in self.events {
122            gpu.destroy_event(e)?;
123        }
124        gpu.free_host_pinned(self.ring, self.ring_bytes)?;
125        gpu.free_arena(self.dev)
126    }
127}
128
129/// Either arena, as the runtime that keeps it alive holds it.
130pub enum ExpertArena {
131    /// Page-locked host memory the GPU reads in place.
132    Pinned(PinnedArena),
133    /// Device memory behind the staging ring.
134    Device(DeviceArena),
135}
136
137impl ExpertArena {
138    /// The arena of `bytes` and the cache over it: device slots with a
139    /// `ring_slots` staging ring when `device`, the page-locked arena otherwise.
140    pub fn alloc(
141        gpu: &dyn GpuBackend,
142        bytes: usize,
143        layout: SlotLayout,
144        device: bool,
145        ring_slots: usize,
146    ) -> Result<(Self, ExpertLru)> {
147        if device {
148            let (a, lru) = DeviceArena::alloc(gpu, bytes, layout, ring_slots)?;
149            return Ok((ExpertArena::Device(a), lru));
150        }
151        let a = PinnedArena::alloc(gpu, bytes)?;
152        let lru = ExpertLru::new(a.host(), a.dev(), a.bytes(), layout)?;
153        Ok((ExpertArena::Pinned(a), lru))
154    }
155
156    pub fn bytes(&self) -> usize {
157        match self {
158            ExpertArena::Pinned(a) => a.bytes(),
159            ExpertArena::Device(a) => a.bytes(),
160        }
161    }
162
163    /// One line for the load log.
164    pub fn describe(&self) -> String {
165        let gib = self.bytes() as f64 / (1024.0 * 1024.0 * 1024.0);
166        match self {
167            ExpertArena::Pinned(_) => format!("{gib:.1} GiB page-locked, read in place"),
168            ExpertArena::Device(a) => format!(
169                "{gib:.1} GiB device memory behind a {:.0} MiB page-locked staging ring",
170                a.ring_bytes as f64 / 1048576.0
171            ),
172        }
173    }
174
175    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
176        match self {
177            ExpertArena::Pinned(a) => a.free(gpu),
178            ExpertArena::Device(a) => a.free(gpu),
179        }
180    }
181}
182
183/// One byte range of a missed expert as a reader thread sees it: `(slot,
184/// key, ring slot within the half, byte offset, byte length)`.
185impl ExpertLru {
186    /// The arena's device address (slot `i` starts at `i * layout.bytes`).
187    pub fn arena_dev(&self) -> DevicePtr {
188        DevicePtr(self.dev)
189    }
190
191    /// The `(layer, expert, slot | -1)` changes since the last call: what a
192    /// device-side slot table has to learn (assignments and evictions).
193    pub fn drain_slot_changes(&mut self) -> Vec<(u32, u32, i32)> {
194        std::mem::take(&mut self.slot_changes)
195    }
196
197    /// The slots are device memory behind a staging ring.
198    pub fn is_device(&self) -> bool {
199        self.staging.is_some()
200    }
201
202    /// One token's gather on whichever path this cache was built for: the
203    /// staging ring (device slots), the reader pool (`predicted` experts
204    /// started in the background), or scoped reader threads. The copies of a
205    /// staged fetch are enqueued on `stream` ahead of the caller's launches.
206    pub fn fetch_many_on<S: ExpertSource + ?Sized>(
207        &mut self,
208        gpu: &dyn GpuBackend,
209        stream: u64,
210        src: &S,
211        keys: &[(u32, u32)],
212        predicted: &[(u32, u32)],
213        threads: usize,
214    ) -> Result<Vec<ExpertSlot>> {
215        if self.staging.is_some() {
216            self.fetch_many_staged(gpu, stream, src, keys, threads)
217        } else if self.has_pool() {
218            self.fetch_many_prefetching(keys, predicted)
219        } else {
220            self.fetch_many(src, keys, threads)
221        }
222    }
223
224    /// The staged gather: misses read into the ring by up to `threads`
225    /// scoped threads (each miss cut into byte ranges as `fetch_many` does),
226    /// then one H2D copy per miss into its device slot on `stream`, in
227    /// half-ring batches.
228    fn fetch_many_staged<S: ExpertSource + ?Sized>(
229        &mut self,
230        gpu: &dyn GpuBackend,
231        stream: u64,
232        src: &S,
233        keys: &[(u32, u32)],
234        threads: usize,
235    ) -> Result<Vec<ExpertSlot>> {
236        let (ring, ring_slots, slot_bytes, events) = match &self.staging {
237            Some(s) => (s.ring, s.ring_slots, s.slot_bytes, s.events),
238            None => bail!("fetch_many_staged on a page-locked expert cache"),
239        };
240        let evictions0 = self.stats.evictions;
241        let mut out: Vec<u32> = Vec::with_capacity(keys.len());
242        let mut misses: Vec<(u32, (u32, u32))> = Vec::new();
243        for &key in keys {
244            if let Some(&i) = self.map.get(&key) {
245                self.touch(i);
246                self.stats.hits += 1;
247                out.push(i);
248                continue;
249            }
250            let i = self.assign(key)?;
251            misses.push((i, key));
252            out.push(i);
253        }
254        let layout = self.layout();
255        let bytes = layout.bytes;
256        let half_slots = ring_slots / 2;
257        let threads = threads.max(1);
258        let mut failures: Vec<(u32, anyhow::Error)> = Vec::new();
259        for batch in misses.chunks(half_slots) {
260            let h = {
261                let st = self.staging.as_mut().expect("staging");
262                let h = st.half;
263                st.half ^= 1;
264                h
265            };
266            // the half's previous copies must have landed before its ring
267            // slots are rewritten
268            if self.staging.as_ref().is_some_and(|s| s.armed[h]) {
269                gpu.event_synchronize(events[h])?;
270                self.staging.as_mut().expect("staging").armed[h] = false;
271            }
272            // SAFETY: `(h + 1) * half_slots * slot_bytes <= ring bytes`.
273            let half_base = SlotPtr(unsafe { ring.add(h * half_slots * slot_bytes) });
274            // the reads: every miss cut into `parts` ranges spread over the
275            // threads (>= 1 MiB a range), exactly as `fetch_many` does; a
276            // source with direct descriptors reads aligned windows instead
277            let parts = (threads / batch.len())
278                .clamp(1, threads)
279                .min((bytes / (1 << 20)).max(1));
280            let per_part = bytes.div_ceil(parts);
281            let mut work: Vec<StagedRead> = Vec::new();
282            // per miss: its segments and their ring offsets (direct) or none
283            let mut direct: Vec<Option<(Vec<DirectSeg>, Vec<usize>)>> =
284                Vec::with_capacity(batch.len());
285            for (j, &(i, (l, e))) in batch.iter().enumerate() {
286                let slot_ring_off = j * slot_bytes;
287                match src.direct_segments(l, e)? {
288                    Some(segs) => {
289                        let (reads, bases) = direct_plan(i, &segs, &layout, slot_ring_off, parts);
290                        work.extend(reads);
291                        direct.push(Some((segs, bases)));
292                    }
293                    None => {
294                        for p in 0..parts {
295                            let off = p * per_part;
296                            let len = per_part.min(bytes.saturating_sub(off));
297                            if len > 0 {
298                                work.push(StagedRead::Buffered {
299                                    slot: i,
300                                    key: (l, e),
301                                    off,
302                                    len,
303                                    ring_off: slot_ring_off + off,
304                                });
305                            }
306                        }
307                        direct.push(None);
308                    }
309                }
310            }
311            let per = work.len().div_ceil(threads.min(work.len()).max(1)).max(1);
312            let results: Vec<Vec<(u32, anyhow::Error)>> = std::thread::scope(|s| {
313                let handles: Vec<_> = work
314                    .chunks(per)
315                    .map(|chunk| {
316                        s.spawn(move || {
317                            let mut errs = Vec::new();
318                            for r in chunk {
319                                if let Err(err) = staged_read(src, half_base, r) {
320                                    errs.push((r.slot(), err));
321                                }
322                            }
323                            errs
324                        })
325                    })
326                    .collect();
327                handles
328                    .into_iter()
329                    .map(|h| h.join().expect("expert reader thread panicked"))
330                    .collect()
331            });
332            let mut seen = std::collections::HashSet::new();
333            let failed_now: Vec<(u32, anyhow::Error)> = results
334                .into_iter()
335                .flatten()
336                .filter(|(i, _)| seen.insert(*i))
337                .collect();
338            // the copies, per miss that read whole, stream-ordered ahead of
339            // this token's launches; the ring stays valid until the event
340            for (j, &(i, _)) in batch.iter().enumerate() {
341                if seen.contains(&i) {
342                    continue;
343                }
344                let slot_dev = self.slot(i).gate.0 - layout.gate_off as u64;
345                // (ring offset, bytes, device offset) of each copy
346                let copies: Vec<(usize, usize, usize)> = match &direct[j] {
347                    Some((segs, bases)) => segs
348                        .iter()
349                        .zip(bases)
350                        .map(|(seg, &base)| (base, seg.len, seg.slot_off))
351                        .collect(),
352                    None => vec![(j * slot_bytes, bytes, 0)],
353                };
354                for (ring_off, len, dev_off) in copies {
355                    // SAFETY: as above; the readers have joined.
356                    let src_bytes =
357                        unsafe { std::slice::from_raw_parts(half_base.get().add(ring_off), len) };
358                    gpu.copy_h2d_async_retained(
359                        src_bytes,
360                        DevicePtr(slot_dev + dev_off as u64),
361                        stream,
362                    )?;
363                }
364            }
365            gpu.record_event(events[h], stream)?;
366            self.staging.as_mut().expect("staging").armed[h] = true;
367            failures.extend(failed_now);
368        }
369        let n_ok = misses.len() - failures.len();
370        self.stats.misses += n_ok as u64;
371        self.stats.bytes_read += (n_ok * bytes) as u64;
372        if self.trace.is_some() {
373            let ev = self.stats.evictions - evictions0;
374            self.trace_line(keys, misses.len(), ev, 0);
375        }
376        if let Some((_, first)) = failures.first() {
377            let msg = format!(
378                "{} of {} expert reads failed; first: {first:#}",
379                failures.len(),
380                misses.len()
381            );
382            for (i, _) in &failures {
383                self.unmap(*i);
384            }
385            bail!(msg);
386        }
387        ensure!(out.len() == keys.len(), "staged fetch lost a key");
388        Ok(out.into_iter().map(|i| self.slot(i)).collect())
389    }
390}
391
392/// One read of the staged miss path into the ring half at `half_base`.
393fn staged_read<S: ExpertSource + ?Sized>(
394    src: &S,
395    half_base: SlotPtr,
396    r: &StagedRead,
397) -> Result<()> {
398    match r {
399        StagedRead::Buffered {
400            key: (l, e),
401            off,
402            len,
403            ring_off,
404            ..
405        } => {
406            // SAFETY: the ring slot of this miss is its alone for the batch
407            // and the ranges of one slot are disjoint.
408            let dst =
409                unsafe { std::slice::from_raw_parts_mut(half_base.get().add(*ring_off), *len) };
410            src.read_expert_range(*l, *e, *off, dst)
411        }
412        StagedRead::Direct {
413            shard,
414            file_off,
415            len,
416            need,
417            ring_off,
418            ..
419        } => {
420            // SAFETY: the windows of one slot are disjoint and inside its
421            // padded regions (`direct_plan`).
422            let dst =
423                unsafe { std::slice::from_raw_parts_mut(half_base.get().add(*ring_off), *len) };
424            let f = src
425                .direct_file(*shard)
426                .with_context(|| format!("shard {shard} has no direct descriptor"))?;
427            pread_at_least(f, *file_off, dst, *need)
428                .with_context(|| format!("direct read, shard {shard}"))
429        }
430    }
431}
spark_runtime/weights/gguf/expert_arena.rs

spark_runtime/weights/gguf/
expert_arena.rs
