spark_runtime/weights/gguf/
engram_rows.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! The engram tables of DeepSeek-V4.1 Flash, read by row on demand. Split
5//! from `expert_stream.rs` (500-LoC cap).
6
7use std::sync::Arc;
8
9use anyhow::{Context, Result, ensure};
10
11use super::container::Q2Group;
12use super::expert_stream::{ShardFiles, block_layer, pread};
13
14/// One engram table: `blk.N.engram_embd.weight`, GGUF dims `[head_dim, rows]`,
15/// one quant block per row.
16#[derive(Clone, Copy, Debug)]
17pub struct EngramTable {
18    pub layer: usize,
19    pub shard: usize,
20    pub base: u64,
21    pub rows: usize,
22    pub row_bytes: usize,
23    pub head_dim: usize,
24    pub ggml_type_id: u32,
25}
26
27/// Row reads from the engram tables. The tables are ~30 GiB each and a token
28/// touches 24 rows per table, so nothing is cached: every call is `ids.len()`
29/// positional reads of one block.
30pub struct EngramRowReader {
31    files: Arc<ShardFiles>,
32    tables: Vec<EngramTable>,
33}
34
35impl EngramRowReader {
36    pub fn new(files: Arc<ShardFiles>) -> Result<Self> {
37        let mut tables = Vec::new();
38        for (shard, t) in files.with_suffix("engram_embd.weight") {
39            let Some(layer) = block_layer(&t.name) else {
40                continue;
41            };
42            ensure!(
43                t.dims.len() == 2,
44                "{}: expected [head_dim, rows], got {:?}",
45                t.name,
46                t.dims
47            );
48            let (head_dim, rows) = (t.dims[0], t.dims[1]);
49            let (qk, bb) = t.ggml_type.block_layout(Q2Group::G128)?;
50            ensure!(
51                qk == head_dim,
52                "{}: head_dim {head_dim} is not one {qk}-element block per row",
53                t.name
54            );
55            tables.push(EngramTable {
56                layer,
57                shard,
58                base: files.abs_offset(shard, t),
59                rows,
60                row_bytes: bb,
61                head_dim,
62                ggml_type_id: t.ggml_type.id(),
63            });
64        }
65        ensure!(!tables.is_empty(), "no engram_embd tensors in this GGUF");
66        tables.sort_by_key(|t| t.layer);
67        Ok(EngramRowReader { files, tables })
68    }
69
70    pub fn tables(&self) -> &[EngramTable] {
71        &self.tables
72    }
73
74    pub fn table(&self, layer: usize) -> Option<&EngramTable> {
75        self.tables.iter().find(|t| t.layer == layer)
76    }
77
78    /// Read rows `ids` of layer `layer`'s table into `dst`
79    /// (`ids.len() * row_bytes` bytes, in `ids` order).
80    pub fn read_rows(&self, layer: usize, ids: &[u64], dst: &mut [u8]) -> Result<()> {
81        let t = self
82            .table(layer)
83            .with_context(|| format!("layer {layer} has no engram table"))?;
84        ensure!(
85            dst.len() == ids.len() * t.row_bytes,
86            "dst is {} bytes for {} rows of {}",
87            dst.len(),
88            ids.len(),
89            t.row_bytes
90        );
91        let file = self.files.file(t.shard);
92        for (i, &id) in ids.iter().enumerate() {
93            ensure!(
94                (id as usize) < t.rows,
95                "engram row {id} >= {} (layer {layer})",
96                t.rows
97            );
98            let off = t.base + id * t.row_bytes as u64;
99            pread(file, off, &mut dst[i * t.row_bytes..(i + 1) * t.row_bytes])
100                .with_context(|| format!("engram layer {layer} row {id}"))?;
101        }
102        Ok(())
103    }
104}
spark_runtime/weights/gguf/engram_rows.rs

spark_runtime/weights/gguf/
engram_rows.rs
