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 avarok_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    KimiK3WeightLoader, LagunaWeightLoader, MinimaxM2WeightLoader, ModelWeightLoader,
19    NemotronHWeightLoader, NllbWeightLoader, Qwen3VLWeightLoader, Qwen3WeightLoader,
20    Qwen35DenseWeightLoader, 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        "kimi_k3" | "kimi_linear" => Ok(Box::new(KimiK3WeightLoader)),
72        "qwen3_next" => Ok(Box::new(Qwen3WeightLoader)),
73        "qwen3_vl_moe" => Ok(Box::new(Qwen3VLWeightLoader)),
74        "qwen3_5_moe" | "qwen3_5" | "qwen35_moe" | "qwen35" => {
75            // Dense check has to come first. Qwen3.6-27B-FP8 is the dense text
76            // sibling of the Qwen3.6 VL family — its config declares the same
77            // `vision_config` block as the MoE-VL siblings (so `is_qwen3_vl()`
78            // returns true), but the checkpoint ships no vision tower and no
79            // MoE router, so the VL loader panics on a missing `mlp.gate`.
80            // `is_qwen35_dense()` requires `num_experts == 0`, which only the
81            // dense text models satisfy — VL-MoE always has experts.
82            if config.is_qwen35_dense() {
83                Ok(Box::new(Qwen35DenseWeightLoader))
84            } else if config.is_qwen3_vl() {
85                Ok(Box::new(Qwen3VLWeightLoader))
86            } else {
87                Ok(Box::new(Qwen35WeightLoader))
88            }
89        }
90        // Qwen3.6: identical architecture to Qwen3.5 MoE at the weight level
91        // (GDN + full-attention + MoE hybrid, same expert layout, same MTP).
92        // Only difference is MRoPE-interleaved layout + attn_output_gate on
93        // full-attention layers — both handled at forward-pass layer time,
94        // not during weight loading.
95        "qwen3_6_moe" | "holo3_1_moe" => Ok(Box::new(Qwen35WeightLoader)),
96        // Nemotron-H family (Mamba-2 + MoE + Attention), including Puzzle
97        // (heterogeneous per-block MoE intermediate / top-k).
98        "nemotron_h" | "nemotron_h_puzzle" => Ok(Box::new(NemotronHWeightLoader)),
99        // NLLB / M2M-100 encoder-decoder translation family.
100        "m2m_100" | "nllb" => Ok(Box::new(NllbWeightLoader)),
101        // Gemma-4 family (pure attention, GeGLU, sliding + full attention)
102        "gemma4" | "gemma_4" => Ok(Box::new(Gemma4WeightLoader)),
103        // Mistral family (MLA + MoE, GQA fallback for initial bring-up)
104        "mistral" => Ok(Box::new(MistralWeightLoader)),
105        // LongCat-Flash(-Lite): MLA dual-sublayer blocks + shortcut MoE with
106        // zero-computation experts (+ n-gram input embeddings).
107        "longcat_flash_ngram" | "longcat_flash" => Ok(Box::new(LongcatWeightLoader)),
108        // Qwen3.8-Flash-Next. `dispatch.rs` normalizes the older
109        // `qwen3_8_flash_next` naming onto `qwen4_exp`, so one arm covers both
110        // published quantizations.
111        "qwen4_exp" => Ok(Box::new(Qwen4ExpWeightLoader)),
112        // MiniMax M2 family (M2.1 / M2.7) — full attention + 256-expert
113        // sigmoid-routed MoE + 3-module MTP.
114        "minimax_m2" => Ok(Box::new(MinimaxM2WeightLoader)),
115        // Step 3.7 Flash — 288-expert sigmoid-routed MoE + shared expert +
116        // mixed full/sliding attention + attention gate + 3 MTP modules.
117        "step3p7" => Ok(Box::new(Step3p7WeightLoader)),
118        "laguna" => Ok(Box::new(LagunaWeightLoader)),
119        // DeepSeek-V4 family (Flash) — MLA + MoE + CSA/HCA hybrid attention + mHC.
120        "deepseek_v4" => Ok(Box::new(DeepSeekV4WeightLoader)),
121        // GLM-5.3-Flash — NoPE MLA behind a DSA kpool indexer + KDA linear attention +
122        // 288-expert sigmoid-routed MoE + mHC. `glm5_next_text` is the inner `model_type`;
123        // the parser canonicalises both onto `glm5_next`.
124        "glm5_next" | "glm5_next_text" => Ok(Box::new(Glm5NextWeightLoader)),
125        // DeepSeek-V4.1 Flash. Ingestion (S1) is complete: the GGUF arch maps,
126        // all seven shards resolve, a full ModelConfig builds from the file's
127        // own metadata, and all 1,046 tensor names translate. The GRAPH is not
128        // built yet — V4.1 adds engram (two ~30 GiB hash tables), shared
129        // compressed attention (`SharedAttentionRuntime`, KV produced by four
130        // layers for forty) and a reworked indexer, none of which V4 has.
131        //
132        // This arm exists so the failure says which of those two things is
133        // missing. Falling through to the catch-all below would report
134        // "Unsupported model type", which is false and sends the reader looking
135        // for a config problem.
136        // DeepSeek-V4.1 Flash: the seven-shard Q2_K GGUF, routed experts and the
137        // engram tables streamed from disk (see weight_loader/deepseek_v41.rs).
138        "deepseek_v41" => Ok(Box::new(
139            crate::weight_loader::deepseek_v41::DeepSeekV41WeightLoader,
140        )),
141        _ => bail!(
142            "Unsupported model type: '{}' (normalized: '{}'). \
143             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",
144            config.model_type,
145            normalized,
146        ),
147    }
148}
149
150mod build;
151mod lm_head_setup;
152mod m2_setup;
153
154pub use build::build_model;
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use crate::layers::mtp_head::MtpQuantization;
160    use spark_runtime::kv_cache::KvCacheDtype;
161    use spark_runtime::prefix_cache::PrefixCache;
162
163    #[test]
164    fn test_unsupported_model_type() {
165        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
166        config.model_type = "llama".to_string();
167
168        let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
169        let store = WeightStore::empty();
170
171        let prefix_cache: Box<dyn PrefixCache> =
172            Box::new(spark_runtime::prefix_cache::NoPrefixCaching);
173        let result = build_model(
174            config,
175            store,
176            Box::new(gpu),
177            1,
178            16,
179            4096,
180            8,
181            MtpQuantization::Nvfp4,
182            false,
183            prefix_cache,
184            0,
185            None,
186            false,
187            1,
188            KvCacheDtype::Fp8,
189            1024 * 1024 * 1024,
190            0.90,
191            0,
192            vec![],
193            0,
194            None,
195            None, // dflash_args
196            None, // lora_args
197            None, // nllb_lang
198            None, // nllb_lora_dir
199        );
200        match result {
201            Err(e) => assert!(e.to_string().contains("Unsupported model type: 'llama'")),
202            Ok(_) => panic!("Expected error for unsupported model type"),
203        }
204    }
205
206    #[test]
207    fn all_declared_model_type_spellings_are_accepted() {
208        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
209        for model_type in [
210            "qwen3_next",
211            "qwen3_vl_moe",
212            "qwen3_5_moe",
213            "qwen3_5",
214            "qwen35_moe",
215            "qwen35",
216            "qwen3_6_moe",
217            "holo3_1_moe",
218            "nemotron_h",
219            "nemotron_h_puzzle",
220            "m2m_100",
221            "nllb",
222            "gemma4",
223            "gemma_4",
224            "mistral",
225            "minimax_m2",
226            "step3p7",
227            "laguna",
228            "deepseek_v4",
229            // Normalization accepts case, hyphens, and dots before dispatch.
230            "QWEN3-NEXT",
231            "nemotron.h.puzzle",
232            "M2M-100",
233        ] {
234            config.model_type = model_type.to_string();
235            assert!(
236                loader_for_config(&config).is_ok(),
237                "declared model type {model_type:?} was rejected"
238            );
239        }
240
241        config.model_type = "unsupported_model".to_string();
242        assert!(loader_for_config(&config).is_err());
243    }
244
245    // The generic `ModelWeightLoader` for NLLB must fail fast: real NLLB serving
246    // goes through the dedicated `NllbGpuModel` encoder-decoder runtime, which
247    // `build_model` selects before this loader is ever consulted. Reaching this
248    // loader is a routing bug, so every mandatory generic tensor-loading
249    // entry point bails. Optional hooks still report their normal absence.
250    #[test]
251    fn nllb_mandatory_generic_loads_fail_fast() {
252        let mut config = ModelConfig::qwen3_next_80b_nvfp4();
253        config.model_type = "nllb".to_string();
254        let loader = loader_for_config(&config).unwrap();
255        let store = WeightStore::empty();
256        let gpu = spark_runtime::gpu::mock::MockGpuBackend::new();
257
258        let errors = [
259            loader
260                .load_layers(&store, &config, &gpu, &[])
261                .err()
262                .expect("generic layer load unexpectedly succeeded"),
263            loader
264                .load_embedding(&store, &config, &gpu)
265                .expect_err("generic embedding load unexpectedly succeeded"),
266            loader
267                .load_final_norm(&store, &config, &gpu)
268                .expect_err("generic final-norm load unexpectedly succeeded"),
269            loader
270                .load_lm_head(&store, &config, &gpu)
271                .expect_err("generic LM-head load unexpectedly succeeded"),
272        ];
273        for err in errors {
274            assert!(
275                err.to_string()
276                    .contains("dedicated GPU encoder-decoder runtime"),
277                "unexpected fail-fast diagnostic: {err}"
278            );
279        }
280    }
281}
spark_model/factory.rs

spark_model/
factory.rs
