ferrox_models/
engine_factory.rs

1//! Load-time engine selection (llama.cpp `llama_model_*` factory analogue).
2//!
3//! GGUF text-generation architectures that the generic [`crate::decoder::Decoder`]
4//! can run are routed there. Dedicated stacks (MLA / DSA / recurrent /
5//! T5 encoder-decoder) are selected here and must not accumulate
6//! `if arch` branches inside matmul / attention kernels.
7
8use crate::capability::{resolve_profile, ArchPath, DecoderFamily};
9use crate::decoder::Decoder;
10use crate::engine::{Engine, Glm52Engine, KimiEngine, MlaEngine};
11use crate::gemma4_engine::Gemma4Engine;
12use crate::gemma4_gguf_loader;
13use crate::glm52_gguf_loader;
14use crate::loader::LoadError;
15use crate::mla_gguf_loader;
16use thiserror::Error;
17
18/// Why a GGUF cannot be served by the currently compiled engine set.
19#[derive(Debug, Error)]
20pub enum EngineSelectError {
21    #[error("architecture {0:?} is outside the text-generation scope ({1})")]
22    OutOfScope(String, &'static str),
23    #[error("architecture {0:?} requires a dedicated engine that is not wired for serve yet: {1}")]
24    DedicatedUnavailable(String, &'static str),
25    #[error("unknown or unsupported architecture {0:?}")]
26    Unknown(String),
27}
28
29/// Result of load-time engine selection for a GGUF architecture string.
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum SelectedEngineKind {
32    /// Standard / Phi / Gemma / Qwen3 GQA path via [`Decoder`].
33    GenericDecoder,
34    /// Kimi / GLM / DeepSeek dedicated stacks (separate loaders).
35    DedicatedStack,
36    /// MLA memory engines (DeepSeek2 / Mistral4) — fail-closed for generic.
37    Mla,
38    /// Gemma-4 dedicated text engine (per-layer emb + shared KV + SWA split).
39    Gemma4,
40    /// Recurrent / hybrid SSM families — fail-closed until wired.
41    RecurrentHybrid,
42    /// T5 encoder-decoder — fail-closed until wired.
43    EncoderDecoder,
44}
45
46/// Resolve which engine kind a GGUF `general.architecture` should use.
47pub fn select_engine_kind(arch: &str) -> Result<SelectedEngineKind, EngineSelectError> {
48    let profile =
49        resolve_profile(arch).ok_or_else(|| EngineSelectError::Unknown(arch.to_string()))?;
50    match profile.path {
51        ArchPath::GenericGqa { .. } | ArchPath::TestFixture { .. } => {
52            Ok(SelectedEngineKind::GenericDecoder)
53        }
54        ArchPath::Deferred { reason } => {
55            Err(EngineSelectError::OutOfScope(arch.to_string(), reason))
56        }
57        ArchPath::DedicatedOnly { reason } => match profile.family {
58            DecoderFamily::Dedicated => Ok(SelectedEngineKind::DedicatedStack),
59            DecoderFamily::Mla => Ok(SelectedEngineKind::Mla),
60            DecoderFamily::GemmaFamily if is_gemma4_arch(arch) => Ok(SelectedEngineKind::Gemma4),
61            DecoderFamily::Hybrid | DecoderFamily::Recurrent => {
62                Ok(SelectedEngineKind::RecurrentHybrid)
63            }
64            DecoderFamily::EncoderDecoder => Ok(SelectedEngineKind::EncoderDecoder),
65            _ => Err(EngineSelectError::DedicatedUnavailable(
66                arch.to_string(),
67                reason,
68            )),
69        },
70    }
71}
72
73/// Fail-closed check used by `ferrox-server` before constructing a
74/// [`Decoder`] for architectures that need another engine.
75pub fn ensure_generic_decoder(arch: &str) -> Result<(), EngineSelectError> {
76    match select_engine_kind(arch)? {
77        SelectedEngineKind::GenericDecoder => Ok(()),
78        SelectedEngineKind::DedicatedStack => Err(EngineSelectError::DedicatedUnavailable(
79            arch.to_string(),
80            if is_glm52_arch(arch) {
81                "use load_glm52_engine_from_path / ServedEngine::Glm52 — not generic Decoder"
82            } else {
83                "use the dedicated Kimi/GLM/DeepSeek loader, not generic Decoder"
84            },
85        )),
86        SelectedEngineKind::Mla => Err(EngineSelectError::DedicatedUnavailable(
87            arch.to_string(),
88            "use load_mla_engine_from_path / ServedEngine::Mla — not generic Decoder",
89        )),
90        SelectedEngineKind::Gemma4 => Err(EngineSelectError::DedicatedUnavailable(
91            arch.to_string(),
92            "use load_gemma4_engine_from_path / ServedEngine::Gemma4 — not generic Decoder",
93        )),
94        SelectedEngineKind::RecurrentHybrid => {
95            let _ = crate::recurrent_engine::RecurrentEngine::reject(arch);
96            let _ = crate::hybrid_engine::HybridEngine::reject(arch);
97            Err(EngineSelectError::DedicatedUnavailable(
98                arch.to_string(),
99                "recurrent/hybrid SSM engine stub present — not yet on the serve path",
100            ))
101        }
102        SelectedEngineKind::EncoderDecoder => {
103            let _ = crate::t5_engine::T5Engine::reject(arch);
104            Err(EngineSelectError::DedicatedUnavailable(
105                arch.to_string(),
106                "T5 encoder-decoder engine stub present — not yet on the serve path",
107            ))
108        }
109    }
110}
111
112/// Type-erased serve handle: today ordinary GGUFs use [`Decoder`];
113/// Kimi/GLM/MLA use dedicated engines once loaders succeed.
114pub enum ServedEngine {
115    Decoder(Box<Decoder>),
116    Kimi(KimiEngine),
117    Glm52(Glm52Engine),
118    Mla(MlaEngine),
119    // Boxed: the Gemma-4 engine is ~200 bytes larger than any other
120    // variant, so inlining it would pad every `ServedEngine` to its size.
121    Gemma4(Box<Gemma4Engine>),
122}
123
124impl ServedEngine {
125    pub fn vocab_size(&self) -> usize {
126        match self {
127            Self::Decoder(d) => Engine::vocab_size(d.as_ref()),
128            Self::Kimi(k) => Engine::vocab_size(k),
129            Self::Glm52(g) => Engine::vocab_size(g),
130            Self::Mla(m) => Engine::vocab_size(m),
131            Self::Gemma4(g) => Engine::vocab_size(g.as_ref()),
132        }
133    }
134}
135
136/// Open a DeepSeek-2 / Mistral-4 GGUF and build [`ServedEngine::Mla`].
137pub fn load_mla_engine_from_path(path: &std::path::Path) -> Result<ServedEngine, LoadError> {
138    let file = ferrox_gguf::ShardedGguf::open(path)?;
139    let arch = file
140        .metadata_str("general.architecture")
141        .unwrap_or("unknown");
142    match select_engine_kind(arch) {
143        Ok(SelectedEngineKind::Mla) => {}
144        Ok(other) => {
145            return Err(LoadError::DedicatedArchitectureRequired(
146                arch.to_string(),
147                match other {
148                    SelectedEngineKind::GenericDecoder => "generic decoder arch, not MLA",
149                    SelectedEngineKind::DedicatedStack => "dedicated non-MLA stack",
150                    SelectedEngineKind::RecurrentHybrid => "hybrid/recurrent, not MLA",
151                    SelectedEngineKind::EncoderDecoder => "encoder-decoder, not MLA",
152                    SelectedEngineKind::Gemma4 => "gemma4 dedicated, not MLA",
153                    SelectedEngineKind::Mla => unreachable!(),
154                },
155            ));
156        }
157        Err(EngineSelectError::Unknown(a)) => {
158            return Err(LoadError::UnsupportedArchitecture(a));
159        }
160        Err(EngineSelectError::OutOfScope(a, r)) => {
161            return Err(LoadError::UnsupportedFeature(a, r.to_string()));
162        }
163        Err(EngineSelectError::DedicatedUnavailable(a, r)) => {
164            return Err(LoadError::DedicatedArchitectureRequired(a, r));
165        }
166    }
167    Ok(ServedEngine::Mla(mla_gguf_loader::load_mla_engine(&file)?))
168}
169
170fn is_glm52_arch(arch: &str) -> bool {
171    matches!(arch, "glm-dsa" | "glm4" | "glm4moe")
172}
173
174fn is_gemma4_arch(arch: &str) -> bool {
175    crate::gemma4_engine::GEMMA4_ARCHES.contains(&arch)
176}
177
178/// Open a GLM-5.2 / GLM4-family GGUF and build [`ServedEngine::Glm52`].
179pub fn load_glm52_engine_from_path(path: &std::path::Path) -> Result<ServedEngine, LoadError> {
180    let file = ferrox_gguf::ShardedGguf::open(path)?;
181    let arch = file
182        .metadata_str("general.architecture")
183        .unwrap_or("unknown");
184    if !is_glm52_arch(arch) {
185        return Err(LoadError::DedicatedArchitectureRequired(
186            arch.to_string(),
187            "not a GLM-5.2 / GLM4-family architecture (expected glm-dsa/glm4/glm4moe)",
188        ));
189    }
190    match select_engine_kind(arch) {
191        Ok(SelectedEngineKind::DedicatedStack) => {}
192        Ok(other) => {
193            return Err(LoadError::DedicatedArchitectureRequired(
194                arch.to_string(),
195                match other {
196                    SelectedEngineKind::GenericDecoder => "generic decoder arch, not GLM DSA",
197                    SelectedEngineKind::Mla => "MLA arch, not GLM DSA",
198                    SelectedEngineKind::RecurrentHybrid => "hybrid/recurrent, not GLM DSA",
199                    SelectedEngineKind::EncoderDecoder => "encoder-decoder, not GLM DSA",
200                    SelectedEngineKind::Gemma4 => "gemma4 dedicated, not GLM DSA",
201                    SelectedEngineKind::DedicatedStack => unreachable!(),
202                },
203            ));
204        }
205        Err(EngineSelectError::Unknown(a)) => {
206            return Err(LoadError::UnsupportedArchitecture(a));
207        }
208        Err(EngineSelectError::OutOfScope(a, r)) => {
209            return Err(LoadError::UnsupportedFeature(a, r.to_string()));
210        }
211        Err(EngineSelectError::DedicatedUnavailable(a, r)) => {
212            return Err(LoadError::DedicatedArchitectureRequired(a, r));
213        }
214    }
215    Ok(ServedEngine::Glm52(glm52_gguf_loader::load_glm52_engine(
216        &file,
217    )?))
218}
219
220/// Open a Gemma-4 GGUF and build [`ServedEngine::Gemma4`].
221pub fn load_gemma4_engine_from_path(path: &std::path::Path) -> Result<ServedEngine, LoadError> {
222    let file = ferrox_gguf::ShardedGguf::open(path)?;
223    let arch = file
224        .metadata_str("general.architecture")
225        .unwrap_or("unknown");
226    if !is_gemma4_arch(arch) {
227        return Err(LoadError::DedicatedArchitectureRequired(
228            arch.to_string(),
229            "not a gemma4 / gemma4-assistant architecture",
230        ));
231    }
232    match select_engine_kind(arch) {
233        Ok(SelectedEngineKind::Gemma4) => {}
234        Ok(other) => {
235            return Err(LoadError::DedicatedArchitectureRequired(
236                arch.to_string(),
237                match other {
238                    SelectedEngineKind::GenericDecoder => "generic decoder arch, not Gemma4",
239                    SelectedEngineKind::Mla => "MLA arch, not Gemma4",
240                    SelectedEngineKind::DedicatedStack => "dedicated non-Gemma4 stack",
241                    SelectedEngineKind::RecurrentHybrid => "hybrid/recurrent, not Gemma4",
242                    SelectedEngineKind::EncoderDecoder => "encoder-decoder, not Gemma4",
243                    SelectedEngineKind::Gemma4 => unreachable!(),
244                },
245            ));
246        }
247        Err(EngineSelectError::Unknown(a)) => {
248            return Err(LoadError::UnsupportedArchitecture(a));
249        }
250        Err(EngineSelectError::OutOfScope(a, r)) => {
251            return Err(LoadError::UnsupportedFeature(a, r.to_string()));
252        }
253        Err(EngineSelectError::DedicatedUnavailable(a, r)) => {
254            return Err(LoadError::DedicatedArchitectureRequired(a, r));
255        }
256    }
257    Ok(ServedEngine::Gemma4(Box::new(
258        gemma4_gguf_loader::load_gemma4_engine(&file)?,
259    )))
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265
266    #[test]
267    fn llama_and_qwen3_select_generic_decoder() {
268        assert_eq!(
269            select_engine_kind("llama").unwrap(),
270            SelectedEngineKind::GenericDecoder
271        );
272        assert_eq!(
273            select_engine_kind("qwen3").unwrap(),
274            SelectedEngineKind::GenericDecoder
275        );
276        assert_eq!(
277            select_engine_kind("gemma3").unwrap(),
278            SelectedEngineKind::GenericDecoder
279        );
280        assert_eq!(
281            select_engine_kind("phi3").unwrap(),
282            SelectedEngineKind::GenericDecoder
283        );
284        assert_eq!(
285            select_engine_kind("mixtral").unwrap(),
286            SelectedEngineKind::GenericDecoder
287        );
288    }
289
290    #[test]
291    fn mamba_is_recurrent_fail_closed_for_generic() {
292        assert!(matches!(
293            select_engine_kind("mamba2").unwrap(),
294            SelectedEngineKind::RecurrentHybrid
295        ));
296        assert!(ensure_generic_decoder("mamba2").is_err());
297    }
298
299    #[test]
300    fn deepseek2_is_mla_not_generic() {
301        assert_eq!(
302            select_engine_kind("deepseek2").unwrap(),
303            SelectedEngineKind::Mla
304        );
305        assert!(ensure_generic_decoder("deepseek2").is_err());
306    }
307
308    #[test]
309    fn glm4_is_dedicated_stack_not_generic() {
310        assert_eq!(
311            select_engine_kind("glm4").unwrap(),
312            SelectedEngineKind::DedicatedStack
313        );
314        assert!(ensure_generic_decoder("glm4").is_err());
315    }
316
317    #[test]
318    fn gemma4_selects_dedicated_engine() {
319        assert_eq!(
320            select_engine_kind("gemma4").unwrap(),
321            SelectedEngineKind::Gemma4
322        );
323        assert_eq!(
324            select_engine_kind("gemma4-assistant").unwrap(),
325            SelectedEngineKind::Gemma4
326        );
327        assert!(ensure_generic_decoder("gemma4").is_err());
328    }
329}
ferrox_models/engine_factory.rs

ferrox_models/
engine_factory.rs
