atlas_core/kimi_k3/
attnres.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Block AttnRes CPU reference.
4//!
5//! Each sublayer mixes a learned softmax over completed block residuals plus
6//! the current intra-block partial sum (Moonshot Block AttnRes):
7//!
8//! ```text
9//! K_i = RMSNorm(V_i)
10//! α   = softmax_i( q · K_i )
11//! h   = Σ α_i V_i
12//! ```
13//!
14//! `q` is the per-layer `*_res_proj` row. Mix=0 is the identity skip (return
15//! the partial / skip source). Mix=1 is the full softmax mixture.
16
17/// Vanilla RMSNorm: `x * w / sqrt(mean(x^2) + eps)`.
18pub fn rms_norm(x: &[f32], w: &[f32], eps: f32) -> Vec<f32> {
19    assert_eq!(x.len(), w.len());
20    let mean_sq = x.iter().map(|v| v * v).sum::<f32>() / x.len() as f32;
21    let inv = 1.0 / (mean_sq + eps).sqrt();
22    x.iter().zip(w).map(|(v, s)| v * s * inv).collect()
23}
24
25/// Softmax mixture over residual sources. `sources[0]` is conventionally the
26/// skip (current partial block). `query` is `[hidden]` (`*_res_proj`).
27pub fn attnres_softmax_mix(
28    sources: &[Vec<f32>],
29    query: &[f32],
30    norm_w: &[f32],
31    eps: f32,
32) -> Vec<f32> {
33    assert!(
34        !sources.is_empty(),
35        "AttnRes needs at least the skip source"
36    );
37    let hidden = query.len();
38    let mut logits = Vec::with_capacity(sources.len());
39    for src in sources {
40        assert_eq!(src.len(), hidden);
41        let k = rms_norm(src, norm_w, eps);
42        let dot = query.iter().zip(&k).map(|(q, kk)| q * kk).sum::<f32>();
43        logits.push(dot);
44    }
45    let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
46    let mut weights: Vec<f32> = logits.iter().map(|l| (l - max).exp()).collect();
47    let z: f32 = weights.iter().sum();
48    for w in &mut weights {
49        *w /= z;
50    }
51    let mut out = vec![0.0f32; hidden];
52    for (src, a) in sources.iter().zip(weights) {
53        for (o, v) in out.iter_mut().zip(src) {
54            *o += a * v;
55        }
56    }
57    out
58}
59
60/// Test / ablation lever: `mix=0` returns `skip`; `mix=1` returns `mixed`.
61pub fn attnres_blend(skip: &[f32], mixed: &[f32], mix: f32) -> Vec<f32> {
62    assert_eq!(skip.len(), mixed.len());
63    skip.iter()
64        .zip(mixed)
65        .map(|(s, m)| (1.0 - mix) * s + mix * m)
66        .collect()
67}
68
69/// Apply AttnRes with an explicit mix lever. `mix=0` is identity skip.
70pub fn attnres_mix(
71    sources: &[Vec<f32>],
72    query: &[f32],
73    norm_w: &[f32],
74    eps: f32,
75    mix: f32,
76) -> Vec<f32> {
77    let skip = &sources[0];
78    let mixed = attnres_softmax_mix(sources, query, norm_w, eps);
79    attnres_blend(skip, &mixed, mix)
80}
81
82/// Stream map keyed by hidden/residual pointer bits (`DevicePtr.0`).
83/// Layer 0 inserts; last layer `remove`s on success. Any `Err` drops the
84/// entry (umbrella `940bd4eeb`).
85#[derive(Debug)]
86pub struct AttnResHub<T> {
87    map: std::collections::HashMap<u64, T>,
88}
89
90impl<T> Default for AttnResHub<T> {
91    fn default() -> Self {
92        Self {
93            map: std::collections::HashMap::new(),
94        }
95    }
96}
97
98impl<T> AttnResHub<T> {
99    pub fn insert(&mut self, key: u64, v: T) {
100        self.map.insert(key, v);
101    }
102
103    pub fn get(&self, key: u64) -> Option<&T> {
104        self.map.get(&key)
105    }
106
107    pub fn remove(&mut self, key: u64) -> Option<T> {
108        self.map.remove(&key)
109    }
110
111    pub fn contains(&self, key: u64) -> bool {
112        self.map.contains_key(&key)
113    }
114
115    pub fn decode<R, E>(
116        &mut self,
117        key: u64,
118        f: impl FnOnce(&mut Self) -> Result<R, E>,
119    ) -> Result<R, E> {
120        let r = f(self);
121        if r.is_err() {
122            self.map.remove(&key);
123        }
124        r
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    fn max_abs(a: &[f32], b: &[f32]) -> f32 {
133        a.iter()
134            .zip(b)
135            .map(|(x, y)| (x - y).abs())
136            .fold(0.0, f32::max)
137    }
138
139    #[test]
140    fn mix_zero_is_identity_skip() {
141        let skip = vec![1.0, 0.0, -0.5, 2.0];
142        let other = vec![0.0, 1.0, 4.0, -3.0];
143        let sources = [skip.clone(), other];
144        let query = vec![0.2, -0.1, 0.4, 0.3];
145        let w = vec![1.0, 1.0, 1.0, 1.0];
146        let id = attnres_mix(&sources, &query, &w, 1e-5, 0.0);
147        assert_eq!(id, skip, "mix=0 must return the skip source");
148    }
149
150    #[test]
151    fn mix_zero_vs_mix_one_diverges() {
152        // Known-bad: a graph that zeros mix weights and still claims mix=1.
153        // The instrument must see mix=0 ≠ mix=1 before a green AttnRes is trusted.
154        let skip = vec![1.0, 0.0, -0.5, 2.0];
155        let other = vec![0.0, 1.0, 4.0, -3.0];
156        let sources = [skip.clone(), other];
157        let query = vec![0.2, -0.1, 0.4, 0.3];
158        let w = vec![1.0, 1.0, 1.0, 1.0];
159        let m0 = attnres_mix(&sources, &query, &w, 1e-5, 0.0);
160        let m1 = attnres_mix(&sources, &query, &w, 1e-5, 1.0);
161        assert!(
162            max_abs(&m0, &m1) > 0.5,
163            "mix=0 ({m0:?}) must diverge from mix=1 ({m1:?})"
164        );
165        assert_eq!(m0, skip);
166        assert_ne!(m1, skip);
167        const RECORDED_MIX1: [f32; 4] = [0.571_599_9, 0.428_400_1, 1.427_800_4, -0.142_000_4];
168        assert!(
169            max_abs(&m1, &RECORDED_MIX1) <= 1e-5,
170            "mix=1 vs recorded fixture max_abs={}",
171            max_abs(&m1, &RECORDED_MIX1)
172        );
173    }
174
175    #[test]
176    fn hub_drops_entry_on_decode_err() {
177        let mut hub = AttnResHub::default();
178        hub.insert(1, vec![1.0f32]);
179        let err: Result<(), &str> = hub.decode(1, |_| Err("cuda fail"));
180        assert!(err.is_err());
181        assert!(
182            !hub.contains(1),
183            "RST: decode Err must drop the stream (940bd4eeb)"
184        );
185        hub.insert(2, vec![2.0]);
186        hub.decode(2, |h| {
187            h.remove(2);
188            Ok::<(), &str>(())
189        })
190        .unwrap();
191        assert!(!hub.contains(2), "last-layer success still removes");
192    }
193}
atlas_core/kimi_k3/attnres.rs

atlas_core/kimi_k3/
attnres.rs
