spark_model/layers/ple/
ids.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! PLE n-gram row ids: EOS-aware right-shift plus the multiply-XOR hash.
4//!
5//! Pure token-id arithmetic — no device, no weights — which is why it is
6//! tested bit-exactly against the reference in `tests.rs` and runs in CI.
7//!
8//! **This does NOT transfer from LongCat (#746).** LongCat accumulates a
9//! polynomial rolling hash (`acc += shift[d] * m[d]`, then `% rows`). Qwen
10//! multiplies each shifted token by a SplitMix64-derived odd multiplier and
11//! **XOR**s them, then takes a per-head prime modulus and adds a per-head
12//! offset into one 320M-row table. The two produce different rows from the
13//! same tokens, and both produce VALID rows — so a mix-up is silent.
14//!
15//! Reference: `Qwen4ExpTextNGramEmbedding.forward` /
16//! `_shift_right_ignore_eos`, `bench/qwen4_exp/ref/modeling_qwen4_exp.py`.
17
18/// Geometry for one PLE site, read from the checkpoint rather than derived.
19///
20/// `multipliers`, `head_vocab_sizes` and `head_offsets` are all SHIPPED
21/// (`ple_embedding.layer_multipliers` / `.ngram_heads_vocab_sizes` /
22/// `.ngram_heads_offsets`). The reference can derive them — from SplitMix64
23/// and a prime search — and `bench/qwen4_exp/ple_golden.py` confirms the
24/// derivation reproduces the shipped values exactly. Reading them is still
25/// right: it cannot drift when the reference does.
26#[derive(Clone, Debug)]
27pub struct PleIdDims {
28    /// `ngram_size` (3). Also the conv dilation, elsewhere.
29    pub ngram_size: usize,
30    /// `heads_per_ngram` (8). Heads are grouped by n-gram order:
31    /// `[0, heads_per_ngram)` uses order 2, the next block order 3, and so on.
32    pub heads_per_ngram: usize,
33    /// `layer_multipliers[ngram_size]`, always odd.
34    pub multipliers: Vec<u64>,
35    /// `ngram_heads_vocab_sizes[ngram_heads]` — a distinct prime per head.
36    pub head_vocab_sizes: Vec<u64>,
37    /// `ngram_heads_offsets[ngram_heads]` — where each head's range starts
38    /// in the single concatenated table.
39    pub head_offsets: Vec<u64>,
40    pub eos_token_id: u32,
41}
42
43impl PleIdDims {
44    /// `(ngram_size - 1) * heads_per_ngram` — 16 here. Times `head_dim`
45    /// (160) this is `ple_embed_dim` (2560): the head slices are
46    /// CONCATENATED, not summed as LongCat's are.
47    pub fn ngram_heads(&self) -> usize {
48        (self.ngram_size - 1) * self.heads_per_ngram
49    }
50
51    /// How many previous tokens a decode step must carry to reproduce
52    /// prefill's ids. `ngram_size - 1` = 2.
53    pub fn context_len(&self) -> usize {
54        self.ngram_size - 1
55    }
56
57    /// Validate against the reference's invariants. Called once at load; a
58    /// mismatch here is a checkpoint we do not understand, not something to
59    /// paper over.
60    pub fn validate(&self) -> anyhow::Result<()> {
61        let heads = self.ngram_heads();
62        anyhow::ensure!(
63            self.multipliers.len() == self.ngram_size,
64            "PLE: layer_multipliers has {} entries, expected ngram_size={}",
65            self.multipliers.len(),
66            self.ngram_size
67        );
68        anyhow::ensure!(
69            self.head_vocab_sizes.len() == heads && self.head_offsets.len() == heads,
70            "PLE: head vocab/offsets are {}/{}, expected ngram_heads={heads}",
71            self.head_vocab_sizes.len(),
72            self.head_offsets.len()
73        );
74        // The reference builds these as `2 * (splitmix64(..) % half) + 1`.
75        // An even multiplier would collapse the low bit of every id.
76        for (i, m) in self.multipliers.iter().enumerate() {
77            anyhow::ensure!(
78                m % 2 == 1,
79                "PLE: layer_multipliers[{i}] = {m} is even; the reference \
80                 derives `2*x + 1`, so this checkpoint is not what we think"
81            );
82        }
83        anyhow::ensure!(
84            self.head_vocab_sizes.iter().all(|v| *v > 0),
85            "PLE: a head vocab size is 0 — modulus would divide by zero"
86        );
87        Ok(())
88    }
89}
90
91/// Right-shift by `shift`, refusing to read across an EOS boundary.
92///
93/// Positions whose source would fall before the current EOS-delimited
94/// segment get EOS instead. Transcribed from `_shift_right_ignore_eos`:
95/// `previous_eos` is the last EOS position STRICTLY BEFORE each index, so a
96/// token sitting on an EOS starts a segment at the previous one.
97fn shift_right_ignore_eos(tokens: &[u32], shift: usize, eos: u32) -> Vec<u32> {
98    if shift == 0 {
99        return tokens.to_vec();
100    }
101    let mut out = vec![eos; tokens.len()];
102    // `prev_eos` tracks `previous_eos_inclusive` shifted by one, i.e. the
103    // cummax over indices strictly less than `pos`.
104    let mut prev_eos: i64 = -1;
105    let mut seen_eos: i64 = -1;
106    for pos in 0..tokens.len() {
107        let segment_start = prev_eos + 1;
108        let position_in_segment = pos as i64 - segment_start;
109        let source = pos as i64 - shift as i64;
110        if position_in_segment >= shift as i64 && source >= 0 {
111            out[pos] = tokens[source as usize];
112        }
113        if tokens[pos] == eos {
114            seen_eos = pos as i64;
115        }
116        prev_eos = seen_eos;
117    }
118    out
119}
120
121/// Row ids for every head, one row per token in `tokens`.
122///
123/// `tokens` must already be `context ++ new`, where `context` is the
124/// `context_len` preceding tokens (EOS-filled at the start of a sequence).
125/// Returns `[tokens.len()][ngram_heads]`; callers slice off the last
126/// `new.len()` rows, exactly as the reference's
127/// `torch.cat(blocks, dim=-1)[:, -input_ids.shape[1]:]` does.
128pub fn ple_ngram_ids(dims: &PleIdDims, tokens: &[u32]) -> Vec<Vec<u64>> {
129    let heads = dims.ngram_heads();
130    let shifted: Vec<Vec<u32>> = (0..dims.ngram_size)
131        .map(|s| shift_right_ignore_eos(tokens, s, dims.eos_token_id))
132        .collect();
133
134    let mut out = vec![vec![0u64; heads]; tokens.len()];
135    for ngram in 2..=dims.ngram_size {
136        let start = (ngram - 2) * dims.heads_per_ngram;
137        for (pos, row) in out.iter_mut().enumerate() {
138            // `mixed = shifted[0]*m[0]`, then XOR `shifted[p]*m[p]` for
139            // p in 1..ngram. Wrapping is the reference's semantics too:
140            // torch int64 multiply wraps, and the multipliers are bounded by
141            // `(2^63 - 1) / vocab_size` precisely so it does not.
142            let mut mixed = (shifted[0][pos] as u64).wrapping_mul(dims.multipliers[0]);
143            for p in 1..ngram {
144                mixed ^= (shifted[p][pos] as u64).wrapping_mul(dims.multipliers[p]);
145            }
146            for h in start..start + dims.heads_per_ngram {
147                row[h] = mixed % dims.head_vocab_sizes[h] + dims.head_offsets[h];
148            }
149        }
150    }
151    out
152}
spark_model/layers/ple/ids.rs

spark_model/layers/ple/
ids.rs
