frink_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 `frink-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 = frink_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
170/// Neither `glm4moe` nor `glm4` is here: GLM-4.5 / 4.5-Air / 4.6 and
171/// GLM-4-0414 are plain GQA and run on the generic decoder
172/// (`tests/glm4moe_graphs.rs`, `tests/glm4_graphs.rs`); both used to be
173/// sent to this loader for MLA keys their graphs never read.
174fn is_glm52_arch(arch: &str) -> bool {
175    arch == "glm-dsa"
176}
177
178fn is_gemma4_arch(arch: &str) -> bool {
179    crate::gemma4_engine::GEMMA4_ARCHES.contains(&arch)
180}
181
182/// Open a GLM-5.2 / GLM4-family GGUF and build [`ServedEngine::Glm52`].
183pub fn load_glm52_engine_from_path(path: &std::path::Path) -> Result<ServedEngine, LoadError> {
184    let file = frink_gguf::ShardedGguf::open(path)?;
185    let arch = file
186        .metadata_str("general.architecture")
187        .unwrap_or("unknown");
188    if !is_glm52_arch(arch) {
189        return Err(LoadError::DedicatedArchitectureRequired(
190            arch.to_string(),
191            "not a GLM-5.2 architecture (expected glm-dsa)",
192        ));
193    }
194    match select_engine_kind(arch) {
195        Ok(SelectedEngineKind::DedicatedStack) => {}
196        Ok(other) => {
197            return Err(LoadError::DedicatedArchitectureRequired(
198                arch.to_string(),
199                match other {
200                    SelectedEngineKind::GenericDecoder => "generic decoder arch, not GLM DSA",
201                    SelectedEngineKind::Mla => "MLA arch, not GLM DSA",
202                    SelectedEngineKind::RecurrentHybrid => "hybrid/recurrent, not GLM DSA",
203                    SelectedEngineKind::EncoderDecoder => "encoder-decoder, not GLM DSA",
204                    SelectedEngineKind::Gemma4 => "gemma4 dedicated, not GLM DSA",
205                    SelectedEngineKind::DedicatedStack => unreachable!(),
206                },
207            ));
208        }
209        Err(EngineSelectError::Unknown(a)) => {
210            return Err(LoadError::UnsupportedArchitecture(a));
211        }
212        Err(EngineSelectError::OutOfScope(a, r)) => {
213            return Err(LoadError::UnsupportedFeature(a, r.to_string()));
214        }
215        Err(EngineSelectError::DedicatedUnavailable(a, r)) => {
216            return Err(LoadError::DedicatedArchitectureRequired(a, r));
217        }
218    }
219    Ok(ServedEngine::Glm52(glm52_gguf_loader::load_glm52_engine(
220        &file,
221    )?))
222}
223
224/// Open a Gemma-4 GGUF and build [`ServedEngine::Gemma4`].
225pub fn load_gemma4_engine_from_path(path: &std::path::Path) -> Result<ServedEngine, LoadError> {
226    let file = frink_gguf::ShardedGguf::open(path)?;
227    let arch = file
228        .metadata_str("general.architecture")
229        .unwrap_or("unknown");
230    if !is_gemma4_arch(arch) {
231        return Err(LoadError::DedicatedArchitectureRequired(
232            arch.to_string(),
233            "not a gemma4 / gemma4-assistant architecture",
234        ));
235    }
236    match select_engine_kind(arch) {
237        Ok(SelectedEngineKind::Gemma4) => {}
238        Ok(other) => {
239            return Err(LoadError::DedicatedArchitectureRequired(
240                arch.to_string(),
241                match other {
242                    SelectedEngineKind::GenericDecoder => "generic decoder arch, not Gemma4",
243                    SelectedEngineKind::Mla => "MLA arch, not Gemma4",
244                    SelectedEngineKind::DedicatedStack => "dedicated non-Gemma4 stack",
245                    SelectedEngineKind::RecurrentHybrid => "hybrid/recurrent, not Gemma4",
246                    SelectedEngineKind::EncoderDecoder => "encoder-decoder, not Gemma4",
247                    SelectedEngineKind::Gemma4 => unreachable!(),
248                },
249            ));
250        }
251        Err(EngineSelectError::Unknown(a)) => {
252            return Err(LoadError::UnsupportedArchitecture(a));
253        }
254        Err(EngineSelectError::OutOfScope(a, r)) => {
255            return Err(LoadError::UnsupportedFeature(a, r.to_string()));
256        }
257        Err(EngineSelectError::DedicatedUnavailable(a, r)) => {
258            return Err(LoadError::DedicatedArchitectureRequired(a, r));
259        }
260    }
261    Ok(ServedEngine::Gemma4(Box::new(
262        gemma4_gguf_loader::load_gemma4_engine(&file)?,
263    )))
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    #[test]
271    fn llama_and_qwen3_select_generic_decoder() {
272        assert_eq!(
273            select_engine_kind("llama").unwrap(),
274            SelectedEngineKind::GenericDecoder
275        );
276        assert_eq!(
277            select_engine_kind("qwen3").unwrap(),
278            SelectedEngineKind::GenericDecoder
279        );
280        assert_eq!(
281            select_engine_kind("gemma3").unwrap(),
282            SelectedEngineKind::GenericDecoder
283        );
284        assert_eq!(
285            select_engine_kind("phi3").unwrap(),
286            SelectedEngineKind::GenericDecoder
287        );
288        // `mixtral` was HERE and is not a generic-decoder row any
289        // more: no converter writes that string, libllama refuses it,
290        // and every real Mixtral checkpoint declares `llama` (which is
291        // the first case in this test). See `capability::NO_UPSTREAM_ARCH`.
292        assert_eq!(
293            select_engine_kind("olmoe").unwrap(),
294            SelectedEngineKind::GenericDecoder
295        );
296    }
297
298    /// `mamba2` reaches the generic decoder since 2026-09-14
299    /// (`layer_shapes::PURE_RECURRENT`); RWKV is what the recurrent
300    /// stub still refuses.
301    #[test]
302    fn mamba_is_generic_and_rwkv_is_recurrent_fail_closed() {
303        assert!(matches!(
304            select_engine_kind("mamba2").unwrap(),
305            SelectedEngineKind::GenericDecoder
306        ));
307        assert!(ensure_generic_decoder("mamba2").is_ok());
308        assert!(matches!(
309            select_engine_kind("rwkv7").unwrap(),
310            SelectedEngineKind::RecurrentHybrid
311        ));
312        assert!(ensure_generic_decoder("rwkv7").is_err());
313    }
314
315    #[test]
316    fn deepseek2_is_mla_not_generic() {
317        assert_eq!(
318            select_engine_kind("deepseek2").unwrap(),
319            SelectedEngineKind::Mla
320        );
321        assert!(ensure_generic_decoder("deepseek2").is_err());
322    }
323
324    #[test]
325    fn glm4_is_generic_and_glm_dsa_is_the_dedicated_stack() {
326        assert_eq!(
327            select_engine_kind("glm4").unwrap(),
328            SelectedEngineKind::GenericDecoder
329        );
330        assert!(ensure_generic_decoder("glm4").is_ok());
331        assert_eq!(
332            select_engine_kind("glm-dsa").unwrap(),
333            SelectedEngineKind::DedicatedStack
334        );
335    }
336
337    #[test]
338    fn gemma4_selects_dedicated_engine() {
339        assert_eq!(
340            select_engine_kind("gemma4").unwrap(),
341            SelectedEngineKind::Gemma4
342        );
343        assert_eq!(
344            select_engine_kind("gemma4-assistant").unwrap(),
345            SelectedEngineKind::Gemma4
346        );
347        assert!(ensure_generic_decoder("gemma4").is_err());
348    }
349}
frink_models/engine_factory.rs

frink_models/
engine_factory.rs
