ferrox_models/
glm_dsa.rs

1//! GLM-5.2's real DSA (DeepSeek Sparse Attention) attention layer:
2//! RoPE-carrying MLA (`ferrox_models::mla`'s math, inlined here rather
3//! than reused directly -- see below) plus the lightning indexer
4//! (`ferrox_core::attention::lightning_indexer_topk`) selecting which
5//! causal positions are visible, then sparse attention restricted to
6//! exactly those (`ferrox_core::attention::causal_mla_attention_sparse`).
7//!
8//! Real tensor names/shapes/dispatch confirmed against llama.cpp PR
9//! #23346 (DeepSeek-V3.2, `src/models/deepseek32.cpp`) and PR #25407
10//! (GLM-5.2's `indexer_types`/interleaved-RoPE diff on top,
11//! `src/models/glm-dsa.cpp`), both fetched live and read line-by-line
12//! (`gh api -H "Accept: application/vnd.github.raw"
13//! repos/ggerganov/llama.cpp/contents/src/models/glm-dsa.cpp`, since
14//! `gh pr diff` alone doesn't show unchanged context for tensor
15//! creation that predates PR #25407) -- see docs/MODELS.md.
16//!
17//! Not the same weight layout as `ferrox_models::mla::MlaAttnWeights`:
18//! GLM-5.2's real GGUF main-attention K/V decompression uses separate
19//! per-head `wk_b`/`wv_b` 3D tensors (`blk.N.attn_k_b`/`attn_v_b`), not
20//! Kimi K3's combined `kv_b_proj` -- see [`Glm52AttnWeights`]'s doc
21//! comment for the "absorbed vs. un-absorbed" direction this matters
22//! for. That's why this is a new set of weight/forward structures
23//! rather than a reuse of `mla::MlaAttnWeights`/`mla_forward_token`,
24//! even though the underlying low-rank-compression math is the same
25//! family.
26//!
27//! Two more real, non-obvious facts from that source, beyond what
28//! `ferrox_models::mla`'s module doc comment already covers for the
29//! shared MLA math:
30//!
31//! 1. Per-layer full/shared indexer dispatch is **not a fixed period**.
32//!    GLM-5.2's real per-layer `indexer_types` array has layers 0-1 as
33//!    "full", then a repeating 1-full+3-shared pattern -- but the *real
34//!    mechanism* a "shared" layer uses is "reuse the top-k from the
35//!    nearest preceding full layer," not "recompute every 4th layer."
36//!    Confirmed directly from `glm-dsa.cpp`'s per-layer loop: a single
37//!    `prev_top_k` local variable is reassigned only when a full layer
38//!    runs, and carried forward unconditionally into every following
39//!    shared layer until the next full layer reassigns it --
40//!    `GGML_ASSERT(prev_top_k != nullptr && "shared indexer layer must
41//!    follow a previous full indexer layer")` on the shared-layer
42//!    branch confirms a shared layer can never be the first layer
43//!    processed. [`glm52_attn_forward_token`]'s `prev_top_k` parameter
44//!    mirrors this exactly: caller-threaded, per-token-forward-pass
45//!    scoped (reset to `None` at the start of each new token, the same
46//!    "one token, all layers in order" scope the real ggml local
47//!    variable has), not persisted across tokens.
48//! 2. The lightning indexer's own q/k split into rope/nope halves is
49//!    **rope-FIRST, nope-second** -- read directly from `glm-dsa.cpp`'s
50//!    `indexer_q_pe`/`indexer_q_nope` `ggml_view_3d` byte offsets
51//!    (`indexer_q_pe` at offset 0, `indexer_q_nope` at
52//!    `ggml_row_size(..., n_embd_indexer_head_nope)`), the **opposite**
53//!    of the main attention's nope-first/rope-second convention.
54//!    **GENUINE DISCLOSED CAVEAT**: GLM-5.2's own real `config.json`
55//!    gives `index_head_dim=128` with the indexer's rope portion
56//!    reusing the main attention's `n_rot()`=64, so
57//!    `nope_dim == rope_dim == 64` for this specific model -- meaning
58//!    this physical-order reading is **not numerically distinguishable
59//!    from its opposite** by inspecting GLM-5.2's own hyperparameters
60//!    alone (the offset expression's value is identical either way).
61//!    This implementation commits to the literal "first view in the
62//!    code is the rope view" reading rather than silently picking
63//!    whichever seemed more consistent with the main attention's
64//!    convention; see docs/MODELS.md for the same caveat recorded
65//!    against the evidence ledger.
66//!
67//! Tested here against synthetic weights, cross-validated against an
68//! independent Python transcription of one "full" indexer layer's
69//! RoPE+indexer+sparse
70//! math across four decode steps, including a step where top-k
71//! sparsity actually excludes a causally-visible position) plus
72//! dedicated Rust-only tests for the full/shared dispatch bookkeeping
73//! itself (not real "math to get subtly wrong" the way RoPE/
74//! indexer-scoring/sparse-selection are, so not re-derived in Python).
75
76use ferrox_core::attention::{
77    apply_rope_interleaved, causal_mla_attention_sparse, lightning_indexer_topk,
78};
79use ferrox_core::matmul::{layer_norm, rms_norm};
80use ferrox_core::weight_matrix::WeightMatrix;
81
82use crate::config::MlaRopeConfig;
83
84/// GLM-5.2's real per-layer MLA hyperparameters. Unlike
85/// `ferrox_models::mla::MlaConfig`, there is no `use_output_gate` (real
86/// GLM-5.2 tensor list has no Kimi-K3-style `attn_gate` equivalent) and
87/// rope is unconditional, not `Option` (every real GLM-5.2 layer's main
88/// attention applies it -- `rope_interleave: true` in its
89/// `config.json`, no per-layer exception unlike the indexer's
90/// full/shared split).
91#[derive(Debug, Clone)]
92pub struct Glm52MlaConfig {
93    pub num_heads: usize,
94    pub q_lora_rank: usize,
95    pub kv_lora_rank: usize,
96    pub qk_nope_head_dim: usize,
97    pub qk_rope_head_dim: usize,
98    pub v_head_dim: usize,
99    pub rope: MlaRopeConfig,
100}
101
102/// GLM-5.2's real lightning-indexer hyperparameters
103/// (`index_head_dim`=128, `index_n_heads`=32, `index_topk`=2048 in the
104/// real published `config.json` -- see docs/MODELS.md; kept
105/// generic here, not hardcoded, so small synthetic tests can use tiny
106/// values).
107#[derive(Debug, Clone)]
108pub struct IndexerConfig {
109    pub n_heads: usize,
110    /// Total per-(shared-)head dimension, rope + nope.
111    pub head_dim: usize,
112    /// The rope portion of `head_dim` -- real GLM-5.2 reuses the main
113    /// attention's `n_rot()` for this (see module doc comment point 2).
114    pub rope_dim: usize,
115    pub top_k: usize,
116    pub rope_theta: f32,
117}
118
119/// The real lightning indexer's weights for one layer (only present on
120/// "full" indexer layers -- see [`glm52_attn_forward_token`]'s
121/// `is_full_indexer_layer` parameter).
122pub struct IndexerWeights {
123    pub k_norm_weight: Vec<f32>,
124    pub k_norm_bias: Vec<f32>,
125    /// `n_embd -> n_heads` (one scalar weight per indexer head).
126    pub proj: WeightMatrix,
127    /// `n_embd -> head_dim`: a single shared (MQA-style) key, not
128    /// per-head -- real `indexer.attn_k` tensor shape confirms this
129    /// (no `n_heads` factor), matching Kimi K3's MLA `k_rot` MQA
130    /// pattern (`ferrox_models::mla`'s module doc comment point 2).
131    pub attn_k: WeightMatrix,
132    /// `q_lora_rank -> n_heads*head_dim`.
133    pub attn_q_b: WeightMatrix,
134}
135
136/// GLM-5.2's real per-layer MLA+indexer attention weights.
137pub struct Glm52AttnWeights {
138    pub q_a_proj: WeightMatrix,
139    pub q_a_layernorm: Vec<f32>,
140    pub q_b_proj: WeightMatrix,
141    pub kv_a_proj_with_mqa: WeightMatrix,
142    pub kv_a_layernorm: Vec<f32>,
143    /// Per-head, in the DECOMPRESSION direction (`kv_lora_rank ->
144    /// qk_nope_head_dim`) -- i.e. **already transposed** from the real
145    /// GGUF's on-disk `attn_k_b` layout, which llama.cpp instead
146    /// applies in the opposite ("absorbed": `qk_nope_head_dim ->
147    /// kv_lora_rank`, pre-multiplying the *query* rather than
148    /// decompressing K) direction as a compute optimization --
149    /// mathematically equivalent, just not the form llama.cpp actually
150    /// executes. `glm52_gguf_loader` performs this transpose at load
151    /// time so this module can reuse the same direct-decompression
152    /// math `ferrox_models::mla` already has, via
153    /// `causal_mla_attention_sparse`, instead of implementing a second,
154    /// absorbed-computation attention primitive. `wk_b[h]` is
155    /// `[qk_nope_head_dim, kv_lora_rank]`.
156    pub wk_b: Vec<WeightMatrix>,
157    /// Per-head, natively in the decompression direction
158    /// (`kv_lora_rank -> v_head_dim`) already -- no transpose needed
159    /// (real GGUF `attn_v_b` is used by llama.cpp to decompress V
160    /// directly too, unlike `wk_b`). `wv_b[h]` is
161    /// `[v_head_dim, kv_lora_rank]`.
162    pub wv_b: Vec<WeightMatrix>,
163    pub o_proj: WeightMatrix,
164    /// Only present on "full" indexer layers -- see
165    /// [`glm52_attn_forward_token`]'s doc comment.
166    pub indexer: Option<IndexerWeights>,
167}
168
169/// Growable per-layer decode state: the main K/V cache (same shape
170/// convention as `ferrox_models::mla`'s `k_cache`/`v_cache`) plus,
171/// for "full" indexer layers only, the indexer's own K cache (a
172/// separate cache since `IndexerConfig::head_dim` generally differs
173/// from the main attention's `qk_nope_head_dim + qk_rope_head_dim`).
174/// "Shared" layers never touch `indexer_k_cache` -- they don't run
175/// their own indexer at all, see module doc comment point 1.
176#[derive(Debug, Clone, Default)]
177pub struct Glm52AttnState {
178    pub k_cache: Vec<f32>,
179    pub v_cache: Vec<f32>,
180    pub indexer_k_cache: Vec<f32>,
181}
182
183impl Glm52AttnState {
184    pub fn new() -> Self {
185        Self::default()
186    }
187}
188
189/// One decode step for one GLM-5.2 DSA attention layer.
190///
191/// `is_full_indexer_layer` selects which of the two real per-layer
192/// behaviors this call uses (see module doc comment point 1):
193/// - `true` ("full"): `weights.indexer` must be `Some`; this call
194///   computes a fresh top-k from the indexer, appending to
195///   `state.indexer_k_cache`, and writes the result into `*prev_top_k`
196///   for later "shared" layers *within this same token's forward
197///   pass* to reuse.
198/// - `false` ("shared"): reuses `prev_top_k` as-is (must already be
199///   `Some` -- the real architecture guarantees the first layer
200///   processed is always "full," matching llama.cpp's
201///   `GGML_ASSERT(prev_top_k != nullptr ...)`); `weights.indexer` is
202///   ignored (may be `None`).
203///
204/// `prev_top_k` is caller-owned and must be reset to `None` at the
205/// start of every new token's forward pass across all layers (not
206/// persisted token-to-token) -- mirroring the real ggml local variable
207/// of the same name and scope in `glm-dsa.cpp`'s per-token graph build.
208#[allow(clippy::too_many_arguments)]
209pub fn glm52_attn_forward_token(
210    weights: &Glm52AttnWeights,
211    cfg: &Glm52MlaConfig,
212    indexer_cfg: &IndexerConfig,
213    hidden: &[f32],
214    rms_norm_eps: f32,
215    is_full_indexer_layer: bool,
216    state: &mut Glm52AttnState,
217    prev_top_k: &mut Option<Vec<usize>>,
218) -> Vec<f32> {
219    let q_head_dim = cfg.qk_nope_head_dim + cfg.qk_rope_head_dim;
220    let pos = state.k_cache.len() / (cfg.num_heads * q_head_dim);
221
222    let q_a = weights.q_a_proj.apply(hidden);
223    let qr = rms_norm(&q_a, &weights.q_a_layernorm, rms_norm_eps);
224
225    let visible: Vec<usize> = if is_full_indexer_layer {
226        let indexer = weights
227            .indexer
228            .as_ref()
229            .expect("a \"full\" indexer layer must carry indexer weights");
230        let idx_head_dim = indexer_cfg.head_dim;
231        let idx_rope_dim = indexer_cfg.rope_dim;
232
233        let indexer_q_full = indexer.attn_q_b.apply(&qr); // [n_heads*head_dim]
234        let indexer_q: Vec<Vec<f32>> = (0..indexer_cfg.n_heads)
235            .map(|h| {
236                let head_slice = &indexer_q_full[h * idx_head_dim..(h + 1) * idx_head_dim];
237                // rope-first, nope-second -- see module doc comment point 2.
238                let mut rope_part = head_slice[..idx_rope_dim].to_vec();
239                apply_rope_interleaved(&mut rope_part, pos, indexer_cfg.rope_theta);
240                rope_part.extend_from_slice(&head_slice[idx_rope_dim..]);
241                rope_part
242            })
243            .collect();
244
245        let indexer_k_raw = indexer.attn_k.apply(hidden); // [head_dim], shared (MQA)
246        let indexer_k_normed = layer_norm(
247            &indexer_k_raw,
248            &indexer.k_norm_weight,
249            &indexer.k_norm_bias,
250            rms_norm_eps,
251        );
252        let mut idx_rope_part = indexer_k_normed[..idx_rope_dim].to_vec();
253        apply_rope_interleaved(&mut idx_rope_part, pos, indexer_cfg.rope_theta);
254        idx_rope_part.extend_from_slice(&indexer_k_normed[idx_rope_dim..]);
255        state.indexer_k_cache.extend_from_slice(&idx_rope_part);
256
257        let idx_seq_len = state.indexer_k_cache.len() / idx_head_dim;
258        let indexer_keys: Vec<Vec<f32>> = (0..idx_seq_len)
259            .map(|t| state.indexer_k_cache[t * idx_head_dim..(t + 1) * idx_head_dim].to_vec())
260            .collect();
261
262        let indexer_weights_vec = indexer.proj.apply(hidden); // [n_heads]
263        let top_k = lightning_indexer_topk(
264            &indexer_q,
265            &indexer_keys,
266            &indexer_weights_vec,
267            indexer_cfg.top_k,
268        );
269        *prev_top_k = Some(top_k.clone());
270        top_k
271    } else {
272        prev_top_k
273            .clone()
274            .expect("a \"shared\" indexer layer must follow a previous \"full\" layer's top-k within this token's forward pass")
275    };
276
277    let mut query = weights.q_b_proj.apply(&qr); // [n_heads*q_head_dim], nope first then rope
278    for h in 0..cfg.num_heads {
279        let q_rot_h = &mut query[h * q_head_dim + cfg.qk_nope_head_dim..(h + 1) * q_head_dim];
280        apply_rope_interleaved(q_rot_h, pos, cfg.rope.theta);
281    }
282
283    let kv_cmpr_pe = weights.kv_a_proj_with_mqa.apply(hidden);
284    let (kv_cmpr_raw, k_pe_raw) = kv_cmpr_pe.split_at(cfg.kv_lora_rank);
285    let mut k_pe = k_pe_raw.to_vec();
286    apply_rope_interleaved(&mut k_pe, pos, cfg.rope.theta);
287    let kv_cmpr = rms_norm(kv_cmpr_raw, &weights.kv_a_layernorm, rms_norm_eps);
288
289    let mut key_step = vec![0f32; cfg.num_heads * q_head_dim];
290    let mut value_step = vec![0f32; cfg.num_heads * cfg.v_head_dim];
291    for h in 0..cfg.num_heads {
292        let k_pass = weights.wk_b[h].apply(&kv_cmpr); // [qk_nope_head_dim]
293        let v_h = weights.wv_b[h].apply(&kv_cmpr); // [v_head_dim]
294
295        let key_h = &mut key_step[h * q_head_dim..(h + 1) * q_head_dim];
296        key_h[..cfg.qk_nope_head_dim].copy_from_slice(&k_pass);
297        key_h[cfg.qk_nope_head_dim..].copy_from_slice(&k_pe);
298
299        value_step[h * cfg.v_head_dim..(h + 1) * cfg.v_head_dim].copy_from_slice(&v_h);
300    }
301
302    state.k_cache.extend_from_slice(&key_step);
303    state.v_cache.extend_from_slice(&value_step);
304    let seq_len = state.k_cache.len() / (cfg.num_heads * q_head_dim);
305
306    let attn_out = causal_mla_attention_sparse(
307        &query,
308        &state.k_cache,
309        &state.v_cache,
310        cfg.num_heads,
311        q_head_dim,
312        cfg.v_head_dim,
313        seq_len,
314        &visible,
315    );
316
317    weights.o_proj.apply(&attn_out)
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323    use crate::config::MlaRopeConfig;
324    use ferrox_core::tensor::Tensor;
325
326    fn wm(data: &[f32], rows: usize, cols: usize) -> WeightMatrix {
327        assert_eq!(data.len(), rows * cols);
328        WeightMatrix::F32(Tensor::new(data.to_vec(), vec![rows, cols]))
329    }
330
331    const HIDDEN_SIZE: usize = 8;
332    const NUM_HEADS: usize = 2;
333    const QK_NOPE_HEAD_DIM: usize = 4;
334    const QK_ROPE_HEAD_DIM: usize = 4;
335    const KV_LORA_RANK: usize = 4;
336    const Q_LORA_RANK: usize = 6;
337    const V_HEAD_DIM: usize = 3;
338    const EPS: f32 = 1e-5;
339    const ROPE_THETA: f32 = 10000.0;
340
341    const IDX_N_HEADS: usize = 2;
342    const IDX_ROPE_DIM: usize = 2;
343    const IDX_NOPE_DIM: usize = 2;
344    const IDX_HEAD_DIM: usize = IDX_ROPE_DIM + IDX_NOPE_DIM;
345    const TOP_K: usize = 2;
346
347    // Generated by an independent Python reference -- do not hand-edit.
348    const GLM_Q_A_PROJ: [f32; 48] = [
349        -0.0282433, -0.200786, 0.229784, 0.162609, 0.161474, 0.190923, -0.28271, -0.259493,
350        -0.299452, -0.445281, 0.438341, 0.223402, -0.114472, 0.108293, 0.138669, -0.0983588,
351        -0.139929, -0.287077, 0.233259, -0.0924404, -0.48316, 0.236404, 0.487659, -0.22182,
352        -0.425881, 0.474637, 0.224366, 0.227792, 0.155548, 0.139447, -0.0826791, -0.435336,
353        0.438776, -0.328727, 0.415205, -0.280294, 0.25502, -0.264693, -0.374199, 0.626933,
354        0.175746, 0.198436, -0.0882081, -0.131098, 0.416386, 0.0938253, -0.086288, -0.173465,
355    ];
356    const GLM_Q_A_LAYERNORM_W: [f32; 6] = [0.983826, 0.879749, 1.11838, 1.05149, 1.1818, 0.834347];
357    const GLM_Q_B_PROJ: [f32; 96] = [
358        0.291994,
359        -0.728247,
360        -0.316577,
361        0.0813452,
362        -0.266261,
363        0.389674,
364        0.0833719,
365        0.351552,
366        0.356938,
367        0.350334,
368        -0.0473987,
369        -0.0266404,
370        -0.169264,
371        0.0701104,
372        0.0207743,
373        0.44759,
374        0.372409,
375        0.283663,
376        0.161893,
377        0.0691206,
378        0.164776,
379        -0.159844,
380        0.244357,
381        0.254148,
382        0.781266,
383        -0.0410461,
384        0.00448047,
385        0.167438,
386        0.134256,
387        -0.117364,
388        0.613949,
389        -0.207111,
390        0.42746,
391        0.45351,
392        0.237126,
393        -0.50974,
394        0.328859,
395        -0.250491,
396        -0.356138,
397        0.122879,
398        0.254109,
399        0.120117,
400        -0.244618,
401        0.090442,
402        0.572282,
403        -0.175117,
404        0.150304,
405        0.127176,
406        -0.230927,
407        -0.181049,
408        0.0503238,
409        -0.252932,
410        -0.00813607,
411        -0.169141,
412        0.178562,
413        -0.172518,
414        -0.163208,
415        -0.286795,
416        0.358209,
417        0.355661,
418        -0.0321808,
419        0.025399,
420        -0.227651,
421        -0.0153813,
422        -0.0254572,
423        -0.364581,
424        -0.450488,
425        0.155816,
426        0.0033226,
427        0.481021,
428        -0.000260049,
429        -0.230117,
430        -0.0422523,
431        0.269254,
432        -0.225551,
433        -0.265757,
434        -0.192519,
435        -0.300859,
436        -0.152023,
437        0.31445,
438        -0.229592,
439        -0.417754,
440        -0.219984,
441        0.0230321,
442        0.162062,
443        -0.162489,
444        -0.504785,
445        0.117479,
446        -0.152083,
447        0.203557,
448        -0.232979,
449        -0.537171,
450        -0.131909,
451        -0.0782392,
452        0.187798,
453        -0.364894,
454    ];
455    const GLM_KV_A_PROJ: [f32; 64] = [
456        -0.0108888,
457        0.122853,
458        -0.388147,
459        -0.320502,
460        0.288834,
461        0.0587081,
462        0.0027565,
463        -0.0303023,
464        0.252204,
465        0.103756,
466        -0.563955,
467        0.539224,
468        -0.515732,
469        0.475067,
470        -0.179422,
471        0.512039,
472        0.564391,
473        -0.309453,
474        0.178157,
475        0.4829,
476        0.304922,
477        -0.327155,
478        0.235531,
479        -0.223351,
480        0.113775,
481        -0.326219,
482        0.129363,
483        0.343847,
484        -0.555633,
485        -0.00371874,
486        -0.480022,
487        0.0622793,
488        0.121396,
489        0.902273,
490        -0.271857,
491        -0.0787809,
492        -0.148056,
493        -0.246381,
494        -0.388923,
495        -0.326308,
496        0.754771,
497        -0.188557,
498        0.157124,
499        -0.242718,
500        -0.196856,
501        0.168396,
502        0.116464,
503        0.406121,
504        -0.0524445,
505        -0.226537,
506        -0.220791,
507        -0.42747,
508        -0.109609,
509        0.327875,
510        0.238249,
511        0.262922,
512        0.0603609,
513        0.259383,
514        -0.125942,
515        0.0253563,
516        -0.672037,
517        -0.0822506,
518        -0.313883,
519        -0.079927,
520    ];
521    const GLM_KV_A_LAYERNORM_W: [f32; 4] = [1.11964, 1.00794, 0.927425, 0.974059];
522    const GLM_WK_B_0: [f32; 16] = [
523        0.152348, 0.223828, 0.104805, 0.146991, 0.167048, 0.0596911, 0.13428, -0.165432, 0.185469,
524        -0.438659, 0.0505363, -0.378193, -0.367254, 0.291605, 0.25605, -0.278039,
525    ];
526    const GLM_WV_B_0: [f32; 12] = [
527        -0.0451786, 0.016892, -0.503596, 0.38556, 0.0672434, -0.345241, 0.261026, 0.113356,
528        0.195666, 0.124068, 0.169155, 0.0241996,
529    ];
530    const GLM_WK_B_1: [f32; 16] = [
531        0.32183, 0.077576, 0.657456, -0.210592, -0.241762, -0.113545, 0.305935, -0.142537,
532        -0.131723, -0.0698308, -0.231153, 0.0832406, -0.184562, -0.395013, -0.434206, 0.643208,
533    ];
534    const GLM_WV_B_1: [f32; 12] = [
535        0.0460725, -0.199187, 0.525018, 0.704311, 0.214609, 0.155865, -0.13945, -0.361349,
536        0.200727, 0.669041, -0.35694, 0.405635,
537    ];
538    const GLM_O_PROJ: [f32; 48] = [
539        0.156636,
540        0.435755,
541        0.254836,
542        -0.28038,
543        -0.00686566,
544        0.254093,
545        0.13879,
546        0.298608,
547        -0.654407,
548        0.544604,
549        -0.40823,
550        0.557235,
551        -0.401607,
552        0.0393622,
553        -0.0108063,
554        -0.425778,
555        -0.0790213,
556        0.183181,
557        0.770074,
558        0.431033,
559        -0.191665,
560        -0.321149,
561        -0.243943,
562        -0.0704616,
563        0.180775,
564        -0.216385,
565        0.0824125,
566        -0.320591,
567        -0.182163,
568        -0.0257085,
569        -0.0184709,
570        0.292862,
571        -0.215734,
572        0.652291,
573        -0.0461593,
574        0.249014,
575        -0.205017,
576        0.0634068,
577        0.087137,
578        0.529326,
579        0.477227,
580        0.171185,
581        0.0539693,
582        0.0189488,
583        -0.138254,
584        -0.173556,
585        0.65771,
586        -0.0616593,
587    ];
588    const GLM_INDEXER_PROJ: [f32; 16] = [
589        -0.087587, -0.0520619, 0.169093, -0.459473, -0.354101, -0.151364, -0.238526, 0.783742,
590        0.0472875, -0.286511, -0.0857637, 0.0416524, 0.383631, 0.309099, -0.0708246, -0.45329,
591    ];
592    const GLM_INDEXER_ATTN_K: [f32; 32] = [
593        -0.0243618, 0.264064, 0.116963, 0.0329414, 0.16249, 0.374711, 0.0555387, 0.011557,
594        0.720572, -0.288658, 0.660557, -0.0520769, -0.398689, 0.10919, -0.21608, -0.272101,
595        0.324867, -0.0406344, 0.436622, -0.353795, 0.30306, -0.256956, 0.381955, 0.346494,
596        0.579363, 0.0492111, 0.0489935, -0.110262, -0.359962, 0.0704509, 0.402797, -0.121248,
597    ];
598    const GLM_INDEXER_ATTN_Q_B: [f32; 48] = [
599        0.423031,
600        0.0939434,
601        0.165178,
602        -0.337296,
603        -0.0761678,
604        -0.0724417,
605        -0.485867,
606        0.693554,
607        0.201884,
608        0.544458,
609        -0.0574976,
610        0.2739,
611        0.174839,
612        0.602076,
613        -0.00910648,
614        0.12034,
615        -0.0199685,
616        0.238296,
617        0.0537028,
618        0.221498,
619        -0.0995168,
620        0.0551425,
621        0.126781,
622        0.14289,
623        0.0625852,
624        0.14437,
625        0.00611794,
626        0.166845,
627        -0.566204,
628        0.132155,
629        -0.292038,
630        -0.203268,
631        -0.226073,
632        0.2511,
633        -0.339645,
634        0.0101247,
635        0.0409607,
636        -0.369437,
637        0.101741,
638        -0.517995,
639        0.150119,
640        0.19386,
641        0.115472,
642        -0.196285,
643        -0.0137043,
644        0.291252,
645        -0.0213699,
646        0.00902874,
647    ];
648    const GLM_INDEXER_K_NORM_W: [f32; 4] = [1.01802, 0.864429, 0.781694, 1.0421];
649    const GLM_INDEXER_K_NORM_B: [f32; 4] = [-0.00865097, -0.0181871, -0.0141045, 0.0431985];
650
651    const GLM_HIDDEN_0: [f32; 8] = [
652        0.173677, 0.403207, -0.819441, 0.0385262, -0.382448, 0.0666326, 0.107605, -0.50578,
653    ];
654    const GLM_HIDDEN_1: [f32; 8] = [
655        0.0557989, -0.774157, 0.586028, -0.0179948, 0.00680277, 0.436953, 0.188515, -0.2977,
656    ];
657    const GLM_HIDDEN_2: [f32; 8] = [
658        -0.718918, 1.18182, -0.484152, -0.465392, 0.221702, -0.438492, -0.388645, -0.164228,
659    ];
660    const GLM_HIDDEN_3: [f32; 8] = [
661        0.650653, -0.256895, -0.94904, 0.929269, -0.585101, 0.0973952, 0.140374, 0.190254,
662    ];
663
664    const GLM_GOLDEN_OUT_0: [f32; 8] = [
665        0.361853, -0.250738, 0.433283, -0.278145, 0.27569, -0.414117, 0.121914, 0.401593,
666    ];
667    const GLM_GOLDEN_OUT_1: [f32; 8] = [
668        -0.013624, 0.187721, -0.050068, -0.116258, -0.0672519, 0.179838, 0.13011, -0.114857,
669    ];
670    const GLM_GOLDEN_OUT_2: [f32; 8] = [
671        0.28988, -0.593942, 0.388018, 0.0189887, 0.317219, -0.64959, -0.143766, 0.524578,
672    ];
673    const GLM_GOLDEN_OUT_3: [f32; 8] = [
674        0.212341, 0.238018, 0.206285, -0.35234, 0.114026, 0.0278299, 0.262389, 0.0396803,
675    ];
676
677    const GLM_GOLDEN_VISIBLE_0: [usize; 1] = [0];
678    const GLM_GOLDEN_VISIBLE_1: [usize; 2] = [0, 1];
679    const GLM_GOLDEN_VISIBLE_2: [usize; 2] = [0, 2];
680    const GLM_GOLDEN_VISIBLE_3: [usize; 2] = [0, 3];
681
682    fn cfg() -> Glm52MlaConfig {
683        Glm52MlaConfig {
684            num_heads: NUM_HEADS,
685            q_lora_rank: Q_LORA_RANK,
686            kv_lora_rank: KV_LORA_RANK,
687            qk_nope_head_dim: QK_NOPE_HEAD_DIM,
688            qk_rope_head_dim: QK_ROPE_HEAD_DIM,
689            v_head_dim: V_HEAD_DIM,
690            rope: MlaRopeConfig { theta: ROPE_THETA },
691        }
692    }
693
694    fn indexer_cfg() -> IndexerConfig {
695        IndexerConfig {
696            n_heads: IDX_N_HEADS,
697            head_dim: IDX_HEAD_DIM,
698            rope_dim: IDX_ROPE_DIM,
699            top_k: TOP_K,
700            rope_theta: ROPE_THETA,
701        }
702    }
703
704    fn make_weights() -> Glm52AttnWeights {
705        let q_head_dim = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM;
706        Glm52AttnWeights {
707            q_a_proj: wm(&GLM_Q_A_PROJ, Q_LORA_RANK, HIDDEN_SIZE),
708            q_a_layernorm: GLM_Q_A_LAYERNORM_W.to_vec(),
709            q_b_proj: wm(&GLM_Q_B_PROJ, NUM_HEADS * q_head_dim, Q_LORA_RANK),
710            kv_a_proj_with_mqa: wm(&GLM_KV_A_PROJ, KV_LORA_RANK + QK_ROPE_HEAD_DIM, HIDDEN_SIZE),
711            kv_a_layernorm: GLM_KV_A_LAYERNORM_W.to_vec(),
712            wk_b: vec![
713                wm(&GLM_WK_B_0, QK_NOPE_HEAD_DIM, KV_LORA_RANK),
714                wm(&GLM_WK_B_1, QK_NOPE_HEAD_DIM, KV_LORA_RANK),
715            ],
716            wv_b: vec![
717                wm(&GLM_WV_B_0, V_HEAD_DIM, KV_LORA_RANK),
718                wm(&GLM_WV_B_1, V_HEAD_DIM, KV_LORA_RANK),
719            ],
720            o_proj: wm(&GLM_O_PROJ, HIDDEN_SIZE, NUM_HEADS * V_HEAD_DIM),
721            indexer: Some(IndexerWeights {
722                k_norm_weight: GLM_INDEXER_K_NORM_W.to_vec(),
723                k_norm_bias: GLM_INDEXER_K_NORM_B.to_vec(),
724                proj: wm(&GLM_INDEXER_PROJ, IDX_N_HEADS, HIDDEN_SIZE),
725                attn_k: wm(&GLM_INDEXER_ATTN_K, IDX_HEAD_DIM, HIDDEN_SIZE),
726                attn_q_b: wm(
727                    &GLM_INDEXER_ATTN_Q_B,
728                    IDX_N_HEADS * IDX_HEAD_DIM,
729                    Q_LORA_RANK,
730                ),
731            }),
732        }
733    }
734
735    #[test]
736    fn full_layer_matches_independent_python_reference_across_four_decode_steps() {
737        let weights = make_weights();
738        let cfg = cfg();
739        let idx_cfg = indexer_cfg();
740        let mut state = Glm52AttnState::new();
741        let mut prev_top_k: Option<Vec<usize>> = None;
742
743        let hiddens = [
744            &GLM_HIDDEN_0[..],
745            &GLM_HIDDEN_1[..],
746            &GLM_HIDDEN_2[..],
747            &GLM_HIDDEN_3[..],
748        ];
749        let goldens = [
750            &GLM_GOLDEN_OUT_0[..],
751            &GLM_GOLDEN_OUT_1[..],
752            &GLM_GOLDEN_OUT_2[..],
753            &GLM_GOLDEN_OUT_3[..],
754        ];
755        let golden_visible: [&[usize]; 4] = [
756            &GLM_GOLDEN_VISIBLE_0,
757            &GLM_GOLDEN_VISIBLE_1,
758            &GLM_GOLDEN_VISIBLE_2,
759            &GLM_GOLDEN_VISIBLE_3,
760        ];
761
762        for (pos, ((hidden, golden), visible)) in hiddens
763            .iter()
764            .zip(goldens.iter())
765            .zip(golden_visible.iter())
766            .enumerate()
767        {
768            // Every layer here is "full" -- this test is specifically
769            // about the RoPE+indexer+sparse math, not the full/shared
770            // dispatch bookkeeping (covered separately below).
771            let out = glm52_attn_forward_token(
772                &weights,
773                &cfg,
774                &idx_cfg,
775                hidden,
776                EPS,
777                true,
778                &mut state,
779                &mut prev_top_k,
780            );
781            assert_eq!(out.len(), golden.len());
782            for (i, (a, b)) in out.iter().zip(golden.iter()).enumerate() {
783                assert!(
784                    (a - b).abs() < 1e-3,
785                    "position {pos} element {i}: rust={a} python={b}"
786                );
787            }
788            // Cross-check the indexer's own top-k selection against the
789            // Python reference too -- this is the sparsity mechanism
790            // itself, not just its downstream numerical effect.
791            assert_eq!(
792                prev_top_k.as_deref(),
793                Some(*visible),
794                "position {pos}: indexer top-k selection mismatch"
795            );
796        }
797    }
798
799    #[test]
800    #[should_panic(expected = "must follow a previous")]
801    fn shared_layer_without_a_preceding_full_layer_panics() {
802        let weights = make_weights();
803        let cfg = cfg();
804        let idx_cfg = indexer_cfg();
805        let mut state = Glm52AttnState::new();
806        let mut prev_top_k: Option<Vec<usize>> = None;
807
808        // First layer processed is "shared" with no prior "full" layer
809        // in this token's forward pass -- the real architecture's
810        // GGML_ASSERT guarantees this can never happen for a real
811        // checkpoint (layer 0 is always "full"), so this must panic
812        // loudly rather than silently produce wrong output.
813        glm52_attn_forward_token(
814            &weights,
815            &cfg,
816            &idx_cfg,
817            &GLM_HIDDEN_0,
818            EPS,
819            false,
820            &mut state,
821            &mut prev_top_k,
822        );
823    }
824
825    #[test]
826    fn shared_layer_reuses_the_nearest_preceding_full_layers_top_k_exactly() {
827        // Three layers processed for the same token, in order:
828        // full (layer A) -> shared (layer B) -> shared (layer C).
829        // Real semantics: B and C must both reuse EXACTLY layer A's
830        // top-k, not recompute their own -- confirmed structurally by
831        // checking `prev_top_k` is untouched by B/C's calls (a "shared"
832        // layer with `weights.indexer: None` would panic on `.expect()`
833        // inside the `is_full_indexer_layer` branch if it ever tried to
834        // compute its own top-k instead of reusing `prev_top_k`).
835        let mut weights = make_weights();
836        let cfg = cfg();
837        let idx_cfg = indexer_cfg();
838        let mut state = Glm52AttnState::new();
839        let mut prev_top_k: Option<Vec<usize>> = None;
840
841        // Layer A: full.
842        glm52_attn_forward_token(
843            &weights,
844            &cfg,
845            &idx_cfg,
846            &GLM_HIDDEN_0,
847            EPS,
848            true,
849            &mut state,
850            &mut prev_top_k,
851        );
852        let top_k_after_full = prev_top_k.clone();
853        assert!(top_k_after_full.is_some());
854
855        // Now strip the indexer weights entirely -- a "shared" layer
856        // must never touch them, so this proves it by construction:
857        // if the shared-layer code path ever dereferenced
858        // `weights.indexer`, this would panic on the `.expect()` inside
859        // the (never-taken, for `is_full_indexer_layer=false`) full
860        // branch's `Option::as_ref().expect(...)` -- but that branch is
861        // gated on `is_full_indexer_layer`, so it's simply never
862        // reached.
863        weights.indexer = None;
864
865        // Layer B: shared, reusing layer A's cached indexer state
866        // implicitly via `prev_top_k` (its own `state.indexer_k_cache`
867        // is a per-attention-instance field in this test setup, reused
868        // across calls since it's the same `state`/`weights` -- a real
869        // decoder would have one `Glm52AttnState` per *layer*, not
870        // shared across layers the way this test reuses one for
871        // brevity, since each layer has its own main K/V cache; what's
872        // being verified here is purely the `prev_top_k` reuse
873        // contract, which is decoder-scoped, not per-layer-state
874        // scoped).
875        glm52_attn_forward_token(
876            &weights,
877            &cfg,
878            &idx_cfg,
879            &GLM_HIDDEN_1,
880            EPS,
881            false,
882            &mut state,
883            &mut prev_top_k,
884        );
885        assert_eq!(
886            prev_top_k, top_k_after_full,
887            "shared layer B must not alter prev_top_k"
888        );
889
890        // Layer C: shared again, still reusing layer A's top-k.
891        glm52_attn_forward_token(
892            &weights,
893            &cfg,
894            &idx_cfg,
895            &GLM_HIDDEN_2,
896            EPS,
897            false,
898            &mut state,
899            &mut prev_top_k,
900        );
901        assert_eq!(
902            prev_top_k, top_k_after_full,
903            "shared layer C must not alter prev_top_k"
904        );
905    }
906
907    #[test]
908    fn full_layer_grows_the_indexer_cache_but_shared_layer_does_not() {
909        let mut weights = make_weights();
910        let cfg = cfg();
911        let idx_cfg = indexer_cfg();
912        let mut state = Glm52AttnState::new();
913        let mut prev_top_k: Option<Vec<usize>> = None;
914
915        glm52_attn_forward_token(
916            &weights,
917            &cfg,
918            &idx_cfg,
919            &GLM_HIDDEN_0,
920            EPS,
921            true,
922            &mut state,
923            &mut prev_top_k,
924        );
925        let len_after_full = state.indexer_k_cache.len();
926        assert_eq!(len_after_full, IDX_HEAD_DIM);
927
928        weights.indexer = None;
929        glm52_attn_forward_token(
930            &weights,
931            &cfg,
932            &idx_cfg,
933            &GLM_HIDDEN_1,
934            EPS,
935            false,
936            &mut state,
937            &mut prev_top_k,
938        );
939        assert_eq!(
940            state.indexer_k_cache.len(),
941            len_after_full,
942            "a \"shared\" layer must not grow the indexer K cache"
943        );
944    }
945}
ferrox_models/glm_dsa.rs

ferrox_models/
glm_dsa.rs
