metrale_model_arch/
engram_v41.rs

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

metrale_model_arch/
engram_v41.rs
