logicaffeine_compile/model/
ngram_hash.rs

1//! The row ids of Qwen3.8-Flash-Next's hashed n-gram embedding (`qwen4_exp`'s per-layer embedding), as transformers'
2//! `Qwen4ExpTextNGramEmbedding` computes them.
3//!
4//! For each position and each n-gram order n from 2 to `ngram_size`, the token and the n − 1 before it are each
5//! multiplied by the layer's multiplier for their place and the products exclusive-ored; the result, reduced by each of
6//! the order's `heads_per_ngram` heads' primes and offset into that head's block, names one table row per head. A
7//! token before the start of its segment (the end-of-sequence token ends a segment) reads as the end-of-sequence
8//! token. The multipliers are odd, drawn by splitmix64 from the layer's seed, and small enough that every product of a
9//! token and a multiplier fits in a signed 64-bit integer; the heads' primes are the consecutive primes after
10//! `ngram_vocab_size_base`, and the table is padded to a multiple of `make_ngram_vocab_size_divisible_by` rows.
11
12use serde_json::Value;
13
14const MASK64: u128 = (1u128 << 64) - 1;
15const SPLITMIX_GAMMA: u64 = 0x9E37_79B9_7F4A_7C15;
16const SPLITMIX_M1: u64 = 0xBF58_476D_1CE4_E5B9;
17const SPLITMIX_M2: u64 = 0x94D0_49BB_1331_11EB;
18const PRIME_1: u64 = 10007;
19
20fn splitmix64(value: u64) -> u64 {
21    let mut v = value.wrapping_add(SPLITMIX_GAMMA);
22    v = (v ^ (v >> 30)).wrapping_mul(SPLITMIX_M1);
23    v = (v ^ (v >> 27)).wrapping_mul(SPLITMIX_M2);
24    v ^ (v >> 31)
25}
26
27fn is_prime(n: u64) -> bool {
28    if n < 2 {
29        return false;
30    }
31    if n % 2 == 0 {
32        return n == 2;
33    }
34    let mut d = 3;
35    while d * d <= n {
36        if n % d == 0 {
37            return false;
38        }
39        d += 2;
40    }
41    true
42}
43
44/// The `count`-th prime after `start`.
45fn nth_prime_after(start: u64, count: u64) -> u64 {
46    let mut p = start;
47    for _ in 0..count {
48        p += 1;
49        while !is_prime(p) {
50            p += 1;
51        }
52    }
53    p
54}
55
56/// One layer's n-gram hash.
57#[derive(Clone, Debug, PartialEq, Eq)]
58pub struct NgramHash {
59    ngram_size: usize,
60    heads_per_ngram: usize,
61    eos: u32,
62    multipliers: Vec<u64>,
63    primes: Vec<u64>,
64    offsets: Vec<u64>,
65    rows_before_padding: u64,
66    rows: u64,
67}
68
69impl NgramHash {
70    /// The hash of the per-layer embedding at place `ple_layer_index` among the layers that carry one.
71    #[allow(clippy::too_many_arguments)]
72    pub fn new(
73        vocab_size: u64,
74        ngram_size: usize,
75        heads_per_ngram: usize,
76        ngram_vocab_size_base: u64,
77        divisible_by: u64,
78        ple_layer_index: u64,
79        seed: u64,
80        eos: u32,
81    ) -> NgramHash {
82        let max_long = (1u64 << 63) - 1;
83        let half_bound = (max_long / vocab_size.max(1) / 2).max(1);
84        let base_seed = seed.wrapping_add(PRIME_1.wrapping_mul(ple_layer_index));
85        let multipliers = (0..ngram_size as u64)
86            .map(|i| {
87                let value = ((base_seed as u128 + SPLITMIX_GAMMA as u128 * (i as u128 + 1)) & MASK64) as u64;
88                2 * (splitmix64(value) % half_bound) + 1
89            })
90            .collect();
91        let heads = (ngram_size - 1) * heads_per_ngram;
92        let mut primes = Vec::with_capacity(heads);
93        let mut offsets = Vec::with_capacity(heads);
94        let mut total = 0u64;
95        for h in 0..heads as u64 {
96            let global = ple_layer_index * heads as u64 + h;
97            let p = nth_prime_after(ngram_vocab_size_base - 1, global + 1);
98            primes.push(p);
99            offsets.push(total);
100            total += p;
101        }
102        let rows = total.div_ceil(divisible_by) * divisible_by;
103        NgramHash { ngram_size, heads_per_ngram, eos, multipliers, primes, offsets, rows_before_padding: total, rows }
104    }
105
106    /// The hash a checkpoint's text configuration describes (its `seed` defaulting to 1234 and its end-of-sequence
107    /// token to the first listed, as the configuration class has them).
108    pub fn from_config(config: &Value, ple_layer_index: u64) -> Option<NgramHash> {
109        let u = |k: &str| config.get(k).and_then(Value::as_u64);
110        let eos = match config.get("eos_token_id")? {
111            Value::Array(a) => a.first()?.as_u64()?,
112            v => v.as_u64()?,
113        };
114        Some(NgramHash::new(
115            u("vocab_size")?,
116            u("ngram_size")? as usize,
117            u("heads_per_ngram")? as usize,
118            u("ngram_vocab_size_base")?,
119            u("make_ngram_vocab_size_divisible_by")?,
120            ple_layer_index,
121            u("seed").unwrap_or(1234),
122            eos as u32,
123        ))
124    }
125
126    pub fn multipliers(&self) -> &[u64] {
127        &self.multipliers
128    }
129
130    pub fn head_primes(&self) -> &[u64] {
131        &self.primes
132    }
133
134    /// Where each head's block starts in the table.
135    pub fn head_offsets(&self) -> &[u64] {
136        &self.offsets
137    }
138
139    pub fn rows_before_padding(&self) -> u64 {
140        self.rows_before_padding
141    }
142
143    /// The table's rows, padded.
144    pub fn rows(&self) -> u64 {
145        self.rows
146    }
147
148    /// The row ids of each of `tokens`, one per head, after the `context` the previous call left (its last
149    /// `ngram_size − 1` tokens; empty on the first call, which reads as end-of-sequence tokens).
150    pub fn ids(&self, context: &[u32], tokens: &[u32]) -> Vec<Vec<u64>> {
151        let keep = self.ngram_size - 1;
152        let mut history: Vec<u32> = if context.is_empty() { vec![self.eos; keep] } else { context.to_vec() };
153        history.extend_from_slice(tokens);
154        // Where each position's segment starts: just after the last end-of-sequence token before it.
155        let mut segment_start = Vec::with_capacity(history.len());
156        let mut start = 0usize;
157        for (p, &t) in history.iter().enumerate() {
158            segment_start.push(start);
159            if t == self.eos {
160                start = p + 1;
161            }
162        }
163        let shifted = |p: usize, s: usize| -> u64 {
164            if p >= s && p - segment_start[p] >= s {
165                u64::from(history[p - s])
166            } else {
167                u64::from(self.eos)
168            }
169        };
170        let first = history.len() - tokens.len();
171        (first..history.len())
172            .map(|p| {
173                let mut row = Vec::with_capacity(self.primes.len());
174                for n in 2..=self.ngram_size {
175                    let mut mixed = shifted(p, 0).wrapping_mul(self.multipliers[0]);
176                    for place in 1..n {
177                        mixed ^= shifted(p, place).wrapping_mul(self.multipliers[place]);
178                    }
179                    for h in 0..self.heads_per_ngram {
180                        let head = (n - 2) * self.heads_per_ngram + h;
181                        row.push(mixed % self.primes[head] + self.offsets[head]);
182                    }
183                }
184                row
185            })
186            .collect()
187    }
188}
logicaffeine_compile/model/ngram_hash.rs

logicaffeine_compile/model/
ngram_hash.rs
