atlas_core/kimi_k3/
mla.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Gated NoPE MLA CPU reference.
4//!
5//! Do not reuse `qwen3_attention` blindly: K3 sets `mla_use_nope=true` while
6//! still allocating `qk_rope_head_dim` slots, and `mla_use_output_gate=true`
7//! applies a full-rank `g_proj` sigmoid gate on the attention output.
8//!
9//! RoPE dims stay in the head (prod 128 nope + 64 rope = 192). NoPE means
10//! those slots are **not rotated**, not that they are dropped.
11//!
12//! CUDA decode is `kernels/gb10/kimi-k3/bf16/mla_decode.cu` (`k3_mla_*`).
13//! BoundLayer FullAttention uses CUDA `mla_decode` unless `K3_CUDA_MLA=0`.
14
15#![allow(clippy::needless_range_loop)]
16
17use super::cache::MlaKv;
18use crate::config::ModelConfig;
19
20#[inline]
21fn sigmoid(x: f32) -> f32 {
22    1.0 / (1.0 + (-x).exp())
23}
24
25#[derive(Clone, Copy, Debug)]
26pub struct MlaConfig {
27    pub heads: usize,
28    pub qk_nope_head_dim: usize,
29    pub qk_rope_head_dim: usize,
30    pub v_head_dim: usize,
31    pub q_lora_rank: usize,
32    pub kv_lora_rank: usize,
33    pub mla_use_nope: bool,
34    pub mla_use_output_gate: bool,
35}
36
37impl MlaConfig {
38    pub fn production() -> Self {
39        Self {
40            heads: 96,
41            qk_nope_head_dim: 128,
42            qk_rope_head_dim: 64,
43            v_head_dim: 128,
44            q_lora_rank: 1536,
45            kv_lora_rank: 512,
46            mla_use_nope: true,
47            mla_use_output_gate: true,
48        }
49    }
50
51    pub fn qk_head_dim(&self) -> usize {
52        self.qk_nope_head_dim + self.qk_rope_head_dim
53    }
54
55    /// `inference-optimization/Kimi-K3-0.40B` text_config MLA dims.
56    pub fn twin_0_40b() -> Self {
57        Self {
58            heads: 8,
59            qk_nope_head_dim: 64,
60            qk_rope_head_dim: 32,
61            v_head_dim: 64,
62            q_lora_rank: 256,
63            kv_lora_rank: 128,
64            mla_use_nope: true,
65            mla_use_output_gate: true,
66        }
67    }
68}
69
70pub fn mla_from(c: &ModelConfig) -> MlaConfig {
71    MlaConfig {
72        heads: c.num_attention_heads,
73        qk_nope_head_dim: c.qk_nope_head_dim,
74        qk_rope_head_dim: c.qk_rope_head_dim,
75        v_head_dim: c.v_head_dim,
76        q_lora_rank: c.q_lora_rank,
77        kv_lora_rank: c.kv_lora_rank,
78        mla_use_nope: c.mla_use_nope,
79        mla_use_output_gate: c.mla_use_output_gate,
80    }
81}
82
83/// FullAttention BoundLayer uses CUDA `mla_decode` unless `K3_CUDA_MLA=0`.
84/// Same polarity as `K3_CUDA_KDA`. Projections stay on the host either way.
85pub fn cuda_mla_enabled() -> bool {
86    !matches!(std::env::var("K3_CUDA_MLA").as_deref(), Ok("0"))
87}
88
89/// Optional RoPE on the rope **slice** of a packed `[nope | rope]` head.
90/// NoPE leaves `q`/`k` unchanged.
91pub fn maybe_rope(x: &mut [f32], nope: usize, rope: usize, pos: usize, theta: f32, use_nope: bool) {
92    if use_nope || rope == 0 {
93        return;
94    }
95    let dim = nope + rope;
96    for head in x.chunks_exact_mut(dim) {
97        let r = &mut head[nope..];
98        for i in 0..rope / 2 {
99            let freq = (pos as f32) / theta.powf(2.0 * i as f32 / rope as f32);
100            let (s, c) = freq.sin_cos();
101            let a = r[i];
102            let b = r[i + rope / 2];
103            r[i] = a * c - b * s;
104            r[i + rope / 2] = a * s + b * c;
105        }
106    }
107}
108
109/// Causal scaled-dot-product attention. `q/k`: `[T, H, dq]`, `v`: `[T, H, dv]`.
110pub fn sdpa(
111    q: &[f32],
112    k: &[f32],
113    v: &[f32],
114    t: usize,
115    heads: usize,
116    dq: usize,
117    dv: usize,
118) -> Vec<f32> {
119    let scale = 1.0 / (dq as f32).sqrt();
120    let mut out = vec![0.0f32; t * heads * dv];
121    for h in 0..heads {
122        for qi in 0..t {
123            let qrow = &q[(qi * heads + h) * dq..(qi * heads + h) * dq + dq];
124            let mut scores = vec![0.0f32; qi + 1];
125            let mut m = f32::NEG_INFINITY;
126            for kj in 0..=qi {
127                let krow = &k[(kj * heads + h) * dq..(kj * heads + h) * dq + dq];
128                let s: f32 = qrow.iter().zip(krow).map(|(a, b)| a * b).sum::<f32>() * scale;
129                scores[kj] = s;
130                if s > m {
131                    m = s;
132                }
133            }
134            let mut z = 0.0f32;
135            for s in &mut scores {
136                *s = (*s - m).exp();
137                z += *s;
138            }
139            let orow = &mut out[(qi * heads + h) * dv..(qi * heads + h) * dv + dv];
140            for kj in 0..=qi {
141                let a = scores[kj] / z;
142                let vrow = &v[(kj * heads + h) * dv..(kj * heads + h) * dv + dv];
143                for d in 0..dv {
144                    orow[d] += a * vrow[d];
145                }
146            }
147        }
148    }
149    out
150}
151
152/// Decode SDPA: one query `[H, dq]` against cached K/V of length `t`.
153pub fn sdpa_one(
154    q: &[f32],
155    k: &[f32],
156    v: &[f32],
157    t: usize,
158    heads: usize,
159    dq: usize,
160    dv: usize,
161) -> Vec<f32> {
162    assert_eq!(q.len(), heads * dq);
163    assert_eq!(k.len(), t * heads * dq);
164    assert_eq!(v.len(), t * heads * dv);
165    let scale = 1.0 / (dq as f32).sqrt();
166    let mut out = vec![0.0f32; heads * dv];
167    for h in 0..heads {
168        let qrow = &q[h * dq..(h + 1) * dq];
169        let mut scores = vec![0.0f32; t];
170        let mut m = f32::NEG_INFINITY;
171        for kj in 0..t {
172            let krow = &k[(kj * heads + h) * dq..(kj * heads + h) * dq + dq];
173            let s: f32 = qrow.iter().zip(krow).map(|(a, b)| a * b).sum::<f32>() * scale;
174            scores[kj] = s;
175            if s > m {
176                m = s;
177            }
178        }
179        let mut z = 0.0f32;
180        for s in &mut scores {
181            *s = (*s - m).exp();
182            z += *s;
183        }
184        let orow = &mut out[h * dv..(h + 1) * dv];
185        for kj in 0..t {
186            let a = scores[kj] / z;
187            let vrow = &v[(kj * heads + h) * dv..(kj * heads + h) * dv + dv];
188            for d in 0..dv {
189                orow[d] += a * vrow[d];
190            }
191        }
192    }
193    out
194}
195
196/// Apply `sigmoid(g) ⊙ attn` when the output gate is on; otherwise identity.
197pub fn apply_output_gate(attn: &[f32], g: &[f32], enabled: bool) -> Vec<f32> {
198    if !enabled {
199        return attn.to_vec();
200    }
201    assert_eq!(attn.len(), g.len());
202    attn.iter().zip(g).map(|(a, gg)| a * sigmoid(*gg)).collect()
203}
204
205/// Gated NoPE attend: optionally skip RoPE, SDPA, optional output gate.
206#[allow(clippy::too_many_arguments)]
207pub fn gated_mla_attend(
208    q: &mut [f32],
209    k: &mut [f32],
210    v: &[f32],
211    g: &[f32],
212    t: usize,
213    cfg: &MlaConfig,
214    pos0: usize,
215    theta: f32,
216) -> Vec<f32> {
217    for p in 0..t {
218        let dim = cfg.qk_head_dim();
219        let qh = &mut q[p * cfg.heads * dim..(p + 1) * cfg.heads * dim];
220        let kh = &mut k[p * cfg.heads * dim..(p + 1) * cfg.heads * dim];
221        maybe_rope(
222            qh,
223            cfg.qk_nope_head_dim,
224            cfg.qk_rope_head_dim,
225            pos0 + p,
226            theta,
227            cfg.mla_use_nope,
228        );
229        maybe_rope(
230            kh,
231            cfg.qk_nope_head_dim,
232            cfg.qk_rope_head_dim,
233            pos0 + p,
234            theta,
235            cfg.mla_use_nope,
236        );
237    }
238    let attn = sdpa(q, k, v, t, cfg.heads, cfg.qk_head_dim(), cfg.v_head_dim);
239    apply_output_gate(&attn, g, cfg.mla_use_output_gate)
240}
241
242/// One decode token: optional RoPE, append K/V, SDPA, optional output gate.
243/// Projections stay in `mla_mixer`. CUDA `k3_mla_*` matches this order.
244#[allow(clippy::too_many_arguments)]
245pub fn mla_decode_token(
246    q: &mut [f32],
247    k: &mut [f32],
248    v: &[f32],
249    g: &[f32],
250    kv: &mut MlaKv,
251    cfg: &MlaConfig,
252    pos: usize,
253    theta: f32,
254) -> Vec<f32> {
255    maybe_rope(
256        q,
257        cfg.qk_nope_head_dim,
258        cfg.qk_rope_head_dim,
259        pos,
260        theta,
261        cfg.mla_use_nope,
262    );
263    maybe_rope(
264        k,
265        cfg.qk_nope_head_dim,
266        cfg.qk_rope_head_dim,
267        pos,
268        theta,
269        cfg.mla_use_nope,
270    );
271    kv.append(k, v);
272    let attn = sdpa_one(
273        q,
274        &kv.k,
275        &kv.v,
276        kv.seq_len,
277        cfg.heads,
278        cfg.qk_head_dim(),
279        cfg.v_head_dim,
280    );
281    apply_output_gate(&attn, g, cfg.mla_use_output_gate)
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    fn tiny_cfg(nope: bool, gate: bool) -> MlaConfig {
289        MlaConfig {
290            heads: 1,
291            qk_nope_head_dim: 2,
292            qk_rope_head_dim: 2,
293            v_head_dim: 2,
294            q_lora_rank: 4,
295            kv_lora_rank: 4,
296            mla_use_nope: nope,
297            mla_use_output_gate: gate,
298        }
299    }
300
301    #[test]
302    fn production_keeps_rope_slots_but_nope() {
303        let p = MlaConfig::production();
304        assert!(p.mla_use_nope);
305        assert!(p.mla_use_output_gate);
306        assert_eq!(p.qk_head_dim(), 192);
307        assert_eq!(p.v_head_dim, 128);
308    }
309
310    #[test]
311    fn nope_does_not_rotate() {
312        let mut x = vec![1.0, 0.0, 1.0, 0.0];
313        let orig = x.clone();
314        maybe_rope(&mut x, 2, 2, 3, 10000.0, true);
315        assert_eq!(x, orig);
316        maybe_rope(&mut x, 2, 2, 3, 10000.0, false);
317        assert_ne!(x, orig, "RoPE on must move the rope slice");
318    }
319
320    #[test]
321    fn output_gate_mutates() {
322        let attn = vec![1.0, 2.0];
323        let g = vec![0.0, 10.0];
324        let off = apply_output_gate(&attn, &g, false);
325        let on = apply_output_gate(&attn, &g, true);
326        assert_eq!(off, attn);
327        assert!((on[0] - 0.5).abs() < 1e-6);
328        assert!(on[1] > 1.9);
329        assert_ne!(on, off);
330    }
331
332    #[test]
333    fn gated_nope_path_runs() {
334        let cfg = tiny_cfg(true, true);
335        let mut q = vec![1.0, 0.0, 0.0, 1.0];
336        let mut k = q.clone();
337        let v = vec![0.5, -0.5];
338        let g = vec![0.0, 0.0];
339        let out = gated_mla_attend(&mut q, &mut k, &v, &g, 1, &cfg, 0, 10000.0);
340        assert_eq!(out.len(), 2);
341    }
342
343    #[test]
344    fn twin_0_40b_geometry() {
345        let t = MlaConfig::twin_0_40b();
346        assert_eq!(t.heads, 8);
347        assert_eq!(t.qk_head_dim(), 96);
348        assert_eq!(t.v_head_dim, 64);
349        assert!(t.mla_use_nope);
350        assert!(t.mla_use_output_gate);
351    }
352
353    #[test]
354    fn decode_token_matches_attend_t1() {
355        let cfg = tiny_cfg(true, true);
356        let mut q = vec![1.0, 0.0, 0.0, 1.0];
357        let mut k = q.clone();
358        let v = vec![0.5, -0.5];
359        let g = vec![0.0, 0.0];
360        let mut q2 = q.clone();
361        let mut k2 = k.clone();
362        let attend = gated_mla_attend(&mut q2, &mut k2, &v, &g, 1, &cfg, 0, 10000.0);
363        let mut kv = MlaKv::default();
364        let dec = mla_decode_token(&mut q, &mut k, &v, &g, &mut kv, &cfg, 0, 10000.0);
365        assert_eq!(attend, dec);
366        assert_eq!(kv.seq_len, 1);
367    }
368
369    #[test]
370    fn cuda_mla_env_default_on() {
371        if std::env::var_os("K3_CUDA_MLA").is_some() {
372            return;
373        }
374        assert!(
375            cuda_mla_enabled(),
376            "FullAttention default is CUDA MLA; K3_CUDA_MLA=0 is the CPU escape"
377        );
378    }
379
380    #[test]
381    fn cuda_mla_env_opt_out() {
382        const THIS: &str = "kimi_k3::mla::tests::cuda_mla_env_opt_out";
383        const MARKER: &str = "K3_CUDA_MLA_CHILD";
384        if std::env::var_os(MARKER).is_some() {
385            assert!(!cuda_mla_enabled(), "K3_CUDA_MLA=0 must keep the CPU mixer");
386            return;
387        }
388        let output = std::process::Command::new(std::env::current_exe().unwrap())
389            .args(["--exact", THIS])
390            .env(MARKER, "1")
391            .env("K3_CUDA_MLA", "0")
392            .output()
393            .unwrap();
394        assert!(
395            output.status.success(),
396            "K3_CUDA_MLA child failed:\n{}",
397            String::from_utf8_lossy(&output.stderr)
398        );
399    }
400}
atlas_core/kimi_k3/mla.rs

atlas_core/kimi_k3/
mla.rs
