frink_models/
engine.rs

1//! A trait-based abstraction over the two structurally different but
2//! text-in/text-out-shaped forward passes this crate has: `Decoder`
3//! (GQA+RoPE, used by GLM-5.2/DeepSeek V4 Pro and every real GGUF
4//! checkpoint) and Kimi K3's dedicated hybrid KDA/Gated-MLA stack
5//! (`kimi_decoder`). This lets `frink-server` share one generic
6//! generation loop across both engines (see
7//! `frink-server::generate::generate_engine`) for the actual
8//! sampling/stop-sequence logic, rather than hand-duplicating it --
9//! while keeping GGUF-only features (the KV block pool, `PrefixCache`)
10//! as `Decoder`-specific code layered on top, not forced into this
11//! trait. The reason: Kimi's KDA state is
12//! a fixed-size recurrent matrix that collapses history irreversibly,
13//! so it cannot support the same restore/truncate operations
14//! `KvCache` can -- unifying those too would mean either a leaky
15//! abstraction or silently pretending Kimi supports something it
16//! doesn't.
17//!
18//! `forward_token`'s `pos` parameter is meaningful for `Decoder` (used
19//! directly for RoPE) but not for `KimiEngine`: Kimi's real forward
20//! pass (`kimi_forward_token`) derives position purely from its own
21//! per-layer state (KDA's recurrent state, MLA's growing K/V buffers)
22//! -- its real signature has no `pos` parameter at all. `KimiEngine`
23//! ignores the argument; this is a real architectural fact about the
24//! model, not an oversight in this trait's design.
25
26use crate::config::{KdaConfig, MlaConfig};
27use crate::decoder::Decoder;
28use crate::deepseek_v4_decoder::{
29    deepseek_v4_forward_token, DeepseekV4DecodeState, DeepseekV4DecoderConfig,
30    DeepseekV4DecoderWeights,
31};
32use crate::glm52_decoder::{
33    glm52_forward_token, Glm52DecodeState, Glm52DecoderConfig, Glm52DecoderWeights,
34};
35use crate::kimi_decoder::{
36    kimi_forward_token, KimiDecodeState, KimiDecoderConfig, KimiDecoderWeights,
37};
38use crate::kimi_tokenizer::KimiTokenizer;
39use crate::tokenizer::SpecialTokens;
40use frink_core::cache::KvCache;
41use frink_core::weight_matrix::WeightMatrix;
42
43mod entry;
44
45pub use entry::Engine;
46
47/// The shared pool-entry assertion each engine's own tests call, and the
48/// probe for whether promotion happens in this process at all. See
49/// `entry.rs` for why each is one function and not one copy per engine.
50#[cfg(test)]
51pub(crate) use entry::{assert_one_pool_entry_per_step, on_workers_promotes_here};
52
53impl Engine for Decoder {
54    type State = Vec<KvCache>;
55
56    fn new_state(&self) -> Vec<KvCache> {
57        self.config.new_kv_caches()
58    }
59
60    fn vocab_size(&self) -> usize {
61        self.config.vocab_size
62    }
63
64    /// Delegates to the inherent [`Decoder::forward_token`], which is
65    /// itself promoted in `decoder/entry.rs`. The two wrappers nest, and
66    /// nesting is free -- the inner one sees a rayon worker and returns
67    /// the body directly -- so this stays a delegation rather than
68    /// reaching past the entry module for a private body.
69    fn forward_token_on_worker(
70        &self,
71        token_id: usize,
72        pos: usize,
73        state: &mut Self::State,
74    ) -> Vec<f32> {
75        Decoder::forward_token(self, token_id, pos, state)
76    }
77}
78
79/// Bundles Kimi K3's weights with the three real config structs
80/// `kimi_forward_token` needs, so `Engine::forward_token`'s three-
81/// argument shape (`token_id`, `pos`, `state`) can wrap Kimi's real
82/// four-config-argument function.
83pub struct KimiEngine {
84    pub weights: KimiDecoderWeights,
85    pub cfg: KimiDecoderConfig,
86    pub mla_cfg: MlaConfig,
87    pub kda_cfg: KdaConfig,
88}
89
90impl Engine for KimiEngine {
91    type State = KimiDecodeState;
92
93    fn new_state(&self) -> KimiDecodeState {
94        KimiDecodeState::new(&self.weights, &self.kda_cfg)
95    }
96
97    fn vocab_size(&self) -> usize {
98        self.weights.output_head.rows()
99    }
100
101    fn forward_token_on_worker(
102        &self,
103        token_id: usize,
104        _pos: usize,
105        state: &mut Self::State,
106    ) -> Vec<f32> {
107        kimi_forward_token(
108            &self.weights,
109            &self.cfg,
110            &self.mla_cfg,
111            &self.kda_cfg,
112            token_id,
113            state,
114        )
115    }
116}
117
118/// GLM-5.2 dedicated DSA stack behind the same [`Engine`] trait as Kimi.
119/// Synthetic / loader-backed weights only — no claim of a full real
120/// ~744B serve path. Lets `generate_engine` exercise GLM without
121/// forcing DSA into the GQA [`Decoder`].
122pub struct Glm52Engine {
123    pub weights: Glm52DecoderWeights,
124    pub cfg: Glm52DecoderConfig,
125}
126
127impl Engine for Glm52Engine {
128    type State = Glm52DecodeState;
129
130    fn new_state(&self) -> Glm52DecodeState {
131        Glm52DecodeState::new(&self.weights)
132    }
133
134    fn vocab_size(&self) -> usize {
135        self.weights.output_head.rows()
136    }
137
138    fn forward_token_on_worker(
139        &self,
140        token_id: usize,
141        _pos: usize,
142        state: &mut Self::State,
143    ) -> Vec<f32> {
144        glm52_forward_token(&self.weights, &self.cfg, token_id, state)
145    }
146}
147
148/// Multi-layer MLA stack for DeepSeek-2 / Mistral-4-style GGUF serve.
149///
150/// Uses [`crate::mla::mla_forward_token`] with asymmetric K/V caches
151/// (plain `Vec<f32>`, not [`KvCache`]). Layers
152/// `[0, leading_dense)` use dense SwiGLU; later layers use MoE
153/// (`frink_moe`) when the GGUF carries experts — fail-closed at load if
154/// expert tensors are missing.
155pub struct MlaEngine {
156    pub embedding: WeightMatrix,
157    pub layers: Vec<MlaLayerWeights>,
158    pub final_norm: Vec<f32>,
159    pub output_head: WeightMatrix,
160    pub mla_cfg: MlaConfig,
161    pub rms_norm_eps: f32,
162    pub hidden_dim: usize,
163    /// Present when any layer uses [`MlaLayerFfn::Moe`].
164    pub moe: Option<MlaMoeRuntime>,
165    /// YaRN, when the file declares it (`crate::mla_yarn`): the `pe`
166    /// frequency rewrite, its magnitude and the softmax scale.
167    pub yarn: Option<crate::mla_yarn::MlaYarn>,
168}
169
170/// MoE routing knobs shared by every MoE layer (DeepSeek-2 / Mistral-4).
171#[derive(Debug, Clone)]
172pub struct MlaMoeRuntime {
173    pub n_experts_active: usize,
174    pub gating: frink_moe::GatingFunction,
175    pub norm_topk_prob: bool,
176    pub expert_weights_scale: f32,
177}
178
179/// One dense layer's FFN: the gate / up / down triple and the
180/// activation that combines the first two, carried TOGETHER so the
181/// forward pass cannot run a `plm` layer's ungated ReLU-squared
182/// weights through SwiGLU. `act` is the architecture's
183/// (`crate::mla_arch::MlaArch::dense_act`); for an ungated one the
184/// loader aliases `gate` to `ffn_up`, as the generic loader does for
185/// `arcee`, and `frink_moe::run_expert` skips the aliased matmul.
186pub struct MlaDenseFfn {
187    pub weights: frink_moe::ExpertWeights,
188    pub act: frink_moe::GluAct,
189}
190
191pub struct MlaMoeFfn {
192    pub router: WeightMatrix,
193    pub experts: Vec<frink_moe::ExpertWeights>,
194    pub shared_expert: frink_moe::ExpertWeights,
195    /// Optional aux-loss-free bias (`blk.N.exp_probs_b.bias`).
196    pub exp_probs_bias: Option<Vec<f32>>,
197}
198
199pub enum MlaLayerFfn {
200    Dense(MlaDenseFfn),
201    Moe(MlaMoeFfn),
202}
203
204pub struct MlaLayerWeights {
205    pub attn_norm: Vec<f32>,
206    pub attn: crate::mla::MlaAttnWeights,
207    pub ffn_norm: Vec<f32>,
208    pub ffn: MlaLayerFfn,
209}
210
211pub struct MlaDecodeState {
212    pub layers: Vec<(Vec<f32>, Vec<f32>)>,
213}
214
215impl MlaEngine {
216    pub fn new_state(&self) -> MlaDecodeState {
217        MlaDecodeState {
218            layers: (0..self.layers.len())
219                .map(|_| (Vec::new(), Vec::new()))
220                .collect(),
221        }
222    }
223
224    fn moe_ffn_forward(&self, ffn: &MlaMoeFfn, x: &[f32]) -> Vec<f32> {
225        use frink_moe::{
226            combine_expert_outputs, route_top_k, route_top_k_sigmoid_with_bias, run_expert,
227            GatingFunction, GluAct,
228        };
229        let moe = self
230            .moe
231            .as_ref()
232            .expect("MlaLayerFfn::Moe requires MlaEngine.moe");
233        let router_logits = ffn.router.apply(x);
234        let decision = match (moe.gating, ffn.exp_probs_bias.as_deref()) {
235            (GatingFunction::Sigmoid, Some(bias)) => route_top_k_sigmoid_with_bias(
236                &router_logits,
237                bias,
238                moe.n_experts_active,
239                moe.norm_topk_prob,
240                moe.expert_weights_scale,
241            ),
242            _ => {
243                let mut d = route_top_k(
244                    &router_logits,
245                    moe.n_experts_active,
246                    moe.gating,
247                    moe.norm_topk_prob,
248                );
249                if (moe.expert_weights_scale - 1.0).abs() > f32::EPSILON {
250                    for w in d.weights.iter_mut() {
251                        *w *= moe.expert_weights_scale;
252                    }
253                }
254                d
255            }
256        };
257        let routed: Vec<(Vec<f32>, f32)> = decision
258            .expert_ids
259            .iter()
260            .zip(decision.weights.iter())
261            .map(|(&e, &w)| (run_expert(x, &ffn.experts[e], GluAct::Swiglu), w))
262            .collect();
263        let shared = run_expert(x, &ffn.shared_expert, GluAct::Swiglu);
264        combine_expert_outputs(&routed, &[shared], x.len())
265    }
266}
267
268impl Engine for MlaEngine {
269    type State = MlaDecodeState;
270
271    fn new_state(&self) -> MlaDecodeState {
272        MlaEngine::new_state(self)
273    }
274
275    fn vocab_size(&self) -> usize {
276        self.output_head.rows()
277    }
278
279    fn forward_token_on_worker(
280        &self,
281        token_id: usize,
282        _pos: usize,
283        state: &mut Self::State,
284    ) -> Vec<f32> {
285        use frink_core::matmul::rms_norm;
286        let mut hidden = self.embedding.dequant_row(token_id);
287        for (layer, (k_cache, v_cache)) in self.layers.iter().zip(state.layers.iter_mut()) {
288            let normed = rms_norm(&hidden, &layer.attn_norm, self.rms_norm_eps);
289            let attn_out = crate::mla::mla_forward_token(
290                &layer.attn,
291                &self.mla_cfg,
292                self.yarn.as_ref(),
293                &normed,
294                self.rms_norm_eps,
295                k_cache,
296                v_cache,
297            );
298            for (h, a) in hidden.iter_mut().zip(attn_out.iter()) {
299                *h += a;
300            }
301            let ffn_in = rms_norm(&hidden, &layer.ffn_norm, self.rms_norm_eps);
302            let down = match &layer.ffn {
303                MlaLayerFfn::Dense(d) => frink_moe::run_expert(&ffn_in, &d.weights, d.act),
304                MlaLayerFfn::Moe(m) => self.moe_ffn_forward(m, &ffn_in),
305            };
306            for (h, d) in hidden.iter_mut().zip(down.iter()) {
307                *h += d;
308            }
309        }
310        let final_normed = rms_norm(&hidden, &self.final_norm, self.rms_norm_eps);
311        self.output_head.apply(&final_normed)
312    }
313}
314
315/// DeepSeek V4 synthetic stack behind [`Engine`]. Preset `deepseek_v4_pro`
316/// remains a sketch until a real GGUF loader + incremental DSV4 KV land.
317pub struct DeepseekV4Engine {
318    pub weights: DeepseekV4DecoderWeights,
319    pub cfg: DeepseekV4DecoderConfig,
320}
321
322impl Engine for DeepseekV4Engine {
323    type State = DeepseekV4DecodeState;
324
325    fn new_state(&self) -> DeepseekV4DecodeState {
326        DeepseekV4DecodeState::new(self.weights.embedding.shape[1])
327    }
328
329    fn vocab_size(&self) -> usize {
330        self.weights.output_head.rows()
331    }
332
333    fn forward_token_on_worker(
334        &self,
335        token_id: usize,
336        _pos: usize,
337        state: &mut Self::State,
338    ) -> Vec<f32> {
339        deepseek_v4_forward_token(&self.weights, &self.cfg, token_id, state)
340    }
341}
342
343/// A minimal text<->token-id interface shared by every real tokenizer
344/// this crate has, regardless of each one's native id width
345/// (`GgufBpeTokenizer`/`GgufSpmTokenizer`/`GgufUnigramTokenizer` use
346/// `u32`, `KimiTokenizer` also uses `u32`) -- lets a generic generation
347/// loop encode/decode without caring which concrete tokenizer it was
348/// given.
349pub trait TextTokenizer {
350    /// `specials` is llama.cpp's `parse_special`: whether a literal
351    /// marker in `text` is the token it names or the characters it is
352    /// written with. See [`SpecialTokens`] for which callers want which.
353    fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize>;
354    fn decode(&self, ids: &[usize]) -> String;
355
356    /// The raw bytes, before any UTF-8 decision is made about them.
357    ///
358    /// A caller decoding ONE token at a time needs these: a character
359    /// split across two tokens is two invalid fragments, and `decode`
360    /// resolves each to U+FFFD separately, losing the bytes (#124).
361    ///
362    /// The default is correct for any tokenizer whose tokens are whole
363    /// text, and no worse than `decode` for one whose tokens are not --
364    /// but such a tokenizer should override this.
365    fn decode_bytes(&self, ids: &[usize]) -> Vec<u8> {
366        self.decode(ids).into_bytes()
367    }
368}
369
370impl TextTokenizer for KimiTokenizer {
371    fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
372        KimiTokenizer::encode(self, text, specials)
373            .into_iter()
374            .map(|id| id as usize)
375            .collect()
376    }
377
378    fn decode(&self, ids: &[usize]) -> String {
379        let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
380        KimiTokenizer::decode(self, &ids32)
381    }
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387    use crate::config::test_dense_fixture;
388
389    /// Locks in the refactor: calling `Decoder` through the generic
390    /// `Engine` trait must be bit-identical to calling its own
391    /// `forward_token`/`forward_batch` directly -- the whole point of
392    /// the trait is that `frink-server`'s generic generation loop can
393    /// use it as a drop-in replacement with zero numeric difference.
394    #[test]
395    fn decoder_via_engine_trait_matches_direct_forward_token_calls() {
396        let decoder = Decoder::new_random_small(test_dense_fixture(), 2, 64);
397        let tokens = [3usize, 7, 1, 9];
398
399        let mut direct_caches: Vec<KvCache> = decoder.config.new_kv_caches();
400        let mut direct_logits = Vec::new();
401        for (pos, &tok) in tokens.iter().enumerate() {
402            direct_logits = decoder.forward_token(tok, pos, &mut direct_caches);
403        }
404
405        let mut engine_state = Engine::new_state(&decoder);
406        let mut engine_logits = Vec::new();
407        for (pos, &tok) in tokens.iter().enumerate() {
408            engine_logits = Engine::forward_token(&decoder, tok, pos, &mut engine_state);
409        }
410
411        assert_eq!(engine_logits, direct_logits);
412        assert_eq!(Engine::vocab_size(&decoder), decoder.config.vocab_size);
413    }
414
415    /// Same equivalence, but against `forward_batch`'s independent
416    /// computation (the ground truth every other test in this
417    /// workspace already uses) rather than a second manual loop.
418    #[test]
419    fn decoder_via_engine_trait_matches_forward_batch_ground_truth() {
420        let decoder = Decoder::new_random_small(test_dense_fixture(), 2, 64);
421        let tokens = vec![2usize, 5, 8];
422
423        let mut batch_caches: Vec<KvCache> = decoder.config.new_kv_caches();
424        let batch_logits = decoder.forward_batch(&tokens, 0, &mut batch_caches);
425        let ground_truth = batch_logits.last().unwrap().clone();
426
427        let mut engine_state = Engine::new_state(&decoder);
428        let mut engine_logits = Vec::new();
429        for (pos, &tok) in tokens.iter().enumerate() {
430            engine_logits = Engine::forward_token(&decoder, tok, pos, &mut engine_state);
431        }
432
433        // Tight tolerance, not bit equality: batched prefill runs the
434        // blocked three-pass softmax while per-token decode keeps the
435        // online accumulator, and the two round differently in the last
436        // ulp (they were bit-identical only while both were online).
437        assert_eq!(engine_logits.len(), ground_truth.len());
438        for (i, (e, g)) in engine_logits.iter().zip(ground_truth.iter()).enumerate() {
439            assert!(
440                (e - g).abs() < 1e-5,
441                "logit {i}: engine {e} vs forward_batch {g}"
442            );
443        }
444    }
445}
frink_models/engine.rs

frink_models/
engine.rs
