spark_model/kimi_k3/
state.rs

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

spark_model/kimi_k3/
state.rs
