avarok_core/kimi_k3/
cpu_forward.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! One-token K3 CPU forward: AttnRes + mixer + MLP, then logits.
4
5use anyhow::Result;
6
7use super::attnres::rms_norm;
8use super::cache::{HybridCache, LayerCache, MlaKv};
9use super::cpu_weights::{
10    Ablation, DenseMlp, K3CpuLayer, K3CpuModel, KdaWeights, MixerW, MlaWeights, MlpW, MoeWeights,
11};
12use super::kda::{KdaConfig, KdaState, bounded_gate, kda_decode_token};
13use super::latent_moe::{LatentMoeConfig, mix_routed_experts, sigmoid_topk};
14use super::mla::{MlaConfig, mla_decode_token};
15use super::ops::{embed_token, matvec, matvec_column_tp};
16use super::situ::sigmoid;
17
18mod dense;
19mod stream;
20pub use stream::AttnResStream;
21
22/// After row-parallel `o_proj` / dense MLP `down`. None at TP=1.
23pub type HiddenReduce<'a> = &'a dyn Fn(&mut [f32]) -> Result<()>;
24
25/// Optional resident dense/shared MLP, with local TP intermediate width.
26pub type DenseMlpCore<'a> =
27    &'a dyn Fn(&DenseMlp, &[f32], usize, usize, f32, f32) -> Result<Vec<f32>>;
28
29/// Per-layer geometry the GPU wrapper and `forward_token` share.
30///
31/// `K3BoundLayer::decode` copies hidden D2H, runs this math, copies H2D.
32/// LinearAttention injects CUDA conv+recurrent unless `K3_CUDA_KDA=0`.
33/// FullAttention injects CUDA gated-NoPE MLA unless `K3_CUDA_MLA=0`.
34pub struct K3LayerCtx<'a> {
35    pub kda: &'a KdaConfig,
36    pub mla: &'a MlaConfig,
37    pub moe: &'a LatentMoeConfig,
38    pub situ_beta: f32,
39    pub situ_linear_beta: f32,
40    pub hidden: usize,
41    pub dense_intermediate: usize,
42    pub eps: f32,
43    pub rope_theta: f32,
44    /// After row-parallel `o_proj` (and dense MLP `down`). None at TP=1.
45    pub reduce_hidden: Option<HiddenReduce<'a>>,
46    pub dense_mlp: Option<DenseMlpCore<'a>>,
47}
48
49impl<'a> K3LayerCtx<'a> {
50    pub fn from_model(model: &'a K3CpuModel) -> Self {
51        Self {
52            kda: &model.kda,
53            mla: &model.mla,
54            moe: &model.moe,
55            situ_beta: model.graph.situ_beta,
56            situ_linear_beta: model.graph.situ_linear_beta,
57            hidden: model.graph.hidden,
58            dense_intermediate: model.dense_intermediate,
59            eps: model.eps,
60            rope_theta: model.rope_theta,
61            reduce_hidden: None,
62            dense_mlp: None,
63        }
64    }
65}
66
67/// Embed one token and run every decoder layer at `pos`. Updates the hybrid cache.
68/// AttnRes is per-token (layer depth), not carried across the sequence.
69pub fn forward_token(
70    model: &K3CpuModel,
71    token: u32,
72    pos: usize,
73    cache: &mut HybridCache,
74    ablation: Ablation,
75) -> Vec<f32> {
76    let embed = embed_token(&model.embed, token, model.graph.hidden, model.vocab);
77    let mut stream = AttnResStream::new(model.graph.hidden, model.graph.attn_res_block_size);
78    stream.partial.clone_from(&embed);
79    let ctx = K3LayerCtx::from_model(model);
80    for layer in &model.layers {
81        if ablation.skip_layer == Some(layer.spec.index) {
82            continue;
83        }
84        let mixer_state = &mut cache.layers[layer.spec.index];
85        forward_one_layer(&ctx, layer, pos, mixer_state, &mut stream, ablation);
86    }
87    let h = stream.mix(
88        &model.output_res_proj,
89        &model.output_res_norm,
90        model.eps,
91        ablation.attnres_mix,
92    );
93    rms_norm(&h, &model.final_norm, model.eps)
94}
95
96/// LM head over a hidden vector. `W` is `[vocab, hidden]`.
97pub fn logits(model: &K3CpuModel, h: &[f32]) -> Vec<f32> {
98    matvec(&model.lm_head, h, model.vocab, model.graph.hidden)
99}
100
101/// One decoder layer: AttnRes + (KDA|MLA) + MLP. Mutates this layer's
102/// [`LayerCache`] and the token's [`AttnResStream`].
103///
104/// Default cores are [`kda_decode_token`] / [`mla_decode_token`] / host
105/// [`mix_routed_experts`]. Serve CUDA injects via [`forward_one_layer_with_cores`].
106pub fn forward_one_layer(
107    ctx: &K3LayerCtx<'_>,
108    layer: &K3CpuLayer,
109    pos: usize,
110    mixer_state: &mut LayerCache,
111    stream: &mut AttnResStream,
112    ablation: Ablation,
113) {
114    forward_one_layer_with_cores(
115        ctx,
116        layer,
117        pos,
118        mixer_state,
119        stream,
120        ablation,
121        cpu_kda_core,
122        cpu_mla_core,
123        cpu_moe_core,
124    )
125    .expect("K3 CPU mixer cores are infallible")
126}
127
128fn cpu_kda_core(
129    x: &[f32],
130    w: &[f32],
131    g: &[f32],
132    b: &[f32],
133    cfg: &KdaConfig,
134    st: &mut KdaState,
135) -> Result<Vec<f32>> {
136    Ok(kda_decode_token(x, w, g, b, cfg, st))
137}
138
139#[allow(clippy::too_many_arguments)]
140fn cpu_mla_core(
141    q: &mut [f32],
142    k: &mut [f32],
143    v: &[f32],
144    g: &[f32],
145    kv: &mut MlaKv,
146    cfg: &MlaConfig,
147    pos: usize,
148    theta: f32,
149) -> Result<Vec<f32>> {
150    Ok(mla_decode_token(q, k, v, g, kv, cfg, pos, theta))
151}
152
153/// Same as [`forward_one_layer`], with a replaceable KDA conv+recurrent core.
154///
155/// Projections, MLA, MLP, and AttnRes stay on this CPU path.
156#[allow(clippy::too_many_arguments)]
157pub fn forward_one_layer_with_kda_decode<F>(
158    ctx: &K3LayerCtx<'_>,
159    layer: &K3CpuLayer,
160    pos: usize,
161    mixer_state: &mut LayerCache,
162    stream: &mut AttnResStream,
163    ablation: Ablation,
164    kda_decode: F,
165) -> Result<()>
166where
167    F: FnMut(&[f32], &[f32], &[f32], &[f32], &KdaConfig, &mut KdaState) -> Result<Vec<f32>>,
168{
169    forward_one_layer_with_cores(
170        ctx,
171        layer,
172        pos,
173        mixer_state,
174        stream,
175        ablation,
176        kda_decode,
177        cpu_mla_core,
178        cpu_moe_core,
179    )
180}
181
182/// Same as [`forward_one_layer`], with a replaceable MLA rope+SDPA+gate core.
183///
184/// Projections, KDA, MLP, and AttnRes stay on this CPU path.
185#[allow(clippy::too_many_arguments)]
186pub fn forward_one_layer_with_mla_decode<F>(
187    ctx: &K3LayerCtx<'_>,
188    layer: &K3CpuLayer,
189    pos: usize,
190    mixer_state: &mut LayerCache,
191    stream: &mut AttnResStream,
192    ablation: Ablation,
193    mla_decode: F,
194) -> Result<()>
195where
196    F: FnMut(
197        &mut [f32],
198        &mut [f32],
199        &[f32],
200        &[f32],
201        &mut MlaKv,
202        &MlaConfig,
203        usize,
204        f32,
205    ) -> Result<Vec<f32>>,
206{
207    forward_one_layer_with_cores(
208        ctx,
209        layer,
210        pos,
211        mixer_state,
212        stream,
213        ablation,
214        cpu_kda_core,
215        mla_decode,
216        cpu_moe_core,
217    )
218}
219
220fn cpu_moe_core(
221    w: &MoeWeights,
222    latent: &[f32],
223    ids: &[usize],
224    mix_w: &[f32],
225    cfg: &LatentMoeConfig,
226) -> Result<Vec<f32>> {
227    Ok(mix_routed_experts(latent, ids, mix_w, &w.experts, cfg))
228}
229
230/// Same as [`forward_one_layer`], with replaceable KDA / MLA / routed-expert cores.
231///
232/// Projections, AttnRes, router, down/up, shared experts, and SiTU mix stay here
233/// unless the MoE callback replaces the expert GEMMs.
234#[allow(clippy::too_many_arguments)]
235pub fn forward_one_layer_with_cores<FK, FM, FE>(
236    ctx: &K3LayerCtx<'_>,
237    layer: &K3CpuLayer,
238    pos: usize,
239    mixer_state: &mut LayerCache,
240    stream: &mut AttnResStream,
241    ablation: Ablation,
242    mut kda_decode: FK,
243    mut mla_decode: FM,
244    mut moe_experts: FE,
245) -> Result<()>
246where
247    FK: FnMut(&[f32], &[f32], &[f32], &[f32], &KdaConfig, &mut KdaState) -> Result<Vec<f32>>,
248    FM: FnMut(
249        &mut [f32],
250        &mut [f32],
251        &[f32],
252        &[f32],
253        &mut MlaKv,
254        &MlaConfig,
255        usize,
256        f32,
257    ) -> Result<Vec<f32>>,
258    FE: FnMut(&MoeWeights, &[f32], &[usize], &[f32], &LatentMoeConfig) -> Result<Vec<f32>>,
259{
260    let eps = ctx.eps;
261    let mix = ablation.attnres_mix;
262    let h = stream.mix(&layer.attn_res_proj, &layer.attn_res_norm, eps, mix);
263    // Archive the *incoming* prefix (HF), not the post-mixer partial.
264    stream.archive_incoming_at_block_start(layer.spec.index);
265    let x = rms_norm(&h, &layer.input_norm, eps);
266    let mut mix_out = match (&layer.mixer, mixer_state) {
267        (MixerW::Kda(w), LayerCache::Kda(state)) => {
268            kda_mixer(w, &x, ctx.kda, state, eps, ablation, &mut kda_decode)?
269        }
270        (MixerW::Mla(w), LayerCache::Mla(kv)) => mla_mixer(
271            w,
272            &x,
273            ctx.mla,
274            kv,
275            pos,
276            ctx.rope_theta,
277            eps,
278            ablation,
279            &mut mla_decode,
280        )?,
281        _ => panic!(
282            "K3 mixer/state mismatch at layer {} (CPU fallback GPU wrapper)",
283            layer.spec.index
284        ),
285    };
286    if let Some(f) = ctx.reduce_hidden {
287        f(&mut mix_out)?;
288    }
289    stream.add(&mix_out);
290    let h = stream.mix(&layer.mlp_res_proj, &layer.mlp_res_norm, eps, mix);
291    let x = rms_norm(&h, &layer.post_norm, eps);
292    let mlp_out = match &layer.mlp {
293        MlpW::Dense(w) => {
294            let mut y = dense::run(ctx, w, &x, ctx.dense_intermediate)?;
295            if let Some(f) = ctx.reduce_hidden {
296                f(&mut y)?;
297            }
298            y
299        }
300        MlpW::Moe(w) => moe_mlp_with(w, &x, ctx, ablation.force_expert, &mut moe_experts)?,
301    };
302    stream.add(&mlp_out);
303    Ok(())
304}
305
306fn apply_o_proj(w: &[f32], x: &[f32], out: usize, inn: usize, ablation: Ablation) -> Vec<f32> {
307    matvec_column_tp(
308        w,
309        x,
310        out,
311        inn,
312        ablation.o_proj_tp,
313        ablation.drop_o_proj_rank,
314    )
315}
316
317fn kda_mixer<F>(
318    w: &KdaWeights,
319    x: &[f32],
320    cfg: &KdaConfig,
321    state: &mut KdaState,
322    eps: f32,
323    ablation: Ablation,
324    kda_decode: &mut F,
325) -> Result<Vec<f32>>
326where
327    F: FnMut(&[f32], &[f32], &[f32], &[f32], &KdaConfig, &mut KdaState) -> Result<Vec<f32>>,
328{
329    let qdim = cfg.qkv_dim();
330    let q = matvec(&w.q_proj, x, qdim, x.len());
331    let k = matvec(&w.k_proj, x, qdim, x.len());
332    let v = matvec(&w.v_proj, x, qdim, x.len());
333    let mut qkv = q;
334    qkv.extend_from_slice(&k);
335    qkv.extend_from_slice(&v);
336    // HF: g = f_b_proj(f_a_proj(x)) — two linears, no SiLU on the bottleneck.
337    let fa = matvec(&w.f_a, x, cfg.head_dim, x.len());
338    let z = matvec(&w.f_b, &fa, qdim, cfg.head_dim);
339    let gate = bounded_gate(
340        &z,
341        &w.dt_bias,
342        &w.a_log,
343        cfg.heads,
344        cfg.head_dim,
345        cfg.gate_lower_bound,
346    );
347    let beta = matvec(&w.b_proj, x, cfg.heads, x.len());
348    let g = matvec(&w.g_proj, x, qdim, x.len());
349    let core = kda_decode(&qkv, &w.conv, &gate, &beta, cfg, state)?;
350    let gated = gated_o_norm(&core, &g, &w.o_norm, cfg.head_dim, eps);
351    Ok(apply_o_proj(&w.o_proj, &gated, x.len(), qdim, ablation))
352}
353
354fn gated_o_norm(core: &[f32], g: &[f32], o_norm: &[f32], head_dim: usize, eps: f32) -> Vec<f32> {
355    let mut out = vec![0.0f32; core.len()];
356    for ((c, gg), o) in core
357        .chunks_exact(head_dim)
358        .zip(g.chunks_exact(head_dim))
359        .zip(out.chunks_exact_mut(head_dim))
360    {
361        let n = rms_norm(c, o_norm, eps);
362        for i in 0..head_dim {
363            o[i] = sigmoid(gg[i]) * n[i];
364        }
365    }
366    out
367}
368
369#[allow(clippy::too_many_arguments)]
370fn mla_mixer<F>(
371    w: &MlaWeights,
372    x: &[f32],
373    cfg: &MlaConfig,
374    kv: &mut MlaKv,
375    pos: usize,
376    theta: f32,
377    eps: f32,
378    ablation: Ablation,
379    mla_decode: &mut F,
380) -> Result<Vec<f32>>
381where
382    F: FnMut(
383        &mut [f32],
384        &mut [f32],
385        &[f32],
386        &[f32],
387        &mut MlaKv,
388        &MlaConfig,
389        usize,
390        f32,
391    ) -> Result<Vec<f32>>,
392{
393    let qk = cfg.qk_head_dim();
394    let qa = matvec(&w.q_a, x, cfg.q_lora_rank, x.len());
395    let qa = rms_norm(&qa, &w.q_a_ln, eps);
396    let mut q = matvec(&w.q_b, &qa, cfg.heads * qk, cfg.q_lora_rank);
397    let kv_in = cfg.kv_lora_rank + cfg.qk_rope_head_dim;
398    let kv_lat = matvec(&w.kv_a, x, kv_in, x.len());
399    let (c, pe) = kv_lat.split_at(cfg.kv_lora_rank);
400    let c = rms_norm(c, &w.kv_a_ln, eps);
401    let kvb_out = cfg.heads * (cfg.qk_nope_head_dim + cfg.v_head_dim);
402    let kvb = matvec(&w.kv_b, &c, kvb_out, cfg.kv_lora_rank);
403    let (k, v) = pack_mla_kv(&kvb, pe, cfg);
404    let mut k = k;
405    let g = matvec(&w.g_proj, x, cfg.heads * cfg.v_head_dim, x.len());
406    let attn = mla_decode(&mut q, &mut k, &v, &g, kv, cfg, pos, theta)?;
407    Ok(apply_o_proj(
408        &w.o_proj,
409        &attn,
410        x.len(),
411        cfg.heads * cfg.v_head_dim,
412        ablation,
413    ))
414}
415
416fn pack_mla_kv(kvb: &[f32], k_pe: &[f32], cfg: &MlaConfig) -> (Vec<f32>, Vec<f32>) {
417    let nope = cfg.qk_nope_head_dim;
418    let rope = cfg.qk_rope_head_dim;
419    let dv = cfg.v_head_dim;
420    let qk = nope + rope;
421    let mut k = vec![0.0f32; cfg.heads * qk];
422    let mut v = vec![0.0f32; cfg.heads * dv];
423    let stride = nope + dv;
424    for h in 0..cfg.heads {
425        let src = &kvb[h * stride..(h + 1) * stride];
426        let kd = &mut k[h * qk..(h + 1) * qk];
427        kd[..nope].copy_from_slice(&src[..nope]);
428        if rope > 0 {
429            kd[nope..].copy_from_slice(k_pe);
430        }
431        v[h * dv..(h + 1) * dv].copy_from_slice(&src[nope..]);
432    }
433    (k, v)
434}
435
436fn moe_mlp_with<F>(
437    w: &MoeWeights,
438    x: &[f32],
439    ctx: &K3LayerCtx<'_>,
440    force: Option<usize>,
441    experts_fn: &mut F,
442) -> Result<Vec<f32>>
443where
444    F: FnMut(&MoeWeights, &[f32], &[usize], &[f32], &LatentMoeConfig) -> Result<Vec<f32>>,
445{
446    let mut logits = matvec(&w.router, x, ctx.moe.n_routed, ctx.moe.hidden);
447    if let Some(e) = force {
448        logits.fill(0.0);
449        logits[e] = 8.0;
450    }
451    let shared = w
452        .shared
453        .as_ref()
454        .map(|s| dense::run(ctx, s, x, ctx.moe.expert_hidden))
455        .transpose()?;
456    let latent = matvec(&w.down, x, ctx.moe.latent, ctx.moe.hidden);
457    let (ids, weights) = sigmoid_topk(&logits, &w.bias, ctx.moe.top_k);
458    let mixed = experts_fn(w, &latent, &ids, &weights, ctx.moe)?;
459    let mixed = if ctx.moe.use_norm {
460        rms_norm(&mixed, &w.norm, ctx.eps)
461    } else {
462        mixed
463    };
464    let mut out = matvec(&w.up, &mixed, ctx.moe.hidden, ctx.moe.latent);
465    if let Some(s) = shared {
466        for (o, ss) in out.iter_mut().zip(&s) {
467            *o += ss;
468        }
469    }
470    Ok(out)
471}
avarok_core/kimi_k3/cpu_forward.rs

avarok_core/kimi_k3/
cpu_forward.rs
