metrale_model_layers/layers/ple/
ids.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: PLE n-gram row ids: the EOS-aware right shift and the multiply-XOR hash.
4//!
5//! Owner: model-layers (PLE).
6//! Invariants:
7//! - Pure token-id arithmetic: no device and no weights, so `tests.rs` checks
8//!   it bit-exactly against ids recorded from the reference forward.
9//! - `ple_ngram_ids` returns one row of `ngram_heads()` ids per input token.
10//!
11//! The LongCat n-gram hash in `ngram_embed/ids.rs` sums `shift * m` and takes
12//! one modulus. This one multiplies each shifted token by an odd multiplier,
13//! XORs the products, then takes a per-head modulus and adds a per-head offset
14//! into one concatenated table. Both give valid rows for the same tokens, so
15//! using the wrong one fails silently.
16//!
17//! Reference: `Qwen4ExpTextNGramEmbedding.forward` and
18//! `_shift_right_ignore_eos` in `bench/qwen4_exp/ref/modeling_qwen4_exp.py`.
19
20/// 2026-09-25: Geometry for one PLE site, read from the checkpoint rather than derived.
21///
22/// `multipliers`, `head_vocab_sizes` and `head_offsets` are the checkpoint's
23/// `ple_embedding.layer_multipliers`, `.ngram_heads_vocab_sizes` and
24/// `.ngram_heads_offsets`. The reference can also derive them (SplitMix64
25/// and a prime search); reading them keeps the ids tied to the checkpoint.
26#[derive(Clone, Debug)]
27pub struct PleIdDims {
28    /// 2026-09-25: `ngram_size`. The loader also uses it as the PLE conv dilation.
29    pub ngram_size: usize,
30    /// 2026-09-25: `heads_per_ngram`. 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    /// 2026-09-25: `layer_multipliers[ngram_size]`; `validate` refuses an even one.
34    pub multipliers: Vec<u64>,
35    /// 2026-09-25: `ngram_heads_vocab_sizes[ngram_heads]`: each head's modulus.
36    pub head_vocab_sizes: Vec<u64>,
37    /// 2026-09-25: `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    /// 2026-09-25: `(ngram_size - 1) * heads_per_ngram`. Times `head_dim` this is
45    /// the PLE embedding width, because the head slices are concatenated
46    /// (`PleLayer::new` checks the product).
47    pub fn ngram_heads(&self) -> usize {
48        (self.ngram_size - 1) * self.heads_per_ngram
49    }
50
51    /// 2026-09-25: How many previous tokens a decode step must carry to reproduce
52    /// prefill's ids.
53    pub fn context_len(&self) -> usize {
54        self.ngram_size - 1
55    }
56
57    /// 2026-09-25: Check the table lengths, that every multiplier is odd and every
58    /// head modulus is non-zero. The loader and `PleLayer::new` call it; an
59    /// error means the checkpoint does not have the reference's geometry.
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        // 2026-09-25: The reference builds these as `2 * (splitmix64(..) % half) + 1`.
75        for (i, m) in self.multipliers.iter().enumerate() {
76            anyhow::ensure!(
77                m % 2 == 1,
78                "PLE: layer_multipliers[{i}] = {m} is even; the reference \
79                 derives `2*x + 1`, so this checkpoint is not what we think"
80            );
81        }
82        anyhow::ensure!(
83            self.head_vocab_sizes.iter().all(|v| *v > 0),
84            "PLE: a head vocab size is 0 — modulus would divide by zero"
85        );
86        Ok(())
87    }
88}
89
90/// 2026-09-25: Right-shift by `shift`, refusing to read across an EOS boundary.
91///
92/// Positions whose source would fall before the current EOS-delimited
93/// segment get EOS instead. Transcribed from `_shift_right_ignore_eos`:
94/// `previous_eos` is the last EOS position strictly before each index, so a
95/// token sitting on an EOS still belongs to the segment the previous EOS started.
96fn shift_right_ignore_eos(tokens: &[u32], shift: usize, eos: u32) -> Vec<u32> {
97    if shift == 0 {
98        return tokens.to_vec();
99    }
100    let mut out = vec![eos; tokens.len()];
101    // 2026-09-25: `prev_eos` is the reference's `previous_eos`: the last EOS index
102    // strictly less than `pos`, or -1.
103    let mut prev_eos: i64 = -1;
104    let mut seen_eos: i64 = -1;
105    for pos in 0..tokens.len() {
106        let segment_start = prev_eos + 1;
107        let position_in_segment = pos as i64 - segment_start;
108        let source = pos as i64 - shift as i64;
109        if position_in_segment >= shift as i64 && source >= 0 {
110            out[pos] = tokens[source as usize];
111        }
112        if tokens[pos] == eos {
113            seen_eos = pos as i64;
114        }
115        prev_eos = seen_eos;
116    }
117    out
118}
119
120/// 2026-09-25: Row ids for every head, one row per token in `tokens`.
121///
122/// `tokens` must already be `context ++ new`, where `context` is the
123/// `context_len` preceding tokens (EOS-filled at the start of a sequence).
124/// Returns `[tokens.len()][ngram_heads]`; callers slice off the last
125/// `new.len()` rows, exactly as the reference's
126/// `torch.cat(blocks, dim=-1)[:, -input_ids.shape[1]:]` does.
127pub fn ple_ngram_ids(dims: &PleIdDims, tokens: &[u32]) -> Vec<Vec<u64>> {
128    let heads = dims.ngram_heads();
129    let shifted: Vec<Vec<u32>> = (0..dims.ngram_size)
130        .map(|s| shift_right_ignore_eos(tokens, s, dims.eos_token_id))
131        .collect();
132
133    let mut out = vec![vec![0u64; heads]; tokens.len()];
134    for ngram in 2..=dims.ngram_size {
135        let start = (ngram - 2) * dims.heads_per_ngram;
136        for (pos, row) in out.iter_mut().enumerate() {
137            // 2026-09-25: `mixed = shifted[0]*m[0]`, then XOR `shifted[p]*m[p]` for
138            // p in 1..ngram. The reference bounds each multiplier by
139            // `(2^63 - 1) / vocab_size`, so the product of a token id below
140            // the vocabulary size and a multiplier does not overflow.
141            let mut mixed = (shifted[0][pos] as u64).wrapping_mul(dims.multipliers[0]);
142            for p in 1..ngram {
143                mixed ^= (shifted[p][pos] as u64).wrapping_mul(dims.multipliers[p]);
144            }
145            for h in start..start + dims.heads_per_ngram {
146                row[h] = mixed % dims.head_vocab_sizes[h] + dims.head_offsets[h];
147            }
148        }
149    }
150    out
151}
metrale_model_layers/layers/ple/ids.rs

metrale_model_layers/layers/ple/
ids.rs
