atlas_core/kimi_k3/
expert_backend.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Expert storage backend for K3 routed MXFP4 packs.
4//!
5//! Official `moonshotai/Kimi-K3` is ~1.56 TB / 96 shards. Two Sparks cannot
6//! hold that resident. Trunk + embed + lm_head + KDA/MLA + shared experts
7//! stay in RAM; routed experts come from per-id NVMe packs.
8//!
9//! `K3_EXPERT_BACKEND=resident | mmap | prefetch` (default resident).
10//! Do not download the 96-shard repo in this module.
11
12use std::collections::HashMap;
13use std::fs::File;
14use std::io::{Read, Write};
15use std::path::{Path, PathBuf};
16
17use anyhow::{Context, Result, bail};
18use memmap2::Mmap;
19
20pub const ENV: &str = "K3_EXPERT_BACKEND";
21const MAGIC: &[u8; 4] = b"K3E1";
22
23#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub enum ExpertBackendKind {
25    /// Twin / dummy experts already in RAM.
26    Resident,
27    /// Per-expert MXFP4 packs on NVMe (`e{id}.mxfp4`).
28    Mmap,
29    /// mmap + next-token expert ids from the last router call.
30    Prefetch,
31}
32
33pub fn kind_from_env() -> ExpertBackendKind {
34    match std::env::var(ENV).as_deref() {
35        Ok("mmap") => ExpertBackendKind::Mmap,
36        Ok("prefetch") => ExpertBackendKind::Prefetch,
37        _ => ExpertBackendKind::Resident,
38    }
39}
40
41/// Dummy pack path: `{dir}/e{id}.mxfp4`.
42pub fn pack_path(dir: &Path, expert_id: usize) -> PathBuf {
43    dir.join(format!("e{expert_id}.mxfp4"))
44}
45
46pub fn write_dummy_pack(dir: &Path, expert_id: usize, payload: &[u8]) -> Result<PathBuf> {
47    std::fs::create_dir_all(dir)?;
48    let path = pack_path(dir, expert_id);
49    let mut f = File::create(&path)?;
50    f.write_all(MAGIC)?;
51    f.write_all(&(expert_id as u32).to_le_bytes())?;
52    f.write_all(&(payload.len() as u32).to_le_bytes())?;
53    f.write_all(payload)?;
54    Ok(path)
55}
56
57fn parse_pack(bytes: &[u8]) -> Result<(u32, &[u8])> {
58    if bytes.len() < 12 || &bytes[..4] != MAGIC {
59        bail!("K3 expert pack: bad magic");
60    }
61    let id = u32::from_le_bytes(bytes[4..8].try_into().unwrap());
62    let n = u32::from_le_bytes(bytes[8..12].try_into().unwrap()) as usize;
63    if bytes.len() != 12 + n {
64        bail!("K3 expert pack: length mismatch");
65    }
66    Ok((id, &bytes[12..]))
67}
68
69/// mmap'd per-expert packs. Missing id fails closed (no silent zero tensor).
70pub struct MmapExpertStore {
71    maps: HashMap<usize, Mmap>,
72}
73
74impl MmapExpertStore {
75    pub fn open(dir: &Path) -> Result<Self> {
76        let mut maps = HashMap::new();
77        if !dir.is_dir() {
78            bail!("K3 expert mmap dir {} missing", dir.display());
79        }
80        for ent in std::fs::read_dir(dir)? {
81            let ent = ent?;
82            let name = ent.file_name();
83            let name = name.to_string_lossy();
84            if !name.starts_with('e') || !name.ends_with(".mxfp4") {
85                continue;
86            }
87            let f = File::open(ent.path())?;
88            // SAFETY: packs are immutable after write_dummy_pack / rental stage.
89            let mmap =
90                unsafe { Mmap::map(&f) }.with_context(|| ent.path().display().to_string())?;
91            let (id, _) = parse_pack(&mmap)?;
92            maps.insert(id as usize, mmap);
93        }
94        Ok(Self { maps })
95    }
96
97    pub fn get(&self, expert_id: usize) -> Result<&[u8]> {
98        let mmap = self
99            .maps
100            .get(&expert_id)
101            .with_context(|| format!("K3 expert {expert_id} pack missing (no silent host F32)"))?;
102        let (id, payload) = parse_pack(mmap)?;
103        if id as usize != expert_id {
104            bail!("K3 expert pack id {id} != requested {expert_id}");
105        }
106        Ok(payload)
107    }
108}
109
110/// Prefetch: remember last router top-k so the next token can hint NVMe.
111#[derive(Clone, Debug, Default)]
112pub struct PrefetchPlanner {
113    last_ids: Vec<usize>,
114}
115
116impl PrefetchPlanner {
117    pub fn note_router(&mut self, ids: &[usize]) {
118        self.last_ids = ids.to_vec();
119    }
120
121    pub fn next_hint(&self) -> &[usize] {
122        &self.last_ids
123    }
124}
125
126/// Read a pack without mmap (tests / tiny payloads).
127pub fn read_pack_file(path: &Path) -> Result<(u32, Vec<u8>)> {
128    let mut bytes = Vec::new();
129    File::open(path)?.read_to_end(&mut bytes)?;
130    let (id, payload) = parse_pack(&bytes)?;
131    Ok((id, payload.to_vec()))
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137
138    fn scratch() -> PathBuf {
139        let nanos = std::time::SystemTime::now()
140            .duration_since(std::time::UNIX_EPOCH)
141            .unwrap()
142            .as_nanos();
143        std::env::temp_dir().join(format!("k3-exp-{}-{nanos}", std::process::id()))
144    }
145
146    #[test]
147    fn default_kind_is_resident() {
148        if std::env::var_os(ENV).is_some() {
149            return;
150        }
151        assert_eq!(kind_from_env(), ExpertBackendKind::Resident);
152    }
153
154    #[test]
155    fn mmap_dummy_pack_roundtrips() {
156        let dir = scratch();
157        write_dummy_pack(&dir, 3, &[1, 2, 3, 4]).unwrap();
158        let store = MmapExpertStore::open(&dir).unwrap();
159        assert_eq!(store.get(3).unwrap(), &[1, 2, 3, 4]);
160        let _ = std::fs::remove_dir_all(&dir);
161    }
162
163    #[test]
164    fn missing_pack_does_not_silent_zero() {
165        let dir = scratch();
166        write_dummy_pack(&dir, 0, &[9]).unwrap();
167        let store = MmapExpertStore::open(&dir).unwrap();
168        let err = store.get(1).unwrap_err().to_string();
169        assert!(
170            err.contains("missing") && err.contains("no silent host F32"),
171            "{err}"
172        );
173        let _ = std::fs::remove_dir_all(&dir);
174    }
175
176    #[test]
177    fn prefetch_notes_router_ids() {
178        let mut p = PrefetchPlanner::default();
179        assert!(p.next_hint().is_empty());
180        p.note_router(&[7, 1]);
181        assert_eq!(p.next_hint(), &[7, 1]);
182    }
183
184    #[test]
185    fn mmap_env_child() {
186        const THIS: &str = "kimi_k3::expert_backend::tests::mmap_env_child";
187        const MARKER: &str = "K3_EXPERT_BACKEND_CHILD";
188        if std::env::var_os(MARKER).is_some() {
189            assert_eq!(kind_from_env(), ExpertBackendKind::Mmap);
190            return;
191        }
192        let output = std::process::Command::new(std::env::current_exe().unwrap())
193            .args(["--exact", THIS])
194            .env(MARKER, "1")
195            .env(ENV, "mmap")
196            .output()
197            .unwrap();
198        assert!(
199            output.status.success(),
200            "{}",
201            String::from_utf8_lossy(&output.stderr)
202        );
203    }
204}
atlas_core/kimi_k3/expert_backend.rs

atlas_core/kimi_k3/
expert_backend.rs
