metrale_model_arch/kimi_k3/
state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `K3CpuFallbackState`, one sequence's Kimi K3 mixer state: the host [`LayerCache`] and the device copies the CUDA mixers use.
4//!
5//! Snapshots and restores use the [`LayerCache`] byte format.
6//!
7//! Owner: model-arch, Kimi K3.
8//! Invariants:
9//! - `restore` rejects a blob that does not parse, or whose mixer kind or
10//!   geometry differs from the current state, before it changes anything.
11//! - `release` drops both device copies even when freeing one fails.
12
13use std::any::Any;
14
15use anyhow::{Result, bail, ensure};
16use metrale_gpu_runtime::gpu::GpuBackend;
17use metrale_model_weights::kimi_k3_host::{LayerCache, MlaConfig};
18
19use super::kda_cuda::KdaDeviceState;
20use super::mla_cuda::MlaDeviceKv;
21
22use metrale_model_layers::layer::LayerState;
23
24/// 2026-09-25: Mixer state owned by one sequence; `K3BoundLayer::release_state`
25/// frees its device copies.
26pub struct K3CpuFallbackState {
27    pub cache: LayerCache,
28    /// 2026-09-25: Once present, this is the current KDA state and `cache` holds
29    /// only what it was seeded from.
30    pub device_kda: Option<KdaDeviceState>,
31    /// 2026-09-25: Once present, this is the current MLA KV and `cache` holds only
32    /// what it was seeded from.
33    pub device_mla: Option<MlaDeviceKv>,
34}
35
36impl K3CpuFallbackState {
37    pub fn new(cache: LayerCache) -> Self {
38        Self {
39            cache,
40            device_kda: None,
41            device_mla: None,
42        }
43    }
44
45    pub fn ensure_device_kda(&mut self, gpu: &dyn GpuBackend) -> Result<()> {
46        if self.device_kda.is_none()
47            && let LayerCache::Kda(host) = &self.cache
48        {
49            self.device_kda = Some(KdaDeviceState::alloc_and_upload(gpu, host)?);
50        }
51        Ok(())
52    }
53
54    pub fn ensure_device_mla(
55        &mut self,
56        gpu: &dyn GpuBackend,
57        cfg: &MlaConfig,
58        cap: usize,
59    ) -> Result<()> {
60        if let Some(device) = &self.device_mla {
61            return device.validate_cfg(cfg);
62        }
63        let LayerCache::Mla(host) = &self.cache else {
64            bail!("K3 CUDA MLA: mixer state is not MLA KV");
65        };
66        let k_row = cfg.heads * cfg.qk_head_dim();
67        let v_row = cfg.heads * cfg.v_head_dim;
68        self.device_mla = Some(MlaDeviceKv::alloc_and_upload(gpu, host, cap, k_row, v_row)?);
69        Ok(())
70    }
71
72    pub fn snapshot(&self, gpu: &dyn GpuBackend, stream: u64) -> Result<Vec<u8>> {
73        let mut cache = self.cache.clone();
74        if let (Some(device), LayerCache::Kda(host)) = (&self.device_kda, &mut cache) {
75            gpu.synchronize(stream)?;
76            device.download(gpu, host)?;
77        }
78        if let (Some(device), LayerCache::Mla(host)) = (&self.device_mla, &mut cache) {
79            gpu.synchronize(stream)?;
80            device.download(gpu, host)?;
81        }
82        Ok(cache.to_bytes())
83    }
84
85    pub fn restore(&mut self, gpu: &dyn GpuBackend, bytes: &[u8], stream: u64) -> Result<()> {
86        let cache = LayerCache::from_bytes(bytes)?;
87        match (&self.cache, &cache) {
88            (LayerCache::Kda(current), LayerCache::Kda(restored)) => ensure!(
89                current.conv.len() == restored.conv.len()
90                    && current.recurrent.len() == restored.recurrent.len(),
91                "K3 restore: KDA state geometry mismatch"
92            ),
93            (LayerCache::Mla(current), LayerCache::Mla(restored)) => {
94                validate_mla_restore(current, restored, self.device_mla.as_ref())?;
95            }
96            _ => bail!("K3 restore: mixer variant mismatch"),
97        }
98        gpu.synchronize(stream)?;
99        self.release(gpu)?;
100        self.cache = cache;
101        Ok(())
102    }
103
104    pub fn release(&mut self, gpu: &dyn GpuBackend) -> Result<()> {
105        let kda = if let Some(device) = self.device_kda.take() {
106            device.free(gpu)
107        } else {
108            Ok(())
109        };
110        let mla = if let Some(device) = self.device_mla.take() {
111            device.free(gpu)
112        } else {
113            Ok(())
114        };
115        kda.and(mla)
116    }
117}
118
119fn validate_mla_restore(
120    current: &metrale_model_weights::kimi_k3_host::MlaKv,
121    restored: &metrale_model_weights::kimi_k3_host::MlaKv,
122    device: Option<&MlaDeviceKv>,
123) -> Result<()> {
124    if restored.seq_len == 0 {
125        ensure!(
126            restored.k.is_empty() && restored.v.is_empty(),
127            "K3 restore: empty MLA seq with payload"
128        );
129        return Ok(());
130    }
131    ensure!(
132        restored.k.len()
133            == restored
134                .seq_len
135                .saturating_mul(restored.k.len() / restored.seq_len)
136            && restored.v.len()
137                == restored
138                    .seq_len
139                    .saturating_mul(restored.v.len() / restored.seq_len)
140            && restored.k.len() / restored.seq_len > 0
141            && restored.v.len() / restored.seq_len > 0,
142        "K3 restore: MLA KV rank"
143    );
144    let k_row = restored.k.len() / restored.seq_len;
145    let v_row = restored.v.len() / restored.seq_len;
146    if current.seq_len > 0 {
147        ensure!(
148            current.k.len() / current.seq_len == k_row
149                && current.v.len() / current.seq_len == v_row,
150            "K3 restore: MLA KV geometry mismatch"
151        );
152    } else if let Some(device) = device {
153        ensure!(
154            device.k_row() == k_row && device.v_row() == v_row,
155            "K3 restore: MLA KV geometry mismatch"
156        );
157    }
158    Ok(())
159}
160
161impl LayerState for K3CpuFallbackState {
162    fn as_any(&self) -> &dyn Any {
163        self
164    }
165    fn as_any_mut(&mut self) -> &mut dyn Any {
166        self
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173    use metrale_gpu_runtime::gpu::GpuBackend;
174    use metrale_gpu_runtime::gpu::mock::MockGpuBackend;
175    use metrale_model_weights::kimi_k3_host::{KdaConfig, KdaState};
176
177    #[test]
178    fn resident_snapshot_restore_and_release_preserve_authoritative_state() {
179        let gpu = MockGpuBackend::new();
180        let mut state =
181            K3CpuFallbackState::new(LayerCache::Kda(KdaState::new(&KdaConfig::twin_0_40b())));
182        state.ensure_device_kda(&gpu).unwrap();
183        state.ensure_device_kda(&gpu).unwrap();
184        assert_eq!(gpu.alloc_count(), 2, "reuse resident state between tokens");
185        assert!(state.restore(&gpu, b"invalid", 0).is_err());
186        assert_eq!(
187            gpu.alloc_count(),
188            2,
189            "invalid snapshot must leave state intact"
190        );
191        let device = state.device_kda.as_ref().unwrap();
192        gpu.copy_h2d(&7.25f32.to_le_bytes(), device.recurrent)
193            .unwrap();
194        let bytes = state.snapshot(&gpu, 0).unwrap();
195        let LayerCache::Kda(saved) = LayerCache::from_bytes(&bytes).unwrap() else {
196            panic!("expected KDA");
197        };
198        assert_eq!(saved.recurrent[0], 7.25);
199        state.restore(&gpu, &bytes, 0).unwrap();
200        assert!(state.device_kda.is_none());
201        assert_eq!(gpu.alloc_count(), 0);
202        state.ensure_device_kda(&gpu).unwrap();
203        assert_eq!(state.snapshot(&gpu, 0).unwrap(), bytes);
204        state.release(&gpu).unwrap();
205        state.release(&gpu).unwrap();
206        assert_eq!(gpu.alloc_count(), 0);
207    }
208
209    #[test]
210    fn failed_resident_allocation_releases_partial_state() {
211        let gpu = MockGpuBackend::new();
212        let cfg = KdaConfig::twin_0_40b();
213        gpu.set_max_allocation_bytes(cfg.conv_elems() * 4);
214        let mut state = K3CpuFallbackState::new(LayerCache::Kda(KdaState::new(&cfg)));
215        assert!(state.ensure_device_kda(&gpu).is_err());
216        assert!(state.device_kda.is_none());
217        assert_eq!(gpu.alloc_count(), 0);
218    }
219
220    #[test]
221    fn resident_mla_snapshot_restore_and_empty_restore() {
222        use metrale_model_weights::kimi_k3_host::{MlaConfig, MlaKv};
223
224        let gpu = MockGpuBackend::new();
225        let cfg = MlaConfig::twin_0_40b();
226        let k_row = cfg.heads * cfg.qk_head_dim();
227        let v_row = cfg.heads * cfg.v_head_dim;
228        let mut seeded = MlaKv::default();
229        seeded.append(&vec![0.5f32; k_row], &vec![0.25f32; v_row]);
230        let mut state = K3CpuFallbackState::new(LayerCache::Mla(seeded));
231        state.ensure_device_mla(&gpu, &cfg, 8).unwrap();
232        state.ensure_device_mla(&gpu, &cfg, 8).unwrap();
233        assert_eq!(gpu.alloc_count(), 2, "reuse resident MLA between tokens");
234        let bytes = state.snapshot(&gpu, 0).unwrap();
235        let LayerCache::Mla(saved) = LayerCache::from_bytes(&bytes).unwrap() else {
236            panic!("expected MLA");
237        };
238        assert_eq!(saved.seq_len, 1);
239        assert_eq!(saved.k[0], 0.5);
240        assert!(state.restore(&gpu, b"invalid", 0).is_err());
241        assert_eq!(
242            gpu.alloc_count(),
243            2,
244            "invalid snapshot must leave state intact"
245        );
246        state.restore(&gpu, &bytes, 0).unwrap();
247        assert!(state.device_mla.is_none());
248        let empty = LayerCache::Mla(MlaKv::default()).to_bytes();
249        state.restore(&gpu, &empty, 0).unwrap();
250        let LayerCache::Mla(cleared) = &state.cache else {
251            panic!("expected MLA");
252        };
253        assert_eq!(cleared.seq_len, 0);
254        state.ensure_device_mla(&gpu, &cfg, 8).unwrap();
255        assert_eq!(state.device_mla.as_ref().unwrap().seq_len, 0);
256        state.release(&gpu).unwrap();
257        assert_eq!(gpu.alloc_count(), 0);
258    }
259}
metrale_model_arch/kimi_k3/state.rs

metrale_model_arch/kimi_k3/
state.rs
