avarok_core/kimi_k3/
cpu_weights.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Host-side 0.40B-shaped (or tiny) K3 weights for the C1 CPU graph.
4
5use super::kda::KdaConfig;
6use super::latent_moe::LatentMoeConfig;
7use super::layer::{K3Graph, K3LayerSpec, MixerKind, MlpKind};
8use super::mla::MlaConfig;
9use super::ops::{fill, ident, ones};
10
11/// Graph mutants. Mix=0 / force-expert-0 are C1; skip-layer is C2;
12/// `o_proj` TP is C7.
13#[derive(Clone, Copy, Debug)]
14pub struct Ablation {
15    pub attnres_mix: f32,
16    pub force_expert: Option<usize>,
17    /// Skip this layer index in `forward_token` (C2 known-bad).
18    pub skip_layer: Option<usize>,
19    /// Column-parallel `o_proj` world size. 1 = unsplit, 2 = in-process TP.
20    pub o_proj_tp: usize,
21    /// Zero this TP rank's `o_proj` shard (C7 known-bad). Requires `o_proj_tp=2`.
22    pub drop_o_proj_rank: Option<usize>,
23}
24
25impl Default for Ablation {
26    fn default() -> Self {
27        Self {
28            attnres_mix: 1.0,
29            force_expert: None,
30            skip_layer: None,
31            o_proj_tp: 1,
32            drop_o_proj_rank: None,
33        }
34    }
35}
36
37impl Ablation {
38    /// GPU-wrapper RST hook. `K3_ATTNRES_MIX=0` is the C5 known-bad (mix=0
39    /// is identity skip). `K3_FORCE_EXPERT=0` is the C6 known-bad.
40    ///
41    /// `spark serve` of the 0.40B twin has no Ablation CLI; these env vars
42    /// are how a planted mutant still changes tokens on the host fallback.
43    pub fn from_env() -> Self {
44        let mut a = Self::default();
45        if let Ok(v) = std::env::var("K3_ATTNRES_MIX")
46            && let Ok(m) = v.parse()
47        {
48            a.attnres_mix = m;
49        }
50        if let Ok(v) = std::env::var("K3_FORCE_EXPERT") {
51            a.force_expert = v.parse().ok();
52        }
53        a
54    }
55}
56
57#[derive(Clone, Debug)]
58pub struct KdaWeights {
59    pub q_proj: Vec<f32>,
60    pub k_proj: Vec<f32>,
61    pub v_proj: Vec<f32>,
62    pub conv: Vec<f32>,
63    pub f_a: Vec<f32>,
64    pub f_b: Vec<f32>,
65    pub dt_bias: Vec<f32>,
66    pub a_log: Vec<f32>,
67    pub b_proj: Vec<f32>,
68    pub g_proj: Vec<f32>,
69    pub o_norm: Vec<f32>,
70    pub o_proj: Vec<f32>,
71}
72
73#[derive(Clone, Debug)]
74pub struct MlaWeights {
75    pub q_a: Vec<f32>,
76    pub q_a_ln: Vec<f32>,
77    pub q_b: Vec<f32>,
78    pub kv_a: Vec<f32>,
79    pub kv_a_ln: Vec<f32>,
80    pub kv_b: Vec<f32>,
81    pub g_proj: Vec<f32>,
82    pub o_proj: Vec<f32>,
83}
84
85#[derive(Clone, Debug)]
86pub struct DenseMlp {
87    pub gate: Vec<f32>,
88    pub up: Vec<f32>,
89    pub down: Vec<f32>,
90}
91
92#[derive(Clone, Debug)]
93pub struct MoeWeights {
94    pub down: Vec<f32>,
95    pub up: Vec<f32>,
96    pub norm: Vec<f32>,
97    pub router: Vec<f32>,
98    pub bias: Vec<f32>,
99    pub experts: Vec<(Vec<f32>, Vec<f32>, Vec<f32>)>,
100    pub shared: Option<DenseMlp>,
101}
102
103#[derive(Clone, Debug)]
104pub enum MixerW {
105    Kda(KdaWeights),
106    Mla(MlaWeights),
107}
108
109#[derive(Clone, Debug)]
110pub enum MlpW {
111    Dense(DenseMlp),
112    Moe(MoeWeights),
113}
114
115#[derive(Clone, Debug)]
116pub struct K3CpuLayer {
117    pub spec: K3LayerSpec,
118    pub input_norm: Vec<f32>,
119    pub post_norm: Vec<f32>,
120    pub attn_res_proj: Vec<f32>,
121    pub attn_res_norm: Vec<f32>,
122    pub mlp_res_proj: Vec<f32>,
123    pub mlp_res_norm: Vec<f32>,
124    pub mixer: MixerW,
125    pub mlp: MlpW,
126}
127
128/// CPU twin (or tiny stand-in) ready for greedy decode.
129#[derive(Clone, Debug)]
130pub struct K3CpuModel {
131    pub graph: K3Graph,
132    pub kda: KdaConfig,
133    pub mla: MlaConfig,
134    pub moe: LatentMoeConfig,
135    pub dense_intermediate: usize,
136    pub eps: f32,
137    pub rope_theta: f32,
138    pub vocab: usize,
139    pub embed: Vec<f32>,
140    pub lm_head: Vec<f32>,
141    pub final_norm: Vec<f32>,
142    pub output_res_proj: Vec<f32>,
143    pub output_res_norm: Vec<f32>,
144    pub layers: Vec<K3CpuLayer>,
145}
146
147impl K3CpuModel {
148    /// 8-layer twin pattern, tiny dims. Deterministic; no HF weights.
149    pub fn synthetic_tiny() -> Self {
150        let hidden = 4;
151        let graph = twin_pattern_graph(hidden);
152        let kda = KdaConfig {
153            heads: 1,
154            head_dim: 2,
155            conv_kernel: 4,
156            gate_lower_bound: Some(-5.0),
157            use_full_rank_gate: true,
158        };
159        let mla = MlaConfig {
160            heads: 1,
161            qk_nope_head_dim: 2,
162            qk_rope_head_dim: 2,
163            v_head_dim: 2,
164            q_lora_rank: 4,
165            kv_lora_rank: 2,
166            mla_use_nope: true,
167            mla_use_output_gate: true,
168        };
169        let moe = LatentMoeConfig {
170            hidden,
171            latent: 4,
172            expert_hidden: 4,
173            n_routed: 2,
174            top_k: 1,
175            n_shared: 1,
176            situ_beta: 4.0,
177            situ_linear_beta: 25.0,
178            use_norm: true,
179            renormalize: true,
180        };
181        Self::build(graph, kda, mla, moe, 4, 8, 1e-5, 10_000.0)
182    }
183
184    /// Same 8-layer twin pattern, slightly wider. Still cheap on CPU.
185    pub fn synthetic_small() -> Self {
186        let hidden = 8;
187        let graph = twin_pattern_graph(hidden);
188        let kda = KdaConfig {
189            heads: 2,
190            head_dim: 4,
191            conv_kernel: 4,
192            gate_lower_bound: Some(-5.0),
193            use_full_rank_gate: true,
194        };
195        let mla = MlaConfig {
196            heads: 2,
197            qk_nope_head_dim: 4,
198            qk_rope_head_dim: 2,
199            v_head_dim: 4,
200            q_lora_rank: 8,
201            kv_lora_rank: 4,
202            mla_use_nope: true,
203            mla_use_output_gate: true,
204        };
205        let moe = LatentMoeConfig {
206            hidden,
207            latent: 8,
208            expert_hidden: 8,
209            n_routed: 2,
210            top_k: 1,
211            n_shared: 1,
212            situ_beta: 4.0,
213            situ_linear_beta: 25.0,
214            use_norm: true,
215            renormalize: true,
216        };
217        Self::build(graph, kda, mla, moe, 8, 16, 1e-5, 10_000.0)
218    }
219
220    /// Production **width**, dummy depth: hidden=7168, heads=96, 2 layers
221    /// (KDA + MLA), vocab=256, 8 routed experts. Head/lora dims stay tiny so
222    /// the CPU test is fast. Not 93 layers / 1.56 TB.
223    pub fn synthetic_prod_width_dummy() -> Self {
224        let hidden = 7168;
225        let graph = prod_width_dummy_graph(hidden);
226        let kda = KdaConfig {
227            heads: 96,
228            head_dim: 2,
229            conv_kernel: 4,
230            gate_lower_bound: Some(-5.0),
231            use_full_rank_gate: true,
232        };
233        let mla = MlaConfig {
234            heads: 96,
235            qk_nope_head_dim: 2,
236            qk_rope_head_dim: 2,
237            v_head_dim: 2,
238            q_lora_rank: 8,
239            kv_lora_rank: 4,
240            mla_use_nope: true,
241            mla_use_output_gate: true,
242        };
243        let moe = LatentMoeConfig {
244            hidden,
245            latent: 64,
246            expert_hidden: 64,
247            n_routed: 8,
248            top_k: 2,
249            n_shared: 1,
250            situ_beta: 4.0,
251            situ_linear_beta: 25.0,
252            use_norm: true,
253            renormalize: true,
254        };
255        let mut model = Self::build(graph, kda, mla, moe, 64, 256, 1e-5, 10_000.0);
256        // ident o_proj is rank-1-empty at this width (inn << hidden). Fill so
257        // both TP column shards contribute to every hidden dim.
258        for (i, layer) in model.layers.iter_mut().enumerate() {
259            let seed = 200 + i as u32;
260            match &mut layer.mixer {
261                MixerW::Kda(w) => w.o_proj = fill(w.o_proj.len(), seed, 0.05),
262                MixerW::Mla(w) => w.o_proj = fill(w.o_proj.len(), seed + 1, 0.05),
263            }
264        }
265        model
266    }
267}
268
269fn twin_pattern_graph(hidden: usize) -> K3Graph {
270    let mixers = [
271        MixerKind::Kda,
272        MixerKind::Kda,
273        MixerKind::Kda,
274        MixerKind::Mla,
275        MixerKind::Kda,
276        MixerKind::Kda,
277        MixerKind::Kda,
278        MixerKind::Mla,
279    ];
280    let layers = mixers
281        .iter()
282        .enumerate()
283        .map(|(i, m)| K3LayerSpec {
284            index: i,
285            mixer: *m,
286            mlp: if i == 0 {
287                MlpKind::Dense
288            } else {
289                MlpKind::LatentMoe
290            },
291        })
292        .collect();
293    K3Graph {
294        layers,
295        hidden,
296        attn_res_block_size: 4,
297        situ_beta: 4.0,
298        situ_linear_beta: 25.0,
299        use_full_rank_gate: true,
300        mla_use_nope: true,
301        mla_use_output_gate: true,
302    }
303}
304
305fn prod_width_dummy_graph(hidden: usize) -> K3Graph {
306    let layers = vec![
307        K3LayerSpec {
308            index: 0,
309            mixer: MixerKind::Kda,
310            mlp: MlpKind::Dense,
311        },
312        K3LayerSpec {
313            index: 1,
314            mixer: MixerKind::Mla,
315            mlp: MlpKind::LatentMoe,
316        },
317    ];
318    K3Graph {
319        layers,
320        hidden,
321        attn_res_block_size: 12,
322        situ_beta: 4.0,
323        situ_linear_beta: 25.0,
324        use_full_rank_gate: true,
325        mla_use_nope: true,
326        mla_use_output_gate: true,
327    }
328}
329
330impl K3CpuModel {
331    #[allow(clippy::too_many_arguments)]
332    fn build(
333        graph: K3Graph,
334        kda: KdaConfig,
335        mla: MlaConfig,
336        moe: LatentMoeConfig,
337        dense_intermediate: usize,
338        vocab: usize,
339        eps: f32,
340        rope_theta: f32,
341    ) -> Self {
342        let h = graph.hidden;
343        let layers = graph
344            .layers
345            .iter()
346            .map(|spec| synth_layer(spec, &kda, &mla, &moe, dense_intermediate))
347            .collect();
348        Self {
349            graph,
350            kda,
351            mla,
352            moe,
353            dense_intermediate,
354            eps,
355            rope_theta,
356            vocab,
357            embed: fill(vocab * h, 1, 0.05),
358            lm_head: fill(vocab * h, 2, 0.05),
359            final_norm: ones(h),
360            output_res_proj: fill(h, 3, 0.05),
361            output_res_norm: ones(h),
362            layers,
363        }
364    }
365}
366
367fn synth_layer(
368    spec: &K3LayerSpec,
369    kda: &KdaConfig,
370    mla: &MlaConfig,
371    moe: &LatentMoeConfig,
372    dense_int: usize,
373) -> K3CpuLayer {
374    let h = moe.hidden;
375    let seed = 10 + spec.index as u32 * 17;
376    let mixer = match spec.mixer {
377        MixerKind::Kda => MixerW::Kda(synth_kda(kda, h, seed)),
378        MixerKind::Mla => MixerW::Mla(synth_mla(mla, h, seed)),
379    };
380    let mlp = match spec.mlp {
381        MlpKind::Dense => MlpW::Dense(synth_dense(h, dense_int, seed + 1)),
382        MlpKind::LatentMoe => MlpW::Moe(synth_moe(moe, seed + 2)),
383    };
384    K3CpuLayer {
385        spec: *spec,
386        input_norm: ones(h),
387        post_norm: ones(h),
388        attn_res_proj: fill(h, seed + 3, 0.05),
389        attn_res_norm: ones(h),
390        mlp_res_proj: fill(h, seed + 4, 0.05),
391        mlp_res_norm: ones(h),
392        mixer,
393        mlp,
394    }
395}
396
397fn synth_kda(kda: &KdaConfig, hidden: usize, seed: u32) -> KdaWeights {
398    let q = kda.qkv_dim();
399    KdaWeights {
400        q_proj: ident(q, hidden),
401        k_proj: ident(q, hidden),
402        v_proj: ident(q, hidden),
403        conv: fill(kda.conv_elems(), seed, 0.05),
404        f_a: fill(kda.head_dim * hidden, seed + 1, 0.05),
405        f_b: fill(q * kda.head_dim, seed + 2, 0.05),
406        dt_bias: vec![0.0; q],
407        a_log: vec![-1.0; kda.heads],
408        b_proj: fill(kda.heads * hidden, seed + 3, 0.05),
409        g_proj: fill(q * hidden, seed + 4, 0.05),
410        o_norm: ones(kda.head_dim),
411        o_proj: ident(hidden, q),
412    }
413}
414
415fn synth_mla(mla: &MlaConfig, hidden: usize, seed: u32) -> MlaWeights {
416    let qk = mla.heads * mla.qk_head_dim();
417    let dv = mla.heads * mla.v_head_dim;
418    let kv_in = mla.kv_lora_rank + mla.qk_rope_head_dim;
419    let kv_b_out = mla.heads * (mla.qk_nope_head_dim + mla.v_head_dim);
420    MlaWeights {
421        q_a: ident(mla.q_lora_rank, hidden),
422        q_a_ln: ones(mla.q_lora_rank),
423        q_b: fill(qk * mla.q_lora_rank, seed, 0.05),
424        kv_a: fill(kv_in * hidden, seed + 1, 0.05),
425        kv_a_ln: ones(mla.kv_lora_rank),
426        kv_b: fill(kv_b_out * mla.kv_lora_rank, seed + 2, 0.05),
427        g_proj: fill(dv * hidden, seed + 3, 0.05),
428        o_proj: ident(hidden, dv),
429    }
430}
431
432fn synth_dense(h: usize, inter: usize, seed: u32) -> DenseMlp {
433    DenseMlp {
434        gate: fill(inter * h, seed, 0.05),
435        up: fill(inter * h, seed + 1, 0.05),
436        down: ident(h, inter),
437    }
438}
439
440fn synth_moe(moe: &LatentMoeConfig, seed: u32) -> MoeWeights {
441    let mut experts = Vec::with_capacity(moe.n_routed);
442    for e in 0..moe.n_routed {
443        let mut w3 = ident(moe.expert_hidden, moe.latent);
444        // Expert 0 stays identity-ish; others scale the up-branch.
445        if e > 0 {
446            w3[0] = 1.0 + e as f32;
447        }
448        experts.push((
449            ident(moe.expert_hidden, moe.latent),
450            ident(moe.latent, moe.expert_hidden),
451            w3,
452        ));
453    }
454    let shared = (moe.n_shared > 0).then(|| synth_dense(moe.hidden, moe.expert_hidden, seed + 9));
455    MoeWeights {
456        down: ident(moe.latent, moe.hidden),
457        up: ident(moe.hidden, moe.latent),
458        norm: ones(moe.latent),
459        router: fill(moe.n_routed * moe.hidden, seed, 0.2),
460        bias: vec![0.0; moe.n_routed],
461        experts,
462        shared,
463    }
464}
avarok_core/kimi_k3/cpu_weights.rs

avarok_core/kimi_k3/
cpu_weights.rs
