metrale_config/
capabilities.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Feature flags derived from a parsed `ModelConfig` by `ModelConfig::capabilities`.
4//!
5//! Owner: config.
6//! Invariants: none beyond the types.
7
8/// 2026-09-26: Recurrent-layer family, chosen in `ModelCapabilities::from_config`.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum SsmArchitecture {
11    None,
12    /// 2026-09-26: Some layer is `LayerType::LinearAttention` and the Mamba-2 head dims are unset.
13    Gdn,
14    /// 2026-09-26: `mamba_num_heads` and `mamba_head_dim` are both non-zero; this wins over `Gdn`.
15    Mamba2,
16}
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum AttentionType {
20    Standard,
21    /// 2026-09-26: Multi-head latent attention, chosen when `kv_lora_rank > 0`.
22    Mla,
23}
24
25#[derive(Debug, Clone)]
26pub struct ModelCapabilities {
27    /// 2026-09-26: Some layer is `LinearAttention`, or the Mamba-2 head dims are set.
28    pub has_ssm_layers: bool,
29    /// 2026-09-26: Some layer is `FullAttention`.
30    pub has_attention_layers: bool,
31    /// 2026-09-26: `num_experts > 0`.
32    pub has_moe_layers: bool,
33    /// 2026-09-26: Equal to `has_ssm_layers`: it follows the architecture, not the tokenizer
34    /// vocabulary.
35    pub supports_thinking: bool,
36    /// 2026-09-26: `ModelConfig::vision` is set.
37    pub supports_vision: bool,
38    /// 2026-09-26: `mtp_num_hidden_layers > 0`.
39    pub has_mtp: bool,
40    pub ssm_architecture: SsmArchitecture,
41    pub attention_type: AttentionType,
42    /// 2026-09-26: Copied from `ModelConfig::nested_config`, which a parser sets when the text
43    /// model's fields sit under a nested key such as `text_config`.
44    pub has_nested_config: bool,
45}
46
47impl ModelCapabilities {
48    pub fn from_config(config: &crate::ModelConfig) -> Self {
49        use crate::LayerType;
50
51        let has_ssm = config
52            .layer_types
53            .iter()
54            .any(|t| matches!(t, LayerType::LinearAttention));
55        let has_mamba2 = config.mamba_num_heads > 0 && config.mamba_head_dim > 0;
56        let has_attention = config
57            .layer_types
58            .iter()
59            .any(|t| matches!(t, LayerType::FullAttention));
60        let has_moe = config.num_experts > 0;
61        let has_vision = config.vision.is_some();
62        let has_mtp = config.mtp_num_hidden_layers > 0;
63        let has_nested = config.nested_config;
64
65        let ssm_arch = if has_mamba2 {
66            SsmArchitecture::Mamba2
67        } else if has_ssm {
68            SsmArchitecture::Gdn
69        } else {
70            SsmArchitecture::None
71        };
72
73        Self {
74            has_ssm_layers: has_ssm || has_mamba2,
75            has_attention_layers: has_attention,
76            has_moe_layers: has_moe,
77            supports_thinking: has_ssm || has_mamba2,
78            supports_vision: has_vision,
79            has_mtp,
80            ssm_architecture: ssm_arch,
81            has_nested_config: has_nested,
82            attention_type: if config.kv_lora_rank > 0 {
83                AttentionType::Mla
84            } else {
85                AttentionType::Standard
86            },
87        }
88    }
89}
metrale_config/capabilities.rs

metrale_config/
capabilities.rs
