aria_inference/
family.rs

1use aria_kernel::EngineError;
2use std::path::Path;
3
4/// Delivery phase for a registered family (requirements §1.1).
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub enum FamilyPhase {
7    A,
8    B,
9    C,
10}
11
12/// Architecture class for graph / loader hooks.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum ArchClass {
15    /// Dense decoder-only LLM (Gemma / Qwen / LFM / …).
16    TextDense,
17    /// MoE text (LFM2-8B-A1B / Inkling); Session top-k router + expert FFN.
18    TextMoE,
19    /// Vision-language.
20    VL,
21    /// Vision-language-action.
22    VLA,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct FamilyEntry {
27    pub path: &'static str,
28    pub base_model: &'static str,
29    pub phase: FamilyPhase,
30    pub arch: ArchClass,
31}
32
33/// Full registry mirroring model/requirements.md §1.1.
34pub const FAMILY_REGISTRY: &[FamilyEntry] = &[
35    FamilyEntry {
36        path: "qwen/qwen3-0.6b",
37        base_model: "Qwen/Qwen3-0.6B",
38        phase: FamilyPhase::B,
39        arch: ArchClass::TextDense,
40    },
41    FamilyEntry {
42        path: "qwen/qwen3-1.7b",
43        base_model: "Qwen/Qwen3-1.7B",
44        phase: FamilyPhase::B,
45        arch: ArchClass::TextDense,
46    },
47    FamilyEntry {
48        path: "qwen/qwen3.5-0.8b",
49        base_model: "Qwen/Qwen3.5-0.8B",
50        phase: FamilyPhase::B,
51        arch: ArchClass::TextDense,
52    },
53    FamilyEntry {
54        path: "qwen/qwen3.5-2b",
55        base_model: "Qwen/Qwen3.5-2B",
56        phase: FamilyPhase::B,
57        arch: ArchClass::TextDense,
58    },
59    FamilyEntry {
60        path: "gemma/gemma-3-270m-it",
61        base_model: "google/gemma-3-270m-it",
62        phase: FamilyPhase::B,
63        arch: ArchClass::TextDense,
64    },
65    FamilyEntry {
66        path: "gemma/gemma-3-1b-it",
67        base_model: "google/gemma-3-1b-it",
68        phase: FamilyPhase::B,
69        arch: ArchClass::TextDense,
70    },
71    FamilyEntry {
72        path: "gemma/gemma-3n-e2b-it",
73        base_model: "google/gemma-3n-E2B-it",
74        phase: FamilyPhase::C,
75        arch: ArchClass::VL,
76    },
77    FamilyEntry {
78        path: "gemma/gemma-3n-e4b-it",
79        base_model: "google/gemma-3n-E4B-it",
80        phase: FamilyPhase::C,
81        arch: ArchClass::VL,
82    },
83    FamilyEntry {
84        path: "gemma/gemma-4-e2b-it",
85        base_model: "google/gemma-4-E2B-it",
86        phase: FamilyPhase::A,
87        arch: ArchClass::VL, // full VL in stage C; tiny text path available from A
88    },
89    FamilyEntry {
90        path: "gemma/gemma-4-e4b-it",
91        base_model: "google/gemma-4-E4B-it",
92        phase: FamilyPhase::C,
93        arch: ArchClass::VL,
94    },
95    FamilyEntry {
96        path: "lfm/lfm2-350m",
97        base_model: "LiquidAI/LFM2-350M",
98        phase: FamilyPhase::B,
99        arch: ArchClass::TextDense,
100    },
101    FamilyEntry {
102        path: "lfm/lfm2-700m",
103        base_model: "LiquidAI/LFM2-700M",
104        phase: FamilyPhase::B,
105        arch: ArchClass::TextDense,
106    },
107    FamilyEntry {
108        path: "lfm/lfm2-1.2b",
109        base_model: "LiquidAI/LFM2-1.2B",
110        phase: FamilyPhase::B,
111        arch: ArchClass::TextDense,
112    },
113    FamilyEntry {
114        path: "lfm/lfm2-2.6b",
115        base_model: "LiquidAI/LFM2-2.6B",
116        phase: FamilyPhase::B,
117        arch: ArchClass::TextDense,
118    },
119    FamilyEntry {
120        path: "lfm/lfm2-8b-a1b",
121        base_model: "LiquidAI/LFM2-8B-A1B",
122        phase: FamilyPhase::B,
123        arch: ArchClass::TextMoE,
124    },
125    FamilyEntry {
126        path: "lfm/lfm2-vl-450m",
127        base_model: "LiquidAI/LFM2-VL-450M",
128        phase: FamilyPhase::C,
129        arch: ArchClass::VL,
130    },
131    FamilyEntry {
132        path: "lfm/lfm2.5-350m",
133        base_model: "LiquidAI/LFM2.5-350M",
134        phase: FamilyPhase::B,
135        arch: ArchClass::TextDense,
136    },
137    FamilyEntry {
138        path: "lfm/lfm2.5-1.2b-instruct",
139        base_model: "LiquidAI/LFM2.5-1.2B-Instruct",
140        phase: FamilyPhase::B,
141        arch: ArchClass::TextDense,
142    },
143    FamilyEntry {
144        path: "lfm/lfm2.5-1.2b-thinking",
145        base_model: "LiquidAI/LFM2.5-1.2B-Thinking",
146        phase: FamilyPhase::B,
147        arch: ArchClass::TextDense,
148    },
149    FamilyEntry {
150        path: "lfm/lfm2.5-2.6b",
151        base_model: "LiquidAI/LFM2.5-2.6B",
152        phase: FamilyPhase::B,
153        arch: ArchClass::TextDense,
154    },
155    FamilyEntry {
156        path: "lfm/lfm2.5-vl-1.6b",
157        base_model: "LiquidAI/LFM2.5-VL-1.6B",
158        phase: FamilyPhase::C,
159        arch: ArchClass::VL,
160    },
161    FamilyEntry {
162        path: "nanbeige/nanbeige4.2-3b",
163        base_model: "Nanbeige/Nanbeige4.2-3B",
164        phase: FamilyPhase::B,
165        arch: ArchClass::TextDense,
166    },
167    FamilyEntry {
168        path: "bonsai/bonsai-27b",
169        base_model: "prism-ml/Bonsai-27B-unpacked",
170        phase: FamilyPhase::B,
171        arch: ArchClass::TextDense,
172    },
173    FamilyEntry {
174        path: "inkling/inkling-small",
175        base_model: "thinkingmachines/Inkling-Small",
176        phase: FamilyPhase::B,
177        arch: ArchClass::TextMoE,
178    },
179    FamilyEntry {
180        path: "openvla/openvla-7b",
181        base_model: "openvla/openvla-7b",
182        phase: FamilyPhase::C,
183        arch: ArchClass::VLA,
184    },
185    FamilyEntry {
186        path: "openpi/openpi-pi0-3b",
187        base_model: "lerobot/pi0_base",
188        phase: FamilyPhase::C,
189        arch: ArchClass::VLA,
190    },
191    FamilyEntry {
192        path: "openpi/openpi-pi0.5-3b",
193        base_model: "lerobot/pi05_base",
194        phase: FamilyPhase::C,
195        arch: ArchClass::VLA,
196    },
197    FamilyEntry {
198        path: "lingbot/lingbot-vla-v2-6b",
199        base_model: "robbyant/lingbot-vla-v2-6b",
200        phase: FamilyPhase::C,
201        arch: ArchClass::VLA,
202    },
203];
204
205/// Runtime family handle used by Session.
206#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
207pub struct Family {
208    pub path: &'static str,
209    pub arch: ArchClass,
210    pub phase: FamilyPhase,
211}
212
213impl Family {
214    pub fn path(self) -> &'static str {
215        self.path
216    }
217
218    pub fn uses_text_decoder(self) -> bool {
219        matches!(
220            self.arch,
221            ArchClass::TextDense | ArchClass::TextMoE | ArchClass::VL | ArchClass::VLA
222        )
223    }
224
225    pub fn is_moe(self) -> bool {
226        self.arch == ArchClass::TextMoE
227    }
228}
229
230pub fn lookup_family(path: &str) -> Result<&'static FamilyEntry, EngineError> {
231    FAMILY_REGISTRY
232        .iter()
233        .find(|e| e.path == path)
234        .ok_or_else(|| EngineError::UnsupportedFamily(path.to_string()))
235}
236
237/// Map a bundle directory or `slug_q4` / `slug_q326_channel` name onto a registry path.
238pub fn infer_family_path(hint: &str) -> Option<&'static str> {
239    let name = Path::new(hint)
240        .file_name()
241        .and_then(|s| s.to_str())
242        .unwrap_or(hint);
243    let slug = strip_quant_suffix(name);
244    if slug.is_empty() {
245        return None;
246    }
247    if let Ok(e) = lookup_family(slug) {
248        return Some(e.path);
249    }
250    let mut best: Option<&'static str> = None;
251    let mut best_len = 0usize;
252    for e in FAMILY_REGISTRY {
253        let tail = e.path.rsplit('/').next().unwrap_or(e.path);
254        if slug == tail && tail.len() > best_len {
255            best = Some(e.path);
256            best_len = tail.len();
257        }
258    }
259    best
260}
261
262fn strip_quant_suffix(name: &str) -> &str {
263    let mut s = name;
264    if let Some(stripped) = s.strip_suffix("_tiny") {
265        s = stripped;
266    }
267    if let Some(idx) = s.rfind("_q") {
268        let suffix = &s[idx + 2..];
269        if is_quant_suffix(suffix) {
270            s = &s[..idx];
271        }
272    }
273    s
274}
275
276/// `4` / `8` / `326` / `3.26`, optionally followed by codebook-share `_channel` / `_group`.
277fn is_quant_suffix(suffix: &str) -> bool {
278    let core = suffix
279        .strip_suffix("_channel")
280        .or_else(|| suffix.strip_suffix("_group"))
281        .unwrap_or(suffix);
282    !core.is_empty() && core.chars().all(|c| c.is_ascii_digit() || c == '.')
283}
284
285/// Qwen3 uses θ=1e6; Qwen3.5 gated attention uses θ=1e7 (`rope_parameters.rope_theta`).
286/// Bundles that omit `rope_theta` hit serde default 10000 and produce empty / garbage
287/// completions (Hello → `content: ""` with a full `max_tokens` budget).
288pub fn effective_rope_theta(family_path: &str, configured: f32) -> f32 {
289    let p = family_path.to_ascii_lowercase();
290    if (configured - 10_000.0).abs() >= 0.5 {
291        return configured;
292    }
293    if p.contains("qwen3.5") {
294        10_000_000.0
295    } else if p.contains("qwen") {
296        1_000_000.0
297    } else {
298        configured
299    }
300}
301
302pub fn family_phase(path: &str) -> Result<FamilyPhase, EngineError> {
303    Ok(lookup_family(path)?.phase)
304}
305
306/// Graph hook id for an architecture class (stage B/C dispatch).
307pub fn graph_hook(arch: ArchClass) -> &'static str {
308    match arch {
309        ArchClass::TextDense => "text_dense_decoder",
310        ArchClass::TextMoE => "text_moe_decoder",
311        ArchClass::VL => "vl_text_plus_vision",
312        ArchClass::VLA => "vla_text_vision_action",
313    }
314}
315
316/// Stage A only: gemma-4-e2b-it.
317pub fn require_stage_a(path: &str) -> Result<Family, EngineError> {
318    let e = lookup_family(path)?;
319    if e.phase != FamilyPhase::A {
320        return Err(EngineError::UnsupportedFamily(format!(
321            "{} is phase {:?}; stage A only runs gemma/gemma-4-e2b-it",
322            path, e.phase
323        )));
324    }
325    Ok(Family {
326        path: e.path,
327        arch: e.arch,
328        phase: e.phase,
329    })
330}
331
332/// Stage B: text / MoE families (+ stage-A golden path).
333pub fn require_stage_b(path: &str) -> Result<Family, EngineError> {
334    let e = lookup_family(path)?;
335    let ok = match (e.phase, e.arch) {
336        (FamilyPhase::A, _) => e.path == "gemma/gemma-4-e2b-it",
337        (FamilyPhase::B, ArchClass::TextDense | ArchClass::TextMoE) => true,
338        _ => false,
339    };
340    if !ok {
341        return Err(EngineError::UnsupportedFamily(format!(
342            "{} (phase {:?}, arch {:?}) is not a stage-B text/MoE family",
343            path, e.phase, e.arch
344        )));
345    }
346    Ok(Family {
347        path: e.path,
348        arch: e.arch,
349        phase: e.phase,
350    })
351}
352
353/// Any registered family that can run the shared text decoder (+ stage C extras).
354pub fn require_runnable(path: &str) -> Result<Family, EngineError> {
355    let e = lookup_family(path)?;
356    Ok(Family {
357        path: e.path,
358        arch: e.arch,
359        phase: e.phase,
360    })
361}
362
363/// Representative path per arch class for tiny E2E tests.
364pub fn arch_class_representatives() -> &'static [(&'static str, ArchClass)] {
365    &[
366        ("gemma/gemma-3-270m-it", ArchClass::TextDense),
367        ("qwen/qwen3.5-2b", ArchClass::TextDense),
368        ("lfm/lfm2-350m", ArchClass::TextDense),
369        ("lfm/lfm2-8b-a1b", ArchClass::TextMoE),
370        ("nanbeige/nanbeige4.2-3b", ArchClass::TextDense),
371        ("bonsai/bonsai-27b", ArchClass::TextDense),
372        ("inkling/inkling-small", ArchClass::TextMoE),
373        ("lfm/lfm2-vl-450m", ArchClass::VL),
374        ("openvla/openvla-7b", ArchClass::VLA),
375    ]
376}
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381
382    /// Mirror model `tests/test_families.py` EXPECTED path → base_model lock table.
383    const EXPECTED_BASE_MODELS: &[(&str, &str)] = &[
384        ("qwen/qwen3-0.6b", "Qwen/Qwen3-0.6B"),
385        ("qwen/qwen3-1.7b", "Qwen/Qwen3-1.7B"),
386        ("qwen/qwen3.5-0.8b", "Qwen/Qwen3.5-0.8B"),
387        ("qwen/qwen3.5-2b", "Qwen/Qwen3.5-2B"),
388        ("gemma/gemma-3-270m-it", "google/gemma-3-270m-it"),
389        ("gemma/gemma-3-1b-it", "google/gemma-3-1b-it"),
390        ("gemma/gemma-3n-e2b-it", "google/gemma-3n-E2B-it"),
391        ("gemma/gemma-3n-e4b-it", "google/gemma-3n-E4B-it"),
392        ("gemma/gemma-4-e2b-it", "google/gemma-4-E2B-it"),
393        ("gemma/gemma-4-e4b-it", "google/gemma-4-E4B-it"),
394        ("lfm/lfm2-350m", "LiquidAI/LFM2-350M"),
395        ("lfm/lfm2-700m", "LiquidAI/LFM2-700M"),
396        ("lfm/lfm2-1.2b", "LiquidAI/LFM2-1.2B"),
397        ("lfm/lfm2-2.6b", "LiquidAI/LFM2-2.6B"),
398        ("lfm/lfm2-8b-a1b", "LiquidAI/LFM2-8B-A1B"),
399        ("lfm/lfm2-vl-450m", "LiquidAI/LFM2-VL-450M"),
400        ("lfm/lfm2.5-350m", "LiquidAI/LFM2.5-350M"),
401        ("lfm/lfm2.5-1.2b-instruct", "LiquidAI/LFM2.5-1.2B-Instruct"),
402        ("lfm/lfm2.5-1.2b-thinking", "LiquidAI/LFM2.5-1.2B-Thinking"),
403        ("lfm/lfm2.5-2.6b", "LiquidAI/LFM2.5-2.6B"),
404        ("lfm/lfm2.5-vl-1.6b", "LiquidAI/LFM2.5-VL-1.6B"),
405        ("nanbeige/nanbeige4.2-3b", "Nanbeige/Nanbeige4.2-3B"),
406        ("bonsai/bonsai-27b", "prism-ml/Bonsai-27B-unpacked"),
407        ("inkling/inkling-small", "thinkingmachines/Inkling-Small"),
408        ("openvla/openvla-7b", "openvla/openvla-7b"),
409        ("openpi/openpi-pi0-3b", "lerobot/pi0_base"),
410        ("openpi/openpi-pi0.5-3b", "lerobot/pi05_base"),
411        ("lingbot/lingbot-vla-v2-6b", "robbyant/lingbot-vla-v2-6b"),
412    ];
413
414    #[test]
415    fn registry_matches_model_expected() {
416        assert_eq!(FAMILY_REGISTRY.len(), EXPECTED_BASE_MODELS.len());
417        for (path, base) in EXPECTED_BASE_MODELS {
418            let e = lookup_family(path).unwrap_or_else(|_| panic!("missing {path}"));
419            assert_eq!(e.base_model, *base, "{path}");
420        }
421        // Every registry row appears in EXPECTED.
422        for e in FAMILY_REGISTRY {
423            assert!(
424                EXPECTED_BASE_MODELS
425                    .iter()
426                    .any(|(p, b)| *p == e.path && *b == e.base_model),
427                "unexpected registry entry {}",
428                e.path
429            );
430        }
431    }
432
433    #[test]
434    fn registry_phase_gates() {
435        assert!(require_stage_b("qwen/qwen3.5-2b").is_ok());
436        assert!(require_stage_b("lfm/lfm2-8b-a1b").is_ok());
437        assert!(matches!(
438            require_stage_b("openvla/openvla-7b"),
439            Err(EngineError::UnsupportedFamily(_))
440        ));
441        assert!(require_runnable("openvla/openvla-7b").is_ok());
442        assert_eq!(graph_hook(ArchClass::TextMoE), "text_moe_decoder");
443        assert_eq!(
444            require_stage_a("gemma/gemma-4-e2b-it").unwrap().path(),
445            "gemma/gemma-4-e2b-it"
446        );
447        assert_eq!(
448            lookup_family("lfm/lfm2-8b-a1b").unwrap().arch,
449            ArchClass::TextMoE
450        );
451        assert_eq!(
452            lookup_family("lfm/lfm2-vl-450m").unwrap().arch,
453            ArchClass::VL
454        );
455        assert_eq!(
456            lookup_family("openvla/openvla-7b").unwrap().arch,
457            ArchClass::VLA
458        );
459    }
460
461    #[test]
462    fn infer_family_from_bundle_dirname() {
463        assert_eq!(infer_family_path("qwen3-0.6b_q4"), Some("qwen/qwen3-0.6b"));
464        assert_eq!(
465            infer_family_path("/home/ubuntu/.ariacompute/models/qwen3-0.6b_q4"),
466            Some("qwen/qwen3-0.6b")
467        );
468        assert_eq!(
469            infer_family_path("gemma-4-e2b-it_q8"),
470            Some("gemma/gemma-4-e2b-it")
471        );
472        assert_eq!(
473            infer_family_path("qwen3-0.6b_q326"),
474            Some("qwen/qwen3-0.6b")
475        );
476        assert_eq!(
477            infer_family_path("qwen3-0.6b_q326_channel"),
478            Some("qwen/qwen3-0.6b")
479        );
480        assert_eq!(
481            infer_family_path("/home/ubuntu/.ariacompute/models/qwen3-0.6b_q326_channel"),
482            Some("qwen/qwen3-0.6b")
483        );
484        assert_eq!(
485            infer_family_path("gemma-3-1b-it_q326"),
486            Some("gemma/gemma-3-1b-it")
487        );
488        assert_eq!(
489            infer_family_path("gemma-3-1b-it_q326_channel"),
490            Some("gemma/gemma-3-1b-it")
491        );
492        assert!(infer_family_path("totally-unknown_q4").is_none());
493    }
494
495    #[test]
496    fn qwen3_rope_theta_not_llama_default() {
497        assert!((effective_rope_theta("qwen/qwen3-0.6b", 10_000.0) - 1_000_000.0).abs() < 1.0);
498        assert!((effective_rope_theta("qwen/qwen3-0.6b", 1_000_000.0) - 1_000_000.0).abs() < 1.0);
499        assert!((effective_rope_theta("gemma/gemma-4-e2b-it", 10_000.0) - 10_000.0).abs() < 1.0);
500        assert!((effective_rope_theta("qwen/qwen3.5-2b", 10_000.0) - 10_000_000.0).abs() < 1.0);
501        assert!(
502            (effective_rope_theta("qwen/qwen3.5-0.8b", 10_000_000.0) - 10_000_000.0).abs() < 1.0
503        );
504    }
505}
aria_inference/family.rs

aria_inference/
family.rs
