metrale_model_engine/
kimi_k3_loader.rs

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

metrale_model_engine/
kimi_k3_loader.rs
