metrale_config/
methods.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Derived getters and small predicates over `ModelConfig` fields.
4//!
5//! Owner: config.
6//! Invariants: none beyond the types.
7
8#![allow(unused_imports)]
9
10use super::{AttnPositionEncoding, LayerType, ModelConfig};
11
12impl ModelConfig {
13    /// 2026-09-26: Every configured stop-token id, primary first; `[eos_token_id]` when
14    /// `eos_token_ids` is empty.
15    pub fn eos_ids(&self) -> Vec<u32> {
16        if self.eos_token_ids.is_empty() {
17            vec![self.eos_token_id]
18        } else {
19            self.eos_token_ids.clone()
20        }
21    }
22
23    pub fn is_eos(&self, id: u32) -> bool {
24        if self.eos_token_ids.is_empty() {
25            id == self.eos_token_id
26        } else {
27            self.eos_token_ids.contains(&id)
28        }
29    }
30
31    /// 2026-09-26: Q heads per KV head; 1 when `num_key_value_heads` is 0.
32    pub fn gqa_ratio(&self) -> usize {
33        self.num_attention_heads
34            .checked_div(self.num_key_value_heads)
35            .unwrap_or(1)
36    }
37
38    /// 2026-09-26: Layer kind for any index in the checkpoint. Text-stack indices go through
39    /// `layer_type`; indices `>= num_hidden_layers` are MTP/NextN layers and resolve through
40    /// `mtp_layer_types`, and past both the answer is `None`.
41    pub fn layer_type_at(&self, layer_idx: usize) -> Option<LayerType> {
42        if layer_idx < self.num_hidden_layers {
43            return Some(self.layer_type(layer_idx));
44        }
45        self.mtp_layer_types
46            .get(layer_idx - self.num_hidden_layers)
47            .copied()
48    }
49
50    /// 2026-09-26: Text-stack layers whose mixer is `LayerType::SparseAttention`.
51    pub fn sparse_attention_layers(&self) -> Vec<usize> {
52        self.layer_types
53            .iter()
54            .enumerate()
55            .filter(|(_, t)| **t == LayerType::SparseAttention)
56            .map(|(i, _)| i)
57            .collect()
58    }
59
60    /// 2026-09-26: True when any layer, in the text stack or the MTP layers, is sparse
61    /// attention.
62    pub fn has_sparse_attention(&self) -> bool {
63        self.layer_types.contains(&LayerType::SparseAttention)
64            || self.mtp_layer_types.contains(&LayerType::SparseAttention)
65    }
66
67    /// 2026-09-26: Layer kind of a text-stack index: `layer_types[layer_idx]` (`FullAttention`
68    /// past its end), or, when `layer_types` is empty, `FullAttention` on every
69    /// `full_attention_interval`-th layer and `LinearAttention` elsewhere.
70    pub fn layer_type(&self, layer_idx: usize) -> LayerType {
71        if !self.layer_types.is_empty() {
72            self.layer_types
73                .get(layer_idx)
74                .cloned()
75                .unwrap_or(LayerType::FullAttention)
76        } else if self.full_attention_interval > 0
77            && (layer_idx + 1).is_multiple_of(self.full_attention_interval)
78        {
79            LayerType::FullAttention
80        } else {
81            LayerType::LinearAttention
82        }
83    }
84
85    /// 2026-09-26: Number of layers that read the paged KV cache: every layer for which
86    /// `LayerType::is_attention` holds (full, sliding and sparse). Without `layer_types`,
87    /// `num_hidden_layers / full_attention_interval`, or every layer when the interval is 0.
88    pub fn num_attention_layers(&self) -> usize {
89        if !self.layer_types.is_empty() {
90            self.layer_types.iter().filter(|t| t.is_attention()).count()
91        } else {
92            self.num_hidden_layers
93                .checked_div(self.full_attention_interval)
94                .unwrap_or(self.num_hidden_layers)
95        }
96    }
97
98    pub fn num_ssm_layers(&self) -> usize {
99        if !self.layer_types.is_empty() {
100            self.layer_types
101                .iter()
102                .filter(|t| **t == LayerType::LinearAttention)
103                .count()
104        } else {
105            self.num_hidden_layers - self.num_attention_layers()
106        }
107    }
108
109    /// 2026-09-26: Whether the model carries recurrent (SSM / linear-attention) state:
110    /// `num_ssm_layers() > 0`. An SSM snapshot tier requested for a model without it is a
111    /// startup error (`ensure_ssm_tier_capability_from`).
112    pub fn has_recurrent_state(&self) -> bool {
113        self.num_ssm_layers() > 0
114    }
115
116    pub fn has_experts(&self) -> bool {
117        self.num_experts > 0
118    }
119
120    /// 2026-09-29: Dims of Q/K that RoPE rotates; 0 when attention has no positional
121    /// encoding, so a RoPE launch that skipped the `attn_position_encoding` check trips the
122    /// `rotary_dim > 0` assert in `ops::rope` instead of rotating silently.
123    pub fn rotary_dim(&self) -> usize {
124        if self.attn_position_encoding == Some(AttnPositionEncoding::None) {
125            0
126        } else if self.rotary_dim > 0 {
127            self.rotary_dim
128        } else {
129            (self.partial_rotary_factor * self.head_dim as f64) as usize
130        }
131    }
132
133    pub fn ssm_qkvz_size(&self) -> usize {
134        let q = self.linear_num_key_heads * self.linear_key_head_dim;
135        let k = self.linear_num_key_heads * self.linear_key_head_dim;
136        let v = self.linear_num_value_heads * self.linear_value_head_dim;
137        let z = self.linear_num_value_heads * self.linear_value_head_dim;
138        q + k + v + z
139    }
140
141    pub fn ssm_qkv_size(&self) -> usize {
142        let q = self.linear_num_key_heads * self.linear_key_head_dim;
143        let k = self.linear_num_key_heads * self.linear_key_head_dim;
144        let v = self.linear_num_value_heads * self.linear_value_head_dim;
145        q + k + v
146    }
147
148    pub fn ssm_z_size(&self) -> usize {
149        self.linear_num_value_heads * self.linear_value_head_dim
150    }
151
152    pub fn ssm_ba_size(&self) -> usize {
153        self.linear_num_value_heads * 2
154    }
155
156    /// 2026-09-26: `(start, end)`, end exclusive, of the experts this EP rank owns; the last
157    /// rank also takes the remainder of `num_experts / ep_world_size`.
158    pub fn local_expert_range(&self) -> (usize, usize) {
159        if self.ep_world_size <= 1 {
160            return (0, self.num_experts);
161        }
162        let per_rank = self.num_experts / self.ep_world_size;
163        let start = self.ep_rank * per_rank;
164        let end = if self.ep_rank == self.ep_world_size - 1 {
165            self.num_experts
166        } else {
167            start + per_rank
168        };
169        (start, end)
170    }
171
172    pub fn is_local_expert(&self, expert_id: usize) -> bool {
173        let (start, end) = self.local_expert_range();
174        expert_id >= start && expert_id < end
175    }
176
177    /// 2026-09-26: Range `[start, end)` of a `total`-sized dimension owned by this TP rank.
178    /// `total` must be divisible by `tp_world_size` (debug-asserted). `(0, total)` when
179    /// `tp_world_size <= 1`.
180    pub fn tp_shard_range(&self, total: usize) -> (usize, usize) {
181        if self.tp_world_size <= 1 {
182            return (0, total);
183        }
184        debug_assert!(
185            total.is_multiple_of(self.tp_world_size),
186            "tp_shard_range: total={} not divisible by tp_world_size={}",
187            total,
188            self.tp_world_size,
189        );
190        let per_rank = total / self.tp_world_size;
191        let start = self.tp_rank * per_rank;
192        (start, start + per_rank)
193    }
194
195    pub fn tp_shard_dim(&self, total: usize) -> usize {
196        if self.tp_world_size <= 1 {
197            return total;
198        }
199        total / self.tp_world_size
200    }
201
202    /// 2026-09-26: `{weight_prefix}.layers.{layer_idx}`, with `model` when `weight_prefix`
203    /// is empty.
204    pub fn layer_prefix(&self, layer_idx: usize) -> String {
205        if self.weight_prefix.is_empty() {
206            format!("model.layers.{layer_idx}")
207        } else {
208            format!("{}.layers.{layer_idx}", self.weight_prefix)
209        }
210    }
211
212    pub fn capabilities(&self) -> crate::capabilities::ModelCapabilities {
213        crate::capabilities::ModelCapabilities::from_config(self)
214    }
215
216    pub fn is_qwen35(&self) -> bool {
217        self.model_type == "qwen3_5_moe"
218    }
219
220    pub fn is_qwen35_dense(&self) -> bool {
221        self.model_type == "qwen3_5" && self.num_experts == 0
222    }
223
224    pub fn is_qwen3_vl(&self) -> bool {
225        if self.model_type == "qwen3_vl_moe" {
226            return true;
227        }
228        if self.model_type == "qwen3_5" && self.vision.is_some() {
229            return true;
230        }
231        false
232    }
233
234    /// 2026-09-26: Whether to keep the LM head in BF16 instead of quantizing it to NVFP4 at
235    /// load. `lm_head_bf16_override` (set by serve's `--lm-head-dtype`) wins when set;
236    /// otherwise yes for MLA models (`kv_lora_rank > 0`), for `laguna`, and for dense
237    /// `gemma4` unless `METRALE_GEMMA4_LMHEAD_NVFP4=1`, and no for the rest.
238    pub fn skip_lm_head_quantization(&self) -> bool {
239        if let Some(force_bf16) = self.lm_head_bf16_override {
240            return force_bf16;
241        }
242        if self.kv_lora_rank > 0 {
243            return true;
244        }
245        if self.model_type == "laguna" {
246            return true;
247        }
248        if self.model_type == "gemma4" && self.num_experts == 0 {
249            return std::env::var("METRALE_GEMMA4_LMHEAD_NVFP4").ok().as_deref() != Some("1");
250        }
251        false
252    }
253
254    pub fn mamba2_d_inner(&self) -> usize {
255        self.mamba_num_heads * self.mamba_head_dim
256    }
257
258    pub fn mamba2_d_xbc(&self) -> usize {
259        self.mamba2_d_inner() + 2 * self.n_groups * self.ssm_state_size
260    }
261
262    pub fn mamba2_in_proj_size(&self) -> usize {
263        self.mamba2_d_inner() + self.mamba2_d_xbc() + self.mamba_num_heads
264    }
265
266    /// 2026-09-29: Whether the SSM geometry is Mamba-2 (`mamba_*` heads) rather than
267    /// GatedDeltaNet (`linear_*` heads).
268    fn mamba2_geometry(&self) -> bool {
269        self.mamba_num_heads > 0 && self.mamba_head_dim > 0
270    }
271
272    /// 2026-09-29: Whether the model has Mamba-2 SSM layers (Nemotron-H).
273    pub fn has_mamba2_layers(&self) -> bool {
274        self.num_ssm_layers() > 0 && self.mamba2_geometry()
275    }
276
277    pub fn ssm_h_state_bytes(&self) -> usize {
278        if self.mamba2_geometry() {
279            self.mamba_num_heads * self.mamba_head_dim * self.ssm_state_size * 4
280        } else {
281            self.linear_num_value_heads * self.linear_value_head_dim * self.linear_key_head_dim * 4
282        }
283    }
284
285    pub fn ssm_conv_state_bytes(&self) -> usize {
286        let d_conv = self.linear_conv_kernel_dim;
287        if self.mamba2_geometry() {
288            self.mamba2_d_xbc() * d_conv * 4
289        } else {
290            let conv_dim = self.linear_num_key_heads * self.linear_key_head_dim * 2
291                + self.linear_num_value_heads * self.linear_value_head_dim;
292            conv_dim * d_conv * 4
293        }
294    }
295
296    pub fn ssm_state_norm_dims(&self) -> (usize, usize, usize) {
297        if self.mamba_num_heads > 0 && self.mamba_head_dim > 0 {
298            (
299                self.mamba_num_heads,
300                self.mamba_head_dim,
301                self.ssm_state_size,
302            )
303        } else {
304            (
305                self.linear_num_value_heads,
306                self.linear_key_head_dim,
307                self.linear_value_head_dim,
308            )
309        }
310    }
311
312    pub fn moe_input_size(&self) -> usize {
313        if self.moe_latent_size > 0 {
314            self.moe_latent_size
315        } else {
316            self.hidden_size
317        }
318    }
319
320    /// 2026-09-26: Routed-expert intermediate size for `layer`: its entry in
321    /// `moe_intermediate_sizes` when present and non-zero (Puzzle schedules), else the scalar.
322    pub fn moe_intermediate_size_for(&self, layer: usize) -> usize {
323        self.moe_intermediate_sizes
324            .get(layer)
325            .copied()
326            .filter(|&s| s > 0)
327            .unwrap_or(self.moe_intermediate_size)
328    }
329
330    /// 2026-09-26: Top-k experts per token for `layer`, by the same rule over
331    /// `num_experts_per_toks`.
332    pub fn num_experts_per_tok_for(&self, layer: usize) -> usize {
333        self.num_experts_per_toks
334            .get(layer)
335            .copied()
336            .filter(|&k| k > 0)
337            .unwrap_or(self.num_experts_per_tok)
338    }
339
340    pub fn max_moe_intermediate_size(&self) -> usize {
341        self.moe_intermediate_sizes
342            .iter()
343            .copied()
344            .max()
345            .unwrap_or(0)
346            .max(self.moe_intermediate_size)
347    }
348
349    pub fn num_moe_layers(&self) -> usize {
350        self.layer_types
351            .iter()
352            .filter(|t| **t == LayerType::Moe)
353            .count()
354    }
355}
metrale_config/methods.rs

metrale_config/
methods.rs
