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

ferrox_models/
engine.rs
