metrale_model_arch/kimi_k3/
bound.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `K3BoundLayer`, the Kimi K3 decoder layer the engine binds, and `K3HostShared`, the state its layers share.
4//!
5//! Decode runs on the host between a copy of the hidden row out and a copy back
6//! (`host_decode.rs`). The KDA mixer uses CUDA `kda_decode` with per-sequence
7//! device state unless `K3_CUDA_KDA=0`, the MLA mixer uses CUDA `mla_decode`
8//! unless `K3_CUDA_MLA=0`, and packed LatentMoE experts use
9//! `moe_w4a16_grouped_gemm_ptrtable_e8m0`.
10//!
11//! Owner: model-arch, Kimi K3.
12//! Invariants:
13//! - `decode_graph_unsupported`, `decode_multi_seq_unsupported`,
14//!   `decode_verify_multi_unsupported` and `decode_rollback_unsupported` are
15//!   all true.
16//! - Layer state is always a [`K3CpuFallbackState`]; the layer never uses the
17//!   SSM pool.
18
19use std::collections::HashMap;
20use std::sync::{Arc, OnceLock};
21
22use anyhow::{Context, Result};
23use metrale_cache::kv_cache::PagedKvCache;
24use metrale_config::ModelConfig;
25use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
26use metrale_model_weights::kimi_k3_host::{
27    AttnResStream, K3CpuLayer, K3Graph, K3LayerSpec, KdaConfig, LatentMoeConfig, LayerCache,
28    MixerKind, MlaConfig,
29};
30use metrale_model_weights::weights::WeightDtype;
31use parking_lot::Mutex;
32
33use super::kda_cuda::K3KdaDecodeKernels;
34use super::mla_cuda::K3MlaDecodeKernels;
35use super::moe_cuda::K3MoeGemmKernels;
36use super::state::K3CpuFallbackState;
37use metrale_model_layers::layer::{ForwardContext, LayerState, TransformerLayer};
38use metrale_model_layers::layer::{
39    LayerAuxState, LayerCapabilities, LayerGraphHooks, LayerSplitPrefill, LayerWeightSetup,
40    LayerWriteOnAccept,
41};
42use metrale_model_layers::weight_map::{DenseWeight, QuantizedWeight};
43
44/// 2026-09-25: Checkpoint name, dtype and element count of one resident weight.
45/// The first decode uses them to copy the weight to host FP32.
46#[derive(Clone, Debug)]
47pub struct WeightMeta {
48    pub name: String,
49    pub dtype: WeightDtype,
50    pub numel: usize,
51}
52
53/// 2026-09-25: One per loaded model; every decoder layer holds the same `Arc`.
54pub struct K3HostShared {
55    pub config: ModelConfig,
56    pub graph: K3Graph,
57    pub kda: KdaConfig,
58    pub mla: MlaConfig,
59    pub moe: LatentMoeConfig,
60    pub output_res_proj: DenseWeight,
61    pub output_res_norm: DenseWeight,
62    pub output_res_proj_meta: (WeightDtype, usize),
63    pub output_res_norm_meta: (WeightDtype, usize),
64    pub output_host: OnceLock<(Vec<f32>, Vec<f32>)>,
65    /// 2026-09-25: Resolved on the first CUDA KDA decode, then reused.
66    pub kda_kernels: OnceLock<K3KdaDecodeKernels>,
67    /// 2026-09-25: Resolved on the first CUDA MLA decode, then reused.
68    pub mla_kernels: OnceLock<K3MlaDecodeKernels>,
69    /// 2026-09-25: Resolved at load for a checkpoint with packed experts, then
70    /// reused.
71    pub moe_kernels: OnceLock<K3MoeGemmKernels>,
72    /// 2026-09-25: One AttnRes stream per token in flight, carried across layers.
73    /// Keyed by the token's `residual` pointer (`hidden` when that is null):
74    /// prefill runs every token through one layer before the next, and each
75    /// token has its own residual row, so each token keeps its own stream.
76    pub attnres: Mutex<HashMap<DevicePtr, AttnResStream>>,
77}
78
79/// 2026-09-25: One decoder layer. `weights` must be BF16 or FP32: the host
80/// bind refuses any other dtype.
81pub struct K3BoundLayer {
82    pub index: usize,
83    pub spec: K3LayerSpec,
84    pub weights: Vec<DenseWeight>,
85    pub weight_meta: Vec<WeightMeta>,
86    /// 2026-09-25: Routed experts stored as `{prefix}.weight_packed`, bound as
87    /// MXFP4 E8M0 through `quantized_mxfp4_e8m0_pair`. Empty when the
88    /// checkpoint stores its experts unpacked.
89    pub mxfp4_experts: Vec<(String, QuantizedWeight)>,
90    pub host: OnceLock<K3CpuLayer>,
91    pub shared: Arc<K3HostShared>,
92}
93
94impl TransformerLayer for K3BoundLayer {
95    #[allow(clippy::too_many_arguments)]
96    fn decode(
97        &self,
98        hidden: DevicePtr,
99        residual: DevicePtr,
100        state: &mut dyn LayerState,
101        _kv_cache: &mut PagedKvCache,
102        seq_len: usize,
103        _block_table: &mut Vec<u32>,
104        _disk_block_ids: &mut Vec<u32>,
105        _disk_last_offloaded_per_layer: &mut Vec<u32>,
106        ctx: &ForwardContext,
107        stream: u64,
108    ) -> Result<()> {
109        self.decode_host(
110            hidden,
111            residual,
112            state,
113            seq_len,
114            ctx,
115            stream,
116            metrale_model_weights::kimi_k3_host::cuda_kda_enabled(),
117            metrale_model_weights::kimi_k3_host::cuda_mla_enabled(),
118        )
119    }
120
121    fn alloc_state(&self, _gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>> {
122        let cache = match self.spec.mixer {
123            MixerKind::Kda => LayerCache::Kda(metrale_model_weights::kimi_k3_host::KdaState::new(
124                &self.shared.kda,
125            )),
126            MixerKind::Mla => {
127                LayerCache::Mla(metrale_model_weights::kimi_k3_host::MlaKv::default())
128            }
129        };
130        Ok(Box::new(K3CpuFallbackState::new(cache)))
131    }
132
133    fn release_state(&self, state: &mut dyn LayerState, gpu: &dyn GpuBackend) -> Result<()> {
134        state
135            .as_any_mut()
136            .downcast_mut::<K3CpuFallbackState>()
137            .context("K3 release_state: expected K3CpuFallbackState")?
138            .release(gpu)
139    }
140}
141
142impl LayerCapabilities for K3BoundLayer {
143    fn decode_graph_unsupported(&self) -> bool {
144        // 2026-09-25: Decode copies to and from the host and synchronises, which
145        // a CUDA graph cannot capture.
146        true
147    }
148
149    fn decode_multi_seq_unsupported(&self) -> bool {
150        true
151    }
152
153    fn decode_verify_multi_unsupported(&self) -> bool {
154        // 2026-09-25: Each call advances one sequence's host AttnRes stream and
155        // mixer state.
156        true
157    }
158
159    fn decode_rollback_unsupported(&self) -> bool {
160        // 2026-09-25: The mixer state is not in the paged KV cache, so lowering its
161        // cursor would not rewind it. Only `restore_aux` sets it back.
162        true
163    }
164
165    fn uses_ssm_pool(&self) -> bool {
166        // 2026-09-25: Mixer state lives in `K3CpuFallbackState` (`alloc_state`).
167        false
168    }
169}
170
171impl LayerWeightSetup for K3BoundLayer {}
172impl LayerWriteOnAccept for K3BoundLayer {}
173impl LayerGraphHooks for K3BoundLayer {}
174
175impl LayerAuxState for K3BoundLayer {
176    fn has_aux_state(&self) -> bool {
177        true
178    }
179
180    fn snapshot_aux(
181        &self,
182        state: &dyn LayerState,
183        gpu: &dyn GpuBackend,
184        stream: u64,
185    ) -> Result<Option<Vec<u8>>> {
186        let st = state
187            .as_any()
188            .downcast_ref::<K3CpuFallbackState>()
189            .context("K3 snapshot_aux: expected K3CpuFallbackState")?;
190        Ok(Some(st.snapshot(gpu, stream)?))
191    }
192
193    fn restore_aux(
194        &self,
195        state: &mut dyn LayerState,
196        blob: &[u8],
197        gpu: &dyn GpuBackend,
198        stream: u64,
199    ) -> Result<()> {
200        let st = state
201            .as_any_mut()
202            .downcast_mut::<K3CpuFallbackState>()
203            .context("K3 restore_aux: expected K3CpuFallbackState")?;
204        st.restore(gpu, blob, stream)
205    }
206}
207
208impl LayerSplitPrefill for K3BoundLayer {}
209impl metrale_model_layers::circuit_exec::CircuitBindings for K3BoundLayer {}
metrale_model_arch/kimi_k3/bound.rs

metrale_model_arch/kimi_k3/
bound.rs
