avarok_core/kimi_k3/
latent_moe.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Stable LatentMoE CPU reference.
4//!
5//! ```text
6//! latent = down_proj(h)                 // hidden → moe_latent (no pre-norm)
7//! scores = sigmoid(router(h)); top-k on scores+bias  // noaux_tc
8//! y = Σ w_i expert_i(latent)            // SiTU-GLU experts
9//! y = RMSNorm(y) if latent_moe_use_norm; y = up_proj(y)
10//! out = y + shared_experts(h)           // shared stay full-width
11//! ```
12
13#![allow(clippy::too_many_arguments)]
14
15use super::attnres::rms_norm;
16use super::situ::{sigmoid, situ_glu_vec};
17use crate::config::ModelConfig;
18
19#[derive(Clone, Copy, Debug)]
20pub struct LatentMoeConfig {
21    pub hidden: usize,
22    pub latent: usize,
23    pub expert_hidden: usize,
24    pub n_routed: usize,
25    pub top_k: usize,
26    pub n_shared: usize,
27    pub situ_beta: f32,
28    pub situ_linear_beta: f32,
29    pub use_norm: bool,
30    pub renormalize: bool,
31}
32
33impl LatentMoeConfig {
34    pub fn production() -> Self {
35        Self {
36            hidden: 7168,
37            latent: 3584,
38            expert_hidden: 3072,
39            n_routed: 896,
40            top_k: 16,
41            n_shared: 2,
42            situ_beta: 4.0,
43            situ_linear_beta: 25.0,
44            use_norm: true,
45            renormalize: true,
46        }
47    }
48}
49
50pub fn moe_from(c: &ModelConfig) -> LatentMoeConfig {
51    LatentMoeConfig {
52        hidden: c.hidden_size,
53        latent: c.moe_latent_size,
54        expert_hidden: c.moe_intermediate_size,
55        n_routed: c.num_experts,
56        top_k: c.num_experts_per_tok,
57        n_shared: c.n_shared_experts,
58        situ_beta: c.activation_situ_beta,
59        situ_linear_beta: c.activation_situ_linear_beta,
60        use_norm: c.latent_moe_use_norm,
61        renormalize: c.norm_topk_prob,
62    }
63}
64
65/// noaux_tc: sigmoid(logits) for mix weights; `scores + bias` only ranks.
66pub fn sigmoid_topk(logits: &[f32], bias: &[f32], k: usize) -> (Vec<usize>, Vec<f32>) {
67    assert_eq!(logits.len(), bias.len());
68    let n = logits.len();
69    let k = k.min(n);
70    let scores: Vec<f32> = logits.iter().map(|l| sigmoid(*l)).collect();
71    let mut idx: Vec<usize> = (0..n).collect();
72    idx.sort_by(|&a, &b| {
73        let ca = scores[a] + bias[a];
74        let cb = scores[b] + bias[b];
75        cb.partial_cmp(&ca)
76            .unwrap_or(std::cmp::Ordering::Equal)
77            .then(a.cmp(&b))
78    });
79    idx.truncate(k);
80    let mut w: Vec<f32> = idx.iter().map(|&i| scores[i]).collect();
81    let z: f32 = w.iter().sum();
82    if z > 0.0 {
83        for ww in &mut w {
84            *ww /= z;
85        }
86    }
87    (idx, w)
88}
89
90/// One SiTU-GLU expert: `down( situ(w1 x, w3 x) )`. Weights are `[out, in]`.
91pub fn expert_situ(
92    x: &[f32],
93    w1: &[f32],
94    w2: &[f32],
95    w3: &[f32],
96    in_dim: usize,
97    hidden: usize,
98    beta: f32,
99    beta_lin: f32,
100) -> Vec<f32> {
101    let gate = matvec(w1, x, hidden, in_dim);
102    let up = matvec(w3, x, hidden, in_dim);
103    let mid = situ_glu_vec(&gate, &up, beta, beta_lin);
104    matvec(w2, &mid, in_dim, hidden)
105}
106
107/// Mix selected SiTU-GLU experts into latent. Empty `w1` means packed-only
108/// (no host dequant) — CUDA grouped GEMM must have run instead.
109pub fn mix_routed_experts(
110    latent: &[f32],
111    ids: &[usize],
112    weights: &[f32],
113    experts: &[(Vec<f32>, Vec<f32>, Vec<f32>)],
114    cfg: &LatentMoeConfig,
115) -> Vec<f32> {
116    let mut mixed = vec![0.0f32; cfg.latent];
117    for (&id, &w) in ids.iter().zip(weights) {
118        assert!(id < experts.len(), "expert id {id} >= {}", experts.len());
119        let (w1, w2, w3) = &experts[id];
120        assert!(
121            !w1.is_empty(),
122            "K3 packed expert {id} has no host w1; CUDA grouped GEMM required"
123        );
124        let eh = if cfg.latent == 0 {
125            cfg.expert_hidden
126        } else {
127            w1.len() / cfg.latent
128        };
129        let y = expert_situ(
130            latent,
131            w1,
132            w2,
133            w3,
134            cfg.latent,
135            eh,
136            cfg.situ_beta,
137            cfg.situ_linear_beta,
138        );
139        for (m, yy) in mixed.iter_mut().zip(y) {
140            *m += w * yy;
141        }
142    }
143    mixed
144}
145
146fn matvec(w: &[f32], x: &[f32], out: usize, inn: usize) -> Vec<f32> {
147    assert_eq!(
148        w.len(),
149        out * inn,
150        "matvec weight {} vs {out}x{inn}",
151        w.len()
152    );
153    assert_eq!(x.len(), inn);
154    let mut y = vec![0.0f32; out];
155    // Interleave independent output rows so the CPU can overlap accumulations.
156    // Each row retains the scalar reduction order: no reassociation or FMA.
157    let grouped = out / 4 * 4;
158    for o in (0..grouped).step_by(4) {
159        let a = &w[o * inn..(o + 1) * inn];
160        let b = &w[(o + 1) * inn..(o + 2) * inn];
161        let c = &w[(o + 2) * inn..(o + 3) * inn];
162        let d = &w[(o + 3) * inn..(o + 4) * inn];
163        let mut acc = [0.0f32; 4];
164        for i in 0..inn {
165            let value = x[i];
166            acc[0] += a[i] * value;
167            acc[1] += b[i] * value;
168            acc[2] += c[i] * value;
169            acc[3] += d[i] * value;
170        }
171        y[o..o + 4].copy_from_slice(&acc);
172    }
173    for o in grouped..out {
174        let row = &w[o * inn..(o + 1) * inn];
175        let mut acc = 0.0f32;
176        for i in 0..inn {
177            acc += row[i] * x[i];
178        }
179        y[o] = acc;
180    }
181    y
182}
183
184/// Routed latent path + optional shared expert (identity-scale for tests).
185pub fn latent_moe_forward(
186    h: &[f32],
187    down: &[f32],
188    up: &[f32],
189    norm_w: &[f32],
190    logits: &[f32],
191    bias: &[f32],
192    experts: &[(Vec<f32>, Vec<f32>, Vec<f32>)],
193    shared: Option<&[f32]>,
194    cfg: &LatentMoeConfig,
195    eps: f32,
196) -> (Vec<f32>, Vec<usize>) {
197    let latent = matvec(down, h, cfg.latent, cfg.hidden);
198    let (ids, weights) = sigmoid_topk(logits, bias, cfg.top_k);
199    let mut mixed = mix_routed_experts(&latent, &ids, &weights, experts, cfg);
200    if cfg.use_norm {
201        mixed = rms_norm(&mixed, norm_w, eps);
202    }
203    let mut out = matvec(up, &mixed, cfg.hidden, cfg.latent);
204    if let Some(s) = shared {
205        for (o, ss) in out.iter_mut().zip(s) {
206            *o += ss;
207        }
208    }
209    (out, ids)
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    #[test]
217    fn sigmoid_topk_picks_highest() {
218        let logits = [0.1, 5.0, 0.2, 4.0];
219        let bias = [0.0, 0.0, 0.0, 0.0];
220        let (ids, w) = sigmoid_topk(&logits, &bias, 2);
221        assert_eq!(ids, vec![1, 3]);
222        assert!((w[0] + w[1] - 1.0).abs() < 1e-6);
223        assert!(w[0] > w[1]);
224    }
225
226    #[test]
227    fn noaux_tc_bias_ranks_not_weighted() {
228        // HF: scores = sigmoid(logits); top-k on scores+bias; mix uses scores.
229        let logits = [0.0f32, 0.0, 0.0];
230        let bias = [0.0, 5.0, 4.0];
231        let (ids, w) = sigmoid_topk(&logits, &bias, 2);
232        assert_eq!(ids, vec![1, 2]);
233        assert!(
234            (w[0] - w[1]).abs() < 1e-5,
235            "equal logits must keep equal mix weights, got {w:?}"
236        );
237    }
238
239    #[test]
240    fn force_expert_zero_mutant_diverges() {
241        // C6 known-bad: forcing expert 0 must change the mix vs true top-k.
242        let cfg = LatentMoeConfig {
243            hidden: 2,
244            latent: 2,
245            expert_hidden: 2,
246            n_routed: 2,
247            top_k: 1,
248            n_shared: 0,
249            situ_beta: 4.0,
250            situ_linear_beta: 25.0,
251            use_norm: false,
252            renormalize: true,
253        };
254        let h = vec![1.0, 0.0];
255        let down = vec![1.0, 0.0, 0.0, 1.0];
256        let up = down.clone();
257        let ident = |out: usize, inn: usize| -> Vec<f32> {
258            let mut w = vec![0.0; out * inn];
259            for i in 0..out.min(inn) {
260                w[i * inn + i] = 1.0;
261            }
262            w
263        };
264        let e0 = (ident(2, 2), ident(2, 2), ident(2, 2));
265        // Expert 1 scales up-branch so SiTU output differs.
266        let mut w3 = ident(2, 2);
267        w3[0] = 3.0;
268        let e1 = (ident(2, 2), ident(2, 2), w3);
269        let experts = [e0, e1];
270        let logits = [0.0, 4.0];
271        let bias = [0.0, 0.0];
272        let (y, ids) = latent_moe_forward(
273            &h,
274            &down,
275            &up,
276            &[1.0, 1.0],
277            &logits,
278            &bias,
279            &experts,
280            None,
281            &cfg,
282            1e-5,
283        );
284        assert_eq!(ids, vec![1]);
285        let (y0, _) = latent_moe_forward(
286            &h,
287            &down,
288            &up,
289            &[1.0, 1.0],
290            &[4.0, 0.0],
291            &bias,
292            &experts,
293            None,
294            &cfg,
295            1e-5,
296        );
297        let err: f32 = y.iter().zip(&y0).map(|(a, b)| (a - b).abs()).sum();
298        assert!(err > 1e-4, "forcing expert 0 must diverge, err={err}");
299    }
300}
301
302#[cfg(test)]
303#[path = "latent_moe_matvec_tests.rs"]
304mod matvec_tests;
avarok_core/kimi_k3/latent_moe.rs

avarok_core/kimi_k3/
latent_moe.rs
