frink_models/
kimi_decoder.rs

1//! A dedicated decoder for Kimi K3's real hybrid architecture, separate
2//! from `frink-models::decoder::Decoder` (which every GQA-only preset
3//! -- GLM-5.2, DeepSeek V4 Pro, the test fixtures -- uses and which
4//! `frink-cli`, `frink-server`, `prefix_cache`, and `speculative` all
5//! depend on): this module composes the already-independently-tested
6//! `kda`, `mla`, `latent_moe`, `block_residual`, and
7//! `frink_core::situ_and_mul` pieces into a real forward pass, without
8//! touching any of that existing, production-quality GQA code path (a
9//! shared, polymorphic `Decoder` was judged too risky to attempt
10//! without a way to verify it end to end).
11//!
12//! Real per-layer flow, transcribed from `KimiDecoderLayer.forward`'s
13//! `_forward_attn_residual` path (the one Kimi K3 actually runs, since
14//! `attn_res_block_size`=12 is set in its real config) in
15//! `modeling_kimi_linear.py`:
16//!
17//! ```text
18//! prefix_sum = hidden
19//! blended = block_residual.is_empty() ? prefix_sum : apply_attn_res(prefix_sum, block_residual, self_attn_res_*)
20//! if layer_idx % attn_res_block_size == 0 { block_residual.push(prefix_sum); prefix_sum = None }
21//! attn_out = kda_or_mla(rms_norm(blended, input_layernorm))
22//! prefix_sum = (prefix_sum is None) ? attn_out : prefix_sum + attn_out
23//! blended2 = apply_attn_res(prefix_sum, block_residual, mlp_res_*)   // block_residual may have just grown above
24//! ffn_out = dense_or_moe(rms_norm(blended2, post_attention_layernorm))
25//! hidden = prefix_sum + ffn_out
26//! ```
27//!
28//! One real, non-obvious fact confirmed by reading
29//! `KimiLinearModel.forward` (not just the per-layer code): `block_residual`
30//! is freshly re-initialized to *empty* inside `forward()`, i.e. once
31//! per forward call -- for single-token incremental decode (what this
32//! module implements), that means every decode step starts with an
33//! empty `block_residual`, not a value carried over from the previous
34//! token. `KimiDecodeState` therefore only needs to carry each layer's
35//! own attention state (`kda::KdaState` or MLA's growable K/V buffers),
36//! not any block-residual bookkeeping across positions.
37
38use frink_core::matmul::{rms_norm, situ_and_mul};
39use frink_core::tensor::Tensor;
40use frink_core::weight_matrix::WeightMatrix;
41
42use crate::block_residual::apply_attn_res;
43use crate::kda::{self, KdaAttnWeights, KdaState};
44use crate::latent_moe::{self, KimiLatentMoeWeights, KimiMoeConfig};
45use crate::mla::{self, MlaAttnWeights};
46
47pub enum KimiLayerAttention {
48    Kda(Box<KdaAttnWeights>),
49    Mla(Box<MlaAttnWeights>),
50}
51
52/// `KimiMLP` used directly on the full hidden dimension -- the sole
53/// dense leading layer's feed-forward block (`n_dense_leading_layers`=1
54/// for Kimi K3).
55pub struct DenseMlpWeights {
56    pub gate_proj: WeightMatrix,
57    pub up_proj: WeightMatrix,
58    pub down_proj: WeightMatrix,
59}
60
61impl DenseMlpWeights {
62    pub(crate) fn forward(&self, x: &[f32], situ_beta: f32, situ_linear_beta: f32) -> Vec<f32> {
63        let gate = self.gate_proj.apply(x);
64        let up = self.up_proj.apply(x);
65        let combined = situ_and_mul(&gate, &up, situ_beta, situ_linear_beta);
66        self.down_proj.apply(&combined)
67    }
68}
69
70pub enum KimiLayerFfn {
71    Dense(Box<DenseMlpWeights>),
72    Moe(Box<KimiLatentMoeWeights>),
73}
74
75pub struct KimiDecoderLayerWeights {
76    pub input_layernorm_weight: Vec<f32>,
77    pub attn: KimiLayerAttention,
78    pub post_attention_layernorm_weight: Vec<f32>,
79    pub ffn: KimiLayerFfn,
80    /// Block-residual weights -- present on every real layer (confirmed
81    /// against a real shard header), used twice per layer (once before
82    /// attention, once before the FFN).
83    pub self_attention_res_norm_weight: Vec<f32>,
84    pub self_attention_res_proj_weight: Vec<f32>,
85    pub mlp_res_norm_weight: Vec<f32>,
86    pub mlp_res_proj_weight: Vec<f32>,
87}
88
89pub struct KimiDecoderWeights {
90    pub embedding: Tensor, // [vocab_size, hidden_dim]
91    pub layers: Vec<KimiDecoderLayerWeights>,
92    pub output_attn_res_norm_weight: Vec<f32>,
93    pub output_attn_res_proj_weight: Vec<f32>,
94    pub final_norm_weight: Vec<f32>,
95    pub output_head: WeightMatrix, // [vocab_size, hidden_dim]
96}
97
98impl KimiDecoderWeights {
99    /// The shared expert store's live counters when this model streams
100    /// routed experts -- `None` when fully resident. Every store-backed
101    /// layer shares one store, so the first found speaks for the model.
102    pub fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
103        self.layers.iter().find_map(|l| match &l.ffn {
104            KimiLayerFfn::Moe(moe) => match &moe.experts {
105                crate::latent_moe::KimiExpertBacking::Stored { store, .. } => Some(store.stats()),
106                crate::latent_moe::KimiExpertBacking::Resident(_) => None,
107            },
108            KimiLayerFfn::Dense(_) => None,
109        })
110    }
111}
112
113pub struct KimiDecoderConfig {
114    pub attn_res_block_size: usize,
115    pub rms_norm_eps: f32,
116    pub situ_beta: f32,
117    pub situ_linear_beta: f32,
118    pub moe: KimiMoeConfig,
119}
120
121pub enum KimiLayerState {
122    Kda(KdaState),
123    Mla {
124        k_cache: Vec<f32>,
125        v_cache: Vec<f32>,
126    },
127}
128
129pub struct KimiDecodeState {
130    layer_states: Vec<KimiLayerState>,
131}
132
133impl KimiDecodeState {
134    /// Builds fresh (all-zero/empty) per-layer state matching each
135    /// layer's real attention kind, as declared by `weights.layers`.
136    pub fn new(weights: &KimiDecoderWeights, kda_cfg: &crate::config::KdaConfig) -> Self {
137        let layer_states = weights
138            .layers
139            .iter()
140            .map(|l| match &l.attn {
141                KimiLayerAttention::Kda(_) => KimiLayerState::Kda(KdaState::new(kda_cfg)),
142                KimiLayerAttention::Mla(_) => KimiLayerState::Mla {
143                    k_cache: Vec::new(),
144                    v_cache: Vec::new(),
145                },
146            })
147            .collect();
148        KimiDecodeState { layer_states }
149    }
150}
151
152/// One decode step across every layer.
153pub fn kimi_forward_token(
154    weights: &KimiDecoderWeights,
155    cfg: &KimiDecoderConfig,
156    mla_cfg: &crate::config::MlaConfig,
157    kda_cfg: &crate::config::KdaConfig,
158    token_id: usize,
159    state: &mut KimiDecodeState,
160) -> Vec<f32> {
161    let hidden_dim = weights.embedding.cols();
162    let mut hidden = weights.embedding.row(token_id).to_vec();
163    // Fresh every call -- see module doc comment for why this is real,
164    // not a simplification: `KimiLinearModel.forward` re-initializes
165    // `block_residual` to empty on every forward call.
166    let mut block_residual: Vec<f32> = Vec::new();
167
168    for (layer_idx, layer) in weights.layers.iter().enumerate() {
169        let prefix_sum_pre_blend = hidden.clone();
170
171        let blended = if block_residual.is_empty() {
172            prefix_sum_pre_blend.clone()
173        } else {
174            apply_attn_res(
175                &prefix_sum_pre_blend,
176                &block_residual,
177                &layer.self_attention_res_norm_weight,
178                &layer.self_attention_res_proj_weight,
179                cfg.rms_norm_eps,
180            )
181        };
182
183        let mut prefix_sum: Option<Vec<f32>> = Some(prefix_sum_pre_blend.clone());
184        if layer_idx % cfg.attn_res_block_size == 0 {
185            block_residual.extend_from_slice(&prefix_sum_pre_blend);
186            prefix_sum = None;
187        }
188
189        let normed = rms_norm(&blended, &layer.input_layernorm_weight, cfg.rms_norm_eps);
190        let attn_out = match (&layer.attn, &mut state.layer_states[layer_idx]) {
191            (KimiLayerAttention::Kda(w), KimiLayerState::Kda(s)) => {
192                kda::kda_forward_token(w, kda_cfg, &normed, cfg.rms_norm_eps, s)
193            }
194            (KimiLayerAttention::Mla(w), KimiLayerState::Mla { k_cache, v_cache }) => {
195                mla::mla_forward_token(
196                    w,
197                    mla_cfg,
198                    None,
199                    &normed,
200                    cfg.rms_norm_eps,
201                    k_cache,
202                    v_cache,
203                )
204            }
205            _ => unreachable!("layer attention kind and decode state kind must always match"),
206        };
207
208        let mut prefix_sum = match prefix_sum {
209            Some(mut ps) => {
210                for (p, a) in ps.iter_mut().zip(attn_out.iter()) {
211                    *p += a;
212                }
213                ps
214            }
215            None => attn_out,
216        };
217
218        let blended2 = apply_attn_res(
219            &prefix_sum,
220            &block_residual,
221            &layer.mlp_res_norm_weight,
222            &layer.mlp_res_proj_weight,
223            cfg.rms_norm_eps,
224        );
225        let normed2 = rms_norm(
226            &blended2,
227            &layer.post_attention_layernorm_weight,
228            cfg.rms_norm_eps,
229        );
230        let ffn_out = match &layer.ffn {
231            KimiLayerFfn::Dense(w) => w.forward(&normed2, cfg.situ_beta, cfg.situ_linear_beta),
232            KimiLayerFfn::Moe(w) => latent_moe::kimi_latent_moe_forward(w, &cfg.moe, &normed2),
233        };
234
235        for (p, f) in prefix_sum.iter_mut().zip(ffn_out.iter()) {
236            *p += f;
237        }
238        hidden = prefix_sum;
239    }
240
241    if !block_residual.is_empty() {
242        hidden = apply_attn_res(
243            &hidden,
244            &block_residual,
245            &weights.output_attn_res_norm_weight,
246            &weights.output_attn_res_proj_weight,
247            cfg.rms_norm_eps,
248        );
249    }
250
251    let final_normed = rms_norm(&hidden, &weights.final_norm_weight, cfg.rms_norm_eps);
252    assert_eq!(final_normed.len(), hidden_dim);
253    weights.output_head.apply(&final_normed)
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use crate::config::{KdaConfig, MlaConfig};
260    use crate::kda::KdaAttnWeights;
261    use crate::latent_moe::KimiExpertBacking;
262    use crate::latent_moe::KimiExpertWeights;
263    use crate::mla::{MlaAttnWeights, MlaKvB, MlaQProj};
264    use frink_core::tensor::Tensor;
265
266    const HIDDEN_DIM: usize = 8;
267    const EPS: f32 = 1e-5;
268    const SITU_BETA: f32 = 4.0;
269    const SITU_LINEAR_BETA: f32 = 25.0;
270    const ATTN_RES_BLOCK_SIZE: usize = 2;
271
272    const KDA_NUM_HEADS: usize = 2;
273    const KDA_HEAD_DIM: usize = 3;
274    const KDA_PROJ: usize = KDA_NUM_HEADS * KDA_HEAD_DIM;
275    const KDA_CONV_SIZE: usize = 4;
276    const KDA_GATE_LOWER_BOUND: f32 = -5.0;
277
278    const DENSE_INTERMEDIATE: usize = 5;
279
280    const MLA_NUM_HEADS: usize = 2;
281    const MLA_QK_NOPE: usize = 3;
282    const MLA_QK_ROPE: usize = 2;
283    const MLA_KV_LORA: usize = 4;
284    const MLA_Q_LORA: usize = 6;
285    const MLA_V_HEAD_DIM: usize = 3;
286    const MLA_PROJ: usize = MLA_NUM_HEADS * MLA_V_HEAD_DIM;
287
288    const MOE_HIDDEN_DIM: usize = 4;
289    const MOE_INTERMEDIATE: usize = 3;
290    const N_EXPERTS: usize = 4;
291    const TOP_K: usize = 2;
292    const SHARED_INTERMEDIATE: usize = 3;
293    const OUTPUT_VOCAB: usize = 5;
294
295    fn wm(data: &[f32], rows: usize, cols: usize) -> WeightMatrix {
296        assert_eq!(data.len(), rows * cols);
297        WeightMatrix::F32(Tensor::new(data.to_vec(), vec![rows, cols]))
298    }
299
300    // Generated by an independent Python reference -- do not hand-edit.
301    const KDA_Q_PROJ: [f32; 48] = [
302        0.0332731, -0.0251273, -0.241248, -0.645647, 0.363556, -0.144589, -0.058427, -0.264839,
303        -0.174983, -0.313658, -0.0222095, 0.0329603, -0.0773659, 0.179789, -0.428436, 0.272643,
304        -0.29953, -0.124363, -0.270533, -0.15906, -0.136326, 0.138484, 0.728104, -0.306624,
305        -0.0415019, -0.62199, -0.0334468, 0.0501355, 0.374942, 0.438693, 0.398001, 0.0660859,
306        0.146161, -0.129336, 0.701423, -0.0954704, -0.37886, -0.0410528, 0.653784, 0.193832,
307        0.451906, 0.367486, 0.0999694, -0.498877, 0.495604, -0.203314, -0.326699, -0.410713,
308    ];
309    const KDA_K_PROJ: [f32; 48] = [
310        -0.421999, -0.0844493, 0.0366557, -0.373529, -0.320309, -0.171037, -0.208623, -0.251102,
311        -0.384313, -0.242582, -0.708695, -0.0410055, -0.163735, 0.045339, 0.0596934, -0.462932,
312        0.344817, 0.314063, 0.183688, -0.215147, 0.0682644, -0.119572, 0.24693, -0.171698,
313        -0.427691, -0.311782, -0.243235, 0.168478, -0.322684, -0.0423794, 0.296734, -0.106622,
314        0.0607117, 0.900444, 0.214488, 0.390483, 0.0470901, -0.084555, -0.509912, 0.273387,
315        -0.0251134, 0.459143, 0.34484, -0.102602, -0.315908, -0.311536, -0.154464, -0.067087,
316    ];
317    const KDA_V_PROJ: [f32; 48] = [
318        0.106298, -0.296877, -0.441024, 0.346085, -0.211878, 0.208265, 0.286445, 0.254616,
319        0.527007, 0.0238059, 0.0520139, -0.0428818, 0.547271, 0.0641891, 0.406677, -0.159912,
320        0.315001, 0.090456, 0.665635, 0.0849947, 0.332725, 0.0468494, -0.340192, -0.57511,
321        -0.199206, 0.00843515, 0.307353, -0.341596, -0.188955, -0.0262392, 0.711621, -0.12957,
322        0.0589482, 0.6676, -0.0487213, 0.318632, -0.129497, 0.45287, 0.312848, 0.404807, 0.300039,
323        -0.111872, 0.00619566, 0.337312, 0.203499, -0.615727, -0.426802, 0.00709124,
324    ];
325    const KDA_Q_CONV_W: [f32; 24] = [
326        0.232949, -0.349125, -0.144423, -1.17421, 0.392014, -0.0551284, -0.188107, 0.34828,
327        0.182368, -1.15816, 0.29745, -0.163951, -0.785389, -0.0799616, 0.0853995, 0.159346,
328        0.438013, 0.16794, 0.11068, 0.222541, 0.27932, 0.584781, 0.0801223, 0.315161,
329    ];
330    const KDA_K_CONV_W: [f32; 24] = [
331        0.0811582, 0.5723, 0.275913, 0.337176, 0.30982, -0.260437, -0.368242, 0.355531, -0.940038,
332        -0.0380444, 0.559686, -0.632208, 0.0297131, -0.246906, 0.435212, -0.65012, -0.207435,
333        -0.259419, -0.219341, 0.426728, -0.0195621, -0.164165, 0.231709, 0.363266,
334    ];
335    const KDA_V_CONV_W: [f32; 24] = [
336        0.174442,
337        0.16448,
338        -0.11751,
339        -0.00693412,
340        0.22015,
341        0.1184,
342        0.59284,
343        -0.100545,
344        0.904578,
345        -0.057162,
346        -0.185486,
347        -0.0155876,
348        0.239106,
349        -0.184512,
350        0.887074,
351        0.319565,
352        0.518244,
353        0.238721,
354        -0.0932803,
355        -0.614529,
356        0.400545,
357        0.458996,
358        -0.217588,
359        0.0540083,
360    ];
361    const KDA_A_LOG: [f32; 2] = [2.57153, 1.68862];
362    const KDA_F_A_PROJ: [f32; 24] = [
363        0.293008, 0.208071, -0.0458425, 0.26919, -0.0604267, -0.414273, -0.347327, 0.435059,
364        -0.0875413, 0.260375, 0.78432, -0.0103574, 0.129559, 0.522936, 0.205997, 0.235916,
365        0.129708, -0.515258, -0.251447, 0.0114761, -0.0608938, -0.142939, 0.0391045, -0.106787,
366    ];
367    const KDA_F_B_PROJ: [f32; 18] = [
368        0.258232, 0.0799449, -0.128173, 0.374093, 0.0249372, 0.564439, 0.291732, -0.673602,
369        0.300492, -0.312056, -0.0782346, -0.560855, 0.319259, -0.194854, -0.109685, -0.295024,
370        -0.0275032, -0.287207,
371    ];
372    const KDA_DT_BIAS: [f32; 6] = [0.406797, 0.44378, -0.118682, -0.101149, 0.144153, 0.100643];
373    const KDA_B_PROJ: [f32; 16] = [
374        -0.557976, 0.101854, -0.529571, 0.437377, 0.282072, 0.0808985, 0.119817, -0.234031,
375        -0.114455, -0.182255, 0.492746, -0.213926, -0.0890158, 0.520648, -0.0288543, 0.378583,
376    ];
377    const KDA_G_PROJ: [f32; 48] = [
378        0.115836, 0.257017, -0.28702, 0.0285552, 0.0241214, -0.586101, -0.284745, -0.1296,
379        0.369078, -0.275476, 0.352622, 0.235832, -0.13984, 0.230448, 0.440173, 0.224343, 0.103446,
380        0.108157, -0.523653, -0.674568, 0.130624, 0.228334, -0.493531, 0.257466, -0.0398756,
381        -0.385554, 0.209178, -0.414277, -0.364987, 0.515843, -0.486846, -0.128941, -0.0565001,
382        0.480469, -0.0731606, 0.279074, -0.16879, -0.589812, 0.200213, -0.3698, 0.355666, 0.269551,
383        0.450168, -0.0689023, -0.234073, -0.81206, 0.250412, 0.0890739,
384    ];
385    const KDA_O_NORM_W: [f32; 3] = [0.970175, 0.902756, 1.03511];
386    const KDA_O_PROJ: [f32; 48] = [
387        -0.148986, -0.157794, 0.225606, -0.302445, -0.0587451, -0.235495, -0.541087, -0.270422,
388        0.604319, 0.371634, 0.144359, 0.409312, 0.108881, -0.352629, -0.129886, -0.258043,
389        -0.168576, -0.241786, 0.0283778, -0.097169, -0.0738384, 0.102603, 0.302548, 0.219944,
390        0.144096, -0.521365, -0.354098, -0.79218, -0.227326, -0.778447, -0.221429, -0.0319402,
391        0.377724, 0.388695, 0.560278, 0.420781, -0.440711, 0.275919, 0.189938, 0.159768, 0.341554,
392        0.594513, 0.333733, -0.0974569, -0.308295, 0.200531, 0.101979, -0.363293,
393    ];
394    const DENSE_GATE_PROJ: [f32; 40] = [
395        0.0530662,
396        -0.0224611,
397        -0.468284,
398        0.394925,
399        -0.0289055,
400        -0.138464,
401        0.302004,
402        -0.0823999,
403        0.295465,
404        0.391958,
405        0.197625,
406        0.236802,
407        -0.000834978,
408        0.539451,
409        0.230436,
410        0.0369415,
411        0.235552,
412        -0.0302823,
413        0.216731,
414        0.106894,
415        -0.0388429,
416        -0.0235654,
417        0.0683206,
418        -0.247831,
419        -0.829944,
420        -0.255354,
421        -0.652798,
422        -0.539486,
423        -0.16364,
424        0.212245,
425        -0.263695,
426        0.266948,
427        0.474532,
428        0.249119,
429        -0.0215431,
430        0.0568342,
431        -0.0727135,
432        -0.267156,
433        0.0426981,
434        -0.921182,
435    ];
436    const DENSE_UP_PROJ: [f32; 40] = [
437        -0.236304,
438        0.180069,
439        -0.2151,
440        0.222938,
441        -0.0463798,
442        -0.31709,
443        -0.00577976,
444        0.254963,
445        -0.152982,
446        0.736692,
447        0.565269,
448        -0.00259989,
449        0.22078,
450        0.0674972,
451        0.193022,
452        0.626185,
453        0.0697913,
454        0.489142,
455        -0.0709509,
456        0.135837,
457        0.0604967,
458        0.134483,
459        0.15545,
460        -0.715912,
461        0.131568,
462        0.447037,
463        0.0964134,
464        -0.568027,
465        -0.539577,
466        -0.385414,
467        -0.308715,
468        0.338456,
469        -0.589399,
470        -0.22471,
471        -0.256634,
472        -0.172552,
473        0.229815,
474        -0.232763,
475        -0.15492,
476        0.112469,
477    ];
478    const DENSE_DOWN_PROJ: [f32; 40] = [
479        0.364035, -0.722506, -0.192076, -0.513796, 0.358074, 0.187812, -0.210479, -0.190417,
480        0.0378089, 0.398016, 0.393775, 0.119131, 0.0944593, 0.160586, -0.302192, -0.0174533,
481        -0.0131564, 0.17226, -0.0443057, -0.354608, 0.221662, 0.333016, -0.0158993, -0.0357403,
482        -0.251197, 0.0164112, 0.210909, 0.292464, 0.201202, -0.121079, 0.407385, 0.275425,
483        -0.537043, -0.0175377, -0.428206, 0.253178, 0.404171, -0.231798, 0.168481, -0.222966,
484    ];
485    const MLA_Q_A_PROJ: [f32; 48] = [
486        0.451391, -0.126339, 0.514525, 0.104224, -0.482352, -0.12592, -0.140897, 0.570421,
487        0.177444, 0.0637828, -0.219946, -0.449056, -0.0832418, -0.0932411, 0.141219, -0.688717,
488        0.311497, -0.0927841, 0.0286042, 0.608332, -0.340392, 0.292946, 0.390898, -0.0359569,
489        -0.157069, -0.194805, 0.595659, -0.129717, 0.435375, 0.707992, -0.267707, -0.446994,
490        0.257943, 0.208935, 0.0124412, 0.0539018, 0.40146, -0.144905, -0.115103, -0.432397,
491        -0.205624, -0.596541, 0.485434, 0.194575, -0.132356, -0.194361, -0.361851, -0.109006,
492    ];
493    const MLA_Q_A_NORM_W: [f32; 6] = [0.969939, 0.841643, 1.00562, 1.08525, 0.926626, 0.990575];
494    const MLA_Q_B_PROJ: [f32; 60] = [
495        0.1374,
496        -0.135647,
497        0.0887308,
498        -0.0979333,
499        -0.259915,
500        -0.0779051,
501        0.115564,
502        0.228723,
503        -0.230383,
504        0.083073,
505        0.192752,
506        -0.0369039,
507        -0.286092,
508        -0.0686448,
509        -0.307404,
510        -0.0636238,
511        0.211306,
512        -0.146429,
513        -0.0687791,
514        0.250252,
515        -0.156639,
516        0.337944,
517        0.178598,
518        -0.616342,
519        -0.0385348,
520        0.382634,
521        0.671192,
522        -0.248981,
523        0.121872,
524        0.275898,
525        0.188452,
526        0.828256,
527        -0.438505,
528        0.621337,
529        0.408834,
530        -0.303079,
531        0.0797437,
532        -0.240968,
533        0.0335185,
534        0.602176,
535        0.461823,
536        -0.23303,
537        0.0937693,
538        -0.0384911,
539        0.173248,
540        -0.515262,
541        0.0811063,
542        0.543123,
543        -0.148281,
544        -0.599126,
545        0.302461,
546        -0.417772,
547        0.064963,
548        -0.000217146,
549        0.188863,
550        0.752821,
551        -0.293606,
552        0.169847,
553        0.739423,
554        0.435163,
555    ];
556    const MLA_KV_A_PROJ: [f32; 48] = [
557        -0.421765, -0.232617, 0.203753, 0.166609, -0.045513, -0.0101532, 0.0955488, -0.165861,
558        -0.0260061, -0.0220384, -0.0330121, -0.158364, -0.0111785, -0.345142, 0.755459, -0.0558389,
559        -0.277054, -0.26961, 0.240628, 0.0947741, -0.167247, -0.23363, -0.157332, 0.345719,
560        0.156487, 0.0865676, -0.192661, -0.214795, -0.567426, -0.0837593, -0.324576, 0.310762,
561        -0.314536, -0.68736, -0.0154466, 0.315627, -0.150268, -0.0287802, 0.356274, -0.246495,
562        -0.0825744, 0.204116, 0.559475, 0.640863, 0.456519, 0.173669, 0.185071, -0.115482,
563    ];
564    const MLA_KV_A_NORM_W: [f32; 4] = [1.10413, 1.0107, 0.871313, 1.16755];
565    const MLA_KV_B_PROJ: [f32; 48] = [
566        0.0866448,
567        -0.423795,
568        0.114522,
569        -0.441771,
570        -0.186988,
571        0.0366587,
572        0.114137,
573        -0.0609123,
574        0.199646,
575        0.722449,
576        0.704719,
577        -0.460906,
578        0.0629041,
579        0.485761,
580        0.328924,
581        0.0319879,
582        -0.298518,
583        0.123369,
584        -0.136789,
585        -0.256657,
586        0.131407,
587        0.435331,
588        0.0737087,
589        -0.130659,
590        0.130558,
591        0.113592,
592        -0.0490255,
593        -0.292942,
594        -0.0925113,
595        -0.412935,
596        -0.222763,
597        -0.189368,
598        0.202786,
599        -0.128307,
600        -0.215928,
601        0.0669919,
602        0.384514,
603        0.637332,
604        -0.0592113,
605        0.0464926,
606        -0.281259,
607        0.107265,
608        0.494055,
609        -0.00978826,
610        -0.398057,
611        0.593258,
612        0.0966843,
613        0.013154,
614    ];
615    const MLA_O_PROJ: [f32; 48] = [
616        -0.384195,
617        0.195792,
618        0.334816,
619        -0.101664,
620        0.103024,
621        -0.142412,
622        0.130369,
623        -0.307274,
624        -0.367605,
625        0.559024,
626        -0.128535,
627        0.0330137,
628        0.241127,
629        0.0644755,
630        -0.0776874,
631        -0.231983,
632        -0.909896,
633        -0.268289,
634        0.205435,
635        -0.268665,
636        -0.663843,
637        0.00206543,
638        -0.454986,
639        0.226469,
640        -0.539808,
641        0.34166,
642        7.24295e-05,
643        0.339034,
644        -0.244381,
645        0.168123,
646        -0.316009,
647        0.00653252,
648        0.107088,
649        0.229998,
650        0.279935,
651        0.318718,
652        -0.262585,
653        0.528549,
654        0.356373,
655        0.217004,
656        0.154669,
657        -0.193219,
658        -0.0690429,
659        0.304651,
660        0.554123,
661        0.158962,
662        0.525434,
663        -0.20413,
664    ];
665    const MLA_G_PROJ: [f32; 48] = [
666        0.378873, -0.362289, -0.172782, -0.305609, -0.15234, 0.13899, -0.358599, -0.786074,
667        -0.168818, 0.168034, -0.27184, -0.103517, 0.34245, 0.0531821, 0.254431, -0.164725,
668        -0.0816316, -0.0860996, -0.0865057, 0.0241667, 0.259761, 0.230357, -0.726831, 0.379549,
669        0.171985, -0.378427, -0.199484, 0.138928, -0.238634, -0.55968, -0.405272, -0.184117,
670        -0.184866, 0.0497178, 0.103438, 0.380376, 0.364498, -0.175124, 0.293164, 0.200305,
671        -0.0883711, -0.155664, -0.011865, 0.304442, -0.0586883, -0.269584, -0.215651, 0.0801283,
672    ];
673    const MOE_ROUTER_WEIGHT: [f32; 32] = [
674        -0.14075, 0.348724, 0.162851, 0.220414, -0.334635, -0.0441166, 0.291079, -0.16397,
675        0.104033, 0.186523, -0.0732253, -0.18498, 0.0548991, 0.0572118, -0.212319, -0.114626,
676        0.0438156, -0.291582, -0.364393, -0.456047, -0.153299, -0.175468, 0.0600189, -0.433395,
677        -0.241496, -0.195144, 0.468446, -0.0704137, 0.0423066, 0.0625186, 0.125581, -0.384111,
678    ];
679    const MOE_BIAS: [f32; 4] = [-0.287968, -0.506168, -0.0823805, -1.16076];
680    const MOE_DOWN_PROJ: [f32; 32] = [
681        0.357856, -0.419139, 0.0951717, 0.486966, 0.138318, -0.321581, 0.023647, -0.287115,
682        -0.218501, 0.466282, -0.674924, 0.358341, -0.651141, -0.190237, 0.154624, -0.131574,
683        -0.0769183, 0.0423567, 0.768159, 0.659497, -0.106854, -0.338984, -0.0425253, -0.326632,
684        -0.157781, -0.0021155, 0.259854, 0.132918, 0.305795, 0.418168, 0.0538999, 0.630383,
685    ];
686    const MOE_UP_PROJ: [f32; 32] = [
687        0.0537259, -0.222431, 0.0517936, -0.328772, -0.218991, 0.0109666, -0.66681, 0.53133,
688        -0.24925, -0.271987, -0.57234, 0.0120699, -0.288328, 0.131623, -0.132867, -0.105812,
689        -0.284854, -0.456721, -0.150868, -0.102689, -0.245856, -0.0891634, 0.389166, 0.518893,
690        -0.0414535, 0.157239, -0.0155488, 0.207036, 0.00431687, -0.261777, -0.541271, 0.211408,
691    ];
692    const MOE_ROUTED_NORM_W: [f32; 4] = [1.14805, 0.951891, 0.971543, 0.908736];
693    const MOE_SHARED_W1: [f32; 24] = [
694        -0.377193,
695        0.27978,
696        -0.0102447,
697        0.229219,
698        -0.380455,
699        0.552381,
700        -0.0980754,
701        0.56176,
702        -0.297,
703        -0.142547,
704        -0.204375,
705        -0.338748,
706        0.388348,
707        0.489957,
708        0.356719,
709        0.0574536,
710        0.110144,
711        -0.00856218,
712        0.426843,
713        0.260044,
714        0.148431,
715        -0.184063,
716        -0.342297,
717        -0.0762749,
718    ];
719    const MOE_SHARED_W2: [f32; 24] = [
720        -0.43575, 0.0521214, -0.10101, 0.072215, 0.337606, 0.0963424, -0.107468, -0.49745,
721        0.364127, 0.217258, 0.359295, -0.0454641, -0.252116, 0.0389652, 0.178223, 0.0836014,
722        0.124572, -0.252946, -0.227852, 0.0663797, 0.110113, 0.458224, 0.520268, -0.0199504,
723    ];
724    const MOE_SHARED_W3: [f32; 24] = [
725        0.368003,
726        -0.483325,
727        -0.379792,
728        0.100306,
729        0.106976,
730        -0.0896776,
731        -0.316544,
732        0.322652,
733        -0.517715,
734        0.249498,
735        0.0612062,
736        -0.601994,
737        -0.193968,
738        0.558544,
739        -0.388051,
740        -0.277637,
741        0.122792,
742        -0.185875,
743        -0.00133169,
744        0.0236871,
745        -0.296105,
746        0.0705485,
747        -0.302294,
748        -0.246982,
749    ];
750    const MOE_E0_W1: [f32; 12] = [
751        0.0440567, 0.174659, -0.355788, -0.234582, 0.238399, 0.600469, -0.130771, 0.811262,
752        0.130333, -0.0412095, -0.714084, 0.158573,
753    ];
754    const MOE_E0_W2: [f32; 12] = [
755        -0.103864, -0.115033, 0.275029, -0.0808199, -0.348501, 0.30532, 0.310035, -0.0757431,
756        -0.132877, 0.711211, 0.226243, 0.163931,
757    ];
758    const MOE_E0_W3: [f32; 12] = [
759        -0.114776, 0.141509, 0.504018, -0.115836, 0.43992, 0.131565, -0.469145, 0.0490231,
760        0.151157, -0.0375621, -0.17218, 0.231705,
761    ];
762    const MOE_E1_W1: [f32; 12] = [
763        0.0488396,
764        -0.341746,
765        0.0418606,
766        0.221871,
767        -0.722022,
768        -0.00220993,
769        -0.0916067,
770        -0.398679,
771        -0.423551,
772        -0.645019,
773        -0.27434,
774        0.0208012,
775    ];
776    const MOE_E1_W2: [f32; 12] = [
777        -0.167076, 0.231248, -0.759916, 0.352397, -0.0142662, -0.24585, -0.0941462, 0.0286856,
778        0.103257, -0.272298, 0.177667, -0.226165,
779    ];
780    const MOE_E1_W3: [f32; 12] = [
781        0.189908, -0.207244, -0.0267815, 0.0907185, 0.191764, 0.105395, 0.0162554, -0.452242,
782        0.430473, -0.291964, -0.261698, 0.0512003,
783    ];
784    const MOE_E2_W1: [f32; 12] = [
785        -0.0189576, 0.0668481, 0.214585, 0.143547, -0.459021, 0.191771, -0.386651, 0.519018,
786        -0.38456, -0.409715, 0.101121, -0.114501,
787    ];
788    const MOE_E2_W2: [f32; 12] = [
789        0.329286, -0.122696, 0.385045, -0.309553, 0.205889, -0.0737322, 0.0664255, 0.243589,
790        -0.101208, 0.615374, 0.791347, -0.288181,
791    ];
792    const MOE_E2_W3: [f32; 12] = [
793        -0.34824, -0.0969774, -0.0404219, -0.142054, -0.140162, 0.168913, 0.329251, 0.568559,
794        0.302453, 0.2389, 0.10173, 0.26265,
795    ];
796    const MOE_E3_W1: [f32; 12] = [
797        -0.270246, -0.148207, 0.301809, 0.317732, 0.216907, 0.245138, 0.390181, -0.120865,
798        0.115518, -0.398155, -0.0685247, -0.263379,
799    ];
800    const MOE_E3_W2: [f32; 12] = [
801        -0.0315493, -0.147451, -0.471437, -0.149953, -0.0759736, -0.42393, 0.500115, 0.50571,
802        -0.126137, -0.0207527, 0.0606406, -0.138175,
803    ];
804    const MOE_E3_W3: [f32; 12] = [
805        0.0539619, -0.604005, -0.229171, 0.432303, 0.240478, 0.157986, 0.482184, -0.0937539,
806        0.255604, 0.304551, -0.0564262, -0.17754,
807    ];
808    const L0_SELF_ATTN_RES_NORM_W: [f32; 8] = [
809        1.09202, 0.89273, 1.09569, 0.903039, 0.86568, 0.985966, 0.904269, 0.901858,
810    ];
811    const L0_SELF_ATTN_RES_PROJ_W: [f32; 8] = [
812        0.181373, -0.268475, -0.278623, -0.0176931, 0.15743, -0.189069, -0.0939655, 0.0588636,
813    ];
814    const L0_MLP_RES_NORM_W: [f32; 8] = [
815        1.01577, 1.04681, 0.991541, 1.00423, 0.89548, 1.09467, 1.05601, 0.968893,
816    ];
817    const L0_MLP_RES_PROJ_W: [f32; 8] = [
818        -0.388332, -0.132232, 0.21217, -0.337705, 0.246598, -0.20951, 0.154862, -0.225175,
819    ];
820    const L0_INPUT_LAYERNORM_W: [f32; 8] = [
821        0.970917, 0.970733, 0.946098, 1.23423, 0.883455, 1.12311, 0.809911, 1.12998,
822    ];
823    const L0_POST_ATTN_LAYERNORM_W: [f32; 8] = [
824        0.986205, 1.04826, 1.00761, 1.04983, 0.890071, 0.817534, 1.17858, 0.915493,
825    ];
826    const L1_SELF_ATTN_RES_NORM_W: [f32; 8] = [
827        0.985561, 0.917139, 1.10889, 1.02053, 0.951296, 1.08218, 1.05087, 0.816958,
828    ];
829    const L1_SELF_ATTN_RES_PROJ_W: [f32; 8] = [
830        -0.13164, -0.436846, -0.109316, 0.0236738, 0.0868917, 0.038431, 0.12599, 0.680195,
831    ];
832    const L1_MLP_RES_NORM_W: [f32; 8] = [
833        1.123, 1.03367, 1.15067, 1.06599, 1.10144, 0.905083, 0.911655, 1.11872,
834    ];
835    const L1_MLP_RES_PROJ_W: [f32; 8] = [
836        -0.0917493, -0.258613, 0.62677, 0.0330113, 0.0841344, -0.332195, -0.355091, 0.152981,
837    ];
838    const L1_INPUT_LAYERNORM_W: [f32; 8] = [
839        0.878093, 0.911769, 1.15303, 1.1774, 0.961136, 1.01693, 0.935918, 1.08085,
840    ];
841    const L1_POST_ATTN_LAYERNORM_W: [f32; 8] = [
842        0.982172, 0.868114, 0.876461, 1.00153, 1.10216, 1.01644, 0.784298, 1.0521,
843    ];
844    const OUTPUT_ATTN_RES_NORM_W: [f32; 8] = [
845        0.978339, 0.909439, 1.00794, 0.900698, 1.11013, 1.13568, 1.00689, 1.06737,
846    ];
847    const OUTPUT_ATTN_RES_PROJ_W: [f32; 8] = [
848        -0.0780474, 0.195513, -0.118552, 0.46788, 0.285518, -0.16824, -0.350209, -0.0767004,
849    ];
850    const FINAL_NORM_W: [f32; 8] = [
851        0.969882, 0.931832, 1.10373, 1.16478, 0.891072, 1.0772, 0.85093, 0.91487,
852    ];
853    const EMBEDDING_ROW: [f32; 8] = [
854        0.309738, 0.0378015, -0.337758, 0.323415, 0.162627, -0.685942, 0.134525, 0.42463,
855    ];
856    const OUTPUT_HEAD: [f32; 40] = [
857        0.04209, 0.261934, -0.329344, 0.0625128, 0.272985, 0.290578, 0.0418593, 0.216545, -0.63005,
858        -0.246383, 0.0668848, -0.179091, -0.277674, 0.285747, 0.0959162, -0.0944821, -0.542151,
859        0.2749, 0.299618, -0.0826765, 0.136924, -0.193359, -0.132659, 0.306145, 0.0721566,
860        0.0321218, -0.360207, -0.152256, -0.779984, 0.286993, -0.0684097, 0.142516, -0.124933,
861        0.180322, 0.00527221, 0.0710325, 0.453557, -0.152982, 0.0701835, 0.314973,
862    ];
863    const GOLDEN_LOGITS: [f32; 5] = [0.487583, -1.04021, 0.695218, -0.214018, 1.12848];
864
865    fn expert(
866        w1: &[f32],
867        w2: &[f32],
868        w3: &[f32],
869        ffn_dim: usize,
870        in_dim: usize,
871    ) -> KimiExpertWeights {
872        KimiExpertWeights {
873            w1: wm(w1, ffn_dim, in_dim),
874            w2: wm(w2, in_dim, ffn_dim),
875            w3: wm(w3, ffn_dim, in_dim),
876        }
877    }
878
879    fn make_weights() -> KimiDecoderWeights {
880        let layer0 = KimiDecoderLayerWeights {
881            input_layernorm_weight: L0_INPUT_LAYERNORM_W.to_vec(),
882            attn: KimiLayerAttention::Kda(Box::new(KdaAttnWeights {
883                q_proj: wm(&KDA_Q_PROJ, KDA_PROJ, HIDDEN_DIM),
884                k_proj: wm(&KDA_K_PROJ, KDA_PROJ, HIDDEN_DIM),
885                v_proj: wm(&KDA_V_PROJ, KDA_PROJ, HIDDEN_DIM),
886                q_conv_weight: KDA_Q_CONV_W.to_vec(),
887                k_conv_weight: KDA_K_CONV_W.to_vec(),
888                v_conv_weight: KDA_V_CONV_W.to_vec(),
889                a_log: KDA_A_LOG.to_vec(),
890                f_a_proj: wm(&KDA_F_A_PROJ, KDA_HEAD_DIM, HIDDEN_DIM),
891                f_b_proj: wm(&KDA_F_B_PROJ, KDA_PROJ, KDA_HEAD_DIM),
892                dt_bias: KDA_DT_BIAS.to_vec(),
893                b_proj: wm(&KDA_B_PROJ, KDA_NUM_HEADS, HIDDEN_DIM),
894                g_proj: wm(&KDA_G_PROJ, KDA_PROJ, HIDDEN_DIM),
895                o_norm_weight: KDA_O_NORM_W.to_vec(),
896                o_proj: wm(&KDA_O_PROJ, HIDDEN_DIM, KDA_PROJ),
897            })),
898            post_attention_layernorm_weight: L0_POST_ATTN_LAYERNORM_W.to_vec(),
899            ffn: KimiLayerFfn::Dense(Box::new(DenseMlpWeights {
900                gate_proj: wm(&DENSE_GATE_PROJ, DENSE_INTERMEDIATE, HIDDEN_DIM),
901                up_proj: wm(&DENSE_UP_PROJ, DENSE_INTERMEDIATE, HIDDEN_DIM),
902                down_proj: wm(&DENSE_DOWN_PROJ, HIDDEN_DIM, DENSE_INTERMEDIATE),
903            })),
904            self_attention_res_norm_weight: L0_SELF_ATTN_RES_NORM_W.to_vec(),
905            self_attention_res_proj_weight: L0_SELF_ATTN_RES_PROJ_W.to_vec(),
906            mlp_res_norm_weight: L0_MLP_RES_NORM_W.to_vec(),
907            mlp_res_proj_weight: L0_MLP_RES_PROJ_W.to_vec(),
908        };
909
910        let layer1 = KimiDecoderLayerWeights {
911            input_layernorm_weight: L1_INPUT_LAYERNORM_W.to_vec(),
912            attn: KimiLayerAttention::Mla(Box::new(MlaAttnWeights {
913                q: MlaQProj::LowRank {
914                    a: wm(&MLA_Q_A_PROJ, MLA_Q_LORA, HIDDEN_DIM),
915                    norm: MLA_Q_A_NORM_W.to_vec(),
916                    b: wm(
917                        &MLA_Q_B_PROJ,
918                        MLA_NUM_HEADS * (MLA_QK_NOPE + MLA_QK_ROPE),
919                        MLA_Q_LORA,
920                    ),
921                },
922                kv_a_proj_with_mqa: wm(&MLA_KV_A_PROJ, MLA_KV_LORA + MLA_QK_ROPE, HIDDEN_DIM),
923                kv_a_layernorm: MLA_KV_A_NORM_W.to_vec(),
924                kv_b: MlaKvB::Combined(wm(
925                    &MLA_KV_B_PROJ,
926                    MLA_NUM_HEADS * (MLA_QK_NOPE + MLA_V_HEAD_DIM),
927                    MLA_KV_LORA,
928                )),
929                o_proj: wm(&MLA_O_PROJ, HIDDEN_DIM, MLA_PROJ),
930                g_proj: Some(wm(&MLA_G_PROJ, MLA_PROJ, HIDDEN_DIM)),
931            })),
932            post_attention_layernorm_weight: L1_POST_ATTN_LAYERNORM_W.to_vec(),
933            ffn: KimiLayerFfn::Moe(Box::new(KimiLatentMoeWeights {
934                router_weight: wm(&MOE_ROUTER_WEIGHT, N_EXPERTS, HIDDEN_DIM),
935                e_score_correction_bias: MOE_BIAS.to_vec(),
936                down_proj: wm(&MOE_DOWN_PROJ, MOE_HIDDEN_DIM, HIDDEN_DIM),
937                up_proj: wm(&MOE_UP_PROJ, HIDDEN_DIM, MOE_HIDDEN_DIM),
938                routed_expert_norm_weight: Some(MOE_ROUTED_NORM_W.to_vec()),
939                experts: KimiExpertBacking::Resident(vec![
940                    expert(
941                        &MOE_E0_W1,
942                        &MOE_E0_W2,
943                        &MOE_E0_W3,
944                        MOE_INTERMEDIATE,
945                        MOE_HIDDEN_DIM,
946                    ),
947                    expert(
948                        &MOE_E1_W1,
949                        &MOE_E1_W2,
950                        &MOE_E1_W3,
951                        MOE_INTERMEDIATE,
952                        MOE_HIDDEN_DIM,
953                    ),
954                    expert(
955                        &MOE_E2_W1,
956                        &MOE_E2_W2,
957                        &MOE_E2_W3,
958                        MOE_INTERMEDIATE,
959                        MOE_HIDDEN_DIM,
960                    ),
961                    expert(
962                        &MOE_E3_W1,
963                        &MOE_E3_W2,
964                        &MOE_E3_W3,
965                        MOE_INTERMEDIATE,
966                        MOE_HIDDEN_DIM,
967                    ),
968                ]),
969                shared_expert: expert(
970                    &MOE_SHARED_W1,
971                    &MOE_SHARED_W2,
972                    &MOE_SHARED_W3,
973                    SHARED_INTERMEDIATE,
974                    HIDDEN_DIM,
975                ),
976            })),
977            self_attention_res_norm_weight: L1_SELF_ATTN_RES_NORM_W.to_vec(),
978            self_attention_res_proj_weight: L1_SELF_ATTN_RES_PROJ_W.to_vec(),
979            mlp_res_norm_weight: L1_MLP_RES_NORM_W.to_vec(),
980            mlp_res_proj_weight: L1_MLP_RES_PROJ_W.to_vec(),
981        };
982
983        KimiDecoderWeights {
984            embedding: Tensor::new(EMBEDDING_ROW.to_vec(), vec![1, HIDDEN_DIM]),
985            layers: vec![layer0, layer1],
986            output_attn_res_norm_weight: OUTPUT_ATTN_RES_NORM_W.to_vec(),
987            output_attn_res_proj_weight: OUTPUT_ATTN_RES_PROJ_W.to_vec(),
988            final_norm_weight: FINAL_NORM_W.to_vec(),
989            output_head: wm(&OUTPUT_HEAD, OUTPUT_VOCAB, HIDDEN_DIM),
990        }
991    }
992
993    fn decoder_cfg() -> KimiDecoderConfig {
994        KimiDecoderConfig {
995            attn_res_block_size: ATTN_RES_BLOCK_SIZE,
996            rms_norm_eps: EPS,
997            situ_beta: SITU_BETA,
998            situ_linear_beta: SITU_LINEAR_BETA,
999            moe: KimiMoeConfig {
1000                n_experts_active: TOP_K,
1001                moe_renormalize: true,
1002                routed_scaling_factor: 1.0,
1003                situ_beta: SITU_BETA,
1004                situ_linear_beta: SITU_LINEAR_BETA,
1005                rms_norm_eps: EPS,
1006            },
1007        }
1008    }
1009
1010    fn mla_cfg() -> MlaConfig {
1011        MlaConfig {
1012            num_heads: MLA_NUM_HEADS,
1013            q_lora_rank: MLA_Q_LORA,
1014            kv_lora_rank: MLA_KV_LORA,
1015            qk_nope_head_dim: MLA_QK_NOPE,
1016            qk_rope_head_dim: MLA_QK_ROPE,
1017            v_head_dim: MLA_V_HEAD_DIM,
1018            use_output_gate: true,
1019            rope: None,
1020        }
1021    }
1022
1023    fn kda_cfg() -> KdaConfig {
1024        KdaConfig {
1025            num_heads: KDA_NUM_HEADS,
1026            head_dim: KDA_HEAD_DIM,
1027            short_conv_kernel_size: KDA_CONV_SIZE,
1028            gate_lower_bound: KDA_GATE_LOWER_BOUND,
1029            use_full_rank_gate: true,
1030        }
1031    }
1032
1033    /// One Kimi K3 decode step enters the CPU worker pool once.
1034    ///
1035    /// Kimi has no checkpoint this machine can load, so the instance is
1036    /// the same synthetic two-layer stack every other test in this file
1037    /// uses. The count being asserted does not depend on the weights:
1038    /// it is how many times the driving thread crossed rayon's cold
1039    /// submission path, which before `engine/entry.rs` was once per
1040    /// parallel region and is now once per step.
1041    ///
1042    /// Sabotage: drop the `par::on_workers` from
1043    /// `Engine::forward_token` and this goes red with the region count.
1044    #[test]
1045    fn one_kimi_decode_step_enters_the_pool_once() {
1046        crate::engine::assert_one_pool_entry_per_step(
1047            &crate::KimiEngine {
1048                weights: make_weights(),
1049                cfg: decoder_cfg(),
1050                mla_cfg: mla_cfg(),
1051                kda_cfg: kda_cfg(),
1052            },
1053            0,
1054        );
1055    }
1056
1057    #[test]
1058    fn two_mixed_layers_match_independent_python_reference() {
1059        let weights = make_weights();
1060        let cfg = decoder_cfg();
1061        let mla_cfg = mla_cfg();
1062        let kda_cfg = kda_cfg();
1063        let mut state = KimiDecodeState::new(&weights, &kda_cfg);
1064
1065        let logits = kimi_forward_token(&weights, &cfg, &mla_cfg, &kda_cfg, 0, &mut state);
1066        assert_eq!(logits.len(), GOLDEN_LOGITS.len());
1067        for (i, (a, b)) in logits.iter().zip(GOLDEN_LOGITS.iter()).enumerate() {
1068            assert!((a - b).abs() < 1e-3, "logit {i}: rust={a} python={b}");
1069        }
1070    }
1071}
frink_models/kimi_decoder.rs

frink_models/
kimi_decoder.rs
