ferrum-models 0.11.0

//! Qwen3.5 / Qwen3.6 safetensors inventory and manifest validation.
//!
//! This module only inspects safetensors metadata. It does not materialize
//! tensor data, so it can run before the W3 executor allocates model weights.

use std::collections::BTreeSet;
use std::fs;
use std::fs::File;
use std::path::{Path, PathBuf};

use ferrum_kernels::backend::Backend;
use ferrum_quantization::{Linear, WeightLoader};
use ferrum_types::{FerrumError, Result as FerrumResult};
use memmap2::Mmap;
use safetensors::SafeTensors;

use crate::qwen35_config::{
    Qwen35LayerType, Qwen35MlpKind, Qwen35TextConfig, Qwen35WeightManifest, Qwen35WeightSpec,
};

const PREFIX_CANDIDATES: &[&str] = &["model.language_model", "model"];
const QWEN35_MOE_EXPERT_GPTQ_PROJECTIONS: &[&str] = &["gate_proj", "up_proj", "down_proj"];

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Qwen35WeightInventory {
    names: BTreeSet<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Qwen35WeightValidation {
    pub prefix: String,
    pub missing_required: Vec<String>,
    pub present_required: Vec<String>,
    pub present_optional: Vec<String>,
    pub missing_optional: Vec<String>,
}

/// Complete, pre-allocation classification of every tensor advertised by a
/// Qwen3.5-family checkpoint.  Production block-quantized sources use this
/// partition to prove that text-only execution did not silently ignore an
/// archive member.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct Qwen35TensorPartition {
    pub execution: Vec<String>,
    pub vision_non_executed: Vec<String>,
    pub mtp_non_executed: Vec<String>,
    pub unknown: Vec<String>,
}

