avarok_core/kimi_k3/
cpu_load.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Load unpacked BF16/F32 K3 safetensors into [`K3CpuModel`].
4//! Packed MXFP4 (`weight_packed`) is refused unless `K3_ALLOW_MXFP4=1`, which
5//! unpacks via the DSV4 E8M0 host path (`crate::mxfp4_e8m0`).
6
7use std::collections::HashMap;
8use std::fs::File;
9use std::path::{Path, PathBuf};
10#[cfg(test)]
11use std::sync::OnceLock;
12
13use anyhow::{Context, Result, bail};
14use memmap2::Mmap;
15use safetensors::SafeTensors;
16use safetensors::tensor::Dtype;
17
18use super::cpu_bind::{assemble, inventory, text_key};
19use super::cpu_weights::K3CpuModel;
20use super::{kda_from, mla_from, moe_from};
21use crate::config::parse_config;
22
23impl K3CpuModel {
24    /// Bind a HuggingFace twin dir (`config.json` + unpacked `.safetensors`).
25    pub fn from_pretrained(dir: &Path) -> Result<Self> {
26        if !dir.exists() {
27            bail!("K3 twin dir {} does not exist", dir.display());
28        }
29        let cfg_path = dir.join("config.json");
30        let raw = std::fs::read_to_string(&cfg_path)
31            .with_context(|| format!("read {}", cfg_path.display()))?;
32        let c = parse_config(&raw).context("parse K3 config.json")?;
33        let graph = super::layer::K3Graph::from_config(&c);
34        let kda = kda_from(&c);
35        let mla = mla_from(&c);
36        let moe = moe_from(&c);
37        let mut store = load_dir(dir)?;
38        let mut got = HashMap::new();
39        for (name, n) in inventory(&c, &graph, &kda, &mla, &moe) {
40            got.insert(name.clone(), take_len(&mut store, &name, n)?);
41        }
42        if c.tie_word_embeddings {
43            let embed_k = text_key(&c.weight_prefix, "model.embed_tokens.weight");
44            let head_k = text_key(&c.weight_prefix, "lm_head.weight");
45            if !got.contains_key(&head_k) {
46                let embed = got
47                    .get(&embed_k)
48                    .with_context(|| format!("tied lm_head needs {embed_k}"))?
49                    .clone();
50                got.insert(head_k, embed);
51            }
52        }
53        let mut unused: Vec<String> = store
54            .keys()
55            .filter(|k| is_text_weight(k))
56            .cloned()
57            .collect();
58        if !unused.is_empty() {
59            unused.sort();
60            eprintln!(
61                "K3CpuModel::from_pretrained unmapped tensors: {}",
62                unused.join(", ")
63            );
64        }
65        assemble(&c, graph, kda, mla, moe, &mut got)
66    }
67}
68
69/// Aviation prompt 0 ids through the comma. C1 first generated id is 1459.
70#[cfg(test)]
71pub(super) const TWIN_PROMPT0: &[u32] = &[18805, 308, 799, 5624, 12524, 318, 57195, 11];
72#[cfg(test)]
73pub(super) const TWIN_PROMPT0_FIRST: u32 = 1459;
74
75/// Explicit checkpoint tests share the expensive host bind. Missing input fails.
76#[cfg(test)]
77pub(super) fn twin_from_env() -> Option<&'static K3CpuModel> {
78    static TWIN: OnceLock<K3CpuModel> = OnceLock::new();
79    Some(TWIN.get_or_init(|| {
80        let p = std::env::var("K3_TWIN").expect("K3_TWIN must name the 0.40B checkpoint");
81        let path = Path::new(&p);
82        assert!(path.is_dir(), "K3_TWIN={p} is not a directory");
83        K3CpuModel::from_pretrained(path)
84            .unwrap_or_else(|e| panic!("K3_TWIN={p} load failed: {e:#}"))
85    }))
86}
87
88fn is_vision(name: &str) -> bool {
89    name.starts_with("vision_tower.") || name.starts_with("mm_projector.")
90}
91
92fn is_text_weight(name: &str) -> bool {
93    !is_vision(name)
94        && (name.starts_with("language_model.")
95            || name.starts_with("model.")
96            || name.starts_with("lm_head."))
97}
98
99fn take_len(store: &mut HashMap<String, Vec<f32>>, key: &str, n: usize) -> Result<Vec<f32>> {
100    let v = take_tensor(store, key)?;
101    if v.len() != n {
102        bail!("{key}: {} elems, expected {n}", v.len());
103    }
104    Ok(v)
105}
106
107fn take_tensor(store: &mut HashMap<String, Vec<f32>>, key: &str) -> Result<Vec<f32>> {
108    if let Some(v) = store.remove(key) {
109        return Ok(v);
110    }
111    if let Some(rest) = key.strip_prefix("language_model.")
112        && let Some(v) = store.remove(rest)
113    {
114        return Ok(v);
115    }
116    if key.ends_with("lm_head.weight")
117        && let Some(v) = store.remove("lm_head.weight")
118    {
119        return Ok(v);
120    }
121    bail!("tensor {key} not found")
122}
123
124fn load_dir(dir: &Path) -> Result<HashMap<String, Vec<f32>>> {
125    let mut out = HashMap::new();
126    let mut packed = super::mxfp4::PackedSink::default();
127    for path in shard_files(dir)? {
128        ingest_shard(&path, &mut out, &mut packed)?;
129    }
130    packed.finish(&mut out)?;
131    if out.is_empty() {
132        bail!("no language-model tensors in {}", dir.display());
133    }
134    Ok(out)
135}
136
137fn shard_files(dir: &Path) -> Result<Vec<PathBuf>> {
138    let index = dir.join("model.safetensors.index.json");
139    if index.exists() {
140        let raw =
141            std::fs::read_to_string(&index).with_context(|| format!("read {}", index.display()))?;
142        let v: serde_json::Value = serde_json::from_str(&raw).context("safetensors index JSON")?;
143        let map = v
144            .get("weight_map")
145            .and_then(|m| m.as_object())
146            .context("index missing weight_map")?;
147        let mut files: Vec<String> = map
148            .values()
149            .filter_map(|x| x.as_str().map(str::to_string))
150            .collect();
151        files.sort();
152        files.dedup();
153        return Ok(files.into_iter().map(|f| dir.join(f)).collect());
154    }
155    let single = dir.join("model.safetensors");
156    if single.exists() {
157        return Ok(vec![single]);
158    }
159    let mut found = Vec::new();
160    for ent in std::fs::read_dir(dir).with_context(|| format!("read_dir {}", dir.display()))? {
161        let p = ent?.path();
162        if p.extension().and_then(|e| e.to_str()) == Some("safetensors") {
163            found.push(p);
164        }
165    }
166    found.sort();
167    if found.is_empty() {
168        bail!(
169            "no safetensors in {} (expected model.safetensors or index.json)",
170            dir.display()
171        );
172    }
173    Ok(found)
174}
175
176fn ingest_shard(
177    path: &Path,
178    out: &mut HashMap<String, Vec<f32>>,
179    packed: &mut super::mxfp4::PackedSink,
180) -> Result<()> {
181    let file = File::open(path).with_context(|| format!("open {}", path.display()))?;
182    // SAFETY: read-only mapping of an immutable checkpoint shard.
183    let mmap = unsafe { Mmap::map(&file) }.with_context(|| format!("mmap {}", path.display()))?;
184    let st = SafeTensors::deserialize(&mmap)
185        .with_context(|| format!("safetensors {}", path.display()))?;
186    for name in st.names() {
187        if is_vision(name) {
188            continue;
189        }
190        let t = st.tensor(name)?;
191        if packed.take(name, t.dtype(), t.shape(), t.data())? {
192            continue;
193        }
194        let data = to_f32(name, t.dtype(), t.data())?;
195        if out.insert(name.to_string(), data).is_some() {
196            bail!("duplicate tensor {name} in {}", path.display());
197        }
198    }
199    Ok(())
200}
201
202fn to_f32(name: &str, dtype: Dtype, data: &[u8]) -> Result<Vec<f32>> {
203    match dtype {
204        Dtype::F32 => {
205            if !data.len().is_multiple_of(4) {
206                bail!(
207                    "tensor {name}: F32 byte length {} not multiple of 4",
208                    data.len()
209                );
210            }
211            Ok(data
212                .chunks_exact(4)
213                .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
214                .collect())
215        }
216        Dtype::BF16 => {
217            if !data.len().is_multiple_of(2) {
218                bail!(
219                    "tensor {name}: BF16 byte length {} not multiple of 2",
220                    data.len()
221                );
222            }
223            Ok(data
224                .chunks_exact(2)
225                .map(|b| {
226                    let u = u16::from_le_bytes([b[0], b[1]]);
227                    f32::from_bits((u as u32) << 16)
228                })
229                .collect())
230        }
231        other => bail!("tensor {name}: unsupported dtype {other:?} (want BF16 or F32)"),
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238    use crate::kimi_k3::cpu_weights::{Ablation, MixerW};
239    use crate::kimi_k3::greedy::greedy_decode;
240    use crate::kimi_k3::layer::K3Graph;
241    use safetensors::tensor::{TensorView, serialize};
242    use std::env;
243
244    const TINY_CFG: &str = r#"{
245      "model_type": "kimi_k3",
246      "text_config": {
247        "activation_situ_beta": 4.0,
248        "activation_situ_linear_beta": 25.0,
249        "attn_res_block_size": 2,
250        "first_k_dense_replace": 1,
251        "head_dim": 2,
252        "hidden_act": "situ",
253        "hidden_size": 4,
254        "intermediate_size": 4,
255        "kv_lora_rank": 2,
256        "latent_moe_use_norm": true,
257        "linear_attn_config": {
258          "full_attn_layers": [2],
259          "head_dim": 2,
260          "kda_layers": [1],
261          "num_heads": 1,
262          "short_conv_kernel_size": 4,
263          "use_full_rank_gate": true
264        },
265        "mla_use_nope": true,
266        "mla_use_output_gate": true,
267        "model_type": "kimi_linear",
268        "moe_intermediate_size": 4,
269        "moe_renormalize": true,
270        "moe_router_activation_func": "sigmoid",
271        "num_attention_heads": 1,
272        "num_experts": 2,
273        "num_experts_per_token": 1,
274        "num_hidden_layers": 2,
275        "num_shared_experts": 1,
276        "q_lora_rank": 4,
277        "qk_nope_head_dim": 2,
278        "qk_rope_head_dim": 2,
279        "rms_norm_eps": 1e-5,
280        "rope_theta": 10000.0,
281        "routed_expert_hidden_size": 4,
282        "topk_method": "noaux_tc",
283        "v_head_dim": 2,
284        "vocab_size": 8
285      }
286    }"#;
287
288    fn write_f32_st(dir: &Path, tensors: &[(String, usize)]) {
289        let backing: Vec<(String, Vec<u8>)> = tensors
290            .iter()
291            .map(|(n, len)| {
292                let bytes: Vec<u8> = (0..*len)
293                    .flat_map(|i| (0.01 * (i as f32 + 1.0)).to_le_bytes())
294                    .collect();
295                (n.clone(), bytes)
296            })
297            .collect();
298        let views: Vec<(String, TensorView)> = tensors
299            .iter()
300            .zip(backing.iter())
301            .map(|((n, len), (_, b))| {
302                (
303                    n.clone(),
304                    TensorView::new(Dtype::F32, vec![*len], b).expect(n),
305                )
306            })
307            .collect();
308        let bytes = serialize(views, None).expect("serialize");
309        std::fs::write(dir.join("model.safetensors"), bytes).unwrap();
310    }
311
312    fn scratch_dir(tag: &str) -> PathBuf {
313        let nanos = std::time::SystemTime::now()
314            .duration_since(std::time::UNIX_EPOCH)
315            .unwrap()
316            .as_nanos();
317        let dir = env::temp_dir().join(format!("k3-{tag}-{}-{nanos}", std::process::id()));
318        std::fs::create_dir_all(&dir).unwrap();
319        dir
320    }
321
322    fn tiny_dir() -> PathBuf {
323        let dir = scratch_dir("cpu-load");
324        std::fs::write(dir.join("config.json"), TINY_CFG).unwrap();
325        let c = parse_config(TINY_CFG).unwrap();
326        let graph = K3Graph::from_config(&c);
327        let inv = inventory(&c, &graph, &kda_from(&c), &mla_from(&c), &moe_from(&c));
328        write_f32_st(&dir, &inv);
329        dir
330    }
331
332    #[test]
333    fn from_pretrained_missing_dir_errors() {
334        let err = K3CpuModel::from_pretrained(Path::new("/no/such/k3-twin"))
335            .unwrap_err()
336            .to_string();
337        assert!(err.contains("does not exist"), "{err}");
338    }
339
340    #[test]
341    fn from_pretrained_tiny_safetensors() {
342        let dir = tiny_dir();
343        let m = K3CpuModel::from_pretrained(&dir).expect("load tiny");
344        assert_eq!(m.layers.len(), 2);
345        assert_eq!(m.vocab, 8);
346        assert_eq!(m.graph.hidden, 4);
347        assert!(matches!(m.layers[0].mixer, MixerW::Kda(_)));
348        assert!(matches!(m.layers[1].mixer, MixerW::Mla(_)));
349        let out = greedy_decode(&m, &[1, 2], 3, Ablation::default());
350        assert_eq!(out.len(), 5);
351        let _ = std::fs::remove_dir_all(&dir);
352    }
353
354    #[test]
355    fn from_pretrained_refuses_weight_packed() {
356        let dir = scratch_dir("packed");
357        std::fs::write(dir.join("config.json"), TINY_CFG).unwrap();
358        let data = [0u8; 4];
359        let view = TensorView::new(Dtype::U8, vec![4], &data).unwrap();
360        let bytes = serialize(
361            [(
362                "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight_packed"
363                    .to_string(),
364                view,
365            )],
366            None,
367        )
368        .unwrap();
369        std::fs::write(dir.join("model.safetensors"), bytes).unwrap();
370        let err = K3CpuModel::from_pretrained(&dir).unwrap_err().to_string();
371        assert!(err.contains("S5 MXFP4 not this slice"), "{err}");
372        let _ = std::fs::remove_dir_all(&dir);
373    }
374
375    #[test]
376    fn bf16_converts_one() {
377        let bits = 0x3f80u16;
378        let v = to_f32("t", Dtype::BF16, &bits.to_le_bytes()).unwrap();
379        assert_eq!(v, vec![1.0]);
380    }
381}
avarok_core/kimi_k3/cpu_load.rs

avarok_core/kimi_k3/
cpu_load.rs
