metrale_model_engine/
factory.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Model factory: [`loader_for_config`] picks the
4//! [`ModelWeightLoader`] for a config's `model_type`, and [`build_model`]
5//! builds the model from the config and the weights.
6//!
7//! Owner: metrale-model-engine.
8//! Invariants: none beyond the types.
9
10use anyhow::{Result, bail};
11use metrale_config::ModelConfig;
12use metrale_model_weights::weights::WeightStore;
13
14use crate::kimi_k3_loader::KimiK3WeightLoader;
15use metrale_model_arch::mistral_loader::MistralWeightLoader;
16use metrale_model_arch::weight_loader::LongcatWeightLoader;
17use metrale_model_arch::weight_loader::Qwen4ExpWeightLoader;
18use metrale_model_arch::weight_loader::{
19    DeepSeekV4WeightLoader, DflashConfig, Gemma4WeightLoader, Glm5NextWeightLoader,
20    LagunaWeightLoader, MinimaxM2WeightLoader, ModelWeightLoader, NemotronHWeightLoader,
21    NllbWeightLoader, Qwen3VLWeightLoader, Qwen3WeightLoader, Qwen35DenseWeightLoader,
22    Qwen35WeightLoader, Step3p7WeightLoader,
23};
24
25/// 2026-09-25: DFlash drafter inputs for [`build_model`]: the drafter's own
26/// [`WeightStore`], its parsed `config.json`, and the CLI overrides for γ and
27/// the sliding-window size.
28///
29/// The server (`serve_phases/weights.rs`) loads the drafter checkpoint and
30/// parses its config with
31/// `metrale_model_arch::weight_loader::dflash_loader::parse_dflash_config`.
32/// [`build_model`] passes them to
33/// `metrale_model_arch::dflash_head::BlockDiffusionDraftHead::from_weights`.
34pub struct DflashBuildArgs<'a> {
35    pub drafter_store: &'a WeightStore,
36    pub drafter_config: DflashConfig,
37    pub gamma: Option<usize>,
38    pub window_size: Option<usize>,
39}
40
41/// 2026-09-25: LoRA inputs for [`build_model`] (`--lora-adapter NAME=PATH`):
42/// the adapters and the pool-shape CLI values.
43///
44/// [`build_model`] allocates the LoRA pool before the buffer arena and the
45/// free-memory sample the KV budget is computed from, so the pool is counted
46/// against the budget.
47pub struct LoraBuildArgs<'a> {
48    /// 2026-09-25: The adapters to pack, one per `--lora-adapter NAME=PATH`;
49    /// slot k is `adapters[k]`.
50    pub adapters: Vec<metrale_model_layers::lora::LoraAdapterInput<'a>>,
51    pub max_lora_rank: usize,
52    pub max_loras: usize,
53}
54
55/// 2026-09-25: The weight loader for `config.model_type`, matched after
56/// lower-casing it and turning `-` and `.` into `_`. An unknown type is an
57/// error.
58pub fn loader_for_config(config: &ModelConfig) -> Result<Box<dyn ModelWeightLoader>> {
59    let normalized = config.model_type.to_lowercase().replace(['-', '.'], "_");
60    match normalized.as_str() {
61        "kimi_k3" | "kimi_linear" => Ok(Box::new(KimiK3WeightLoader)),
62        "qwen3_next" => Ok(Box::new(Qwen3WeightLoader)),
63        "qwen3_vl_moe" => Ok(Box::new(Qwen3VLWeightLoader)),
64        "qwen3_5_moe" | "qwen3_5" | "qwen35_moe" | "qwen35" => {
65            // 2026-09-25: `is_qwen35_dense()` (model_type `qwen3_5` with no
66            // experts) is tested first: a dense config can carry a
67            // `vision_config`, which `is_qwen3_vl()` also accepts.
68            if config.is_qwen35_dense() {
69                Ok(Box::new(Qwen35DenseWeightLoader))
70            } else if config.is_qwen3_vl() {
71                Ok(Box::new(Qwen3VLWeightLoader))
72            } else {
73                Ok(Box::new(Qwen35WeightLoader))
74            }
75        }
76        "qwen3_6_moe" | "holo3_1_moe" => Ok(Box::new(Qwen35WeightLoader)),
77        "nemotron_h" | "nemotron_h_puzzle" => Ok(Box::new(NemotronHWeightLoader)),
78        "m2m_100" | "nllb" => Ok(Box::new(NllbWeightLoader)),
79        "gemma4" | "gemma_4" => Ok(Box::new(Gemma4WeightLoader)),
80        "mistral" => Ok(Box::new(MistralWeightLoader)),
81        "longcat_flash_ngram" | "longcat_flash" => Ok(Box::new(LongcatWeightLoader)),
82        // 2026-09-25: The config parser also maps `qwen3_8_flash_next` to
83        // `qwen4_exp` (metrale-config `dispatch.rs`, `parsers/qwen4_exp.rs`).
84        "qwen4_exp" => Ok(Box::new(Qwen4ExpWeightLoader)),
85        "minimax_m2" => Ok(Box::new(MinimaxM2WeightLoader)),
86        "step3p7" => Ok(Box::new(Step3p7WeightLoader)),
87        "laguna" => Ok(Box::new(LagunaWeightLoader)),
88        "deepseek_v4" => Ok(Box::new(DeepSeekV4WeightLoader)),
89        // 2026-09-25: The config parser sets `glm5_next` for both
90        // (metrale-config `parsers/glm5_next/parse.rs`).
91        "glm5_next" | "glm5_next_text" => Ok(Box::new(Glm5NextWeightLoader)),
92        "deepseek_v41" => Ok(Box::new(
93            metrale_model_arch::weight_loader::deepseek_v41::DeepSeekV41WeightLoader,
94        )),
95        _ => bail!(
96            "Unsupported model type: '{}' (normalized: '{}'). \
97             Supported: qwen3_next, glm5_next, qwen3_5_moe, qwen3_5, qwen3_6_moe, holo3_1_moe, qwen3_vl_moe, nemotron_h, nemotron_h_puzzle, gemma4, mistral, minimax_m2, step3p7, laguna, deepseek_v4, qwen4_exp, m2m_100, deepseek_v41",
98            config.model_type,
99            normalized,
100        ),
101    }
102}
103
104mod build;
105mod lm_head_setup;
106mod m2_setup;
107
108pub use build::build_model;
109pub use build::slots::{AUTO_MAX_SLOTS, BuiltModel, SlotPlan, SlotRequest, balance_slots};
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use metrale_cache::kv_cache::KvCacheDtype;
115    use metrale_model_layers::layers::mtp_head::MtpQuantization;
116    use metrale_telemetry::prefix_cache::PrefixCache;
117
118    #[test]
119    fn test_unsupported_model_type() {
120        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
121        config.model_type = "llama".to_string();
122
123        let gpu = metrale_gpu_runtime::gpu::mock::MockGpuBackend::new();
124        let store = WeightStore::empty();
125
126        let prefix_cache: Box<dyn PrefixCache> =
127            Box::new(metrale_telemetry::prefix_cache::NoPrefixCaching);
128        let result = build_model(
129            config,
130            store,
131            Box::new(gpu),
132            1,
133            16,
134            4096,
135            SlotPlan {
136                request: SlotRequest::Count(8),
137                reserve_for: &|_| Ok(1 << 30),
138            },
139            MtpQuantization::Nvfp4,
140            false,
141            prefix_cache,
142            0,
143            None,
144            false,
145            1,
146            KvCacheDtype::Fp8,
147            0.90,
148            0,
149            vec![],
150            0,
151            None,
152            None,
153            None,
154            None,
155            None,
156        );
157        match result {
158            Err(e) => assert!(e.to_string().contains("Unsupported model type: 'llama'")),
159            Ok(_) => panic!("Expected error for unsupported model type"),
160        }
161    }
162
163    #[test]
164    fn all_declared_model_type_spellings_are_accepted() {
165        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
166        for model_type in [
167            "qwen3_next",
168            "qwen3_vl_moe",
169            "qwen3_5_moe",
170            "qwen3_5",
171            "qwen35_moe",
172            "qwen35",
173            "qwen3_6_moe",
174            "holo3_1_moe",
175            "nemotron_h",
176            "nemotron_h_puzzle",
177            "m2m_100",
178            "nllb",
179            "gemma4",
180            "gemma_4",
181            "mistral",
182            "minimax_m2",
183            "step3p7",
184            "laguna",
185            "deepseek_v4",
186            // 2026-09-25: Case, hyphens and dots are normalized before the match.
187            "QWEN3-NEXT",
188            "nemotron.h.puzzle",
189            "M2M-100",
190        ] {
191            config.model_type = model_type.to_string();
192            assert!(
193                loader_for_config(&config).is_ok(),
194                "declared model type {model_type:?} was rejected"
195            );
196        }
197
198        config.model_type = "unsupported_model".to_string();
199        assert!(loader_for_config(&config).is_err());
200    }
201
202    // 2026-09-25: `build_model` serves NLLB with `NllbGpuModel` before it asks
203    // for a loader, so NLLB's generic loader fails every mandatory load; its
204    // optional hooks return their defaults.
205    #[test]
206    fn nllb_mandatory_generic_loads_fail_fast() {
207        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
208        config.model_type = "nllb".to_string();
209        let loader = loader_for_config(&config).unwrap();
210        let store = WeightStore::empty();
211        let gpu = metrale_gpu_runtime::gpu::mock::MockGpuBackend::new();
212
213        let errors = [
214            loader
215                .load_layers(&store, &config, &gpu, &[])
216                .err()
217                .expect("generic layer load unexpectedly succeeded"),
218            loader
219                .load_embedding(&store, &config, &gpu)
220                .expect_err("generic embedding load unexpectedly succeeded"),
221            loader
222                .load_final_norm(&store, &config, &gpu)
223                .expect_err("generic final-norm load unexpectedly succeeded"),
224            loader
225                .load_lm_head(&store, &config, &gpu)
226                .expect_err("generic LM-head load unexpectedly succeeded"),
227        ];
228        for err in errors {
229            assert!(
230                err.to_string()
231                    .contains("dedicated GPU encoder-decoder runtime"),
232                "unexpected fail-fast diagnostic: {err}"
233            );
234        }
235    }
236}
metrale_model_engine/factory.rs

metrale_model_engine/
factory.rs
