spark_model/
factory.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Model factory: builds the right model from config + weights.
4//!
5//! Weight loader selection is registry-driven — add new models by implementing
6//! [`ModelWeightLoader`] and registering in [`loader_for_config`]. No other
7//! code changes needed.
8
9use anyhow::{Result, bail};
10use atlas_core::config::ModelConfig;
11use spark_runtime::weights::WeightStore;
12
13use crate::mistral_loader::MistralWeightLoader;
14use crate::weight_loader::LongcatWeightLoader;
15use crate::weight_loader::Qwen4ExpWeightLoader;
16use crate::weight_loader::{
17    DeepSeekV4WeightLoader, DflashConfig, Gemma4WeightLoader, Glm5NextWeightLoader,
18    LagunaWeightLoader, MinimaxM2WeightLoader, ModelWeightLoader, NemotronHWeightLoader,
19    NllbWeightLoader, Qwen3VLWeightLoader, Qwen3WeightLoader, Qwen35DenseWeightLoader,
20    Qwen35WeightLoader, Step3p7WeightLoader,
21};
22
23/// DFlash speculative-decoding build arguments. `None` for non-DFlash runs;
24/// `Some(...)` carries the drafter's separate [`WeightStore`], parsed
25/// `config.json`, and CLI overrides for γ and the sliding-window size.
26///
27/// Construction order: the caller (`spark-server::main`) loads the drafter
28/// checkpoint into a fresh [`WeightStore`] via the same `WeightStore::load`
29/// path used for the target, then parses `config.json` via
30/// [`crate::weight_loader::dflash_loader::parse_dflash_config`]. Both inputs
31/// flow through to [`build_model`] which validates dimensions against the
32/// target before constructing [`crate::layers::BlockDiffusionDraftHead`].
33pub struct DflashBuildArgs<'a> {
34    pub drafter_store: &'a WeightStore,
35    pub drafter_config: DflashConfig,
36    pub gamma: Option<usize>,
37    pub window_size: Option<usize>,
38}
39
40/// LoRA adapter build arguments (`--lora-adapter NAME=PATH`). `None` for
41/// base-only runs; `Some(...)` carries the adapter's separate on-device
42/// [`WeightStore`] (loaded via
43/// `spark_runtime::weights::adapter::load_adapter_safetensors`), the parsed
44/// `adapter_config.json`, and the pool-shape CLI knobs.
45///
46/// Unlike DFlash (loaded post-construction), the LoRA pool is allocated at
47/// the TOP of `build_model` — before the buffer arena and the free-memory
48/// snapshot — so its bytes are automatically debited from the KV budget.
49pub struct LoraBuildArgs<'a> {
50    /// One or more adapters to pack (repeated `--lora-adapter NAME=PATH`),
51    /// each carrying its NAME, its on-device `WeightStore`, and its parsed
52    /// `adapter_config.json`. Slot k = `adapters[k]`. A single element is
53    /// byte-identical to the pre-multi-adapter path.
54    pub adapters: Vec<crate::lora::LoraAdapterInput<'a>>,
55    pub max_lora_rank: usize,
56    pub max_loras: usize,
57}
58
59// ── Loader registry ─────────────────────────────────────────────────────────
60// Adding a new model: implement ModelWeightLoader and add a match arm below.
61// Everything else (KV cache, buffers, TransformerModel) is model-agnostic.
62
63/// Select the weight loader for a given model config.
64///
65/// This is the ONLY place model_type strings are matched. All downstream
66/// code is model-agnostic.
67pub fn loader_for_config(config: &ModelConfig) -> Result<Box<dyn ModelWeightLoader>> {
68    let normalized = config.model_type.to_lowercase().replace(['-', '.'], "_");
69    match normalized.as_str() {
70        // Qwen3 family: sub-dispatch by config predicates
71        "qwen3_next" => Ok(Box::new(Qwen3WeightLoader)),
72        "qwen3_vl_moe" => Ok(Box::new(Qwen3VLWeightLoader)),
73        "qwen3_5_moe" | "qwen3_5" | "qwen35_moe" | "qwen35" => {
74            // Dense check has to come first. Qwen3.6-27B-FP8 is the dense text
75            // sibling of the Qwen3.6 VL family — its config declares the same
76            // `vision_config` block as the MoE-VL siblings (so `is_qwen3_vl()`
77            // returns true), but the checkpoint ships no vision tower and no
78            // MoE router, so the VL loader panics on a missing `mlp.gate`.
79            // `is_qwen35_dense()` requires `num_experts == 0`, which only the
80            // dense text models satisfy — VL-MoE always has experts.
81            if config.is_qwen35_dense() {
82                Ok(Box::new(Qwen35DenseWeightLoader))
83            } else if config.is_qwen3_vl() {
84                Ok(Box::new(Qwen3VLWeightLoader))
85            } else {
86                Ok(Box::new(Qwen35WeightLoader))
87            }
88        }
89        // Qwen3.6: identical architecture to Qwen3.5 MoE at the weight level
90        // (GDN + full-attention + MoE hybrid, same expert layout, same MTP).
91        // Only difference is MRoPE-interleaved layout + attn_output_gate on
92        // full-attention layers — both handled at forward-pass layer time,
93        // not during weight loading.
94        "qwen3_6_moe" | "holo3_1_moe" => Ok(Box::new(Qwen35WeightLoader)),
95        // Nemotron-H family (Mamba-2 + MoE + Attention), including Puzzle
96        // (heterogeneous per-block MoE intermediate / top-k).
97        "nemotron_h" | "nemotron_h_puzzle" => Ok(Box::new(NemotronHWeightLoader)),
98        // NLLB / M2M-100 encoder-decoder translation family.
99        "m2m_100" | "nllb" => Ok(Box::new(NllbWeightLoader)),
100        // Gemma-4 family (pure attention, GeGLU, sliding + full attention)
101        "gemma4" | "gemma_4" => Ok(Box::new(Gemma4WeightLoader)),
102        // Mistral family (MLA + MoE, GQA fallback for initial bring-up)
103        "mistral" => Ok(Box::new(MistralWeightLoader)),
104        // LongCat-Flash(-Lite): MLA dual-sublayer blocks + shortcut MoE with
105        // zero-computation experts (+ n-gram input embeddings).
106        "longcat_flash_ngram" | "longcat_flash" => Ok(Box::new(LongcatWeightLoader)),
107        // Qwen3.8-Flash-Next. `dispatch.rs` normalizes the older
108        // `qwen3_8_flash_next` naming onto `qwen4_exp`, so one arm covers both
109        // published quantizations.
110        "qwen4_exp" => Ok(Box::new(Qwen4ExpWeightLoader)),
111        // MiniMax M2 family (M2.1 / M2.7) — full attention + 256-expert
112        // sigmoid-routed MoE + 3-module MTP.
113        "minimax_m2" => Ok(Box::new(MinimaxM2WeightLoader)),
114        // Step 3.7 Flash — 288-expert sigmoid-routed MoE + shared expert +
115        // mixed full/sliding attention + attention gate + 3 MTP modules.
116        "step3p7" => Ok(Box::new(Step3p7WeightLoader)),
117        "laguna" => Ok(Box::new(LagunaWeightLoader)),
118        // DeepSeek-V4 family (Flash) — MLA + MoE + CSA/HCA hybrid attention + mHC.
119        "deepseek_v4" => Ok(Box::new(DeepSeekV4WeightLoader)),
120        // GLM-5.3-Flash — NoPE MLA behind a DSA kpool indexer + KDA linear attention +
121        // 288-expert sigmoid-routed MoE + mHC. `glm5_next_text` is the inner `model_type`;
122        // the parser canonicalises both onto `glm5_next`.
123        "glm5_next" | "glm5_next_text" => Ok(Box::new(Glm5NextWeightLoader)),
124        _ => bail!(
125            "Unsupported model type: '{}' (normalized: '{}'). \
126             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",
127            config.model_type,
128            normalized,
129        ),
130    }
131}
132
133mod build;
134mod lm_head_setup;
135mod m2_setup;
136
137pub use build::build_model;
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use crate::layers::mtp_head::MtpQuantization;
143    use spark_runtime::kv_cache::KvCacheDtype;
144    use spark_runtime::prefix_cache::PrefixCache;
145
146    #[test]
147    fn test_unsupported_model_type() {
148        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
149        config.model_type = "llama".to_string();
150
151        let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
152        let store = WeightStore::empty();
153
154        let prefix_cache: Box<dyn PrefixCache> =
155            Box::new(spark_runtime::prefix_cache::NoPrefixCaching);
156        let result = build_model(
157            config,
158            store,
159            Box::new(gpu),
160            1,
161            16,
162            4096,
163            8,
164            MtpQuantization::Nvfp4,
165            false,
166            prefix_cache,
167            0,
168            None,
169            false,
170            1,
171            KvCacheDtype::Fp8,
172            1024 * 1024 * 1024,
173            0.90,
174            0,
175            vec![],
176            0,
177            None,
178            None, // dflash_args
179            None, // lora_args
180            None, // nllb_lang
181            None, // nllb_lora_dir
182        );
183        match result {
184            Err(e) => assert!(e.to_string().contains("Unsupported model type: 'llama'")),
185            Ok(_) => panic!("Expected error for unsupported model type"),
186        }
187    }
188
189    #[test]
190    fn all_declared_model_type_spellings_are_accepted() {
191        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
192        for model_type in [
193            "qwen3_next",
194            "qwen3_vl_moe",
195            "qwen3_5_moe",
196            "qwen3_5",
197            "qwen35_moe",
198            "qwen35",
199            "qwen3_6_moe",
200            "holo3_1_moe",
201            "nemotron_h",
202            "nemotron_h_puzzle",
203            "m2m_100",
204            "nllb",
205            "gemma4",
206            "gemma_4",
207            "mistral",
208            "minimax_m2",
209            "step3p7",
210            "laguna",
211            "deepseek_v4",
212            // Normalization accepts case, hyphens, and dots before dispatch.
213            "QWEN3-NEXT",
214            "nemotron.h.puzzle",
215            "M2M-100",
216        ] {
217            config.model_type = model_type.to_string();
218            assert!(
219                loader_for_config(&config).is_ok(),
220                "declared model type {model_type:?} was rejected"
221            );
222        }
223
224        config.model_type = "unsupported_model".to_string();
225        assert!(loader_for_config(&config).is_err());
226    }
227
228    // The generic `ModelWeightLoader` for NLLB must fail fast: real NLLB serving
229    // goes through the dedicated `NllbGpuModel` encoder-decoder runtime, which
230    // `build_model` selects before this loader is ever consulted. Reaching this
231    // loader is a routing bug, so every mandatory generic tensor-loading
232    // entry point bails. Optional hooks still report their normal absence.
233    #[test]
234    fn nllb_mandatory_generic_loads_fail_fast() {
235        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
236        config.model_type = "nllb".to_string();
237        let loader = loader_for_config(&config).unwrap();
238        let store = WeightStore::empty();
239        let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
240
241        let errors = [
242            loader
243                .load_layers(&store, &config, &gpu, &[])
244                .err()
245                .expect("generic layer load unexpectedly succeeded"),
246            loader
247                .load_embedding(&store, &config, &gpu)
248                .expect_err("generic embedding load unexpectedly succeeded"),
249            loader
250                .load_final_norm(&store, &config, &gpu)
251                .expect_err("generic final-norm load unexpectedly succeeded"),
252            loader
253                .load_lm_head(&store, &config, &gpu)
254                .expect_err("generic LM-head load unexpectedly succeeded"),
255        ];
256        for err in errors {
257            assert!(
258                err.to_string()
259                    .contains("dedicated GPU encoder-decoder runtime"),
260                "unexpected fail-fast diagnostic: {err}"
261            );
262        }
263    }
264}
spark_model/factory.rs

spark_model/
factory.rs
