spark_model/layers/
ngram_embed.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! N-gram embedding id math — LongCat-Flash-Lite / Qwen3.8-Flash-Next family.
4//!
5//! Host-side, kernel-independent core of the n-gram embedding path (arxiv
6//! 2601.21204): the polynomial-rolling-hash ids that select rows in the
7//! `K * (N-1)` giant lookup tables. Split from the (future) GPU module so
8//! the integer math has a pure unit-test surface — the Rust ids are checked
9//! BIT-EXACT against the Python reference via
10//! `bench/ngram_ref/ngram_id_fixtures.json` (generated by
11//! `bench/ngram_ref/make_fixtures.py` from the line-faithful numpy port of
12//! the HF `modeling_longcat_ngram.py`).
13//!
14//! Mechanism (table `index = (i-2)*K + j` for n-gram size `i`, split `j`):
15//!
16//! ```text
17//!   T      = ratio * vocab_size + 2*index + 1        (table row count)
18//!   mods   = [V^1 mod T, V^2 mod T, ..., V^(i-1) mod T]
19//!   id_t   = ( x_t + Σ_{d=1..i-1} shift_d(x)_t * mods[d-1] ) mod T
20//! ```
21//!
22//! where `shift_d` is a right-shift by `d` that RESETS at document
23//! boundaries: a position within `d` tokens of a segment start (segments
24//! end at an EOS token, inclusive) contributes token id 0 instead of
25//! crossing the boundary. Ids depend ONLY on token ids — never on hidden
26//! state — which is what makes the lookups deterministic, prefetchable and
27//! speculative-decode-friendly.
28//!
29//! Overflow contract: `x * mod < 2^17 * 2^24 < 2^41` and at most N-1 terms
30//! accumulate, so the running sum fits comfortably in i64/u64 WITHOUT
31//! intermediate reduction at LongCat scale (V=131072, T≈10.2M). The
32//! per-term products must still be computed in 64-bit — 32-bit would
33//! overflow — and `mods` themselves must be built with 64-bit modmul.
34
35/// The n-gram trio + derived dims, extracted from `ModelConfig`.
36#[derive(Debug, Clone, Copy)]
37pub struct NgramDims {
38    pub vocab_size: u64,
39    pub ratio: u64,
40    /// Largest n-gram size N (>= 2).
41    pub neighbor_num: usize,
42    /// Hash splits K per n-gram size.
43    pub split_num: usize,
44    pub eos_token_id: u32,
45    pub hidden_size: usize,
46}
47
48impl NgramDims {
49    pub fn from_config(c: &atlas_core::config::ModelConfig) -> Option<Self> {
50        if c.ngram_vocab_size_ratio == 0 {
51            return None;
52        }
53        Some(Self {
54            vocab_size: c.vocab_size as u64,
55            ratio: c.ngram_vocab_size_ratio as u64,
56            neighbor_num: c.emb_neighbor_num,
57            split_num: c.emb_split_num,
58            eos_token_id: c.eos_token_id,
59            hidden_size: c.hidden_size,
60        })
61    }
62
63    pub fn num_tables(&self) -> usize {
64        self.split_num * (self.neighbor_num - 1)
65    }
66
67    /// Per-table embedding dim = hidden / num_tables (validated at parse).
68    pub fn table_dim(&self) -> usize {
69        self.hidden_size / self.num_tables()
70    }
71
72    /// Row count of table `index`: `ratio*vocab + 2*index + 1` — the
73    /// consecutive odd offsets give the K tables of one n-gram size
74    /// distinct (near-coprime) sizes so a collision in one split is
75    /// independent of the others.
76    pub fn table_rows(&self, index: usize) -> u64 {
77        self.ratio * self.vocab_size + 2 * index as u64 + 1
78    }
79
80    /// `[V^1 mod T, ..., V^(i-1) mod T]` for table (i, j).
81    pub fn vocab_mods(&self, ngram: usize, split: usize) -> Vec<u64> {
82        let index = (ngram - 2) * self.split_num + split;
83        let t = self.table_rows(index);
84        let mut mods = Vec::with_capacity(ngram - 1);
85        let mut power: u64 = 1;
86        for _ in 0..ngram - 1 {
87            // 64-bit modmul: V < 2^18, power < T < 2^25 → product < 2^43. OK.
88            power = (power * self.vocab_size) % t;
89            mods.push(power);
90        }
91        mods
92    }
93}
94
95/// `shift_right_ignore_eos` over one sequence: out[t] = ctx[t-n] unless the
96/// span [t-n, t] crosses a document boundary (EOS-inclusive segment end),
97/// in which case 0. Mirrors the reference implementation exactly, including
98/// the quirk that a segment shorter than n contributes nothing.
99mod embed;
100mod ids;
101mod table;
102
103pub use embed::NgramEmbedding;
104pub use ids::ngram_ids;
105pub use table::NgramTable;
106
107#[cfg(test)]
108mod tests;
spark_model/layers/ngram_embed.rs

spark_model/layers/
ngram_embed.rs
