atlas_core/kimi_k3/
layer.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! One K3 decoder layer as host graph data: mixer + MLP + AttnRes sites.
4//!
5//! ```text
6//! h = AttnRes(blocks, partial, attn_res_*)
7//! mix_out = KDA | gated-NoPE-MLA (on RMSNorm(h))
8//! partial += mix_out
9//! h = AttnRes(blocks, partial, mlp_res_*)
10//! mlp_out = dense SiTU-GLU | LatentMoE (on RMSNorm(h))
11//! partial += mlp_out
12//! at layer_idx % block_size == 0: archive incoming prefix, reset partial
13//! ```
14
15use crate::config::{LayerType, ModelConfig};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum MixerKind {
19    Kda,
20    Mla,
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum MlpKind {
25    Dense,
26    LatentMoe,
27}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub struct K3LayerSpec {
31    pub index: usize,
32    pub mixer: MixerKind,
33    pub mlp: MlpKind,
34}
35
36#[derive(Debug, Clone)]
37pub struct K3Graph {
38    pub layers: Vec<K3LayerSpec>,
39    pub hidden: usize,
40    pub attn_res_block_size: usize,
41    pub situ_beta: f32,
42    pub situ_linear_beta: f32,
43    pub use_full_rank_gate: bool,
44    pub mla_use_nope: bool,
45    pub mla_use_output_gate: bool,
46}
47
48impl K3Graph {
49    pub fn from_config(c: &ModelConfig) -> Self {
50        let dense_end = c.mlp_only_layers.len();
51        let layers = c
52            .layer_types
53            .iter()
54            .enumerate()
55            .map(|(i, t)| K3LayerSpec {
56                index: i,
57                mixer: match t {
58                    LayerType::LinearAttention => MixerKind::Kda,
59                    _ => MixerKind::Mla,
60                },
61                mlp: if i < dense_end {
62                    MlpKind::Dense
63                } else {
64                    MlpKind::LatentMoe
65                },
66            })
67            .collect();
68        Self {
69            layers,
70            hidden: c.hidden_size,
71            attn_res_block_size: c.attn_res_block_size,
72            situ_beta: c.activation_situ_beta,
73            situ_linear_beta: c.activation_situ_linear_beta,
74            use_full_rank_gate: c.use_full_rank_gate,
75            mla_use_nope: c.mla_use_nope,
76            mla_use_output_gate: c.mla_use_output_gate,
77        }
78    }
79
80    pub fn kda_count(&self) -> usize {
81        self.layers
82            .iter()
83            .filter(|l| l.mixer == MixerKind::Kda)
84            .count()
85    }
86
87    pub fn mla_count(&self) -> usize {
88        self.layers
89            .iter()
90            .filter(|l| l.mixer == MixerKind::Mla)
91            .count()
92    }
93
94    pub fn last_is_mla(&self) -> bool {
95        self.layers
96            .last()
97            .is_some_and(|l| l.mixer == MixerKind::Mla)
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use crate::config::parse_config;
105
106    #[test]
107    fn twin_0_40b_graph_is_six_kda_last_mla() {
108        const TWIN: &str = include_str!("../../../../docs/k3/fixtures/Kimi-K3-0.40B-config.json");
109        let c = parse_config(TWIN).expect("0.40B twin");
110        let g = K3Graph::from_config(&c);
111        assert_eq!(g.layers.len(), 8);
112        assert_eq!(g.kda_count(), 6);
113        assert_eq!(g.mla_count(), 2);
114        assert!(g.last_is_mla());
115        // 3:1 then trailing MLA: 0-2 KDA, 3 MLA, 4-6 KDA, 7 MLA.
116        for i in [0, 1, 2, 4, 5, 6] {
117            assert_eq!(g.layers[i].mixer, MixerKind::Kda, "layer {i}");
118        }
119        assert_eq!(g.layers[3].mixer, MixerKind::Mla);
120        assert_eq!(g.layers[7].mixer, MixerKind::Mla);
121        assert_eq!(g.layers[0].mlp, MlpKind::Dense);
122        assert!(g.layers[1..].iter().all(|l| l.mlp == MlpKind::LatentMoe));
123        assert_eq!(g.attn_res_block_size, 4);
124        assert_eq!(g.hidden, 1024);
125        assert!(g.mla_use_nope);
126        assert!(g.mla_use_output_gate);
127        assert!(g.use_full_rank_gate);
128        assert_eq!(g.situ_beta, 4.0);
129        assert_eq!(g.situ_linear_beta, 25.0);
130    }
131
132    #[test]
133    fn official_graph_census() {
134        const OFFICIAL: &str =
135            include_str!("../../../../docs/k3/fixtures/moonshotai-Kimi-K3-config.json");
136        let c = parse_config(OFFICIAL).expect("official");
137        let g = K3Graph::from_config(&c);
138        assert_eq!(g.layers.len(), 93);
139        assert_eq!(g.kda_count(), 69);
140        assert_eq!(g.mla_count(), 24);
141        assert!(g.last_is_mla());
142        // HF 1-based 92 and 93 are both MLA (0-based 91, 92).
143        assert_eq!(g.layers[91].mixer, MixerKind::Mla);
144        assert_eq!(g.layers[92].mixer, MixerKind::Mla);
145        assert_eq!(g.attn_res_block_size, 12);
146        assert_eq!(g.layers[0].mlp, MlpKind::Dense);
147    }
148}
atlas_core/kimi_k3/layer.rs

atlas_core/kimi_k3/
layer.rs
