spark_model/weight_loader/
kimi_k3.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Kimi K3 weight loader — C1: BF16 0.40B twin bind.
4//! Packed experts bind to DSV4 E8M0 (`quantized_k3_mxfp4_e8m0`) only when
5//! `K3_ALLOW_MXFP4=1`. Default refuse so an accidental official download
6//! cannot silently land.
7
8use anyhow::{Result, bail};
9use avarok_core::config::ModelConfig;
10use spark_runtime::gpu::GpuBackend;
11use spark_runtime::kv_cache::KvCacheDtype;
12use spark_runtime::weights::WeightStore;
13
14use super::ModelWeightLoader;
15use crate::layer::TransformerLayer;
16use crate::weight_map::DenseWeight;
17
18mod bf16;
19mod classes;
20mod dry_run;
21mod mxfp4;
22mod tp;
23
24pub use dry_run::{KimiK3DryRun, dry_run_index_json, dry_run_weight_map};
25
26pub struct KimiK3WeightLoader;
27
28impl KimiK3WeightLoader {
29    pub fn dry_run(index_json: &str, language_model_only: bool) -> Result<KimiK3DryRun> {
30        dry_run::dry_run_index_json(index_json, language_model_only)
31    }
32}
33
34/// Packed experts without `K3_ALLOW_MXFP4=1` refuse. Opt-in lands DSV4 E8M0.
35pub fn refuse_mxfp4(store: &WeightStore) -> Result<()> {
36    if store.names().any(|n| n.contains("weight_packed"))
37        && !avarok_core::kimi_k3::mxfp4::allow_mxfp4()
38    {
39        bail!("S5 MXFP4 not this slice");
40    }
41    Ok(())
42}
43
44impl ModelWeightLoader for KimiK3WeightLoader {
45    fn supports_tp(&self) -> bool {
46        // Q/K/V col-parallel, O/g row/col per mixer, KDA head/channel
47        // companions, dense gate/up/down, expert w1/w3/w2. Embed / RMSNorm /
48        // router stay replicated. `lm_head` stays full — vocab-parallel is
49        // `lmhead_vocab_shard` at GEMV (minimax/glm5_next).
50        true
51    }
52
53    fn binds_vision_encoder(&self) -> bool {
54        false
55    }
56
57    fn load_layers(
58        &self,
59        store: &WeightStore,
60        config: &ModelConfig,
61        gpu: &dyn GpuBackend,
62        _layer_kv_dtypes: &[KvCacheDtype],
63    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
64        refuse_mxfp4(store)?;
65        bf16::load_layers(store, config, gpu)
66    }
67
68    fn load_embedding(
69        &self,
70        store: &WeightStore,
71        config: &ModelConfig,
72        gpu: &dyn GpuBackend,
73    ) -> Result<DenseWeight> {
74        refuse_mxfp4(store)?;
75        bf16::load_embedding(store, config, gpu)
76    }
77
78    fn load_final_norm(
79        &self,
80        store: &WeightStore,
81        config: &ModelConfig,
82        gpu: &dyn GpuBackend,
83    ) -> Result<DenseWeight> {
84        refuse_mxfp4(store)?;
85        bf16::load_final_norm(store, config, gpu)
86    }
87
88    fn load_lm_head(
89        &self,
90        store: &WeightStore,
91        config: &ModelConfig,
92        gpu: &dyn GpuBackend,
93    ) -> Result<DenseWeight> {
94        refuse_mxfp4(store)?;
95        bf16::load_lm_head(store, config, gpu)
96    }
97
98    fn load_mtp_weights(
99        &self,
100        _store: &WeightStore,
101        _config: &ModelConfig,
102        _gpu: &dyn GpuBackend,
103    ) -> Result<Option<crate::weight_loader::MtpWeights>> {
104        Ok(None)
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use crate::factory::loader_for_config;
112    use avarok_core::config::parse_config;
113    use spark_runtime::gpu::GpuBackend;
114    use spark_runtime::weights::{WeightDtype, WeightTensor};
115    use std::collections::HashMap;
116    use std::process::Command;
117
118    #[test]
119    fn kimi_k3_supports_tp() {
120        assert!(KimiK3WeightLoader.supports_tp());
121    }
122
123    #[test]
124    fn kimi_k3_does_not_probe_nvfp4_tgemm() {
125        assert!(
126            !crate::layers::tgemm_probe_ok("kimi_k3"),
127            "BF16/FP32 twin must not look up w4a16_gemm_t_p3"
128        );
129        assert!(crate::layers::tgemm_probe_ok("qwen3_5_moe"));
130    }
131
132    #[test]
133    fn loader_for_config_kimi_k3() {
134        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
135        for ty in ["kimi_k3", "kimi_linear", "Kimi-K3"] {
136            config.model_type = ty.to_string();
137            let loader = loader_for_config(&config).expect(ty);
138            let store = WeightStore::empty();
139            let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
140            let err = match loader.load_layers(&store, &config, &gpu, &[]) {
141                Ok(_) => panic!("{ty}: empty store must not bind"),
142                Err(e) => e.to_string(),
143            };
144            assert!(
145                err.contains("not found") || err.contains("S5 MXFP4"),
146                "{ty}: {err}"
147            );
148            assert!(!err.contains("K3-WIP"), "{ty}: stale WIP bail: {err}");
149        }
150    }
151
152    #[test]
153    fn load_bails_on_mxfp4_packed() {
154        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
155        config.model_type = "kimi_k3".into();
156        let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
157        let ptr = gpu.alloc(4).unwrap();
158        let store = WeightStore::from_map(HashMap::from([(
159            "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight_packed".to_string(),
160            WeightTensor {
161                ptr,
162                shape: vec![2, 2],
163                dtype: WeightDtype::UInt8,
164            },
165        )]));
166        let loader = KimiK3WeightLoader;
167        let err = match loader.load_layers(&store, &config, &gpu, &[]) {
168            Ok(_) => panic!("packed MXFP4 must not bind layers"),
169            Err(e) => e.to_string(),
170        };
171        assert!(err.contains("S5 MXFP4 not this slice"), "{err}");
172        let err = match loader.load_embedding(&store, &config, &gpu) {
173            Ok(_) => panic!("packed MXFP4 must not bind embed"),
174            Err(e) => e.to_string(),
175        };
176        assert!(err.contains("S5 MXFP4 not this slice"), "{err}");
177    }
178
179    #[test]
180    fn packed_tp_refuses_before_weight_binding() {
181        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
182        config.tp_world_size = 8;
183        let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
184        let store = WeightStore::from_map(HashMap::from([(
185            "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight_packed".into(),
186            WeightTensor {
187                ptr: gpu.alloc(16).unwrap(),
188                shape: vec![1, 16],
189                dtype: WeightDtype::UInt8,
190            },
191        )]));
192        let err = match bf16::load_layers(&store, &config, &gpu) {
193            Ok(_) => panic!("packed TP must refuse an unmarked store"),
194            Err(err) => err.to_string(),
195        };
196        assert!(err.contains("does not slice packed MXFP4"), "{err}");
197    }
198
199    #[test]
200    fn load_allow_mxfp4_lands_packed_experts() {
201        const THIS: &str = "weight_loader::kimi_k3::tests::load_allow_mxfp4_lands_packed_experts";
202        const MARKER: &str = "K3_MXFP4_GPU_ALLOW_CHILD";
203        if std::env::var_os(MARKER).is_some() {
204            const TWIN: &str =
205                include_str!("../../../../docs/k3/fixtures/Kimi-K3-0.40B-config.json");
206            let mut config = parse_config(TWIN).expect("0.40B twin");
207            config.tp_world_size = 1;
208            let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
209            let graph = avarok_core::kimi_k3::K3Graph::from_config(&config);
210            let mut map = HashMap::new();
211            let mut put = |name: String, shape: Vec<usize>, dtype: WeightDtype| {
212                let n = shape.iter().product::<usize>().max(1);
213                let ptr = gpu.alloc(n.max(4)).unwrap();
214                map.insert(name, WeightTensor { ptr, shape, dtype });
215            };
216            put(
217                bf16::text_key(&config, "model.embed_tokens.weight"),
218                vec![2],
219                WeightDtype::BF16,
220            );
221            put(
222                bf16::text_key(&config, "model.norm.weight"),
223                vec![2],
224                WeightDtype::BF16,
225            );
226            put(
227                bf16::text_key(&config, "model.output_attn_res_proj.weight"),
228                vec![2],
229                WeightDtype::BF16,
230            );
231            put(
232                bf16::text_key(&config, "model.output_attn_res_norm.weight"),
233                vec![2],
234                WeightDtype::BF16,
235            );
236            put(
237                bf16::text_key(&config, "lm_head.weight"),
238                vec![2],
239                WeightDtype::BF16,
240            );
241            let packed_w1 = bf16::text_key(
242                &config,
243                "model.layers.1.block_sparse_moe.experts.0.w1.weight",
244            );
245            for spec in &graph.layers {
246                for k in bf16::layer_keys(
247                    &config,
248                    spec.index,
249                    spec.mixer,
250                    spec.mlp,
251                    config.num_experts,
252                ) {
253                    if k == packed_w1 {
254                        continue;
255                    }
256                    put(k, vec![2], WeightDtype::BF16);
257                }
258            }
259            let prefix = packed_w1.trim_end_matches(".weight");
260            put(
261                format!("{prefix}.weight_packed"),
262                vec![config.moe_intermediate_size, config.moe_latent_size / 2],
263                WeightDtype::UInt8,
264            );
265            put(
266                format!("{prefix}.weight_scale"),
267                vec![config.moe_intermediate_size, config.moe_latent_size / 32],
268                WeightDtype::UInt8,
269            );
270            let store = WeightStore::from_map(map);
271            refuse_mxfp4(&store).expect("K3_ALLOW_MXFP4=1 must skip S5 refuse");
272            let loader = KimiK3WeightLoader;
273            let layers = loader
274                .load_layers(&store, &config, &gpu, &[])
275                .expect("packed expert must land on DSV4 E8M0");
276            assert_eq!(layers.len(), 8);
277            return;
278        }
279        let output = Command::new(std::env::current_exe().unwrap())
280            .args(["--exact", THIS])
281            .env(MARKER, "1")
282            .env("K3_ALLOW_MXFP4", "1")
283            .output()
284            .unwrap();
285        assert!(
286            output.status.success(),
287            "ALLOW child failed:\n{}\n{}",
288            String::from_utf8_lossy(&output.stdout),
289            String::from_utf8_lossy(&output.stderr)
290        );
291    }
292
293    #[test]
294    fn load_bf16_twin_binds_layers() {
295        const TWIN: &str = include_str!("../../../../docs/k3/fixtures/Kimi-K3-0.40B-config.json");
296        let config = parse_config(TWIN).expect("0.40B twin");
297        let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
298        let graph = avarok_core::kimi_k3::K3Graph::from_config(&config);
299        let mut map = HashMap::new();
300        let mut put = |name: String| {
301            let ptr = gpu.alloc(4).unwrap();
302            map.insert(
303                name,
304                WeightTensor {
305                    ptr,
306                    shape: vec![2],
307                    dtype: WeightDtype::BF16,
308                },
309            );
310        };
311        put(bf16::text_key(&config, "model.embed_tokens.weight"));
312        put(bf16::text_key(&config, "model.norm.weight"));
313        put(bf16::text_key(&config, "model.output_attn_res_proj.weight"));
314        put(bf16::text_key(&config, "model.output_attn_res_norm.weight"));
315        put(bf16::text_key(&config, "lm_head.weight"));
316        for spec in &graph.layers {
317            for k in bf16::layer_keys(
318                &config,
319                spec.index,
320                spec.mixer,
321                spec.mlp,
322                config.num_experts,
323            ) {
324                put(k);
325            }
326        }
327        let store = WeightStore::from_map(map);
328        let loader = KimiK3WeightLoader;
329        let layers = loader.load_layers(&store, &config, &gpu, &[]).unwrap();
330        assert_eq!(layers.len(), 8);
331        for layer in &layers {
332            assert!(layer.decode_graph_unsupported());
333            assert!(layer.decode_multi_seq_unsupported());
334            assert!(layer.decode_verify_multi_unsupported());
335            assert!(layer.decode_rollback_unsupported());
336            assert!(layer.has_aux_state());
337        }
338        loader.load_embedding(&store, &config, &gpu).unwrap();
339        loader.load_final_norm(&store, &config, &gpu).unwrap();
340        loader.load_lm_head(&store, &config, &gpu).unwrap();
341    }
342}
spark_model/weight_loader/kimi_k3.rs

spark_model/weight_loader/
kimi_k3.rs