impl Qwen35TensorPartition {
    pub fn require_no_unknown(&self) -> Result<(), String> {
        if self.unknown.is_empty() {
            return Ok(());
        }
        Err(format!(
            "Qwen3.5 checkpoint tensor partition contains {} unknown tensors; first: {}",
            self.unknown.len(),
            self.unknown
                .iter()
                .take(12)
                .cloned()
                .collect::<Vec<_>>()
                .join(", ")
        ))
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Qwen35ResolvedWeightSpec {
    pub role: String,
    pub name: String,
    /// Numeric position captured from an expert wildcard. String ordering is
    /// not semantic here: expert 10 must never precede expert 2.
    pub expert_index: Option<u32>,
    pub source: Option<Qwen35ResolvedWeightSource>,
    pub required: bool,
    pub present: bool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Qwen35ResolvedWeightSource {
    Dense {
        values: String,
    },
    BlockFp8 {
        values: String,
        scale_inv: String,
    },
    Gptq {
        qweight: String,
        scales: String,
        qzeros: String,
        g_idx: Option<String>,
    },
    CompressedTensors {
        weight_packed: String,
        weight_scale: String,
        weight_zero_point: String,
        weight_shape: String,
    },
}

impl Qwen35ResolvedWeightSource {
    pub fn primary_name(&self) -> &str {
        match self {
            Self::Dense { values } => values,
            Self::BlockFp8 { values, .. } => values,
            Self::Gptq { qweight, .. } => qweight,
            Self::CompressedTensors { weight_packed, .. } => weight_packed,
        }
    }

    pub fn component_names(&self) -> Vec<&str> {
        match self {
            Self::Dense { values } => vec![values],
            Self::BlockFp8 { values, scale_inv } => vec![values, scale_inv],
            Self::Gptq {
                qweight,
                scales,
                qzeros,
                g_idx,
            } => {
                let mut names = vec![qweight.as_str(), scales.as_str(), qzeros.as_str()];
                names.extend(g_idx.as_deref());
                names
            }
            Self::CompressedTensors {
                weight_packed,
                weight_scale,
                weight_zero_point,
                weight_shape,
            } => vec![weight_packed, weight_scale, weight_zero_point, weight_shape],
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Qwen35ResolvedLayerWeights {
    pub layer_index: usize,
    pub attention: Qwen35LayerType,
    pub mlp: Qwen35MlpKind,
    pub tensors: Vec<Qwen35ResolvedWeightSpec>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Qwen35ResolvedWeightPlan {
    pub prefix: String,
    pub global_tensors: Vec<Qwen35ResolvedWeightSpec>,
    pub layers: Vec<Qwen35ResolvedLayerWeights>,
}

pub struct Qwen35WeightPlanLoader<'a, B: Backend> {
    plan: &'a Qwen35ResolvedWeightPlan,
    inner: &'a dyn WeightLoader<B>,
}

impl<'a, B: Backend> Qwen35WeightPlanLoader<'a, B> {
    pub fn new(plan: &'a Qwen35ResolvedWeightPlan, inner: &'a dyn WeightLoader<B>) -> Self {
        Self { plan, inner }
    }

    pub fn plan(&self) -> &'a Qwen35ResolvedWeightPlan {
        self.plan
    }

    pub fn has_global_tensor(&self, role: &str) -> bool {
        self.plan
            .global_tensor(role)
            .is_some_and(|tensor| tensor.present && self.inner.has_tensor(&tensor.name))
    }

    pub fn has_layer_tensor(&self, layer_index: usize, role: &str) -> bool {
        self.plan
            .layer_tensor(layer_index, role)
            .is_some_and(|tensor| self.inner.has_tensor(&tensor.name))
    }

    pub fn load_global_tensor(&self, role: &str) -> FerrumResult<B::Buffer> {
        let tensor = self.required_global_tensor(role)?;
        self.inner.load_tensor(&tensor.name)
    }

    pub fn load_layer_tensor(&self, layer_index: usize, role: &str) -> FerrumResult<B::Buffer> {
        let tensor = self.required_layer_tensor(layer_index, role)?;
        self.inner.load_tensor(&tensor.name)
    }

    pub fn load_global_linear(&self, role: &str) -> FerrumResult<Box<dyn Linear<B>>> {
        let tensor = self.required_global_tensor(role)?;
        self.inner.load_linear(&linear_module_name(&tensor.name))
    }

    pub fn load_layer_linear(
        &self,
        layer_index: usize,
        role: &str,
    ) -> FerrumResult<Box<dyn Linear<B>>> {
        let tensor = self.required_layer_tensor(layer_index, role)?;
        self.inner.load_linear(&linear_module_name(&tensor.name))
    }

    pub fn load_layer_dense_gate_up_linear(
        &self,
        layer_index: usize,
    ) -> FerrumResult<Box<dyn Linear<B>>> {
        let gate = self.required_layer_tensor(layer_index, "mlp_gate")?;
        let up = self.required_layer_tensor(layer_index, "mlp_up")?;
        let gate_module = linear_module_name(&gate.name);
        let up_module = linear_module_name(&up.name);
        let prefix = gate_module.strip_suffix("gate_proj").ok_or_else(|| {
            FerrumError::model(format!(
                "Qwen3.5 dense MLP gate tensor for layer {layer_index} does not end with \
                 gate_proj: {gate_module}"
            ))
        })?;
        let expected_up = format!("{prefix}up_proj");
        if up_module != expected_up {
            return Err(FerrumError::model(format!(
                "Qwen3.5 dense MLP gate/up tensors for layer {layer_index} do not share a \
                 fusion prefix: gate={gate_module} up={up_module}"
            )));
        }

        self.inner.load_linear(&format!("{prefix}gate_up_proj"))
    }

    pub fn load_layer_shared_expert_gate_up_linear(
        &self,
        layer_index: usize,
    ) -> FerrumResult<Box<dyn Linear<B>>> {
        let gate = self.required_layer_tensor(layer_index, "moe_shared_expert_gate_proj")?;
        let up = self.required_layer_tensor(layer_index, "moe_shared_expert_up_proj")?;
        let gate_module = linear_module_name(&gate.name);
        let up_module = linear_module_name(&up.name);
        let prefix = gate_module.strip_suffix("gate_proj").ok_or_else(|| {
            FerrumError::model(format!(
                "Qwen3.5 shared expert gate tensor for layer {layer_index} does not end with \
                 gate_proj: {gate_module}"
            ))
        })?;
        let expected_up = format!("{prefix}up_proj");
        if up_module != expected_up {
            return Err(FerrumError::model(format!(
                "Qwen3.5 shared expert gate/up tensors for layer {layer_index} do not share a \
                 fusion prefix: gate={gate_module} up={up_module}"
            )));
        }

        self.inner.load_linear(&format!("{prefix}gate_up_proj"))
    }

    pub fn load_layer_linear_attention_qkvz(
        &self,
        layer_index: usize,
    ) -> FerrumResult<Box<dyn Linear<B>>> {
        self.load_fused_layer_linear(
            layer_index,
            "linear_attn_qkv",
            "linear_attn_z",
            "in_proj_qkv",
            "in_proj_z",
            "in_proj_qkvz",
        )
    }

    pub fn load_layer_linear_attention_ba(
        &self,
        layer_index: usize,
    ) -> FerrumResult<Box<dyn Linear<B>>> {
        self.load_fused_layer_linear(
            layer_index,
            "linear_attn_b",
            "linear_attn_a",
            "in_proj_b",
            "in_proj_a",
            "in_proj_ba",
        )
    }

    pub fn load_layer_stacked_gptq_experts(
        &self,
        layer_index: usize,
        num_experts: usize,
        proj_names: &[&str],
    ) -> FerrumResult<(
        std::sync::Arc<dyn ferrum_kernels::MarlinExpertStack<B>>,
        usize,
        usize,
    )> {
        let expert_prefix = format!(
            "{}.layers.{layer_index}.mlp.experts.{{e}}.",
            self.plan.prefix
        );
        self.inner
            .load_stacked_gptq_experts(&expert_prefix, num_experts, proj_names)
    }

    fn required_global_tensor(&self, role: &str) -> FerrumResult<&Qwen35ResolvedWeightSpec> {
        let tensor = self.plan.global_tensor(role).ok_or_else(|| {
            FerrumError::model(format!(
                "Qwen3.5 resolved weight plan has no global tensor role {role:?}"
            ))
        })?;
        if !tensor.present {
            return Err(FerrumError::model(format!(
                "Qwen3.5 global tensor role {role:?} is absent: {}",
                tensor.name
            )));
        }
        Ok(tensor)
    }

    fn required_layer_tensor(
        &self,
        layer_index: usize,
        role: &str,
    ) -> FerrumResult<&Qwen35ResolvedWeightSpec> {
        self.plan.layer_tensor(layer_index, role).ok_or_else(|| {
            FerrumError::model(format!(
                "Qwen3.5 resolved weight plan has no present layer tensor role {role:?} at layer {layer_index}"
            ))
        })
    }

    fn load_fused_layer_linear(
        &self,
        layer_index: usize,
        first_role: &str,
        second_role: &str,
        first_suffix: &str,
        second_suffix: &str,
        fused_suffix: &str,
    ) -> FerrumResult<Box<dyn Linear<B>>> {
        let first = self.required_layer_tensor(layer_index, first_role)?;
        let second = self.required_layer_tensor(layer_index, second_role)?;
        let first_module = linear_module_name(&first.name);
        let second_module = linear_module_name(&second.name);
        let prefix = first_module.strip_suffix(first_suffix).ok_or_else(|| {
            FerrumError::model(format!(
                "Qwen3.5 layer {layer_index} tensor for role {first_role:?} does not end with \
                 {first_suffix}: {first_module}"
            ))
        })?;
        let expected_second = format!("{prefix}{second_suffix}");
        if second_module != expected_second {
            return Err(FerrumError::model(format!(
                "Qwen3.5 layer {layer_index} fused projection roles {first_role:?}/{second_role:?} \
                 do not share a fusion prefix: first={first_module} second={second_module}"
            )));
        }

        self.inner.load_linear(&format!("{prefix}{fused_suffix}"))
    }
}

impl Qwen35WeightInventory {
    pub fn from_names(names: impl IntoIterator<Item = impl Into<String>>) -> Self {
        Self {
            names: names.into_iter().map(Into::into).collect(),
        }
    }

    pub fn from_safetensors_dir(model_dir: &Path) -> Result<Self, String> {
        let single = model_dir.join("model.safetensors");
        if single.exists() {
            return Self::from_safetensors_files([single]);
        }

        let index = model_dir.join("model.safetensors.index.json");
        if index.exists() {
            return Self::from_safetensors_index(model_dir, &index);
        }

        let mut files = fs::read_dir(model_dir)
            .map_err(|err| format!("read_dir {model_dir:?}: {err}"))?
            .filter_map(|entry| entry.ok().map(|entry| entry.path()))
            .filter(|path| path.extension().is_some_and(|ext| ext == "safetensors"))
            .collect::<Vec<_>>();
        files.sort();
        if files.is_empty() {
            Err(format!("no safetensors files found in {model_dir:?}"))
        } else {
            Self::from_safetensors_files(files)
        }
    }

    pub fn tensor_names(&self) -> impl Iterator<Item = &str> {
        self.names.iter().map(String::as_str)
    }

    pub fn contains(&self, name: &str) -> bool {
        self.names.contains(name)
    }

    pub fn validate_manifest(&self, manifest: &Qwen35WeightManifest) -> Qwen35WeightValidation {
        let mut missing_required = Vec::new();
        let mut present_required = Vec::new();
        let mut present_optional = Vec::new();
        let mut missing_optional = Vec::new();

        for tensor in manifest.global_tensors.iter().chain(
            manifest
                .layers
                .iter()
                .flat_map(|layer| layer.tensors.iter()),
        ) {
            let present = self.contains_weight_spec(tensor);
            match (tensor.required, present) {
                (true, true) => present_required.push(tensor.name.clone()),
                (true, false) => missing_required.push(tensor.name.clone()),
                (false, true) => present_optional.push(tensor.name.clone()),
                (false, false) => missing_optional.push(tensor.name.clone()),
            }
        }

        Qwen35WeightValidation {
            prefix: manifest.prefix.clone(),
            missing_required,
            present_required,
            present_optional,
            missing_optional,
        }
    }

    pub fn resolve_manifest(
        &self,
        manifest: &Qwen35WeightManifest,
    ) -> Result<Qwen35ResolvedWeightPlan, String> {
        let validation = self.validate_manifest(manifest);
        if !validation.missing_required.is_empty() {
            return Err(format!(
                "missing {} required Qwen3.5/Qwen3.6 tensors for prefix {}: {}",
                validation.missing_required.len(),
                validation.prefix,
                validation
                    .missing_required
                    .iter()
                    .take(12)
                    .cloned()
                    .collect::<Vec<_>>()
                    .join(", ")
            ));
        }

        Ok(Qwen35ResolvedWeightPlan {
            prefix: manifest.prefix.clone(),
            global_tensors: self.resolve_weight_specs(&manifest.global_tensors)?,
            layers: manifest
                .layers
                .iter()
                .map(|layer| {
                    Ok(Qwen35ResolvedLayerWeights {
                        layer_index: layer.layer_index,
                        attention: layer.attention,
                        mlp: layer.mlp,
                        tensors: self.resolve_weight_specs(&layer.tensors)?,
                    })
                })
                .collect::<Result<Vec<_>, String>>()?,
        })
    }

    pub fn detect_prefix_and_validate(
        &self,
        config: &Qwen35TextConfig,
    ) -> Result<Qwen35WeightValidation, String> {
        let mut best: Option<Qwen35WeightValidation> = None;
        for prefix in PREFIX_CANDIDATES {
            let manifest = config.weight_manifest(*prefix)?;
            let validation = self.validate_manifest(&manifest);
            if validation.missing_required.is_empty() {
                self.validate_moe_gptq_expert_sidecars(config, &validation.prefix)?;
                return Ok(validation);
            }
            if best.as_ref().is_none_or(|current| {
                validation.missing_required.len() < current.missing_required.len()
            }) {
                best = Some(validation);
            }
        }
        let best = best.expect("PREFIX_CANDIDATES is non-empty");
        Err(format!(
            "missing {} required Qwen3.5/Qwen3.6 tensors for prefix {}: {}",
            best.missing_required.len(),
            best.prefix,
            best.missing_required
                .iter()
                .take(12)
                .cloned()
                .collect::<Vec<_>>()
                .join(", ")
        ))
    }

    pub fn detect_prefix_and_resolve(
        &self,
        config: &Qwen35TextConfig,
    ) -> Result<Qwen35ResolvedWeightPlan, String> {
        let mut best: Option<Qwen35WeightValidation> = None;
        for prefix in PREFIX_CANDIDATES {
            let manifest = config.weight_manifest(*prefix)?;
            let validation = self.validate_manifest(&manifest);
            if validation.missing_required.is_empty() {
                self.validate_moe_gptq_expert_sidecars(config, &validation.prefix)?;
                return self.resolve_manifest(&manifest);
            }
            if best.as_ref().is_none_or(|current| {
                validation.missing_required.len() < current.missing_required.len()
            }) {
                best = Some(validation);
            }
        }
        let best = best.expect("PREFIX_CANDIDATES is non-empty");
        Err(format!(
            "missing {} required Qwen3.5/Qwen3.6 tensors for prefix {}: {}",
            best.missing_required.len(),
            best.prefix,
            best.missing_required
                .iter()
                .take(12)
                .cloned()
                .collect::<Vec<_>>()
                .join(", ")
        ))
    }

    /// Partitions the *entire* archive inventory after a manifest has been
    /// resolved. This intentionally does not infer text-only support from a
    /// repository id or from missing shards: only the two explicit Qwen3.5
    /// auxiliary namespaces are accepted as typed, non-executed components.
    /// Everything else must already be referenced by the resolved execution
    /// plan or remains unknown and therefore fail-closed.
    pub fn partition_resolved_plan(
        &self,
        plan: &Qwen35ResolvedWeightPlan,
    ) -> Result<Qwen35TensorPartition, String> {
        let validation = plan.validation();
        if !validation.is_pass() {
            return Err(format!(
                "cannot partition unresolved Qwen3.5 plan for prefix {:?}: {} required tensors are missing",
                plan.prefix,
                validation.missing_required.len()
            ));
        }

        let execution_names = plan
            .global_tensors
            .iter()
            .chain(plan.layers.iter().flat_map(|layer| layer.tensors.iter()))
            .filter(|tensor| tensor.present)
            .filter_map(|tensor| tensor.source.as_ref())
            .flat_map(Qwen35ResolvedWeightSource::component_names)
            .collect::<BTreeSet<_>>();
        let mut partition = Qwen35TensorPartition {
            execution: Vec::new(),
            vision_non_executed: Vec::new(),
            mtp_non_executed: Vec::new(),
            unknown: Vec::new(),
        };
        for name in &self.names {
            if execution_names.contains(name.as_str()) {
                partition.execution.push(name.clone());
            } else if name.starts_with("model.visual.") {
                partition.vision_non_executed.push(name.clone());
            } else if name.starts_with("mtp.") {
                partition.mtp_non_executed.push(name.clone());
            } else {
                partition.unknown.push(name.clone());
            }
        }
        Ok(partition)
    }

    fn from_safetensors_files(files: impl IntoIterator<Item = PathBuf>) -> Result<Self, String> {
        let mut names = BTreeSet::new();
        for path in files {
            let file = File::open(&path).map_err(|err| format!("open {path:?}: {err}"))?;
            let mmap = unsafe { Mmap::map(&file).map_err(|err| format!("mmap {path:?}: {err}"))? };
            let safetensors =
                SafeTensors::deserialize(&mmap).map_err(|err| format!("parse {path:?}: {err}"))?;
            names.extend(safetensors.names().into_iter().map(|name| name.to_string()));
        }
        Ok(Self { names })
    }

    fn from_safetensors_index(model_dir: &Path, index: &Path) -> Result<Self, String> {
        let raw = fs::read_to_string(index).map_err(|err| format!("read {index:?}: {err}"))?;
        let value: serde_json::Value =
            serde_json::from_str(&raw).map_err(|err| format!("parse {index:?}: {err}"))?;
        let weight_map = value
            .get("weight_map")
            .and_then(serde_json::Value::as_object)
            .ok_or_else(|| format!("{index:?} missing weight_map"))?;
        let shard_files = weight_map
            .values()
            .filter_map(serde_json::Value::as_str)
            .collect::<BTreeSet<_>>();
        for shard in shard_files {
            let path = model_dir.join(shard);
            if !path.exists() {
                return Err(format!("missing safetensors shard {path:?}"));
            }
        }
        Ok(Self {
            names: weight_map.keys().cloned().collect(),
        })
    }

    fn contains_weight_spec(&self, tensor: &Qwen35WeightSpec) -> bool {
        !self.matching_names(tensor).is_empty()
    }

    fn validate_moe_gptq_expert_sidecars(
        &self,
        config: &Qwen35TextConfig,
        prefix: &str,
    ) -> Result<(), String> {
        let Some(moe) = &config.moe else {
            return Ok(());
        };
        let mut any_gptq_expert_tensor = false;
        let mut expected_sidecars = Vec::new();
        let mut gate_up_g_idx = Vec::new();
        let mut down_g_idx = Vec::new();
        for layer_index in config.sparse_moe_layers() {
            for expert in 0..moe.num_experts {
                for proj in QWEN35_MOE_EXPERT_GPTQ_PROJECTIONS {
                    let stem = qwen35_moe_expert_gptq_stem(prefix, layer_index, expert, proj);
                    let qweight = format!("{stem}.qweight");
                    let scales = format!("{stem}.scales");
                    let qzeros = format!("{stem}.qzeros");
                    let g_idx = format!("{stem}.g_idx");
                    let has_qweight = self.contains(&qweight);
                    let has_scales = self.contains(&scales);
                    let has_qzeros = self.contains(&qzeros);
                    let has_g_idx = self.contains(&g_idx);
                    let has_any = has_qweight || has_scales || has_qzeros || has_g_idx;
                    any_gptq_expert_tensor |= has_any;
                    expected_sidecars.push((qweight, has_qweight));
                    expected_sidecars.push((scales, has_scales));
                    expected_sidecars.push((qzeros, has_qzeros));
                    if *proj == "down_proj" {
                        down_g_idx.push((g_idx, has_g_idx));
                    } else {
                        gate_up_g_idx.push((g_idx, has_g_idx));
                    }
                }
            }
        }
        if !any_gptq_expert_tensor {
            return Ok(());
        }
        let missing = expected_sidecars
            .into_iter()
            .filter_map(|(name, present)| (!present).then_some(name))
            .collect::<Vec<_>>();
        if !missing.is_empty() {
            return Err(format!(
                "incomplete Qwen3.5 per-expert GPTQ sidecars for prefix {prefix}: missing {} tensors, first: {}",
                missing.len(),
                missing
                    .iter()
                    .take(12)
                    .cloned()
                    .collect::<Vec<_>>()
                    .join(", ")
            ));
        }
        for (label, entries) in [("gate/up", gate_up_g_idx), ("down", down_g_idx)] {
            let present = entries.iter().filter(|(_, present)| *present).count();
            if present != 0 && present != entries.len() {
                let first_missing = entries
                    .iter()
                    .filter_map(|(name, present)| (!*present).then_some(name.as_str()))
                    .take(8)
                    .collect::<Vec<_>>()
                    .join(", ");
                return Err(format!(
                    "inconsistent Qwen3.5 per-expert GPTQ g_idx sidecars for {label} stack under prefix {prefix}: present {present}/{}; first missing: {first_missing}",
                    entries.len()
                ));
            }
        }
        Ok(())
    }

    fn matching_names(&self, tensor: &Qwen35WeightSpec) -> Vec<(Option<u32>, String)> {
        if !tensor.name.contains('*') {
            let mut candidates = vec![tensor.name.clone()];
            if tensor.name.ends_with(".lm_head.weight") {
                candidates.push("lm_head.weight".to_string());
            }
            for candidate in candidates {
                if self.contains(&candidate) {
                    return vec![(None, candidate)];
                }
                if role_accepts_quantized_linear_alias(&tensor.role) {
                    if let Some(primary) = self.quantized_linear_primary_name(&candidate) {
                        return vec![(None, primary)];
                    }
                }
            }
            return Vec::new();
        }
        let mut pieces = tensor.name.splitn(2, '*');
        let prefix = pieces.next().unwrap_or("");
        let suffix = pieces.next().unwrap_or("");
        let mut matches = self
            .names
            .iter()
            .filter_map(|name| {
                let captured = name.strip_prefix(prefix)?.strip_suffix(suffix)?;
                let expert_index = captured.parse::<u32>().ok()?;
                (captured == expert_index.to_string()).then(|| (expert_index, name.clone()))
            })
            .collect::<Vec<_>>();
        matches.sort_by_key(|(expert_index, _)| *expert_index);
        matches
            .into_iter()
            .map(|(expert_index, name)| (Some(expert_index), name))
            .collect()
    }

    fn quantized_linear_primary_name(&self, dense_weight_name: &str) -> Option<String> {
        let module = dense_weight_name.strip_suffix(".weight")?;
        let qweight = format!("{module}.qweight");
        let scales = format!("{module}.scales");
        let qzeros = format!("{module}.qzeros");
        if self.contains(&qweight) && self.contains(&scales) && self.contains(&qzeros) {
            return Some(qweight);
        }
        let weight_packed = format!("{module}.weight_packed");
        let weight_scale = format!("{module}.weight_scale");
        let weight_zero_point = format!("{module}.weight_zero_point");
        let weight_shape = format!("{module}.weight_shape");
        // Preserve an observed primary sidecar even when its bundle is
        // incomplete. `resolved_source` can then report the exact missing
        // compressed-tensors tensor instead of degrading to a generic
        // missing dense-weight error.
        self.contains(&weight_packed).then_some(weight_packed)
    }

    fn resolve_weight_specs(
        &self,
        specs: &[Qwen35WeightSpec],
    ) -> Result<Vec<Qwen35ResolvedWeightSpec>, String> {
        specs
            .iter()
            .flat_map(|spec| {
                let matches = self.matching_names(spec);
                if matches.is_empty() {
                    return vec![Ok(Qwen35ResolvedWeightSpec {
                        role: spec.role.clone(),
                        name: spec.name.clone(),
                        expert_index: None,
                        source: None,
                        required: spec.required,
                        present: false,
                    })];
                }
                matches
                    .into_iter()
                    .map(|(expert_index, name)| {
                        let source = self.resolved_source(&name)?;
                        Ok(Qwen35ResolvedWeightSpec {
                            role: spec.role.clone(),
                            name: name.clone(),
                            expert_index,
                            source: Some(source),
                            required: spec.required,
                            present: true,
                        })
                    })
                    .collect()
            })
            .collect()
    }

    fn resolved_source(&self, name: &str) -> Result<Qwen35ResolvedWeightSource, String> {
        if let Some(module) = name.strip_suffix(".weight_packed") {
            let weight_scale = format!("{module}.weight_scale");
            let weight_zero_point = format!("{module}.weight_zero_point");
            let weight_shape = format!("{module}.weight_shape");
            let missing = [&weight_scale, &weight_zero_point, &weight_shape]
                .into_iter()
                .filter(|sidecar| !self.contains(sidecar))
                .cloned()
                .collect::<Vec<_>>();
            if !missing.is_empty() {
                return Err(format!(
                    "incomplete compressed-tensors source bundle for {name:?}; missing {}",
                    missing.join(", ")
                ));
            }
            return Ok(Qwen35ResolvedWeightSource::CompressedTensors {
                weight_packed: name.to_owned(),
                weight_scale,
                weight_zero_point,
                weight_shape,
            });
        }
        if let Some(module) = name.strip_suffix(".weight") {
            let scale_inv = format!("{module}.weight_scale_inv");
            if self.contains(&scale_inv) {
                return Ok(Qwen35ResolvedWeightSource::BlockFp8 {
                    values: name.to_owned(),
                    scale_inv,
                });
            }
        }
        let Some(module) = name.strip_suffix(".qweight") else {
            return Ok(Qwen35ResolvedWeightSource::Dense {
                values: name.to_owned(),
            });
        };
        let scales = format!("{module}.scales");
        let qzeros = format!("{module}.qzeros");
        let missing = [&scales, &qzeros]
            .into_iter()
            .filter(|sidecar| !self.contains(sidecar))
            .cloned()
            .collect::<Vec<_>>();
        if !missing.is_empty() {
            return Err(format!(
                "incomplete GPTQ source bundle for {name:?}; missing {}",
                missing.join(", ")
            ));
        }
        let g_idx = format!("{module}.g_idx");
        Ok(Qwen35ResolvedWeightSource::Gptq {
            qweight: name.to_owned(),
            scales,
            qzeros,
            g_idx: self.contains(&g_idx).then_some(g_idx),
        })
    }
}

fn qwen35_moe_expert_gptq_stem(
    prefix: &str,
    layer_index: usize,
    expert: usize,
    proj: &str,
) -> String {
    format!("{prefix}.layers.{layer_index}.mlp.experts.{expert}.{proj}")
}

impl Qwen35WeightValidation {
    pub fn is_pass(&self) -> bool {
        self.missing_required.is_empty()
    }
}

impl Qwen35ResolvedWeightPlan {
    pub fn validation(&self) -> Qwen35WeightValidation {
        let mut missing_required = Vec::new();
        let mut present_required = Vec::new();
        let mut present_optional = Vec::new();
        let mut missing_optional = Vec::new();

        for tensor in self
            .global_tensors
            .iter()
            .chain(self.layers.iter().flat_map(|layer| layer.tensors.iter()))
        {
            match (tensor.required, tensor.present) {
                (true, true) => present_required.push(tensor.name.clone()),
                (true, false) => missing_required.push(tensor.name.clone()),
                (false, true) => present_optional.push(tensor.name.clone()),
                (false, false) => missing_optional.push(tensor.name.clone()),
            }
        }

        Qwen35WeightValidation {
            prefix: self.prefix.clone(),
            missing_required,
            present_required,
            present_optional,
            missing_optional,
        }
    }

    pub fn layer_tensor(
        &self,
        layer_index: usize,
        role: &str,
    ) -> Option<&Qwen35ResolvedWeightSpec> {
        self.layers
            .iter()
            .find(|layer| layer.layer_index == layer_index)
            .and_then(|layer| {
                layer
                    .tensors
                    .iter()
                    .find(|tensor| tensor.role == role && tensor.present)
            })
    }

    pub fn layer_tensors<'plan>(
        &'plan self,
        layer_index: usize,
        role: &'plan str,
    ) -> impl Iterator<Item = &'plan Qwen35ResolvedWeightSpec> + 'plan {
        self.layers
            .iter()
            .find(|layer| layer.layer_index == layer_index)
            .into_iter()
            .flat_map(|layer| layer.tensors.iter())
            .filter(move |tensor| tensor.role == role && tensor.present)
    }

    pub fn global_tensor(&self, role: &str) -> Option<&Qwen35ResolvedWeightSpec> {
        self.global_tensors
            .iter()
            .find(|tensor| tensor.role == role)
    }
}

fn linear_module_name(tensor_name: &str) -> String {
    tensor_name
        .strip_suffix(".weight")
        .or_else(|| tensor_name.strip_suffix(".qweight"))
        .or_else(|| tensor_name.strip_suffix(".weight_packed"))
        .unwrap_or(tensor_name)
        .to_string()
}

fn role_accepts_quantized_linear_alias(role: &str) -> bool {
    matches!(
        role,
        "lm_head"
            | "linear_attn_qkv"
            | "linear_attn_z"
            | "linear_attn_b"
            | "linear_attn_a"
            | "linear_attn_out"
            | "self_attn_q"
            | "self_attn_k"
            | "self_attn_v"
            | "self_attn_o"
            | "mlp_gate"
            | "mlp_up"
            | "mlp_down"
            | "moe_router"
            | "moe_shared_expert_gate"
            | "moe_shared_expert_gate_proj"
            | "moe_shared_expert_up_proj"
            | "moe_shared_expert_down_proj"
    )
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::sync::Mutex;

    use ferrum_kernels::backend::cpu::CpuBackend;
    use ferrum_quantization::{DenseLinear, QuantConfig, WeightLoader};
    use ferrum_types::{FerrumError, Result as FerrumResult};
    use safetensors::tensor::{serialize_to_file, Dtype, TensorView};
    use tempfile::TempDir;

    use super::*;

    fn dense_config() -> Qwen35TextConfig {
        Qwen35TextConfig::from_hf_config_str(
            r#"{
              "model_type": "qwen3_5",
              "text_config": {
                "model_type": "qwen3_5_text",
                "hidden_size": 16,
                "num_hidden_layers": 4,
                "layer_types": ["linear_attention", "linear_attention", "linear_attention", "full_attention"],
                "linear_num_key_heads": 2,
                "linear_num_value_heads": 2,
                "linear_key_head_dim": 4,
                "linear_value_head_dim": 4,
                "linear_conv_kernel_dim": 4,
                "mamba_ssm_dtype": "float32",
                "head_dim": 4,
                "num_attention_heads": 2,
                "num_key_value_heads": 1,
                "intermediate_size": 32,
                "tie_word_embeddings": true
              }
            }"#,
        )
        .unwrap()
    }

    fn moe_config() -> Qwen35TextConfig {
        Qwen35TextConfig::from_hf_config_str(
            r#"{
              "model_type": "qwen3_5_moe",
              "text_config": {
                "model_type": "qwen3_5_moe_text",
                "hidden_size": 16,
                "num_hidden_layers": 4,
                "layer_types": ["linear_attention", "linear_attention", "linear_attention", "full_attention"],
                "linear_num_key_heads": 2,
                "linear_num_value_heads": 4,
                "linear_key_head_dim": 4,
                "linear_value_head_dim": 4,
                "linear_conv_kernel_dim": 4,
                "mamba_ssm_dtype": "float32",
                "head_dim": 4,
                "num_attention_heads": 2,
                "num_key_value_heads": 1,
                "num_experts": 8,
                "num_experts_per_tok": 2,
                "moe_intermediate_size": 8,
                "shared_expert_intermediate_size": 8,
                "tie_word_embeddings": false
              }
            }"#,
        )
        .unwrap()
    }

    fn write_safetensors(dir: &Path, names: &[String]) {
        let tensors = names
            .iter()
            .map(|name| {
                let bytes: &'static [u8] =
                    Box::leak(0.0f32.to_le_bytes().to_vec().into_boxed_slice());
                (
                    name.clone(),
                    TensorView::new(Dtype::F32, vec![1], bytes).unwrap(),
                )
            })
            .collect::<Vec<_>>();
        serialize_to_file(
            tensors,
            &None::<HashMap<String, String>>,
            &dir.join("model.safetensors"),
        )
        .unwrap();
    }

    struct RecordingLoader {
        tensors: HashMap<String, Vec<f32>>,
        linear_names: Mutex<Vec<String>>,
    }

    impl RecordingLoader {
        fn from_names(names: impl IntoIterator<Item = impl Into<String>>) -> Self {
            Self {
                tensors: names
                    .into_iter()
                    .map(|name| (name.into(), vec![1.0]))
                    .collect(),
                linear_names: Mutex::new(Vec::new()),
            }
        }

        fn linear_names(&self) -> Vec<String> {
            self.linear_names.lock().unwrap().clone()
        }
    }

    impl WeightLoader<CpuBackend> for RecordingLoader {
        fn load_tensor(&self, name: &str) -> FerrumResult<Vec<f32>> {
            self.tensors
                .get(name)
                .cloned()
                .ok_or_else(|| FerrumError::model(format!("missing tensor {name}")))
        }

        fn load_linear(
            &self,
            name: &str,
        ) -> FerrumResult<Box<dyn ferrum_quantization::Linear<CpuBackend>>> {
            self.linear_names.lock().unwrap().push(name.to_string());
            Ok(Box::new(DenseLinear::<CpuBackend>::from_rows(
                &[0.0, 0.0],
                1,
                2,
            )))
        }

        fn has_tensor(&self, name: &str) -> bool {
            self.tensors.contains_key(name)
        }

        fn quant_config(&self) -> Option<&QuantConfig> {
            None
        }
    }

    fn required_names_with_gptq_linear_aliases(
        config: &Qwen35TextConfig,
        prefix: &str,
    ) -> Vec<String> {
        let manifest = config.weight_manifest(prefix).unwrap();
        manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .flat_map(|tensor| {
                if role_accepts_quantized_linear_alias(&tensor.role) {
                    let module = tensor
                        .name
                        .strip_suffix(".weight")
                        .expect("quantized linear aliases are declared as .weight specs");
                    vec![
                        format!("{module}.qweight"),
                        format!("{module}.scales"),
                        format!("{module}.qzeros"),
                    ]
                } else {
                    vec![tensor.name.clone()]
                }
            })
            .collect()
    }

    fn append_complete_moe_expert_gptq_sidecars(
        names: &mut Vec<String>,
        config: &Qwen35TextConfig,
        prefix: &str,
        include_g_idx: bool,
    ) {
        let moe = config.moe.as_ref().unwrap();
        for layer_index in config.sparse_moe_layers() {
            for expert in 0..moe.num_experts {
                for proj in QWEN35_MOE_EXPERT_GPTQ_PROJECTIONS {
                    let stem = qwen35_moe_expert_gptq_stem(prefix, layer_index, expert, proj);
                    names.push(format!("{stem}.qweight"));
                    names.push(format!("{stem}.scales"));
                    names.push(format!("{stem}.qzeros"));
                    if include_g_idx {
                        names.push(format!("{stem}.g_idx"));
                    }
                }
            }
        }
    }

    #[test]
    fn validates_single_file_safetensors_against_manifest() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        let tmp = TempDir::new().unwrap();
        write_safetensors(tmp.path(), &names);

        let inventory = Qwen35WeightInventory::from_safetensors_dir(tmp.path()).unwrap();
        let validation = inventory.detect_prefix_and_validate(&config).unwrap();
        assert_eq!(validation.prefix, "model");
        assert!(validation.is_pass());
        assert!(validation
            .missing_optional
            .contains(&"model.lm_head.weight".to_string()));
    }

    #[test]
    fn reports_missing_required_tensor() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        names.retain(|name| name != "model.layers.3.self_attn.q_proj.weight");
        let inventory = Qwen35WeightInventory::from_names(names);
        let err = inventory
            .detect_prefix_and_validate(&config)
            .expect_err("missing full-attention q_proj should fail");
        assert!(err.contains("self_attn.q_proj.weight"), "{err}");
    }

