spark_model/layers/ngram_embed/
ids.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The n-gram id core: EOS-aware right-shift and the polynomial rolling
4//! hash. Pure token-id arithmetic — no device, no weights — which is why
5//! it can be tested bit-exactly against the Python reference.
6
7use super::NgramDims;
8
9fn shift_right_ignore_eos(ctx: &[u32], n: usize, eos: u32) -> Vec<u32> {
10    let len = ctx.len();
11    let mut out = vec![0u32; len];
12    let mut prev = 0usize;
13    for (pos, &tok) in ctx.iter().enumerate() {
14        if tok == eos {
15            let end = pos + 1;
16            if end - prev > n {
17                out[prev + n..end].copy_from_slice(&ctx[prev..end - n]);
18            }
19            prev = end;
20        }
21    }
22    if prev < len && len - prev > n {
23        out[prev + n..len].copy_from_slice(&ctx[prev..len - n]);
24    }
25    out
26}
27
28/// Compute the row ids for EVERY table over `ctx` (the n-1 cached context
29/// tokens followed by the new tokens). Returns `num_tables` vectors of
30/// `ctx.len()` ids each, table-major in reference index order
31/// (`(ngram-2)*K + split`); callers slice the last `seq_len` entries.
32pub fn ngram_ids(dims: &NgramDims, ctx: &[u32]) -> Vec<Vec<u64>> {
33    let mut out = Vec::with_capacity(dims.num_tables());
34    // shift_d computed once per d, shared across splits (reference computes
35    // shifted_ids once per n-gram size; d ranges 1..=N-1).
36    let shifts: Vec<Vec<u32>> = (1..dims.neighbor_num)
37        .map(|d| shift_right_ignore_eos(ctx, d, dims.eos_token_id))
38        .collect();
39    for ngram in 2..=dims.neighbor_num {
40        for split in 0..dims.split_num {
41            let index = (ngram - 2) * dims.split_num + split;
42            let t = dims.table_rows(index);
43            let mods = dims.vocab_mods(ngram, split);
44            let ids = ctx
45                .iter()
46                .enumerate()
47                .map(|(pos, &x)| {
48                    let mut acc = x as u64;
49                    for (d, &m) in mods.iter().enumerate() {
50                        acc += shifts[d][pos] as u64 * m;
51                    }
52                    acc % t
53                })
54                .collect();
55            out.push(ids);
56        }
57    }
58    out
59}
60
61// ── GPU module ───────────────────────────────────────────────────────────
spark_model/layers/ngram_embed/ids.rs

spark_model/layers/ngram_embed/
ids.rs
