metrale_config/
position_encoding.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: The positional encoding a model's full-attention layers apply to Q and K.
4//!
5//! Owner: config.
6//! Invariants:
7//! - A config returned by `parse_config` or `config_from_gguf` has
8//!   `attn_position_encoding` set (`resolve_attn_position_encoding`).
9//! - A config that declares no positional encoding carries no RoPE variant parameter
10//!   (explicit `rotary_dim`, MRoPE, YaRN); such a config is refused, not guessed at.
11
12use anyhow::{Context, Result, bail};
13
14use super::ModelConfig;
15
16/// 2026-09-29: What full attention does to Q and K before the dot product.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum AttnPositionEncoding {
19    /// 2026-09-29: Rotary embedding over `rotary_dim()` dims at `rope_theta`, or the layer's
20    /// RoPE variant (YaRN, proportional, MRoPE) where the loader configures one.
21    Rope,
22    /// 2026-09-29: No positional encoding: Q and K enter attention as projected. Nemotron-H
23    /// attention has none; position reaches it through the Mamba-2 layers only (HF
24    /// `NemotronHAttention.forward` never calls `apply_rotary_pos_emb`).
25    None,
26}
27
28impl ModelConfig {
29    /// 2026-09-29: The resolved positional encoding of full attention. An error for a config
30    /// that never went through `parse_config` / `config_from_gguf` and was not built by a
31    /// factory that sets it.
32    pub fn attn_position_encoding(&self) -> Result<AttnPositionEncoding> {
33        self.attn_position_encoding.with_context(|| {
34            format!(
35                "model_type `{}`: attention positional encoding was never resolved \
36                 (config not built by parse_config/config_from_gguf)",
37                self.model_type
38            )
39        })
40    }
41}
42
43/// 2026-09-29: Settle `attn_position_encoding` after the family parse. A family parser that
44/// knows its attention has no positional encoding declares `None`; any other family rotates
45/// Q/K by its rope fields, which is what `Rope` records. A `None` declaration next to a RoPE
46/// variant parameter is contradictory and refused.
47pub(crate) fn resolve_attn_position_encoding(config: &mut ModelConfig) -> Result<()> {
48    match config.attn_position_encoding {
49        Some(AttnPositionEncoding::None) => {
50            let conflicts: Vec<&str> = [
51                (config.rotary_dim > 0, "rotary_dim"),
52                (config.mrope_interleaved, "mrope_interleaved"),
53                (config.mrope_section.iter().any(|&s| s > 0), "mrope_section"),
54                (config.yarn_factor > 0.0, "yarn_factor"),
55            ]
56            .into_iter()
57            .filter_map(|(set, name)| set.then_some(name))
58            .collect();
59            if !conflicts.is_empty() {
60                bail!(
61                    "model_type `{}` declares no positional encoding for attention, but the \
62                     config also sets RoPE parameters: {}. Refusing an ambiguous config.",
63                    config.model_type,
64                    conflicts.join(", ")
65                );
66            }
67        }
68        Some(AttnPositionEncoding::Rope) => {}
69        None => config.attn_position_encoding = Some(AttnPositionEncoding::Rope),
70    }
71    Ok(())
72}
73
74#[cfg(test)]
75#[path = "position_encoding_tests.rs"]
76mod tests;
metrale_config/position_encoding.rs

metrale_config/
position_encoding.rs
