metrale_model_layers/layers/
ngram_embed.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: N-gram embedding (arXiv 2601.21204): the dimensions, the
4//! host-side polynomial-rolling-hash ids that select rows in the
5//! `K * (N-1)` lookup tables (`ids`), the tables (`table`) and the GPU
6//! embedding (`embed`). `ids_match_python_reference` checks the ids against
7//! `bench/ngram_ref/ngram_id_fixtures.json`, written by
8//! `bench/ngram_ref/make_fixtures.py`.
9//!
10//! Table `index = (i-2)*K + j` for n-gram size `i` and split `j`:
11//!
12//! ```text
13//!   T      = ratio * vocab_size + 2*index + 1        (table row count)
14//!   mods   = [V^1 mod T, V^2 mod T, ..., V^(i-1) mod T]
15//!   id_t   = ( x_t + Σ_{d=1..i-1} shift_d(x)_t * mods[d-1] ) mod T
16//! ```
17//!
18//! where `shift_d` is a right-shift by `d` that resets at document
19//! boundaries: a position within `d` tokens of a segment start (a segment
20//! ends at an EOS token, inclusive) contributes token id 0. The ids depend
21//! only on token ids.
22//!
23//! Owner: model-layers (n-gram embedding).
24//! Invariants:
25//! - `ngram_ids` sums in u64 and reduces once per id, so the ids are exact
26//!   only while `neighbor_num * vocab_size * T < 2^64`; the code does not
27//!   check this.
28
29/// 2026-09-25: The n-gram dimensions, read from `ModelConfig`.
30#[derive(Debug, Clone, Copy)]
31pub struct NgramDims {
32    pub vocab_size: u64,
33    pub ratio: u64,
34    /// 2026-09-25: Largest n-gram size N.
35    pub neighbor_num: usize,
36    /// 2026-09-25: Hash splits K per n-gram size.
37    pub split_num: usize,
38    pub eos_token_id: u32,
39    pub hidden_size: usize,
40}
41
42impl NgramDims {
43    pub fn from_config(c: &metrale_config::ModelConfig) -> Option<Self> {
44        if c.ngram_vocab_size_ratio == 0 {
45            return None;
46        }
47        Some(Self {
48            vocab_size: c.vocab_size as u64,
49            ratio: c.ngram_vocab_size_ratio as u64,
50            neighbor_num: c.emb_neighbor_num,
51            split_num: c.emb_split_num,
52            eos_token_id: c.eos_token_id,
53            hidden_size: c.hidden_size,
54        })
55    }
56
57    pub fn num_tables(&self) -> usize {
58        self.split_num * (self.neighbor_num - 1)
59    }
60
61    /// 2026-09-25: Per-table embedding dim, `hidden_size / num_tables()`
62    /// (integer division).
63    pub fn table_dim(&self) -> usize {
64        self.hidden_size / self.num_tables()
65    }
66
67    /// 2026-09-25: Row count of table `index`: `ratio*vocab + 2*index + 1`,
68    /// so the K tables of one n-gram size have distinct sizes.
69    pub fn table_rows(&self, index: usize) -> u64 {
70        self.ratio * self.vocab_size + 2 * index as u64 + 1
71    }
72
73    /// 2026-09-25: `V^1 mod T, ..., V^(i-1) mod T` for table (i, j).
74    pub fn vocab_mods(&self, ngram: usize, split: usize) -> Vec<u64> {
75        let index = (ngram - 2) * self.split_num + split;
76        let t = self.table_rows(index);
77        let mut mods = Vec::with_capacity(ngram - 1);
78        let mut power: u64 = 1;
79        for _ in 0..ngram - 1 {
80            power = (power * self.vocab_size) % t;
81            mods.push(power);
82        }
83        mods
84    }
85}
86
87mod embed;
88mod ids;
89mod table;
90
91pub use embed::NgramEmbedding;
92pub use ids::ngram_ids;
93pub use table::NgramTable;
94
95#[cfg(test)]
96mod tests;
metrale_model_layers/layers/ngram_embed.rs

metrale_model_layers/layers/
ngram_embed.rs
