spark_model/kimi_k3/
device_cache.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GPU-resident hybrid cache. Host [`HybridCache`] bytes are prefix snapshot
4//! only. Decode must not D2H/H2D conv, recurrent, or the growing MLA KV.
5
6use anyhow::Result;
7use atlas_core::kimi_k3::{HybridCache, KdaConfig, LayerCache};
8use spark_runtime::gpu::{DevicePtr, GpuBackend};
9
10use super::kda_cuda::KdaDeviceState;
11
12/// One layer's device buffers. MLA `k`/`v` are pre-allocated to `cap` tokens
13/// so append is a strided H2D of the new row only.
14pub struct DeviceLayerCache {
15    pub kda: Option<KdaDeviceState>,
16    pub mla_k: Option<DevicePtr>,
17    pub mla_v: Option<DevicePtr>,
18    pub seq_len: usize,
19    pub cap: usize,
20}
21
22pub struct DeviceHybridCache {
23    pub layers: Vec<DeviceLayerCache>,
24}
25
26impl DeviceHybridCache {
27    /// Allocate device buffers from a host snapshot (prefix restore / cold start).
28    pub fn from_host(
29        gpu: &dyn GpuBackend,
30        host: &HybridCache,
31        kda_cfg: &KdaConfig,
32        mla_cap: usize,
33        mla_k_row: usize,
34        mla_v_row: usize,
35    ) -> Result<Self> {
36        let mut layers = Vec::with_capacity(host.layers.len());
37        for slot in &host.layers {
38            layers.push(match slot {
39                LayerCache::Kda(s) => DeviceLayerCache {
40                    kda: Some(KdaDeviceState::alloc_and_upload(gpu, s)?),
41                    mla_k: None,
42                    mla_v: None,
43                    seq_len: 0,
44                    cap: 0,
45                },
46                LayerCache::Mla(kv) => {
47                    let cap = mla_cap.max(kv.seq_len.max(1));
48                    let k = gpu.alloc((cap * mla_k_row * 4).max(1))?;
49                    let v = gpu.alloc((cap * mla_v_row * 4).max(1))?;
50                    if !kv.k.is_empty() {
51                        let kb: Vec<u8> = kv.k.iter().flat_map(|x| x.to_le_bytes()).collect();
52                        gpu.copy_h2d(&kb, k)?;
53                    }
54                    if !kv.v.is_empty() {
55                        let vb: Vec<u8> = kv.v.iter().flat_map(|x| x.to_le_bytes()).collect();
56                        gpu.copy_h2d(&vb, v)?;
57                    }
58                    DeviceLayerCache {
59                        kda: None,
60                        mla_k: Some(k),
61                        mla_v: Some(v),
62                        seq_len: kv.seq_len,
63                        cap,
64                    }
65                }
66            });
67        }
68        let _ = kda_cfg;
69        Ok(Self { layers })
70    }
71
72    pub fn kda(&self, layer: usize) -> Option<&KdaDeviceState> {
73        self.layers.get(layer).and_then(|l| l.kda.as_ref())
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80    use atlas_core::config::parse_config;
81    use atlas_core::kimi_k3::{K3Graph, MixerKind};
82    use spark_runtime::gpu::mock::MockGpuBackend;
83
84    #[test]
85    fn twin_device_cache_kda_slots_are_resident() {
86        const TWIN: &str = include_str!("../../../../docs/k3/fixtures/Kimi-K3-0.40B-config.json");
87        let c = parse_config(TWIN).unwrap();
88        let g = K3Graph::from_config(&c);
89        let kda = KdaConfig::twin_0_40b();
90        let host = HybridCache::from_graph(&g, &kda);
91        let gpu = MockGpuBackend::new();
92        let before = gpu.d2h_blocking_count();
93        let dev = DeviceHybridCache::from_host(&gpu, &host, &kda, 16, 8 * 96, 8 * 64).unwrap();
94        assert_eq!(gpu.d2h_blocking_count(), before, "seed is H2D, not D2H");
95        for i in [0, 1, 2, 4, 5, 6] {
96            assert!(dev.kda(i).is_some(), "KDA layer {i}");
97            assert!(matches!(g.layers[i].mixer, MixerKind::Kda));
98        }
99        for i in [3, 7] {
100            assert!(dev.layers[i].mla_k.is_some());
101            assert!(dev.kda(i).is_none());
102        }
103    }
104}
spark_model/kimi_k3/device_cache.rs

spark_model/kimi_k3/
device_cache.rs
