spark_model/kimi_k3/
bound.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! BF16/FP32 twin layer bind. Decode copies hidden D2H, runs mixer+MLP+AttnRes,
4//! copies H2D. LinearAttention / KDA conv+recurrent uses CUDA `kda_decode`
5//! with per-sequence resident recurrence unless `K3_CUDA_KDA=0`.
6//! FullAttention / MLA uses CUDA `mla_decode`
7//! unless `K3_CUDA_MLA=0`. Packed LatentMoE experts launch DSV4
8//! `moe_w4a16_grouped_gemm_ptrtable_e8m0`. Router / down / up / shared / SiTU
9//! stay on the host.
10
11use std::collections::HashMap;
12use std::sync::{Arc, OnceLock};
13
14use anyhow::{Context, Result};
15use avarok_core::config::ModelConfig;
16use avarok_core::kimi_k3::{
17    AttnResStream, K3CpuLayer, K3Graph, K3LayerSpec, KdaConfig, LatentMoeConfig, LayerCache,
18    MixerKind, MlaConfig,
19};
20use parking_lot::Mutex;
21use spark_runtime::gpu::{DevicePtr, GpuBackend};
22use spark_runtime::kv_cache::PagedKvCache;
23use spark_runtime::weights::WeightDtype;
24
25use super::kda_cuda::K3KdaDecodeKernels;
26use super::mla_cuda::K3MlaDecodeKernels;
27use super::moe_cuda::K3MoeGemmKernels;
28use super::state::K3CpuFallbackState;
29use crate::layer::{ForwardContext, LayerState, TransformerLayer};
30use crate::weight_map::{DenseWeight, QuantizedWeight};
31
32/// Name + device dtype/numel for a lazy host bind (copy from GPU if needed).
33#[derive(Clone, Debug)]
34pub struct WeightMeta {
35    pub name: String,
36    pub dtype: WeightDtype,
37    pub numel: usize,
38}
39
40/// Shared across every decoder layer of one loaded K3 model.
41pub struct K3HostShared {
42    pub config: ModelConfig,
43    pub graph: K3Graph,
44    pub kda: KdaConfig,
45    pub mla: MlaConfig,
46    pub moe: LatentMoeConfig,
47    pub output_res_proj: DenseWeight,
48    pub output_res_norm: DenseWeight,
49    pub output_res_proj_meta: (WeightDtype, usize),
50    pub output_res_norm_meta: (WeightDtype, usize),
51    pub output_host: OnceLock<(Vec<f32>, Vec<f32>)>,
52    /// Resolved once per loaded model. LinearAttention decode launches these.
53    pub kda_kernels: OnceLock<K3KdaDecodeKernels>,
54    /// Resolved once per loaded model. FullAttention decode launches these
55    /// unless `K3_CUDA_MLA=0`.
56    pub mla_kernels: OnceLock<K3MlaDecodeKernels>,
57    /// Resolved once per loaded model. Packed LatentMoE launches these.
58    pub moe_kernels: OnceLock<K3MoeGemmKernels>,
59    /// AttnRes is per-token across layers. Keyed by this step's `residual`
60    /// pointer so prefill (layer-outer, token-inner) still sees the same
61    /// stream as CPU `forward_token` (token-outer, layer-inner).
62    pub attnres: Mutex<HashMap<DevicePtr, AttnResStream>>,
63}
64
65/// One decoder layer whose BF16 (or FP32) tensors were bound from the store.
66pub struct K3BoundLayer {
67    pub index: usize,
68    pub spec: K3LayerSpec,
69    pub weights: Vec<DenseWeight>,
70    pub weight_meta: Vec<WeightMeta>,
71    /// Packed routed experts landed on DSV4 `quantized_mxfp4_e8m0_pair`.
72    /// Empty for the BF16 twin.
73    pub mxfp4_experts: Vec<(String, QuantizedWeight)>,
74    pub host: OnceLock<K3CpuLayer>,
75    pub shared: Arc<K3HostShared>,
76}
77
78impl TransformerLayer for K3BoundLayer {
79    fn decode_graph_unsupported(&self) -> bool {
80        // Host round-trip cannot live in a CUDA graph.
81        true
82    }
83
84    fn decode_multi_seq_unsupported(&self) -> bool {
85        true
86    }
87
88    fn decode_verify_multi_unsupported(&self) -> bool {
89        // Host AttnRes streams and mixer caches are advanced one sequence at a time.
90        true
91    }
92
93    fn decode_rollback_unsupported(&self) -> bool {
94        // Lowering a paged-KV cursor does not rewind the host KDA recurrence
95        // or MLA cache. Prefix snapshots are a separate explicit restore path.
96        true
97    }
98
99    fn uses_ssm_pool(&self) -> bool {
100        // KDA is `linear_attention` in config. Own [`K3CpuFallbackState`], not GDN pool.
101        false
102    }
103
104    fn has_aux_state(&self) -> bool {
105        true
106    }
107
108    fn snapshot_aux(
109        &self,
110        state: &dyn LayerState,
111        gpu: &dyn GpuBackend,
112        stream: u64,
113    ) -> Result<Option<Vec<u8>>> {
114        let st = state
115            .as_any()
116            .downcast_ref::<K3CpuFallbackState>()
117            .context("K3 snapshot_aux: expected K3CpuFallbackState")?;
118        Ok(Some(st.snapshot(gpu, stream)?))
119    }
120
121    fn restore_aux(
122        &self,
123        state: &mut dyn LayerState,
124        blob: &[u8],
125        gpu: &dyn GpuBackend,
126        stream: u64,
127    ) -> Result<()> {
128        let st = state
129            .as_any_mut()
130            .downcast_mut::<K3CpuFallbackState>()
131            .context("K3 restore_aux: expected K3CpuFallbackState")?;
132        st.restore(gpu, blob, stream)
133    }
134
135    #[allow(clippy::too_many_arguments)]
136    fn decode(
137        &self,
138        hidden: DevicePtr,
139        residual: DevicePtr,
140        state: &mut dyn LayerState,
141        _kv_cache: &mut PagedKvCache,
142        seq_len: usize,
143        _block_table: &mut Vec<u32>,
144        _disk_block_ids: &mut Vec<u32>,
145        _disk_last_offloaded_per_layer: &mut Vec<u32>,
146        ctx: &ForwardContext,
147        stream: u64,
148    ) -> Result<()> {
149        self.decode_host(
150            hidden,
151            residual,
152            state,
153            seq_len,
154            ctx,
155            stream,
156            avarok_core::kimi_k3::cuda_kda_enabled(),
157            avarok_core::kimi_k3::cuda_mla_enabled(),
158        )
159    }
160
161    fn alloc_state(&self, _gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>> {
162        let cache = match self.spec.mixer {
163            MixerKind::Kda => {
164                LayerCache::Kda(avarok_core::kimi_k3::KdaState::new(&self.shared.kda))
165            }
166            MixerKind::Mla => LayerCache::Mla(avarok_core::kimi_k3::MlaKv::default()),
167        };
168        Ok(Box::new(K3CpuFallbackState::new(cache)))
169    }
170
171    fn release_state(&self, state: &mut dyn LayerState, gpu: &dyn GpuBackend) -> Result<()> {
172        state
173            .as_any_mut()
174            .downcast_mut::<K3CpuFallbackState>()
175            .context("K3 release_state: expected K3CpuFallbackState")?
176            .release(gpu)
177    }
178}
spark_model/kimi_k3/bound.rs

spark_model/kimi_k3/
bound.rs
