metrale_model_arch/kimi_k3/
device_cache.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `DeviceHybridCache`: device copies of every layer's Kimi K3 mixer state, built from a host [`HybridCache`].
4//!
5//! Owner: model-arch, Kimi K3.
6//! Invariants: none beyond the types.
7
8use anyhow::Result;
9use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
10use metrale_model_weights::kimi_k3_host::{HybridCache, KdaConfig, LayerCache};
11
12use super::kda_cuda::KdaDeviceState;
13
14/// 2026-09-25: One layer's device state: KDA state for a KDA layer, or MLA
15/// `k`/`v` buffers sized for `cap` tokens for an MLA layer.
16pub struct DeviceLayerCache {
17    pub kda: Option<KdaDeviceState>,
18    pub mla_k: Option<DevicePtr>,
19    pub mla_v: Option<DevicePtr>,
20    pub seq_len: usize,
21    pub cap: usize,
22}
23
24pub struct DeviceHybridCache {
25    pub layers: Vec<DeviceLayerCache>,
26}
27
28impl DeviceHybridCache {
29    /// 2026-09-25: Allocates each layer's device state and uploads the host
30    /// state into it. An MLA layer's `cap` is `mla_cap`, raised to the host
31    /// history length.
32    pub fn from_host(
33        gpu: &dyn GpuBackend,
34        host: &HybridCache,
35        kda_cfg: &KdaConfig,
36        mla_cap: usize,
37        mla_k_row: usize,
38        mla_v_row: usize,
39    ) -> Result<Self> {
40        let mut layers = Vec::with_capacity(host.layers.len());
41        for slot in &host.layers {
42            layers.push(match slot {
43                LayerCache::Kda(s) => DeviceLayerCache {
44                    kda: Some(KdaDeviceState::alloc_and_upload(gpu, s)?),
45                    mla_k: None,
46                    mla_v: None,
47                    seq_len: 0,
48                    cap: 0,
49                },
50                LayerCache::Mla(kv) => {
51                    let cap = mla_cap.max(kv.seq_len.max(1));
52                    let k = gpu.alloc((cap * mla_k_row * 4).max(1))?;
53                    let v = gpu.alloc((cap * mla_v_row * 4).max(1))?;
54                    if !kv.k.is_empty() {
55                        let kb: Vec<u8> = kv.k.iter().flat_map(|x| x.to_le_bytes()).collect();
56                        gpu.copy_h2d(&kb, k)?;
57                    }
58                    if !kv.v.is_empty() {
59                        let vb: Vec<u8> = kv.v.iter().flat_map(|x| x.to_le_bytes()).collect();
60                        gpu.copy_h2d(&vb, v)?;
61                    }
62                    DeviceLayerCache {
63                        kda: None,
64                        mla_k: Some(k),
65                        mla_v: Some(v),
66                        seq_len: kv.seq_len,
67                        cap,
68                    }
69                }
70            });
71        }
72        let _ = kda_cfg;
73        Ok(Self { layers })
74    }
75
76    pub fn kda(&self, layer: usize) -> Option<&KdaDeviceState> {
77        self.layers.get(layer).and_then(|l| l.kda.as_ref())
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use metrale_config::parse_config;
85    use metrale_gpu_runtime::gpu::mock::MockGpuBackend;
86    use metrale_model_weights::kimi_k3_host::{K3Graph, MixerKind};
87
88    #[test]
89    fn twin_device_cache_kda_slots_are_resident() {
90        const TWIN: &str = include_str!("../../../../docs/k3/fixtures/Kimi-K3-0.40B-config.json");
91        let c = parse_config(TWIN).unwrap();
92        let g = K3Graph::from_config(&c);
93        let kda = KdaConfig::twin_0_40b();
94        let host = HybridCache::from_graph(&g, &kda);
95        let gpu = MockGpuBackend::new();
96        let before = gpu.d2h_blocking_count();
97        let dev = DeviceHybridCache::from_host(&gpu, &host, &kda, 16, 8 * 96, 8 * 64).unwrap();
98        assert_eq!(gpu.d2h_blocking_count(), before, "seed is H2D, not D2H");
99        for i in [0, 1, 2, 4, 5, 6] {
100            assert!(dev.kda(i).is_some(), "KDA layer {i}");
101            assert!(matches!(g.layers[i].mixer, MixerKind::Kda));
102        }
103        for i in [3, 7] {
104            assert!(dev.layers[i].mla_k.is_some());
105            assert!(dev.kda(i).is_none());
106        }
107    }
108}
metrale_model_arch/kimi_k3/device_cache.rs

metrale_model_arch/kimi_k3/
device_cache.rs
