metrale_model_arch/glm5next_dsa_ref/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: CPU reference for the GLM-5.3-Flash DSA kpool indexer and NoPE MLA.
4//!
5//! Owner: model-arch (GLM-5.3-Flash DSA reference).
6//! Invariants:
7//! - Nothing here runs on a GPU or reads a checkpoint; production code does not call it.
8//! - Every index row `expand_selection` returns is written for its full `out_width`, with
9//!   `INVALID` in every slot that holds no token.
10//!
11//! `examples/dsa_indexer_microtest.rs` checks it against goldens from `gen_dsa_indexer_golden.py`
12//! and `gen_dsa_mla_golden.py`, which run the `transformers` `glm5_next` indexer and attention.
13//!
14//! * `k_norm` is a LayerNorm: it subtracts the mean and adds a bias
15//!   (`indexer.k_norm.bias`). The DSA block's other norms (`q_a_layernorm`, `kv_a_layernorm`)
16//!   have no bias.
17//! * The pool softmax runs over the slot axis, separately for each `head_dim` channel.
18//! * Pooling starts at the first valid token, not at slot 0: with left padding
19//!   `[P, P, A, B, C, D]`, pool 0 is `[A, B, C, D]`.
20//! * A pool is valid only if all `kpool` slots are valid. A trailing partial pool is not a pool;
21//!   the tail append covers it, so a 7-token sequence has one pool.
22//! * NoPE: `qk_rope_head_dim = 0`, so `kv_a_proj_with_mqa` emits `kv_lora_rank` values and there
23//!   is no rope section.
24
25/// 2026-09-25: Indexer and MLA geometry. No field has a default.
26#[derive(Clone, Copy, Debug)]
27pub struct DsaDims {
28    pub hidden: usize,
29    /// 2026-09-25: Indexer heads (`index_n_heads`), not the MLA head count.
30    pub index_heads: usize,
31    /// 2026-09-25: Indexer head dim (`index_head_dim`), not the MLA head dim.
32    pub index_head_dim: usize,
33    pub index_kpool: usize,
34    pub index_topk: usize,
35    pub always_select_tail: bool,
36    pub q_lora_rank: usize,
37    pub heads: usize,
38    pub kv_lora_rank: usize,
39    pub qk_nope_head_dim: usize,
40    /// 2026-09-25: Zero on GLM-5.3-Flash (MODEL.toml); [`expand_kv`] asserts it is zero.
41    pub qk_rope_head_dim: usize,
42    pub v_head_dim: usize,
43}
44
45impl DsaDims {
46    pub fn qk_head_dim(&self) -> usize {
47        self.qk_nope_head_dim + self.qk_rope_head_dim
48    }
49    /// 2026-09-25: Pools selected per query: `index_topk / index_kpool`, capped by how many pools
50    /// exist.
51    pub fn select_k(&self, n_pools: usize) -> usize {
52        (self.index_topk / self.index_kpool).min(n_pools)
53    }
54    /// 2026-09-25: Width of the emitted index row. The tail adds `kpool - 1` slots.
55    pub fn out_width(&self) -> usize {
56        self.index_topk
57            + if self.always_select_tail {
58                self.index_kpool - 1
59            } else {
60                0
61            }
62    }
63    pub fn is_nope(&self) -> bool {
64        self.qk_rope_head_dim == 0
65    }
66}
67
68/// 2026-09-25: The invalid-index sentinel, the `-1` the golden generator writes.
69pub const INVALID: i32 = -1;
70
71#[inline]
72fn sigmoid(x: f32) -> f32 {
73    1.0 / (1.0 + (-x).exp())
74}
75
76/// 2026-09-25: LayerNorm over the trailing `d`: subtract the mean, divide by the standard
77/// deviation, then `w * x + b`.
78pub fn layer_norm(x: &[f32], w: &[f32], b: &[f32], d: usize, eps: f32) -> Vec<f32> {
79    let mut out = vec![0.0f32; x.len()];
80    for (row_in, row_out) in x.chunks_exact(d).zip(out.chunks_exact_mut(d)) {
81        let mean = row_in.iter().sum::<f32>() / d as f32;
82        let var = row_in.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / d as f32;
83        let inv = 1.0 / (var + eps).sqrt();
84        for i in 0..d {
85            row_out[i] = (row_in[i] - mean) * inv * w[i] + b[i];
86        }
87    }
88    out
89}
90
91/// 2026-09-25: RMSNorm without bias, the contrast to [`layer_norm`] in the tests.
92pub fn rms_norm(x: &[f32], w: &[f32], d: usize, eps: f32) -> Vec<f32> {
93    let mut out = vec![0.0f32; x.len()];
94    for (row_in, row_out) in x.chunks_exact(d).zip(out.chunks_exact_mut(d)) {
95        let inv = 1.0 / (row_in.iter().map(|v| v * v).sum::<f32>() / d as f32 + eps).sqrt();
96        for i in 0..d {
97            row_out[i] = row_in[i] * inv * w[i];
98        }
99    }
100    out
101}
102
103/// 2026-09-25: `y = x @ w^T` for `x: [m, k]`, `w: [n, k]` (torch `Linear` layout), no bias.
104pub fn linear(x: &[f32], m: usize, k: usize, w: &[f32], n: usize) -> Vec<f32> {
105    let mut out = vec![0.0f32; m * n];
106    for row in 0..m {
107        for col in 0..n {
108            let mut acc = 0.0f32;
109            for i in 0..k {
110                acc += x[row * k + i] * w[col * k + i];
111            }
112            out[row * n + col] = acc;
113        }
114    }
115    out
116}
117
118/// 2026-09-25: The compressed k-pool candidates.
119pub struct Pools {
120    /// 2026-09-25: `[n_pools, index_head_dim]`: softmax-weighted average of the pool's keys.
121    pub keys: Vec<f32>,
122    /// 2026-09-25: `[n_pools, index_kpool]`: raw token index per slot, [`INVALID`] where the slot
123    /// is not a valid token.
124    pub indices: Vec<i32>,
125    /// 2026-09-25: `[n_pools]`: 1 only when every slot is valid.
126    pub valid: Vec<u8>,
127    pub n_pools: usize,
128}
129
130/// 2026-09-25: Build the pools. `k`/`gate` are `[seq, index_head_dim]`, `valid` is `[seq]`. The
131/// indexer golden records the output of `transformers`' `get_pooled_states` for comparison.
132pub fn pool_states(
133    k: &[f32],
134    gate: &[f32],
135    valid: &[u8],
136    ape: &[f32],
137    dims: DsaDims,
138    seq: usize,
139) -> Pools {
140    let (d, kp) = (dims.index_head_dim, dims.index_kpool);
141    let n_pools = seq.div_ceil(kp);
142    // 2026-09-25: Pooling starts at the first valid token, so left padding is skipped.
143    let first_key = valid.iter().position(|v| *v != 0).unwrap_or(seq) as i64;
144
145    let mut keys = vec![0.0f32; n_pools * d];
146    let mut indices = vec![INVALID; n_pools * kp];
147    let mut pvalid = vec![0u8; n_pools];
148    let mut logits = vec![0.0f32; kp];
149
150    for p in 0..n_pools {
151        let mut slot_valid = [false; 64];
152        let mut slot_idx = [0usize; 64];
153        let mut all = true;
154        for s in 0..kp {
155            let raw = first_key + (p * kp + s) as i64;
156            let in_range = raw >= 0 && (raw as usize) < seq;
157            let ok = in_range && valid[raw as usize] != 0;
158            slot_valid[s] = ok;
159            slot_idx[s] = if in_range { raw as usize } else { 0 };
160            all &= ok;
161            indices[p * kp + s] = if ok { raw as i32 } else { INVALID };
162        }
163        pvalid[p] = all as u8;
164
165        // 2026-09-25: Softmax over the slot axis, independently per channel.
166        for dd in 0..d {
167            let mut mx = f32::NEG_INFINITY;
168            for s in 0..kp {
169                logits[s] = if slot_valid[s] {
170                    gate[slot_idx[s] * d + dd] + ape[s * d + dd]
171                } else {
172                    f32::NEG_INFINITY
173                };
174                mx = mx.max(logits[s]);
175            }
176            let mut sum = 0.0f32;
177            for s in 0..kp {
178                let e = if logits[s] == f32::NEG_INFINITY {
179                    0.0
180                } else {
181                    (logits[s] - mx).exp()
182                };
183                logits[s] = e;
184                sum += e;
185            }
186            // 2026-09-25: A pool with no valid slot gets weight 0, so its key is 0, not NaN.
187            let inv = if sum > 0.0 { 1.0 / sum } else { 0.0 };
188            let mut acc = 0.0f32;
189            for s in 0..kp {
190                if slot_valid[s] {
191                    acc += logits[s] * inv * k[slot_idx[s] * d + dd];
192                }
193            }
194            keys[p * d + dd] = acc;
195        }
196    }
197    // 2026-09-25: Invalid pools are dropped from the axis. `n_pools`, and the `select_k` derived
198    // from it, therefore depend on the padding, not only on the sequence length: a 7-token
199    // sequence has one pool and a budget of one.
200    let keep: Vec<usize> = (0..n_pools).filter(|p| pvalid[*p] != 0).collect();
201    if keep.len() == n_pools {
202        return Pools {
203            keys,
204            indices,
205            valid: pvalid,
206            n_pools,
207        };
208    }
209    let mut ck = vec![0.0f32; keep.len() * d];
210    let mut ci = vec![INVALID; keep.len() * kp];
211    let mut cv = vec![0u8; keep.len()];
212    for (j, p) in keep.iter().enumerate() {
213        ck[j * d..(j + 1) * d].copy_from_slice(&keys[p * d..(p + 1) * d]);
214        ci[j * kp..(j + 1) * kp].copy_from_slice(&indices[p * kp..(p + 1) * kp]);
215        cv[j] = pvalid[*p];
216    }
217    Pools {
218        keys: ck,
219        indices: ci,
220        valid: cv,
221        n_pools: keep.len(),
222    }
223}
224
225/// 2026-09-25: The original ids of the pools [`pool_states`] keeps. Computed from padding and
226/// sequence length alone.
227pub fn kept_pools(valid: &[u8], dims: DsaDims, seq: usize) -> Vec<i32> {
228    let kp = dims.index_kpool;
229    let n_pools = seq.div_ceil(kp);
230    let first_key = valid.iter().position(|v| *v != 0).unwrap_or(seq) as i64;
231    (0..n_pools)
232        .filter(|p| {
233            (0..kp).all(|s| {
234                let raw = first_key + (p * kp + s) as i64;
235                raw >= 0 && (raw as usize) < seq && valid[raw as usize] != 0
236            })
237        })
238        .map(|p| p as i32)
239        .collect()
240}
241
242/// 2026-09-25: Per-(query, pool) index score: `sum_h weights[h] * relu(scale * q_h . key)`.
243///
244/// `q` is `[q_rows, index_heads, index_head_dim]`; `weights` is `[q_rows, index_heads]` and must
245/// already carry the `index_heads^-0.5` factor, as the golden generator applies it.
246pub fn index_scores(
247    q: &[f32],
248    weights: &[f32],
249    pools: &Pools,
250    dims: DsaDims,
251    q_rows: usize,
252) -> Vec<f32> {
253    let (h, d, p) = (dims.index_heads, dims.index_head_dim, pools.n_pools);
254    let scale = (d as f32).powf(-0.5);
255    let mut out = vec![0.0f32; q_rows * p];
256    for r in 0..q_rows {
257        for pp in 0..p {
258            let mut acc = 0.0f32;
259            for hh in 0..h {
260                let mut dot = 0.0f32;
261                for dd in 0..d {
262                    dot += q[(r * h + hh) * d + dd] * pools.keys[pp * d + dd];
263                }
264                // 2026-09-25: ReLU after the scale. `relu(s * x) == s * relu(x)` only because the
265                // scale is positive.
266                acc += weights[r * h + hh] * (scale * dot).max(0.0);
267            }
268            out[r * p + pp] = acc;
269        }
270    }
271    out
272}
273
274/// 2026-09-25: Which keys a query at `q_pos` may see: causal and not padding.
275pub fn visible(valid_keys: &[u8], q_pos: usize, key_idx: usize) -> bool {
276    key_idx <= q_pos && valid_keys[key_idx] != 0
277}
278
279/// 2026-09-25: Select up to `select_k` pools per query, ordered by higher score, then smaller
280/// pool index, so ties resolve the same way on every run. Invalid candidates score `f32::MIN`.
281pub fn topk_pools(
282    scores: &[f32],
283    valid_candidates: &[u8],
284    n_pools: usize,
285    q_rows: usize,
286    select_k: usize,
287) -> Vec<i32> {
288    let mut out = vec![INVALID; q_rows * select_k];
289    let mut buf: Vec<(f32, usize)> = Vec::with_capacity(n_pools);
290    for r in 0..q_rows {
291        buf.clear();
292        for p in 0..n_pools {
293            let s = if valid_candidates[r * n_pools + p] != 0 {
294                scores[r * n_pools + p]
295            } else {
296                f32::MIN
297            };
298            buf.push((s, p));
299        }
300        buf.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap().then(a.1.cmp(&b.1)));
301        for (j, (_, p)) in buf.iter().take(select_k).enumerate() {
302            out[r * select_k + j] = *p as i32;
303        }
304    }
305    out
306}
307
308/// 2026-09-25: Expand selected pools into raw token indices, append the visible tail, and pad with
309/// [`INVALID`]. Every row is written for all `out_width` entries: the output starts filled with
310/// [`INVALID`], and a padded query row stays that way.
311#[allow(clippy::too_many_arguments)]
312pub fn expand_selection(
313    selected: &[i32],
314    pools: &Pools,
315    valid_candidates: &[u8],
316    valid_keys: &[u8],
317    q_positions: &[usize],
318    q_mask: &[u8],
319    dims: DsaDims,
320    seq: usize,
321    select_k: usize,
322) -> Vec<i32> {
323    let (kp, width) = (dims.index_kpool, dims.out_width());
324    let q_rows = q_positions.len();
325    let mut out = vec![INVALID; q_rows * width];
326
327    let first_key = valid_keys.iter().position(|v| *v != 0).unwrap_or(seq) as i64;
328    for r in 0..q_rows {
329        let row = &mut out[r * width..(r + 1) * width];
330        if q_mask[r] == 0 {
331            continue;
332        }
333        let mut w = 0usize;
334        for j in 0..select_k {
335            let p = selected[r * select_k + j];
336            let ok = p >= 0 && valid_candidates[r * pools.n_pools + p as usize] != 0;
337            for s in 0..kp {
338                row[w] = if ok {
339                    pools.indices[p as usize * kp + s]
340                } else {
341                    INVALID
342                };
343                w += 1;
344            }
345        }
346        if dims.always_select_tail {
347            // 2026-09-25: The incomplete pool, as raw token indices.
348            let vis_count = (0..seq)
349                .filter(|k| visible(valid_keys, q_positions[r], *k))
350                .count();
351            let tail_count = vis_count % kp;
352            let tail_start = first_key + vis_count as i64 - tail_count as i64;
353            for t in 0..kp - 1 {
354                let idx = tail_start + t as i64;
355                let ok = t < tail_count
356                    && idx >= 0
357                    && (idx as usize) < seq
358                    && visible(valid_keys, q_positions[r], idx as usize);
359                row[w] = if ok { idx as i32 } else { INVALID };
360                w += 1;
361            }
362        }
363        debug_assert_eq!(
364            w,
365            width.min(select_k * kp + if dims.always_select_tail { kp - 1 } else { 0 })
366        );
367    }
368    out
369}
370
371/// 2026-09-25: Turn an index row into the boolean visibility mask the attention reads. Duplicate
372/// indices collapse, so a repeated token is attended once; out-of-range and [`INVALID`] entries
373/// are dropped.
374pub fn topk_to_mask(topk: &[i32], q_rows: usize, width: usize, kv_len: usize) -> Vec<u8> {
375    let mut mask = vec![0u8; q_rows * kv_len];
376    for r in 0..q_rows {
377        for j in 0..width {
378            let i = topk[r * width + j];
379            if i >= 0 && (i as usize) < kv_len {
380                mask[r * kv_len + i as usize] = 1;
381            }
382        }
383    }
384    mask
385}
386
387/// 2026-09-25: Public sigmoid; nothing in the DSA reference calls it.
388pub fn sigmoid_f32(x: f32) -> f32 {
389    sigmoid(x)
390}
391
392mod mla;
393pub use mla::{expand_kv, mla_masked_attention};
394
395#[cfg(test)]
396mod tests;
metrale_model_arch/glm5next_dsa_ref/mod.rs

metrale_model_arch/glm5next_dsa_ref/
mod.rs
