avarok_core/kimi_k3/
weights.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Text decoder inventory shared by checkpoint admission and layer binding.
4use super::{K3Graph, MixerKind, MlpKind};
5use crate::config::ModelConfig;
6
7pub fn text_key(config: &ModelConfig, rest: &str) -> String {
8    let p = config.weight_prefix.trim_end_matches('.');
9    if p.is_empty() {
10        rest.to_string()
11    } else {
12        format!("{p}.{rest}")
13    }
14}
15
16pub fn layer_keys(
17    config: &ModelConfig,
18    i: usize,
19    mixer: MixerKind,
20    mlp: MlpKind,
21    n_experts: usize,
22) -> Vec<String> {
23    let lp = text_key(config, &format!("model.layers.{i}"));
24    let mut k = vec![
25        format!("{lp}.input_layernorm.weight"),
26        format!("{lp}.post_attention_layernorm.weight"),
27        format!("{lp}.mlp_res_norm.weight"),
28        format!("{lp}.mlp_res_proj.weight"),
29        format!("{lp}.self_attention_res_norm.weight"),
30        format!("{lp}.self_attention_res_proj.weight"),
31        format!("{lp}.self_attn.g_proj.weight"),
32        format!("{lp}.self_attn.o_proj.weight"),
33    ];
34    match mixer {
35        MixerKind::Kda => {
36            k.extend([
37                format!("{lp}.self_attn.A_log"),
38                format!("{lp}.self_attn.b_proj.weight"),
39                format!("{lp}.self_attn.dt_bias"),
40                format!("{lp}.self_attn.f_a_proj.weight"),
41                format!("{lp}.self_attn.f_b_proj.weight"),
42                format!("{lp}.self_attn.k_conv1d.weight"),
43                format!("{lp}.self_attn.k_proj.weight"),
44                format!("{lp}.self_attn.o_norm.weight"),
45                format!("{lp}.self_attn.q_conv1d.weight"),
46                format!("{lp}.self_attn.q_proj.weight"),
47                format!("{lp}.self_attn.v_conv1d.weight"),
48                format!("{lp}.self_attn.v_proj.weight"),
49            ]);
50        }
51        MixerKind::Mla => {
52            k.extend([
53                format!("{lp}.self_attn.kv_a_layernorm.weight"),
54                format!("{lp}.self_attn.kv_a_proj_with_mqa.weight"),
55                format!("{lp}.self_attn.kv_b_proj.weight"),
56                format!("{lp}.self_attn.q_a_layernorm.weight"),
57                format!("{lp}.self_attn.q_a_proj.weight"),
58                format!("{lp}.self_attn.q_b_proj.weight"),
59            ]);
60        }
61    }
62    match mlp {
63        MlpKind::Dense => {
64            k.extend([
65                format!("{lp}.mlp.down_proj.weight"),
66                format!("{lp}.mlp.gate_proj.weight"),
67                format!("{lp}.mlp.up_proj.weight"),
68            ]);
69        }
70        MlpKind::LatentMoe => {
71            k.extend([
72                format!("{lp}.block_sparse_moe.gate.e_score_correction_bias"),
73                format!("{lp}.block_sparse_moe.gate.weight"),
74                format!("{lp}.block_sparse_moe.routed_expert_down_proj.weight"),
75                format!("{lp}.block_sparse_moe.routed_expert_norm.weight"),
76                format!("{lp}.block_sparse_moe.routed_expert_up_proj.weight"),
77                format!("{lp}.block_sparse_moe.shared_experts.down_proj.weight"),
78                format!("{lp}.block_sparse_moe.shared_experts.gate_proj.weight"),
79                format!("{lp}.block_sparse_moe.shared_experts.up_proj.weight"),
80            ]);
81            for e in 0..n_experts {
82                k.extend([
83                    format!("{lp}.block_sparse_moe.experts.{e}.w1.weight"),
84                    format!("{lp}.block_sparse_moe.experts.{e}.w2.weight"),
85                    format!("{lp}.block_sparse_moe.experts.{e}.w3.weight"),
86                ]);
87            }
88        }
89    }
90    k
91}
92
93/// Required logical keys. Routed expert `.weight` may instead be supplied by
94/// both `.weight_packed` and `.weight_scale`; the loader validates those pairs.
95/// A prefixed `lm_head.weight` may use the unprefixed checkpoint fallback.
96pub fn required_names(config: &ModelConfig) -> Vec<String> {
97    let mut names: Vec<_> = [
98        "model.embed_tokens.weight",
99        "model.norm.weight",
100        "model.output_attn_res_proj.weight",
101        "model.output_attn_res_norm.weight",
102        "lm_head.weight",
103    ]
104    .into_iter()
105    .map(|name| text_key(config, name))
106    .collect();
107    for layer in K3Graph::from_config(config).layers {
108        names.extend(layer_keys(
109            config,
110            layer.index,
111            layer.mixer,
112            layer.mlp,
113            config.num_experts,
114        ));
115    }
116    names
117}
avarok_core/kimi_k3/weights.rs

avarok_core/kimi_k3/
weights.rs
