spark_model/layers/
engram_v41.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! DeepSeek-V4.1 Flash engram on the GPU.
5//!
6//! Engram sits before layers 1 and 14 (`engram_layer_ids`). Per token it hashes
7//! the n-gram window of compressed token ids into `n_hash_cols` row ids of a
8//! per-layer table, gathers those rows, projects their concatenation through
9//! `wkv` into one key per hyper-connection stream plus one shared value, and
10//! adds `gate * value` to each stream on the mHC highway. Four stages, each
11//! held to the CPU reference [`crate::layers::deepseek_v41_ref::engram`]:
12//!
13//! 1. **Hash** ([`EngramHasher`]): integer-exact on the CPU; 48 ids per token
14//!    is not a GPU problem. The tables (`token_map`, multipliers, primes,
15//!    offsets) come from the GGUF metadata, proven identical to `engram.py`.
16//! 2. **Rows**: the checkpoint's tables are Q2_K, one 84-byte block per
17//!    256-wide row, read by id through
18//!    `spark_runtime::weights::expert_stream::EngramRowReader` and dequantised
19//!    on the device by the Q2_K arm of `dequant_gguf_bf16` straight into the
20//!    `[tokens * cols, head_dim]` bf16 input ([`EngramV41::rows_from_q2k`]).
21//! 3. **Projection**: the dense bf16 GEMM, `[tokens, cols * head_dim]` x
22//!    `wkv^T` -> `[tokens, dim * (hc + 1)]`.
23//! 4. **Gate** (`engram_v41_gate`, `kernels/gb10/deepseek-v4-flash/nvfp4/engram_v41.cu`):
24//!    normalised dot of stream against key, signed-sqrt sigmoid, `h += gate * value`,
25//!    in place on the FP32 highway `[T, hc, H]` the mHC kernels keep.
26//!
27//! Oracle: `engram_v41_tests.rs` (the hash exact against the reference; the
28//! GPU path against the golden's `engram_out` captures; real-table rows
29//! bit-identical to the CPU decoder).
30
31use std::sync::Arc;
32
33use anyhow::{Context, Result, ensure};
34use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
35use spark_runtime::kernel_args::KernelLaunch;
36
37/// Bytes of one Q2_K super-block = one engram row.
38pub const ENGRAM_ROW_BYTES: usize = 84;
39mod apply;
40
41const DEQUANT_MODULE: &str = "dequant_gguf_bf16";
42const GATE_MODULE: &str = "engram_v41";
43const GEMM_MODULE: &str = "gemm";
44
45/// The hash tables, shaped as the reference indexes them.
46pub struct EngramHashTables {
47    pub layer_ids: Vec<usize>,
48    pub max_ngram: usize,
49    pub n_heads: usize,
50    /// `token_map[engram_pad_token_id]`: the pad in the compressed space.
51    pub pad_id: i64,
52    /// `[vocab]`, -1 = unmapped.
53    pub token_map: Vec<i64>,
54    /// `[layer][lookback]`
55    pub multipliers: Vec<Vec<i64>>,
56    /// `[layer][ngram_size - 2][head]`
57    pub primes: Vec<Vec<Vec<i64>>>,
58    /// `[layer][(ngram_size - 2) * n_heads + head]`
59    pub offsets: Vec<Vec<i64>>,
60}
61
62impl EngramHashTables {
63    /// From the flat GGUF metadata arrays (`deepseek41.engram.*`, as
64    /// `ModelConfig` carries them): multipliers `[layers * max_ngram]`, primes
65    /// `[layers * (max_ngram - 1) * n_heads]`, offsets the same length.
66    pub fn from_flat(
67        layer_ids: Vec<usize>,
68        max_ngram: usize,
69        n_heads: usize,
70        pad_token_id: u32,
71        token_map: Vec<i64>,
72        multipliers: &[u64],
73        primes: &[u64],
74        offsets: &[u64],
75    ) -> Result<Self> {
76        let nl = layer_ids.len();
77        ensure!(
78            nl > 0 && max_ngram >= 2 && n_heads > 0,
79            "engram geometry: {nl} layers, ngram {max_ngram}, heads {n_heads}"
80        );
81        let cols = (max_ngram - 1) * n_heads;
82        ensure!(
83            multipliers.len() == nl * max_ngram,
84            "engram.multipliers: {} entries, expected {}",
85            multipliers.len(),
86            nl * max_ngram
87        );
88        ensure!(
89            primes.len() == nl * cols,
90            "engram.primes: {} entries, expected {}",
91            primes.len(),
92            nl * cols
93        );
94        ensure!(
95            offsets.len() == nl * cols,
96            "engram.offsets: {} entries, expected {}",
97            offsets.len(),
98            nl * cols
99        );
100        let pad_id = *token_map
101            .get(pad_token_id as usize)
102            .with_context(|| format!("engram pad token {pad_token_id} outside token_map"))?;
103        let to_i64 = |v: &u64| i64::try_from(*v).context("engram table value exceeds i64");
104        let multipliers = multipliers
105            .chunks(max_ngram)
106            .map(|c| c.iter().map(to_i64).collect())
107            .collect::<Result<_>>()?;
108        let primes = primes
109            .chunks(cols)
110            .map(|layer| {
111                layer
112                    .chunks(n_heads)
113                    .map(|g| g.iter().map(to_i64).collect())
114                    .collect()
115            })
116            .collect::<Result<_>>()?;
117        let offsets = offsets
118            .chunks(cols)
119            .map(|c| c.iter().map(to_i64).collect())
120            .collect::<Result<_>>()?;
121        Ok(EngramHashTables {
122            layer_ids,
123            max_ngram,
124            n_heads,
125            pad_id,
126            token_map,
127            multipliers,
128            primes,
129            offsets,
130        })
131    }
132
133    pub fn n_hash_cols(&self) -> usize {
134        (self.max_ngram - 1) * self.n_heads
135    }
136
137    pub fn n_layers(&self) -> usize {
138        self.layer_ids.len()
139    }
140
141    /// Position of model layer `layer` in `layer_ids`.
142    pub fn hash_index(&self, layer: usize) -> Option<usize> {
143        self.layer_ids.iter().position(|&l| l == layer)
144    }
145}
146
147/// The n-gram hash for one sequence: keeps the compressed ids of every
148/// position so decode steps see the prefill's look-back, exactly as the
149/// reference's cache does.
150pub struct EngramHasher {
151    t: Arc<EngramHashTables>,
152    cache: Vec<i64>,
153}
154
155impl EngramHasher {
156    pub fn new(t: Arc<EngramHashTables>, max_seq: usize) -> Self {
157        EngramHasher {
158            t,
159            cache: vec![0; max_seq],
160        }
161    }
162
163    pub fn tables(&self) -> &Arc<EngramHashTables> {
164        &self.t
165    }
166
167    pub fn reset(&mut self) {
168        self.cache.iter_mut().for_each(|v| *v = 0);
169    }
170
171    /// Hash `input_ids` placed at `start_pos..`. Returns
172    /// `[seqlen, n_layers, n_hash_cols]` row ids.
173    pub fn hash(&mut self, input_ids: &[u32], start_pos: usize) -> Result<Vec<i64>> {
174        let t = &self.t;
175        let seqlen = input_ids.len();
176        ensure!(
177            start_pos + seqlen <= self.cache.len(),
178            "engram hasher: {} positions, {} needed",
179            self.cache.len(),
180            start_pos + seqlen
181        );
182        for (s, &id) in input_ids.iter().enumerate() {
183            let m = *t
184                .token_map
185                .get(id as usize)
186                .with_context(|| format!("token {id} outside engram token_map"))?;
187            self.cache[start_pos + s] = m;
188        }
189        let (ng, nl, cols) = (t.max_ngram, t.n_layers(), t.n_hash_cols());
190        let mut out = vec![0i64; seqlen * nl * cols];
191        let mut tokens = vec![0i64; ng];
192        for s in 0..seqlen {
193            let pos = start_pos + s;
194            let mut blocked = false;
195            for (shift, tok) in tokens.iter_mut().enumerate() {
196                let src = self.cache[pos.saturating_sub(shift)];
197                blocked = blocked || pos < shift || src == -1;
198                *tok = if blocked { t.pad_id } else { src };
199            }
200            for (li, mults) in t.multipliers.iter().enumerate() {
201                let mut rolling = tokens[0].wrapping_mul(mults[0]);
202                for i in 1..ng {
203                    rolling ^= tokens[i].wrapping_mul(mults[i]);
204                    for h in 0..t.n_heads {
205                        let col = (i - 1) * t.n_heads + h;
206                        out[(s * nl + li) * cols + col] =
207                            rolling.rem_euclid(t.primes[li][i - 1][h]) + t.offsets[li][col];
208                    }
209                }
210            }
211        }
212        Ok(out)
213    }
214
215    /// The row ids one engram layer needs, in `[token][col]` order, from a
216    /// `hash()` result of `tokens` positions.
217    pub fn layer_row_ids(&self, hashes: &[i64], tokens: usize, hash_index: usize) -> Vec<u64> {
218        let (nl, cols) = (self.t.n_layers(), self.t.n_hash_cols());
219        (0..tokens)
220            .flat_map(|s| (0..cols).map(move |c| hashes[(s * nl + hash_index) * cols + c] as u64))
221            .collect()
222    }
223}
224
225/// One engram layer's resident weights on the device.
226pub struct EngramLayerWeights {
227    pub layer: usize,
228    /// `[dim * (hc + 1), cols * head_dim]` bf16.
229    pub wkv: DevicePtr,
230    /// `[hc, dim]` f32, the elementwise product `q_weight * k_weight`.
231    pub qk: DevicePtr,
232    /// This layer's row workspace (`add_layer` allocates): the raw Q2_K rows
233    /// and their bf16 expansion, so every engram layer's rows for a step can
234    /// be resident at once (the whole-step graph uploads them all up front).
235    pub raw: DevicePtr,
236    pub rows: DevicePtr,
237    /// `wkv` as the GGUF ships it, `[out, in / 256]` raw Q2_K blocks (84 B),
238    /// read by the single-token projection in place of the bf16 expansion
239    /// (6x fewer bytes a token, byte-identical results); null = bf16 only.
240    pub wkv_q2k: DevicePtr,
241}
242
243/// The GPU side: kernels, per-layer weights, and workspaces for up to
244/// `max_tokens` positions per call.
245pub struct EngramV41 {
246    pub dim: usize,
247    pub hc: usize,
248    pub head_dim: usize,
249    pub cols: usize,
250    pub eps: f32,
251    pub max_tokens: usize,
252    gemm_k: KernelHandle,
253    /// the projection at one token: bandwidth-bound GEMV, not the 16x16 tile
254    gemv_k: KernelHandle,
255    /// the same projection off the raw Q2_K blocks (`wkv_q2k`)
256    gemv_q2k_k: KernelHandle,
257    gate_k: KernelHandle,
258    dequant_k: KernelHandle,
259    layers: Vec<EngramLayerWeights>,
260    /// `[max_tokens * cols * head_dim]` bf16.
261    rows: DevicePtr,
262    /// `[max_tokens * cols * ENGRAM_ROW_BYTES]` raw Q2_K blocks.
263    raw: DevicePtr,
264    /// `[max_tokens * dim * (hc + 1)]` bf16.
265    kv: DevicePtr,
266}
267
268impl EngramV41 {
269    pub fn new(
270        gpu: &dyn GpuBackend,
271        dim: usize,
272        hc: usize,
273        head_dim: usize,
274        cols: usize,
275        eps: f32,
276        max_tokens: usize,
277    ) -> Result<Self> {
278        ensure!(
279            (1..=4).contains(&hc),
280            "engram gate supports hc_mult 1..=4, got {hc}"
281        );
282        let n_rows = max_tokens * cols;
283        Ok(EngramV41 {
284            dim,
285            hc,
286            head_dim,
287            cols,
288            eps,
289            max_tokens,
290            gemm_k: gpu.kernel(GEMM_MODULE, "dense_gemm_bf16")?,
291            gemv_k: gpu.kernel("gemv", "dense_gemv_bf16")?,
292            gemv_q2k_k: gpu.kernel(GATE_MODULE, "engram_v41_wkv_q2k_gemv")?,
293            gate_k: gpu.kernel(GATE_MODULE, "engram_v41_gate")?,
294            dequant_k: gpu.kernel(DEQUANT_MODULE, "dequant_q2_k_to_bf16")?,
295            layers: Vec::new(),
296            rows: gpu.alloc(n_rows * head_dim * 2)?,
297            raw: gpu.alloc(n_rows * ENGRAM_ROW_BYTES)?,
298            kv: gpu.alloc(max_tokens * dim * (hc + 1) * 2)?,
299        })
300    }
301
302    pub fn in_features(&self) -> usize {
303        self.cols * self.head_dim
304    }
305
306    pub fn out_features(&self) -> usize {
307        self.dim * (self.hc + 1)
308    }
309
310    /// Register a layer's weights; its row workspace is allocated here.
311    pub fn add_layer(&mut self, gpu: &dyn GpuBackend, mut w: EngramLayerWeights) -> Result<()> {
312        let n_rows = self.max_tokens * self.cols;
313        w.raw = gpu.alloc(n_rows * ENGRAM_ROW_BYTES)?;
314        w.rows = gpu.alloc(n_rows * self.head_dim * 2)?;
315        self.layers.push(w);
316        Ok(())
317    }
318
319    pub fn layer(&self, layer: usize) -> Option<&EngramLayerWeights> {
320        self.layers.iter().find(|l| l.layer == layer)
321    }
322
323    /// `q_weight * k_weight` (both `[hc, dim]`, given as f32 values of the
324    /// stored bf16) uploaded as the gate's `qk`.
325    pub fn upload_qk(gpu: &dyn GpuBackend, q: &[f32], k: &[f32]) -> Result<DevicePtr> {
326        ensure!(q.len() == k.len(), "engram q/k weights differ in length");
327        let prod: Vec<u8> = q
328            .iter()
329            .zip(k)
330            .flat_map(|(a, b)| (a * b).to_le_bytes())
331            .collect();
332        let p = gpu.alloc(prod.len())?;
333        gpu.copy_h2d(&prod, p)?;
334        Ok(p)
335    }
336
337    /// Fill the row input from raw Q2_K blocks (`n_rows * 84` bytes, one block
338    /// per row in `[token][col]` order): upload, dequantise on the device.
339    pub fn rows_from_q2k(
340        &self,
341        gpu: &dyn GpuBackend,
342        layer: Option<usize>,
343        raw_blocks: &[u8],
344        n_rows: usize,
345        stream: u64,
346    ) -> Result<()> {
347        let (raw, rows) = self.row_bufs(layer)?;
348        ensure!(
349            n_rows <= self.max_tokens * self.cols,
350            "engram: {n_rows} rows exceeds the {} workspace",
351            self.max_tokens * self.cols
352        );
353        ensure!(
354            self.head_dim == 256,
355            "engram rows are one Q2_K block: head_dim must be 256, got {}",
356            self.head_dim
357        );
358        ensure!(
359            raw_blocks.len() == n_rows * ENGRAM_ROW_BYTES,
360            "engram: {} raw bytes for {n_rows} rows",
361            raw_blocks.len()
362        );
363        gpu.copy_h2d_async(raw_blocks, raw, stream)?;
364        KernelLaunch::new(gpu, self.dequant_k)
365            .grid([n_rows as u32, 1, 1])
366            .block([256, 1, 1])
367            .arg_ptr(raw)
368            .arg_ptr(rows)
369            .arg_u32(n_rows as u32)
370            .arg_u32(ENGRAM_ROW_BYTES as u32)
371            .launch(stream)
372    }
373
374    /// Fill the row input from already-dequantised bf16 rows (bits), for
375    /// oracles and for tables held resident in bf16.
376    pub fn rows_from_bf16(
377        &self,
378        gpu: &dyn GpuBackend,
379        layer: Option<usize>,
380        rows_bf16: &[u16],
381        n_rows: usize,
382    ) -> Result<()> {
383        ensure!(
384            n_rows <= self.max_tokens * self.cols,
385            "engram: {n_rows} rows exceeds the workspace"
386        );
387        ensure!(
388            rows_bf16.len() == n_rows * self.head_dim,
389            "engram: {} values for {n_rows} rows",
390            rows_bf16.len()
391        );
392        let bytes: Vec<u8> = rows_bf16.iter().flat_map(|v| v.to_le_bytes()).collect();
393        gpu.copy_h2d(&bytes, self.row_bufs(layer)?.1)
394    }
395
396    /// The dequantised rows as the GEMM sees them (`[n_rows, head_dim]` bf16):
397    /// the shared workspace (`layer` = `None`).
398    pub fn rows_ptr(&self) -> DevicePtr {
399        self.rows
400    }
401
402    /// `(raw, rows)` of layer `layer`'s own workspace, or the shared one.
403    fn row_bufs(&self, layer: Option<usize>) -> Result<(DevicePtr, DevicePtr)> {
404        match layer {
405            None => Ok((self.raw, self.rows)),
406            Some(l) => {
407                let w = self
408                    .layer(l)
409                    .with_context(|| format!("engram: layer {l} has no weights"))?;
410                Ok((w.raw, w.rows))
411            }
412        }
413    }
414
415    /// The rows `apply` reads for `w`: its own workspace (a null `rows` is a
416    /// layer registered without one: the shared workspace).
417    fn rows_of(&self, w: &EngramLayerWeights) -> DevicePtr {
418        if w.rows.is_null() { self.rows } else { w.rows }
419    }
420
421    /// The projected `[tokens, dim * (hc + 1)]` bf16 buffer after `apply`.
422    pub fn kv_ptr(&self) -> DevicePtr {
423        self.kv
424    }
425
426    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
427        for l in &self.layers {
428            gpu.free(l.wkv)?;
429            gpu.free(l.qk)?;
430            gpu.free(l.raw)?;
431            gpu.free(l.rows)?;
432            if !l.wkv_q2k.is_null() {
433                gpu.free(l.wkv_q2k)?;
434            }
435        }
436        gpu.free(self.rows)?;
437        gpu.free(self.raw)?;
438        gpu.free(self.kv)
439    }
440}
441
442#[cfg(test)]
443#[path = "engram_v41_tests.rs"]
444mod tests;
spark_model/layers/engram_v41.rs

spark_model/layers/
engram_v41.rs