    #[test]
    fn resolves_manifest_to_executor_tensor_plan() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        let inventory = Qwen35WeightInventory::from_names(names);

        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();
        let validation = plan.validation();

        assert_eq!(plan.prefix, "model");
        assert!(validation.is_pass());
        assert!(plan
            .global_tensors
            .iter()
            .any(|tensor| tensor.role == "embed_tokens"
                && tensor.name == "model.embed_tokens.weight"
                && tensor.present));
        assert!(plan
            .global_tensors
            .iter()
            .any(|tensor| tensor.role == "lm_head"
                && tensor.name == "model.lm_head.weight"
                && !tensor.present
                && !tensor.required));
        assert_eq!(
            plan.layer_tensor(0, "linear_attn_qkv")
                .map(|tensor| tensor.name.as_str()),
            Some("model.layers.0.linear_attn.in_proj_qkv.weight")
        );
    }

    #[test]
    fn resolves_required_gptq_linear_aliases_without_dense_weight_tensors() {
        let config = moe_config();
        let names = required_names_with_gptq_linear_aliases(&config, "model");
        assert!(!names.contains(&"model.layers.3.self_attn.q_proj.weight".to_string()));
        assert!(names.contains(&"model.layers.3.self_attn.q_proj.qweight".to_string()));
        assert!(names.contains(&"model.layers.3.self_attn.q_proj.scales".to_string()));
        assert!(names.contains(&"model.layers.3.self_attn.q_proj.qzeros".to_string()));
        let inventory = Qwen35WeightInventory::from_names(names.clone());

        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();
        let validation = plan.validation();

        assert_eq!(plan.prefix, "model");
        assert!(validation.is_pass());
        assert_eq!(
            plan.global_tensor("lm_head")
                .map(|tensor| tensor.name.as_str()),
            Some("model.lm_head.qweight")
        );
        assert_eq!(
            plan.layer_tensor(3, "self_attn_q")
                .map(|tensor| tensor.name.as_str()),
            Some("model.layers.3.self_attn.q_proj.qweight")
        );
        assert_eq!(
            plan.layer_tensor(0, "linear_attn_qkv")
                .map(|tensor| tensor.name.as_str()),
            Some("model.layers.0.linear_attn.in_proj_qkv.qweight")
        );
        assert_eq!(
            plan.layer_tensor(0, "moe_router")
                .map(|tensor| tensor.name.as_str()),
            Some("model.layers.0.mlp.gate.qweight")
        );
        assert_eq!(
            plan.layer_tensor(0, "moe_shared_expert_gate")
                .map(|tensor| tensor.name.as_str()),
            Some("model.layers.0.mlp.shared_expert_gate.qweight")
        );

        let loader = RecordingLoader::from_names(names);
        let planned = Qwen35WeightPlanLoader::<CpuBackend>::new(&plan, &loader);
        assert!(planned.has_global_tensor("lm_head"));
        assert!(planned.has_layer_tensor(3, "self_attn_q"));
        planned.load_global_linear("lm_head").unwrap();
        planned.load_layer_linear(3, "self_attn_q").unwrap();
        planned
            .load_layer_linear(0, "moe_shared_expert_gate")
            .unwrap();
        planned.load_layer_shared_expert_gate_up_linear(0).unwrap();

        assert_eq!(
            loader.linear_names(),
            vec![
                "model.lm_head".to_string(),
                "model.layers.3.self_attn.q_proj".to_string(),
                "model.layers.0.mlp.shared_expert_gate".to_string(),
                "model.layers.0.mlp.shared_expert.gate_up_proj".to_string(),
            ]
        );
    }

    #[test]
    fn rejects_incomplete_gptq_linear_alias() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        names.retain(|name| name != "model.layers.3.self_attn.q_proj.weight");
        names.push("model.layers.3.self_attn.q_proj.qweight".to_string());
        names.push("model.layers.3.self_attn.q_proj.scales".to_string());
        let inventory = Qwen35WeightInventory::from_names(names);

        let err = inventory
            .detect_prefix_and_resolve(&config)
            .expect_err("incomplete GPTQ q_proj alias must not satisfy manifest");

        assert!(err.contains("self_attn.q_proj.weight"), "{err}");
    }

    #[test]
    fn resolves_complete_compressed_tensors_linear_alias() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        let dense = "model.layers.3.self_attn.q_proj.weight";
        let stem = dense.strip_suffix(".weight").unwrap();
        names.retain(|name| name != dense);
        names.extend([
            format!("{stem}.weight_packed"),
            format!("{stem}.weight_scale"),
            format!("{stem}.weight_zero_point"),
            format!("{stem}.weight_shape"),
        ]);

        let plan = Qwen35WeightInventory::from_names(names)
            .detect_prefix_and_resolve(&config)
            .unwrap();
        let source = &plan.layer_tensor(3, "self_attn_q").unwrap().source;
        assert!(matches!(
            source,
            Some(Qwen35ResolvedWeightSource::CompressedTensors {
                weight_packed,
                weight_scale,
                weight_zero_point,
                weight_shape,
            }) if weight_packed == &format!("{stem}.weight_packed")
                && weight_scale == &format!("{stem}.weight_scale")
                && weight_zero_point == &format!("{stem}.weight_zero_point")
                && weight_shape == &format!("{stem}.weight_shape")
        ));
    }

    #[test]
    fn rejects_incomplete_compressed_tensors_bundle_with_exact_missing_sidecar() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        let dense = "model.layers.3.self_attn.q_proj.weight";
        let stem = dense.strip_suffix(".weight").unwrap();
        names.retain(|name| name != dense);
        names.extend([
            format!("{stem}.weight_packed"),
            format!("{stem}.weight_scale"),
            format!("{stem}.weight_zero_point"),
        ]);

        let error = Qwen35WeightInventory::from_names(names)
            .detect_prefix_and_resolve(&config)
            .expect_err("missing weight_shape must fail before weight allocation");
        assert!(
            error.contains("incomplete compressed-tensors source bundle"),
            "{error}"
        );
        assert!(error.contains("weight_shape"), "{error}");
    }

    #[test]
    fn resolves_optional_wildcard_expert_aliases_when_present() {
        let config = moe_config();
        let manifest = config.weight_manifest("model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        names.push("model.layers.0.mlp.experts.0.gate_proj.weight".to_string());
        let inventory = Qwen35WeightInventory::from_names(names);

        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();

        assert_eq!(
            plan.layer_tensor(0, "moe_per_expert_gate_proj")
                .map(|tensor| tensor.name.as_str()),
            Some("model.layers.0.mlp.experts.0.gate_proj.weight")
        );
        assert!(plan
            .validation()
            .present_optional
            .contains(&"model.layers.0.mlp.experts.0.gate_proj.weight".to_string()));
    }

    #[test]
    fn resolves_expert_wildcards_in_numeric_order() {
        let mut names = [10, 2, 1, 0]
            .into_iter()
            .flat_map(|expert| {
                ["qweight", "scales", "qzeros"].map(|component| {
                    format!("model.layers.0.mlp.experts.{expert}.gate_proj.{component}")
                })
            })
            .collect::<Vec<_>>();
        names.extend([
            "model.layers.0.mlp.experts.01.gate_proj.qweight".to_owned(),
            "model.layers.0.mlp.experts.invalid.gate_proj.qweight".to_owned(),
        ]);
        let inventory = Qwen35WeightInventory::from_names(names);
        let resolved = inventory
            .resolve_weight_specs(&[Qwen35WeightSpec {
                role: "moe_per_expert_gate_proj_qweight".to_owned(),
                name: "model.layers.0.mlp.experts.*.gate_proj.qweight".to_owned(),
                required: false,
            }])
            .unwrap();

        assert_eq!(
            resolved
                .iter()
                .map(|weight| weight.expert_index)
                .collect::<Vec<_>>(),
            [Some(0), Some(1), Some(2), Some(10)]
        );
        assert!(resolved
            .iter()
            .all(|weight| matches!(weight.source, Some(Qwen35ResolvedWeightSource::Gptq { .. }))));
        assert_eq!(
            resolved
                .iter()
                .map(|weight| weight.name.as_str())
                .collect::<Vec<_>>(),
            [
                "model.layers.0.mlp.experts.0.gate_proj.qweight",
                "model.layers.0.mlp.experts.1.gate_proj.qweight",
                "model.layers.0.mlp.experts.2.gate_proj.qweight",
                "model.layers.0.mlp.experts.10.gate_proj.qweight",
            ]
        );
    }

    #[test]
    fn validates_complete_per_expert_gptq_sidecars() {
        let config = moe_config();
        let mut names = required_names_with_gptq_linear_aliases(&config, "model");
        append_complete_moe_expert_gptq_sidecars(&mut names, &config, "model", true);
        let inventory = Qwen35WeightInventory::from_names(names);

        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();

        assert_eq!(
            plan.layer_tensor(0, "moe_per_expert_gate_proj_qweight")
                .map(|tensor| tensor.name.as_str()),
            Some("model.layers.0.mlp.experts.0.gate_proj.qweight")
        );
        let source = &plan
            .layer_tensor(0, "moe_per_expert_gate_proj_qweight")
            .unwrap()
            .source;
        assert!(matches!(
            source,
            Some(Qwen35ResolvedWeightSource::Gptq {
                qweight,
                scales,
                qzeros,
                g_idx: Some(g_idx),
            }) if qweight.ends_with(".qweight")
                && scales.ends_with(".scales")
                && qzeros.ends_with(".qzeros")
                && g_idx.ends_with(".g_idx")
        ));
        assert!(plan.validation().is_pass());
    }

    #[test]
    fn rejects_incomplete_per_expert_gptq_sidecars() {
        let config = moe_config();
        let mut names = required_names_with_gptq_linear_aliases(&config, "model");
        names.push("model.layers.0.mlp.experts.0.gate_proj.qweight".to_string());
        let inventory = Qwen35WeightInventory::from_names(names);

        let err = inventory
            .detect_prefix_and_resolve(&config)
            .expect_err("partial per-expert GPTQ sidecars should fail preflight");

        assert!(
            err.contains("incomplete Qwen3.5 per-expert GPTQ sidecars"),
            "{err}"
        );
        assert!(err.contains("scales"), "{err}");
        assert!(err.contains("qzeros"), "{err}");
    }

    #[test]
    fn rejects_inconsistent_per_expert_gptq_g_idx_sidecars() {
        let config = moe_config();
        let mut names = required_names_with_gptq_linear_aliases(&config, "model");
        append_complete_moe_expert_gptq_sidecars(&mut names, &config, "model", false);
        names.push("model.layers.0.mlp.experts.0.gate_proj.g_idx".to_string());
        let inventory = Qwen35WeightInventory::from_names(names);

        let err = inventory
            .detect_prefix_and_resolve(&config)
            .expect_err("partial g_idx sidecars should fail preflight");

        assert!(
            err.contains("inconsistent Qwen3.5 per-expert GPTQ g_idx sidecars"),
            "{err}"
        );
        assert!(err.contains("gate/up"), "{err}");
    }

    #[test]
    fn resolves_top_level_lm_head_alias_for_language_model_prefix() {
        let config = moe_config();
        let manifest = config.weight_manifest("model.language_model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        names.retain(|name| name != "model.language_model.lm_head.weight");
        names.push("lm_head.weight".to_string());
        let inventory = Qwen35WeightInventory::from_names(names);

        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();

        assert_eq!(plan.prefix, "model.language_model");
        assert_eq!(
            plan.global_tensor("lm_head")
                .map(|tensor| (tensor.name.as_str(), tensor.present)),
            Some(("lm_head.weight", true))
        );
        assert!(plan.validation().is_pass());
    }

    #[test]
    fn planned_loader_loads_by_role_and_strips_linear_weight_suffix() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        let inventory = Qwen35WeightInventory::from_names(names.clone());
        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();
        let loader = RecordingLoader::from_names(names);
        let planned = Qwen35WeightPlanLoader::<CpuBackend>::new(&plan, &loader);

        assert_eq!(planned.plan().prefix, "model");
        assert!(planned.has_global_tensor("embed_tokens"));
        assert_eq!(
            planned.load_global_tensor("embed_tokens").unwrap(),
            vec![1.0]
        );
        assert!(planned.has_layer_tensor(0, "linear_attn_qkv"));
        let linear = planned.load_layer_linear(0, "linear_attn_qkv").unwrap();

        assert_eq!(linear.in_features(), 2);
        assert_eq!(
            loader.linear_names(),
            vec!["model.layers.0.linear_attn.in_proj_qkv".to_string()]
        );
        let err = planned
            .load_global_tensor("lm_head")
            .expect_err("tied lm_head is optional and absent");
        assert!(err.to_string().contains("lm_head"), "{err}");
    }

    #[test]
    fn planned_loader_loads_packed_linear_attention_projection_names() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        let inventory = Qwen35WeightInventory::from_names(names.clone());
        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();
        let loader = RecordingLoader::from_names(names);
        let planned = Qwen35WeightPlanLoader::<CpuBackend>::new(&plan, &loader);

        let qkvz = planned.load_layer_linear_attention_qkvz(0).unwrap();
        let ba = planned.load_layer_linear_attention_ba(0).unwrap();

        assert_eq!(qkvz.in_features(), 2);
        assert_eq!(ba.in_features(), 2);
        assert_eq!(
            loader.linear_names(),
            vec![
                "model.layers.0.linear_attn.in_proj_qkvz".to_string(),
                "model.layers.0.linear_attn.in_proj_ba".to_string(),
            ]
        );
    }

    #[test]
    fn partitions_complete_checkpoint_inventory_without_ignoring_auxiliary_tensors() {
        let config = dense_config();
        let manifest = config.weight_manifest("model.language_model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        names.extend([
            "model.visual.blocks.0.attn.qkv.weight".to_string(),
            "model.visual.patch_embed.proj.weight".to_string(),
            "mtp.layers.0.self_attn.q_proj.weight".to_string(),
            "mtp.layers.0.self_attn.q_proj.weight_scale_inv".to_string(),
        ]);
        let inventory = Qwen35WeightInventory::from_names(names);
        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();

        let partition = inventory.partition_resolved_plan(&plan).unwrap();

        partition.require_no_unknown().unwrap();
        assert_eq!(partition.vision_non_executed.len(), 2);
        assert_eq!(partition.mtp_non_executed.len(), 2);
        assert!(!partition.execution.is_empty());
        assert_eq!(
            partition.execution.len()
                + partition.vision_non_executed.len()
                + partition.mtp_non_executed.len(),
            inventory.tensor_names().count()
        );
    }

    #[test]
    fn rejects_unclassified_checkpoint_tensor_before_execution() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        names.push("model.language_model.unexpected.weight".to_string());
        names.push("model.visuality.looks_like_vision.weight".to_string());
        let inventory = Qwen35WeightInventory::from_names(names);
        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();

        let partition = inventory.partition_resolved_plan(&plan).unwrap();
        let error = partition
            .require_no_unknown()
            .expect_err("unknown archive members must fail closed");

        assert_eq!(partition.unknown.len(), 2);
        assert!(error.contains("2 unknown tensors"), "{error}");
        assert!(error.contains("unexpected.weight"), "{error}");
    }

    #[test]
    fn partition_counts_all_quantized_sidecars_as_execution_components() {
        let config = dense_config();
        let names = required_names_with_gptq_linear_aliases(&config, "model");
        let inventory = Qwen35WeightInventory::from_names(names);
        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();

        let partition = inventory.partition_resolved_plan(&plan).unwrap();

        partition.require_no_unknown().unwrap();
        assert!(partition
            .execution
            .iter()
            .any(|name| name.ends_with(".qweight")));
        assert!(partition
            .execution
            .iter()
            .any(|name| name.ends_with(".scales")));
        assert!(partition
            .execution
            .iter()
            .any(|name| name.ends_with(".qzeros")));
        assert_eq!(partition.execution.len(), inventory.tensor_names().count());
    }

    #[test]
    fn resolves_and_partitions_block_fp8_value_scale_pairs() {
        let config = dense_config();
        let manifest = config.weight_manifest("model").unwrap();
        let mut names = manifest
            .global_tensors
            .iter()
            .chain(
                manifest
                    .layers
                    .iter()
                    .flat_map(|layer| layer.tensors.iter()),
            )
            .filter(|tensor| tensor.required)
            .map(|tensor| tensor.name.clone())
            .collect::<Vec<_>>();
        let values = "model.layers.0.linear_attn.in_proj_qkv.weight";
        let scale_inv = "model.layers.0.linear_attn.in_proj_qkv.weight_scale_inv";
        names.push(scale_inv.to_owned());
        let inventory = Qwen35WeightInventory::from_names(names);
        let plan = inventory.detect_prefix_and_resolve(&config).unwrap();

        let source = plan
            .layer_tensor(0, "linear_attn_qkv")
            .and_then(|tensor| tensor.source.as_ref())
            .expect("typed block-FP8 source");
        assert_eq!(
            source,
            &Qwen35ResolvedWeightSource::BlockFp8 {
                values: values.to_owned(),
                scale_inv: scale_inv.to_owned(),
            }
        );

        let partition = inventory.partition_resolved_plan(&plan).unwrap();
        partition.require_no_unknown().unwrap();
        assert!(partition.execution.contains(&values.to_owned()));
        assert!(partition.execution.contains(&scale_inv.to_owned()));
    }

    #[test]
    fn reads_sharded_index_weight_map_without_loading_tensor_data() {
        let tmp = TempDir::new().unwrap();
        fs::write(tmp.path().join("shard-00001.safetensors"), []).unwrap();
        fs::write(
            tmp.path().join("model.safetensors.index.json"),
            serde_json::json!({
                "metadata": {},
                "weight_map": {
                    "model.embed_tokens.weight": "shard-00001.safetensors",
                    "model.norm.weight": "shard-00001.safetensors"
                }
            })
            .to_string(),
        )
        .unwrap();

        let inventory = Qwen35WeightInventory::from_safetensors_dir(tmp.path()).unwrap();
        assert!(inventory.contains("model.embed_tokens.weight"));
        assert!(inventory.contains("model.norm.weight"));
    }
}