spark_runtime/weights/gguf/
expert_stream.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! DeepSeek-V4.1 Flash expert streaming from the split GGUF, by positional read.
5//!
6//! The routed experts of the Q2_K checkpoint are 40 layers x 384 experts of
7//! 12.22 MiB each; they never fit and they are never expanded. Measured on the
8//! real shards (`deepseek_v41_stream_bench_test.rs`): a CPU dequant of one expert
9//! is 851 ms, and page-faulting an expert slice in through the mmap runs at
10//! 0.33 GB/s cold and cached alike, while `pread` of the same 3.7 MiB slices
11//! runs at 4.5 GB/s cold. So the loader's mmap is used for the header only, and
12//! the expert bytes are read by [`pread`] straight into a page-locked arena the
13//! GPU addresses in place (see [`super::expert_lru::ExpertLru`]); the K-quant
14//! kernels consume the raw blocks from there.
15//!
16//! Three pieces:
17//!   * [`ShardFiles`]: every shard opened once, header parsed, mmap dropped.
18//!   * [`ExpertSliceMap`]: where expert `e` of layer `l` lives for each of the
19//!     three stacked tensors (`blk.N.ffn_{gate,up,down}_exps.weight`, GGUF dims
20//!     `[k, n, experts]`), and the read of one expert into one slot.
21//!   * [`EngramRowReader`] (`engram_rows.rs`): the ~30 GiB engram tables, one
22//!     Q2_K block (84 B) per row, read by row id on demand. 48 rows per token.
23//!
24//! Oracle: bytes identical to the mmap path, held by
25//! `deepseek_v41_stream_oracle_test.rs` on the real shards.
26
27use std::collections::HashMap;
28use std::fs::File;
29use std::path::{Path, PathBuf};
30use std::sync::Arc;
31
32use anyhow::{Context, Result, bail, ensure};
33
34use super::container::{GgufFile, Q2Group, TensorInfo};
35use super::sidecar;
36use crate::weights::{find_gguf, find_gguf_shards};
37
38pub use super::engram_rows::{EngramRowReader, EngramTable};
39pub use super::expert_arena::{DeviceArena, ExpertArena};
40pub use super::expert_lru::{ExpertLru, ExpertSlot, LruStats, PinnedArena};
41use super::expert_reads::open_direct;
42pub use super::expert_reads::{DirectSeg, pread, pread_at_least, pread_uncached};
43
44struct Shard {
45    path: PathBuf,
46    file: File,
47    /// The same file opened `O_DIRECT` for the expert miss path (`None`
48    /// when the platform or `ATLAS_DS41_DIRECT_READS=0` says no; the
49    /// buffered `file` then serves every read).
50    direct: Option<File>,
51    gguf: GgufFile,
52}
53
54/// Every shard of a split GGUF, opened once for the process lifetime. The
55/// header is parsed through the loader's own reader and the mmap is dropped:
56/// nothing here goes through the page cache but the reads we ask for.
57pub struct ShardFiles {
58    shards: Vec<Shard>,
59}
60
61impl ShardFiles {
62    /// Open the shard set that `first` (the `split.no == 0` file) names.
63    pub fn open(first: &Path) -> Result<Self> {
64        let set = find_gguf_shards(first)?;
65        let mut shards = Vec::with_capacity(set.paths.len());
66        for path in set.paths {
67            let (file, mmap, gguf) = sidecar::open_gguf(&path)?;
68            drop(mmap);
69            let direct = open_direct(&path);
70            shards.push(Shard {
71                path,
72                file,
73                direct,
74                gguf,
75            });
76        }
77        let n_direct = shards.iter().filter(|s| s.direct.is_some()).count();
78        tracing::info!(
79            "GGUF shards: {} open, {n_direct} with an O_DIRECT descriptor for expert misses",
80            shards.len()
81        );
82        Ok(ShardFiles { shards })
83    }
84
85    /// Open the model in `dir` (shard 0 found as the loader finds it).
86    pub fn open_dir(dir: &Path) -> Result<Self> {
87        let first = find_gguf(dir).with_context(|| format!("no GGUF in {}", dir.display()))?;
88        Self::open(&first)
89    }
90
91    pub fn len(&self) -> usize {
92        self.shards.len()
93    }
94
95    pub fn is_empty(&self) -> bool {
96        self.shards.is_empty()
97    }
98
99    pub fn path(&self, shard: usize) -> &Path {
100        &self.shards[shard].path
101    }
102
103    pub fn file(&self, shard: usize) -> &File {
104        &self.shards[shard].file
105    }
106
107    /// The `O_DIRECT` descriptor of `shard`, if it opened.
108    pub fn direct_file(&self, shard: usize) -> Option<&File> {
109        self.shards[shard].direct.as_ref()
110    }
111
112    /// `(shard, absolute offset, bytes)` of tensor `name` if it is stored as
113    /// Q2_K (84 B super-blocks of 256): the raw blocks a kernel can read in
114    /// place of the loader's bf16 expansion.
115    pub fn locate_q2k(&self, name: &str) -> Option<(usize, u64, usize)> {
116        let (shard, t, off) = self.locate(name)?;
117        let n = t.num_elements();
118        (format!("{:?}", t.ggml_type) == "Q2_K" && n.is_multiple_of(256)).then_some((
119            shard,
120            off,
121            n / 256 * 84,
122        ))
123    }
124
125    /// The parsed header of `shard` (shard 0 carries the model metadata).
126    pub fn header(&self, shard: usize) -> &GgufFile {
127        &self.shards[shard].gguf
128    }
129
130    /// `(shard, tensor, absolute byte offset)` of a tensor by GGUF name.
131    pub fn locate(&self, name: &str) -> Option<(usize, &TensorInfo, u64)> {
132        self.shards.iter().enumerate().find_map(|(i, s)| {
133            s.gguf
134                .tensor(name)
135                .map(|t| (i, t, s.gguf.tensor_abs_offset(t) as u64))
136        })
137    }
138
139    /// Every `(shard, tensor)` whose name ends with `suffix`.
140    pub(super) fn with_suffix<'a>(
141        &'a self,
142        suffix: &'a str,
143    ) -> impl Iterator<Item = (usize, &'a TensorInfo)> + 'a {
144        self.shards.iter().enumerate().flat_map(move |(i, s)| {
145            s.gguf
146                .tensors
147                .iter()
148                .filter(move |t| t.name.ends_with(suffix))
149                .map(move |t| (i, t))
150        })
151    }
152
153    pub(super) fn abs_offset(&self, shard: usize, t: &TensorInfo) -> u64 {
154        self.shards[shard].gguf.tensor_abs_offset(t) as u64
155    }
156}
157
158/// `blk.N.<rest>` -> `N`.
159pub(super) fn block_layer(name: &str) -> Option<usize> {
160    let rest = name.strip_prefix("blk.")?;
161    let end = rest.find('.')?;
162    rest[..end].parse().ok()
163}
164
165/// Where one expert of one stacked projection lives.
166#[derive(Clone, Copy, Debug, PartialEq, Eq)]
167pub struct SliceLoc {
168    pub shard: usize,
169    /// Absolute byte offset of expert 0; expert `e` is at `base + e * bytes`.
170    pub base: u64,
171    /// Bytes per expert on disk.
172    pub bytes: usize,
173    /// Elements per expert (`k * n`).
174    pub elems: usize,
175    pub ggml_type_id: u32,
176}
177
178/// The three projections of one MoE layer.
179#[derive(Clone, Copy, Debug)]
180pub struct ExpertLayerLoc {
181    pub layer: usize,
182    pub gate: SliceLoc,
183    pub up: SliceLoc,
184    pub down: SliceLoc,
185}
186
187/// Byte layout of one expert inside an LRU slot: `[gate | up | down]`.
188#[derive(Clone, Copy, Debug, PartialEq, Eq)]
189pub struct SlotLayout {
190    pub gate_off: usize,
191    pub up_off: usize,
192    pub down_off: usize,
193    pub gate_bytes: usize,
194    pub up_bytes: usize,
195    pub down_bytes: usize,
196    /// Total bytes per expert = one slot.
197    pub bytes: usize,
198}
199
200impl SlotLayout {
201    pub fn new(gate_bytes: usize, up_bytes: usize, down_bytes: usize) -> Self {
202        SlotLayout {
203            gate_off: 0,
204            up_off: gate_bytes,
205            down_off: gate_bytes + up_bytes,
206            gate_bytes,
207            up_bytes,
208            down_bytes,
209            bytes: gate_bytes + up_bytes + down_bytes,
210        }
211    }
212}
213
214/// Anything the LRU can fill a slot from. The on-disk map implements it; the
215/// unit tests implement it in memory.
216pub trait ExpertSource: Sync {
217    fn slot_layout(&self) -> SlotLayout;
218    fn num_experts(&self) -> usize;
219    /// Fill `dst` (exactly `slot_layout().bytes` long) with expert `expert` of
220    /// layer `layer`, in slot layout.
221    fn read_expert(&self, layer: u32, expert: u32, dst: &mut [u8]) -> Result<()>;
222    /// Fill `dst` with bytes `off..off + dst.len()` of the expert's slot
223    /// image, so one miss can be read by several threads at once. The
224    /// default reads the whole expert into scratch; the on-disk map seeks.
225    fn read_expert_range(&self, layer: u32, expert: u32, off: usize, dst: &mut [u8]) -> Result<()> {
226        let mut whole = vec![0u8; self.slot_layout().bytes];
227        self.read_expert(layer, expert, &mut whole)?;
228        dst.copy_from_slice(&whole[off..off + dst.len()]);
229        Ok(())
230    }
231    /// The expert's slot image as file ranges for `O_DIRECT` reads, or
232    /// `None` when this source has no direct descriptor (the default): the
233    /// staged miss path then reads through `read_expert_range`.
234    fn direct_segments(&self, layer: u32, expert: u32) -> Result<Option<Vec<DirectSeg>>> {
235        let _ = (layer, expert);
236        Ok(None)
237    }
238    /// The `O_DIRECT` descriptor a `DirectSeg` names.
239    fn direct_file(&self, shard: usize) -> Option<&File> {
240        let _ = shard;
241        None
242    }
243}
244
245/// The stacked expert tensors of every MoE layer, located once.
246pub struct ExpertSliceMap {
247    files: Arc<ShardFiles>,
248    layers: Vec<ExpertLayerLoc>,
249    by_layer: HashMap<usize, usize>,
250    num_experts: usize,
251    layout: SlotLayout,
252}
253
254impl ExpertSliceMap {
255    pub fn new(files: Arc<ShardFiles>) -> Result<Self> {
256        let mut parts: HashMap<usize, [Option<SliceLoc>; 3]> = HashMap::new();
257        let mut num_experts: Option<usize> = None;
258        for (which, suffix) in [
259            "ffn_gate_exps.weight",
260            "ffn_up_exps.weight",
261            "ffn_down_exps.weight",
262        ]
263        .iter()
264        .enumerate()
265        {
266            for (shard, t) in files.with_suffix(suffix) {
267                let Some(layer) = block_layer(&t.name) else {
268                    continue;
269                };
270                ensure!(
271                    t.dims.len() == 3,
272                    "{}: expected a stacked 3-D expert tensor, got {:?}",
273                    t.name,
274                    t.dims
275                );
276                let (k, n, experts) = (t.dims[0], t.dims[1], t.dims[2]);
277                match num_experts {
278                    None => num_experts = Some(experts),
279                    Some(e) => ensure!(
280                        e == experts,
281                        "{}: {experts} experts, others have {e}",
282                        t.name
283                    ),
284                }
285                let (qk, bb) = t.ggml_type.block_layout(Q2Group::G128)?;
286                let elems = k * n;
287                ensure!(
288                    elems.is_multiple_of(qk),
289                    "{}: {elems} elements not a multiple of the {qk}-block",
290                    t.name
291                );
292                let loc = SliceLoc {
293                    shard,
294                    base: files.abs_offset(shard, t),
295                    bytes: elems / qk * bb,
296                    elems,
297                    ggml_type_id: t.ggml_type.id(),
298                };
299                let slot = parts.entry(layer).or_default();
300                ensure!(slot[which].is_none(), "{}: duplicate", t.name);
301                slot[which] = Some(loc);
302            }
303        }
304        let Some(num_experts) = num_experts else {
305            bail!("no ffn_*_exps tensors in this GGUF")
306        };
307        let mut layers: Vec<ExpertLayerLoc> = parts
308            .into_iter()
309            .map(|(layer, [g, u, d])| {
310                let need = |x: Option<SliceLoc>, w: &str| {
311                    x.with_context(|| format!("layer {layer}: ffn_{w}_exps missing"))
312                };
313                Ok(ExpertLayerLoc {
314                    layer,
315                    gate: need(g, "gate")?,
316                    up: need(u, "up")?,
317                    down: need(d, "down")?,
318                })
319            })
320            .collect::<Result<_>>()?;
321        layers.sort_by_key(|l| l.layer);
322        let first = layers[0];
323        let layout = SlotLayout::new(first.gate.bytes, first.up.bytes, first.down.bytes);
324        for l in &layers {
325            ensure!(
326                (l.gate.bytes, l.up.bytes, l.down.bytes)
327                    == (layout.gate_bytes, layout.up_bytes, layout.down_bytes),
328                "layer {}: expert byte sizes differ from layer {}",
329                l.layer,
330                first.layer
331            );
332        }
333        let by_layer = layers
334            .iter()
335            .enumerate()
336            .map(|(i, l)| (l.layer, i))
337            .collect();
338        Ok(ExpertSliceMap {
339            files,
340            layers,
341            by_layer,
342            num_experts,
343            layout,
344        })
345    }
346
347    /// MoE layers in ascending order.
348    pub fn layers(&self) -> &[ExpertLayerLoc] {
349        &self.layers
350    }
351
352    pub fn layer(&self, layer: usize) -> Option<&ExpertLayerLoc> {
353        self.by_layer.get(&layer).map(|&i| &self.layers[i])
354    }
355
356    pub fn files(&self) -> &Arc<ShardFiles> {
357        &self.files
358    }
359
360    /// Absolute `(shard, offset)` of one projection of one expert.
361    pub fn slice_at(&self, loc: &SliceLoc, expert: usize) -> (usize, u64) {
362        (loc.shard, loc.base + (expert * loc.bytes) as u64)
363    }
364}
365
366impl ExpertSource for ExpertSliceMap {
367    fn slot_layout(&self) -> SlotLayout {
368        self.layout
369    }
370
371    fn num_experts(&self) -> usize {
372        self.num_experts
373    }
374
375    fn read_expert(&self, layer: u32, expert: u32, dst: &mut [u8]) -> Result<()> {
376        let l = self
377            .layer(layer as usize)
378            .with_context(|| format!("layer {layer} has no routed experts"))?;
379        ensure!(
380            (expert as usize) < self.num_experts,
381            "expert {expert} >= {}",
382            self.num_experts
383        );
384        ensure!(
385            dst.len() == self.layout.bytes,
386            "slot is {} bytes, expert is {}",
387            dst.len(),
388            self.layout.bytes
389        );
390        let lay = self.layout;
391        for (loc, off, len) in [
392            (&l.gate, lay.gate_off, lay.gate_bytes),
393            (&l.up, lay.up_off, lay.up_bytes),
394            (&l.down, lay.down_off, lay.down_bytes),
395        ] {
396            let (shard, at) = self.slice_at(loc, expert as usize);
397            pread_uncached(self.files.file(shard), at, &mut dst[off..off + len])
398                .with_context(|| format!("layer {layer} expert {expert} shard {shard}"))?;
399        }
400        Ok(())
401    }
402
403    fn read_expert_range(&self, layer: u32, expert: u32, off: usize, dst: &mut [u8]) -> Result<()> {
404        let l = self
405            .layer(layer as usize)
406            .with_context(|| format!("layer {layer} has no routed experts"))?;
407        ensure!(
408            (expert as usize) < self.num_experts,
409            "expert {expert} >= {}",
410            self.num_experts
411        );
412        let lay = self.layout;
413        let (lo, hi) = (off, off + dst.len());
414        ensure!(
415            hi <= lay.bytes,
416            "range {lo}..{hi} exceeds the {} byte slot",
417            lay.bytes
418        );
419        for (loc, s_off, s_len) in [
420            (&l.gate, lay.gate_off, lay.gate_bytes),
421            (&l.up, lay.up_off, lay.up_bytes),
422            (&l.down, lay.down_off, lay.down_bytes),
423        ] {
424            let (a, b) = (lo.max(s_off), hi.min(s_off + s_len));
425            if a >= b {
426                continue;
427            }
428            let (shard, at) = self.slice_at(loc, expert as usize);
429            pread_uncached(
430                self.files.file(shard),
431                at + (a - s_off) as u64,
432                &mut dst[a - lo..b - lo],
433            )
434            .with_context(|| {
435                format!("layer {layer} expert {expert} shard {shard} range {a}..{b}")
436            })?;
437        }
438        Ok(())
439    }
440
441    fn direct_segments(&self, layer: u32, expert: u32) -> Result<Option<Vec<DirectSeg>>> {
442        let l = self
443            .layer(layer as usize)
444            .with_context(|| format!("layer {layer} has no routed experts"))?;
445        ensure!(
446            (expert as usize) < self.num_experts,
447            "expert {expert} >= {}",
448            self.num_experts
449        );
450        let lay = self.layout;
451        let mut segs = Vec::with_capacity(3);
452        for (loc, slot_off, len) in [
453            (&l.gate, lay.gate_off, lay.gate_bytes),
454            (&l.up, lay.up_off, lay.up_bytes),
455            (&l.down, lay.down_off, lay.down_bytes),
456        ] {
457            let (shard, file_off) = self.slice_at(loc, expert as usize);
458            if self.files.direct_file(shard).is_none() {
459                return Ok(None);
460            }
461            segs.push(DirectSeg {
462                shard,
463                file_off,
464                slot_off,
465                len,
466            });
467        }
468        Ok(Some(segs))
469    }
470
471    fn direct_file(&self, shard: usize) -> Option<&File> {
472        self.files.direct_file(shard)
473    }
474}
spark_runtime/weights/gguf/expert_stream.rs

spark_runtime/weights/gguf/
expert_stream.rs
