use std::num::NonZeroUsize;
use std::sync::Arc;
use frink_gguf::{GgufValue, TensorSource};
use crate::capability::SwaPattern;
use crate::loader::LoadError;
use crate::mtp_blocks::TrunkLayers;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SwaLayers {
All,
Period {
period: NonZeroUsize,
dense_first: bool,
},
PerLayer(Arc<[bool]>),
}
impl SwaLayers {
pub fn period(period: usize, dense_first: bool) -> Self {
match NonZeroUsize::new(period) {
Some(period) => Self::Period {
period,
dense_first,
},
None => Self::All,
}
}
pub fn from_default(layout: Option<SwaPattern>) -> Self {
match layout {
Some(p) => Self::period(p.period, p.dense_first),
None => Self::All,
}
}
#[inline]
pub fn slides(&self, layer_idx: usize) -> bool {
match self {
Self::All => true,
Self::Period {
period,
dense_first,
} => {
let period = period.get();
if *dense_first {
!layer_idx.is_multiple_of(period)
} else {
layer_idx % period < period - 1
}
}
Self::PerLayer(layers) => layers.get(layer_idx).copied().unwrap_or(false),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PatternKeyRead {
ScalarPeriod,
PerLayerBool,
ScalarThenArray,
}
pub const PER_LAYER_ARRAY_READERS: &[(&str, &str)] = &[
("gemma4", "src/models/gemma4.cpp:5"),
("gemma4-assistant", "src/models/gemma4-assistant.cpp:7"),
("dflash", "src/models/dflash.cpp:69"),
("step35", "src/models/step35.cpp:26"),
("mimo2", "src/models/mimo2.cpp:12"),
("spark2_5", "src/models/spark2-5.cpp:8"),
("maple", "src/models/maple.cpp:8"),
("granite_swa", "src/models/granite-swa.cpp:17"),
];
pub const ARRAY_AT_TRUNK_LENGTH: &[(&str, &str)] =
&[("cohere2moe", "src/models/cohere2moe.cpp:23,35")];
pub fn array_length(arch: &str, trunk: &TrunkLayers) -> usize {
if ARRAY_AT_TRUNK_LENGTH.iter().any(|(a, _)| *a == arch) {
trunk.n_layers
} else {
trunk.block_count
}
}
pub const SCALAR_THEN_ARRAY_READERS: &[(&str, &str)] = &[
("mellum", "src/models/mellum.cpp:12-17"),
("cohere2moe", "src/models/cohere2moe.cpp:32-36"),
("muse-glimmer", "src/models/muse-glimmer.cpp:26"),
];
pub fn pattern_key_read(arch: &str) -> PatternKeyRead {
if PER_LAYER_ARRAY_READERS.iter().any(|(a, _)| *a == arch) {
PatternKeyRead::PerLayerBool
} else if SCALAR_THEN_ARRAY_READERS.iter().any(|(a, _)| *a == arch) {
PatternKeyRead::ScalarThenArray
} else {
PatternKeyRead::ScalarPeriod
}
}
pub fn read_swa_layers(
file: &impl TensorSource,
arch: &str,
key: &str,
trunk: &TrunkLayers,
seeded: Option<SwaPattern>,
) -> Result<SwaLayers, LoadError> {
let value = file.metadata(key);
let period_from_scalar = |v: &GgufValue| -> Result<SwaLayers, LoadError> {
let period = v.as_u64().ok_or_else(|| {
LoadError::UnsupportedFeature(
arch.to_string(),
format!("{key} is neither an unsigned integer nor an array: {v:?}"),
)
})?;
Ok(SwaLayers::period(
period as usize,
seeded.is_some_and(|p| p.dense_first),
))
};
let per_layer = |items: &[GgufValue]| -> Result<SwaLayers, LoadError> {
let want = array_length(arch, trunk);
if items.len() != want {
return Err(LoadError::UnsupportedFeature(
arch.to_string(),
format!(
"{key} has {} entries where llama.cpp's read passes n_layer() = {want} \
(block_count {}, trunk {}); llama.cpp refuses this too (`key has wrong \
array length`, llama-model-loader.cpp:461-465)",
items.len(),
trunk.block_count,
trunk.n_layers
),
));
}
let mut out = Vec::with_capacity(trunk.n_layers);
for (il, item) in items.iter().take(trunk.n_layers).enumerate() {
out.push(item.as_bool().ok_or_else(|| {
LoadError::UnsupportedFeature(
arch.to_string(),
format!("{key} entry {il} is not a bool or integer: {item:?}"),
)
})?);
}
Ok(SwaLayers::PerLayer(out.into()))
};
match pattern_key_read(arch) {
PatternKeyRead::ScalarPeriod => match value {
None | Some(GgufValue::Array(_)) => Ok(SwaLayers::from_default(seeded)),
Some(scalar) => period_from_scalar(scalar),
},
PatternKeyRead::ScalarThenArray => match value {
None => Err(LoadError::MissingHparam(key.to_string())),
Some(GgufValue::Array(items)) => per_layer(items),
Some(scalar) => period_from_scalar(scalar),
},
PatternKeyRead::PerLayerBool => match value {
None => Err(LoadError::MissingHparam(key.to_string())),
Some(GgufValue::Array(items)) => per_layer(items),
Some(scalar) => {
let every = scalar.as_bool().ok_or_else(|| {
LoadError::UnsupportedFeature(
arch.to_string(),
format!("{key} is neither a bool, an integer nor an array: {scalar:?}"),
)
})?;
Ok(SwaLayers::PerLayer(vec![every; trunk.n_layers].into()))
}
},
}
}
#[cfg(test)]
mod tests {
use super::*;
fn trunk(block_count: usize, n_layers: usize) -> TrunkLayers {
TrunkLayers {
block_count,
n_layers,
n_mtp_blocks: block_count - n_layers,
}
}
struct Meta(Vec<(String, GgufValue)>);
impl TensorSource for Meta {
fn metadata(&self, key: &str) -> Option<&GgufValue> {
self.0.iter().find(|(k, _)| k == key).map(|(_, v)| v)
}
fn find_tensor(&self, _: &str) -> Option<&frink_gguf::TensorInfo> {
None
}
fn tensor_bytes(&self, name: &str) -> Result<&[u8], frink_gguf::GgufError> {
Err(frink_gguf::GgufError::TensorNotFound(name.to_string()))
}
fn tensor_mapped_range(
&self,
name: &str,
) -> Result<(Arc<frink_gguf::MmapHandle>, std::ops::Range<usize>), frink_gguf::GgufError>
{
Err(frink_gguf::GgufError::TensorNotFound(name.to_string()))
}
}
const KEY: &str = "x.attention.sliding_window_pattern";
fn with(value: Option<GgufValue>) -> Meta {
Meta(value.into_iter().map(|v| (KEY.to_string(), v)).collect())
}
fn bools(v: &[bool]) -> GgufValue {
GgufValue::Array(v.iter().map(|&b| GgufValue::Bool(b)).collect())
}
const LAST_DENSE_4: Option<SwaPattern> = Some(SwaPattern {
period: 4,
dense_first: false,
});
const DENSE_FIRST_4: Option<SwaPattern> = Some(SwaPattern {
period: 4,
dense_first: true,
});
#[test]
fn period_matches_set_swa_pattern_in_both_phases() {
let last_dense = SwaLayers::period(4, false);
let dense_first = SwaLayers::period(4, true);
let got: Vec<(bool, bool)> = (0..8)
.map(|il| (last_dense.slides(il), dense_first.slides(il)))
.collect();
let want: Vec<(bool, bool)> = (0..8u32).map(|il| (il % 4 < 3, il % 4 != 0)).collect();
assert_eq!(got, want);
assert_eq!(SwaLayers::period(0, false), SwaLayers::All);
assert!((0..8).all(|il| SwaLayers::All.slides(il)));
assert!((0..8).all(|il| !SwaLayers::period(1, false).slides(il)));
assert!((0..8).all(|il| !SwaLayers::period(1, true).slides(il)));
}
#[test]
fn per_layer_answers_false_past_the_trunk() {
let layers = SwaLayers::PerLayer(vec![true, false].into());
assert!(layers.slides(0));
assert!(!layers.slides(1));
assert!(!layers.slides(2));
}
#[test]
fn every_listed_reader_answers_its_mode() {
for (arch, _) in PER_LAYER_ARRAY_READERS {
assert_eq!(
pattern_key_read(arch),
PatternKeyRead::PerLayerBool,
"{arch}"
);
}
for (arch, _) in SCALAR_THEN_ARRAY_READERS {
assert_eq!(
pattern_key_read(arch),
PatternKeyRead::ScalarThenArray,
"{arch}"
);
}
for arch in ["exaone4", "exaone-moe", "olmo2", "gemma3", "llama"] {
assert_eq!(
pattern_key_read(arch),
PatternKeyRead::ScalarPeriod,
"{arch}"
);
}
}
#[test]
fn a_scalar_mode_architecture_ignores_the_array_and_keeps_its_seed() {
let file = with(Some(bools(&[false, false, false, true])));
let got = read_swa_layers(&file, "exaone-moe", KEY, &trunk(4, 4), LAST_DENSE_4).unwrap();
assert_eq!(got, SwaLayers::period(4, false));
let got =
read_swa_layers(&with(None), "exaone-moe", KEY, &trunk(4, 4), LAST_DENSE_4).unwrap();
assert_eq!(got, SwaLayers::period(4, false));
let file = with(Some(GgufValue::U32(2)));
let got = read_swa_layers(&file, "exaone-moe", KEY, &trunk(4, 4), LAST_DENSE_4).unwrap();
assert_eq!(got, SwaLayers::period(2, false));
}
#[test]
fn an_array_mode_architecture_honours_the_array_at_block_count_length() {
let file = with(Some(bools(&[false, true, true, false, true])));
let got = read_swa_layers(&file, "mimo2", KEY, &trunk(5, 4), None).unwrap();
assert_eq!(
got,
SwaLayers::PerLayer(vec![false, true, true, false].into())
);
let file = with(Some(bools(&[false, true, true, false])));
assert!(matches!(
read_swa_layers(&file, "mimo2", KEY, &trunk(5, 4), None),
Err(LoadError::UnsupportedFeature(a, m)) if a == "mimo2" && m.contains("4 entries where llama.cpp's read passes n_layer() = 5")
));
assert!(matches!(
read_swa_layers(&with(None), "mimo2", KEY, &trunk(4, 4), None),
Err(LoadError::MissingHparam(k)) if k == KEY
));
}
#[test]
fn cohere2moe_s_array_is_trunk_length() {
assert_eq!(array_length("cohere2moe", &trunk(5, 4)), 4);
assert_eq!(array_length("mimo2", &trunk(5, 4)), 5);
let file = with(Some(bools(&[false, true, true, false])));
let got = read_swa_layers(&file, "cohere2moe", KEY, &trunk(5, 4), DENSE_FIRST_4).unwrap();
assert_eq!(
got,
SwaLayers::PerLayer(vec![false, true, true, false].into())
);
let file = with(Some(bools(&[false, true, true, false, true])));
assert!(matches!(
read_swa_layers(&file, "cohere2moe", KEY, &trunk(5, 4), DENSE_FIRST_4),
Err(LoadError::UnsupportedFeature(a, m)) if a == "cohere2moe" && m.contains("n_layer() = 4")
));
}
#[test]
fn an_array_mode_architecture_broadcasts_a_scalar_as_a_bool() {
let got = read_swa_layers(
&with(Some(GgufValue::U32(6))),
"step35",
KEY,
&trunk(3, 3),
None,
)
.unwrap();
assert_eq!(got, SwaLayers::PerLayer(vec![true; 3].into()));
let got = read_swa_layers(
&with(Some(GgufValue::U32(0))),
"step35",
KEY,
&trunk(3, 3),
None,
)
.unwrap();
assert_eq!(got, SwaLayers::PerLayer(vec![false; 3].into()));
}
#[test]
fn a_scalar_then_array_architecture_takes_either_shape_and_requires_one() {
let got = read_swa_layers(
&with(Some(GgufValue::U32(2))),
"mellum",
KEY,
&trunk(4, 4),
LAST_DENSE_4,
)
.unwrap();
assert_eq!(got, SwaLayers::period(2, false));
let got = read_swa_layers(
&with(Some(bools(&[true, true, false, true]))),
"mellum",
KEY,
&trunk(4, 4),
LAST_DENSE_4,
)
.unwrap();
assert_eq!(
got,
SwaLayers::PerLayer(vec![true, true, false, true].into())
);
assert!(matches!(
read_swa_layers(&with(None), "mellum", KEY, &trunk(4, 4), LAST_DENSE_4),
Err(LoadError::MissingHparam(k)) if k == KEY
));
}
}