metrale_model_arch/glm5next_dsa_ref/
mla.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The NoPE MLA half of the DSA reference: `expand_kv` and `mla_masked_attention`.
4//!
5//! Owner: model-arch (GLM-5.3-Flash DSA reference).
6//! Invariants: none beyond the types.
7
8use super::*;
9
10/// 2026-09-25: NoPE MLA over a per-query selected key set.
11///
12/// `q`: `[q_rows, heads, qk_head_dim]`, `k`: `[kv_len, heads, qk_head_dim]`,
13/// `v`: `[kv_len, heads, v_head_dim]`, `mask`: `[q_rows, kv_len]`.
14///
15/// The scale is `qk_head_dim^-0.5` over the full head dim; it equals `qk_nope_head_dim^-0.5` only
16/// while the rope part is zero-width.
17pub fn mla_masked_attention(
18    q: &[f32],
19    k: &[f32],
20    v: &[f32],
21    mask: &[u8],
22    dims: DsaDims,
23    q_rows: usize,
24    kv_len: usize,
25) -> Vec<f32> {
26    let (h, qd, vd) = (dims.heads, dims.qk_head_dim(), dims.v_head_dim);
27    let scale = (qd as f32).powf(-0.5);
28    let mut out = vec![0.0f32; q_rows * h * vd];
29    for r in 0..q_rows {
30        for hh in 0..h {
31            // 2026-09-25: Online softmax: one pass, no `[kv_len]` score buffer.
32            let mut m = f32::NEG_INFINITY;
33            let mut l = 0.0f32;
34            let mut acc = vec![0.0f32; vd];
35            for kk in 0..kv_len {
36                if mask[r * kv_len + kk] == 0 {
37                    continue;
38                }
39                let mut dot = 0.0f32;
40                for dd in 0..qd {
41                    dot += q[(r * h + hh) * qd + dd] * k[(kk * h + hh) * qd + dd];
42                }
43                let s = dot * scale;
44                let m_new = m.max(s);
45                let corr = if m == f32::NEG_INFINITY {
46                    0.0
47                } else {
48                    (m - m_new).exp()
49                };
50                let p = (s - m_new).exp();
51                l = l * corr + p;
52                for dd in 0..vd {
53                    acc[dd] = acc[dd] * corr + p * v[(kk * h + hh) * vd + dd];
54                }
55                m = m_new;
56            }
57            let inv = if l > 0.0 { 1.0 / l } else { 0.0 };
58            for dd in 0..vd {
59                out[(r * h + hh) * vd + dd] = acc[dd] * inv;
60            }
61        }
62    }
63    out
64}
65
66/// 2026-09-25: Expand the compressed latent into per-head K and V. NoPE only: it asserts
67/// `qk_rope_head_dim == 0`, so K is exactly `qk_nope_head_dim` wide.
68pub fn expand_kv(kv_c: &[f32], w_kv_b: &[f32], dims: DsaDims, seq: usize) -> (Vec<f32>, Vec<f32>) {
69    let (h, nope, vd, r) = (
70        dims.heads,
71        dims.qk_nope_head_dim,
72        dims.v_head_dim,
73        dims.kv_lora_rank,
74    );
75    assert_eq!(dims.qk_rope_head_dim, 0, "expand_kv is the NoPE path");
76    let wide = linear(kv_c, seq, r, w_kv_b, h * (nope + vd));
77    let mut k = vec![0.0f32; seq * h * nope];
78    let mut v = vec![0.0f32; seq * h * vd];
79    for t in 0..seq {
80        for hh in 0..h {
81            let src = t * h * (nope + vd) + hh * (nope + vd);
82            k[(t * h + hh) * nope..(t * h + hh) * nope + nope]
83                .copy_from_slice(&wide[src..src + nope]);
84            v[(t * h + hh) * vd..(t * h + hh) * vd + vd]
85                .copy_from_slice(&wide[src + nope..src + nope + vd]);
86        }
87    }
88    (k, v)
89}
metrale_model_arch/glm5next_dsa_ref/mla.rs

metrale_model_arch/glm5next_dsa_ref/
mla.rs
