atlas_core/
capabilities.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Model capabilities — config-derived feature flags.
4//!
5//! Instead of checking `config.model_type == "qwen3_5_moe"` throughout the codebase,
6//! call sites use `config.capabilities().has_moe_layers` etc. Adding a new model
7//! only requires implementing the capability derivation, not updating every call site.
8
9/// SSM (State Space Model) architecture family.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum SsmArchitecture {
12    /// No SSM layers — pure attention (Mistral, Llama, etc.)
13    None,
14    /// Gated Delta Networks (Qwen3.5 family)
15    Gdn,
16    /// Mamba-2 (Nemotron-H family)
17    Mamba2,
18}
19
20/// Attention architecture family.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum AttentionType {
23    /// Standard multi-head or grouped-query attention (Qwen3.5, Llama, etc.)
24    Standard,
25    /// Multi-head Latent Attention — compressed KV via low-rank projection.
26    /// (Mistral Small 4, DeepSeek-V2/V3)
27    Mla,
28}
29
30/// Feature flags derived from model config at parse time.
31///
32/// These replace scattered `is_nemotron_h()`, `is_qwen35()` checks with
33/// model-agnostic predicates. New models get capabilities automatically
34/// from their config — no code changes needed.
35#[derive(Debug, Clone)]
36pub struct ModelCapabilities {
37    /// Model has SSM/recurrent layers (GDN or Mamba-2).
38    pub has_ssm_layers: bool,
39    /// Model has full attention layers.
40    pub has_attention_layers: bool,
41    /// Model has MoE (Mixture of Experts) layers.
42    pub has_moe_layers: bool,
43    /// Model supports `<think>` reasoning tokens.
44    /// Derived from tokenizer vocabulary, not model type name.
45    pub supports_thinking: bool,
46    /// Model has vision encoder (multimodal).
47    pub supports_vision: bool,
48    /// Model has MTP (Multi-Token Prediction) draft head.
49    pub has_mtp: bool,
50    /// SSM architecture family (determines state layout).
51    pub ssm_architecture: SsmArchitecture,
52    /// Attention architecture (standard GQA vs MLA compressed latent).
53    pub attention_type: AttentionType,
54    /// Model wraps language model config in a nested field (e.g., `text_config`).
55    pub has_nested_config: bool,
56}
57
58impl ModelCapabilities {
59    /// Derive capabilities from a ModelConfig.
60    ///
61    /// This is the SSOT for model feature detection. When adding a new model,
62    /// ensure its config fields populate the right capabilities here — then
63    /// all downstream code works automatically.
64    pub fn from_config(config: &super::config::ModelConfig) -> Self {
65        use super::config::LayerType;
66
67        let has_ssm = config
68            .layer_types
69            .iter()
70            .any(|t| matches!(t, LayerType::LinearAttention));
71        let has_mamba2 = config.mamba_num_heads > 0 && config.mamba_head_dim > 0;
72        let has_attention = config
73            .layer_types
74            .iter()
75            .any(|t| matches!(t, LayerType::FullAttention));
76        let has_moe = config.num_experts > 0;
77        let has_vision = config.vision.is_some();
78        let has_mtp = config.mtp_num_hidden_layers > 0;
79        let has_nested = config.nested_config;
80
81        let ssm_arch = if has_mamba2 {
82            SsmArchitecture::Mamba2
83        } else if has_ssm {
84            SsmArchitecture::Gdn
85        } else {
86            SsmArchitecture::None
87        };
88
89        Self {
90            has_ssm_layers: has_ssm || has_mamba2,
91            has_attention_layers: has_attention,
92            has_moe_layers: has_moe,
93            // Models with SSM layers support <think> tokens. Long-term this should
94            // be derived from tokenizer vocabulary, not architecture.
95            supports_thinking: has_ssm || has_mamba2,
96            supports_vision: has_vision,
97            has_mtp,
98            ssm_architecture: ssm_arch,
99            has_nested_config: has_nested,
100            attention_type: if config.kv_lora_rank > 0 {
101                AttentionType::Mla
102            } else {
103                AttentionType::Standard
104            },
105        }
106    }
107}
atlas_core/capabilities.rs

atlas_core/
capabilities.rs
