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

ferrox_models/
engine.rs
