hf2q 0.1.1

//! Qwen3.5 / Qwen3.5-MoE inference support.
//!
//! Entry point for **both** variants:
//!
//! * **Dense** (`general.architecture = "qwen35"`) — 27B dense Qwen3.5.
//! * **MoE** (`general.architecture = "qwen35moe"`) — 35B-A3B mixture of experts.
//!
//! Both variants share ≥90% of the forward surface (linear-attention,
//! gated full-attention, MROPE, MTP, tokenizer, hybrid KV cache, every
//! mlx-native kernel). Only the FFN block differs.
//!
//! Owned by **ADR-013**. Companion conversion spec is ADR-012.
//!
//! # Module layout
//!
//! * `mod.rs` (this file) — shared: [`Qwen35Config`], [`Qwen35LayerKind`],
//!   [`Qwen35Variant`], metadata parser, `ARCH_QWEN35*` constants.
//! * `dense.rs` — 27B dense-specific: SwiGLU FFN, tensor-name resolution,
//!   dense forward entry point.
//! * `moe.rs` — MoE-specific: 256-expert dispatch, shared-expert-gate,
//!   MoE tensor-name resolution, MoE forward entry point.
//! * `kernels.rs` — thin wrappers around mlx-native's new ops.
//! * `kv_cache.rs` — hybrid cache (full-attention KV + linear-attention SSM state).

use anyhow::{anyhow, bail, Result};
use mlx_native::gguf::{GgufFile, MetadataValue};

pub mod activation_capture_real;
pub mod chunk_allocs_arena;
pub mod decode_pool;
pub mod delta_net;
pub mod dense;
pub mod dense_ffn_arena;
pub mod dn_prefill_arena;
pub mod dump_bisect;
pub(super) mod encoder_stage;
pub mod fa_prefill_arena;
pub mod fa_projections_arena;
pub mod ffn;
pub mod forward_cpu;
pub mod forward_gpu;
pub mod full_attn;
pub mod gpu_delta_net;
pub mod gpu_ffn;
pub mod gpu_full_attn;
pub mod in_memory_loader;
pub use chunk_allocs_arena::ChunkAllocsArena;
pub use dense_ffn_arena::{
    DenseFfnArena, DenseFfnOutputRingBuffer, LayerBoundaryArena, MoeFfnArena,
    MoeFfnOutputRingBuffer,
};
pub use dn_prefill_arena::DnPrefillArena;
pub use fa_prefill_arena::FaPrefillArena;
pub use fa_projections_arena::FaProjectionsArena;
pub mod io_heads;
pub mod kernels;
pub mod kv_cache;
pub mod model;
pub mod moe;
pub mod mtp;
pub mod mtp_weights_load;
pub mod spec_decode;
pub mod tokenizer;
pub mod wave5b8_profile;
pub mod weight_loader;
pub mod weight_pool;

/// `general.architecture` value for the dense variant.
pub const ARCH_QWEN35: &str = "qwen35";
/// `general.architecture` value for the MoE variant.
pub const ARCH_QWEN35MOE: &str = "qwen35moe";

/// `general.architecture` value emitted by hf2q's Wedge-4f convert pipeline
/// for **Qwen3-VL text** models (`Qwen/Qwen3-VL-2B-Instruct`,
/// `Qwen/Qwen3-VL-4B-Instruct`, etc.).
///
/// This is the dense Qwen3-VL variant. Note the underscore: the value is
/// `qwen3_vl`, not the llama.cpp upstream's `qwen3vl`. hf2q's convert pipeline
/// (`src/convert/...`) preserves HF's `model_type = "qwen3_vl_text"` family
/// stem when it stamps `general.architecture`. Both the underscored and
/// non-underscored variants are recognized by [`is_qwen3_vl_arch`] so we
/// stay forward-compatible with future convert-pipeline alignment.
///
/// Wedge-4 / iter-227 (2026-05-02) introduces this constant solely so the
/// runtime arch dispatch in `serve::cmd_generate` and
/// `serve::api::engine::LoadedModel::load` can detect Qwen3-VL GGUFs and
/// route them with an operator-actionable error rather than falling
/// through to the Gemma 4 path and dying inside the per-layer MoE expert
/// loader with `missing blk.0.ffn_gate_up_exps.weight`.
///
/// **Why not route to `Qwen35Model::load_from_gguf`?** Qwen3-VL text is a
/// plain dense transformer with `attn_{q,k,v,o}` + `ffn_{gate,up,down}`
/// tensors and zero SSM / DeltaNet structure. [`Qwen35Config::from_gguf`]
/// requires `{prefix}.ssm.{state_size,group_count,inner_size,conv_kernel}`
/// and `{prefix}.full_attention_interval` keys, none of which are emitted
/// by Wedge-4f convert (correctly: Qwen3-VL has none of those features).
/// A separate `Qwen3VlModel` load path is the structurally honest answer
/// and is iter-228+ scope; iter-227 closes only the dispatch gap.
pub const ARCH_QWEN3_VL: &str = "qwen3_vl";

/// llama.cpp upstream's `general.architecture` string for the same family
/// (no underscore). Recognized so a future convert-pipeline alignment to
/// the upstream string doesn't silently re-break dispatch.
pub const ARCH_QWEN3VL_UPSTREAM: &str = "qwen3vl";

/// llama.cpp upstream's `general.architecture` string for the **MoE**
/// Qwen3-VL variant (e.g. `Qwen3-VL-30B-A3B`). hf2q does not yet emit
/// or load this variant, but it is recognized by [`is_qwen3_vl_arch`]
/// so the dispatch error message tells the operator we know what they
/// have but cannot serve it yet (rather than falling through silently).
pub const ARCH_QWEN3VLMOE_UPSTREAM: &str = "qwen3vlmoe";

/// True iff `arch` is any Qwen3-VL `general.architecture` string we
/// recognize — covers both hf2q's underscored convention (`qwen3_vl`) and
/// llama.cpp upstream's no-underscore convention (`qwen3vl`,
/// `qwen3vlmoe`). Used at the runtime dispatch sites in
/// `serve::cmd_generate` and `serve::api::engine::LoadedModel::load` so
/// either arch-string spelling lands on the same operator-actionable
/// error path.
///
/// Returns `false` for `qwen35` / `qwen35moe` (those route through the
/// existing Qwen3.5 dense + MoE paths) and for unrelated arches like
/// `gemma4`, `bert`, `nomic-bert`, etc.
pub fn is_qwen3_vl_arch(arch: &str) -> bool {
    arch == ARCH_QWEN3_VL || arch == ARCH_QWEN3VL_UPSTREAM || arch == ARCH_QWEN3VLMOE_UPSTREAM
}

/// True iff a Qwen3-VL `general.architecture` string identifies the
/// MoE variant (Qwen3-VL-30B-A3B etc.). Today only the upstream
/// `qwen3vlmoe` value triggers this; hf2q's convert pipeline does not
/// yet emit a Qwen3-VL-MoE GGUF, but the predicate is wired so the
/// dispatch error message can distinguish dense vs MoE Qwen3-VL.
pub fn is_qwen3_vl_moe_arch(arch: &str) -> bool {
    arch == ARCH_QWEN3VLMOE_UPSTREAM
}

/// Dense vs MoE flavor.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Qwen35Variant {
    Dense,
    Moe,
}

impl Qwen35Variant {
    /// Resolve from a `general.architecture` metadata string.
    pub fn from_arch(arch: &str) -> Option<Self> {
        match arch {
            ARCH_QWEN35 => Some(Qwen35Variant::Dense),
            ARCH_QWEN35MOE => Some(Qwen35Variant::Moe),
            _ => None,
        }
    }

    /// The metadata-key prefix this variant emits (e.g. `qwen35moe.`).
    pub fn key_prefix(&self) -> &'static str {
        match self {
            Qwen35Variant::Dense => ARCH_QWEN35,
            Qwen35Variant::Moe => ARCH_QWEN35MOE,
        }
    }
}

/// Per-layer kind for the Qwen3.5 interleaved hybrid stack.
///
/// **Distinct from Gemma-4's `LayerType::{Sliding, Full}`** (different
/// semantic axis). Per `project_model_class_split.md` and ADR-013
/// Decision 2, Qwen3.5 owns its own enum in its own module.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Qwen35LayerKind {
    /// Gated DeltaNet linear-attention block.
    LinearAttention,
    /// Gated full-attention block (standard SDPA with output gate).
    FullAttention,
}

/// MoE-only config fields. `None` on the dense variant.
#[derive(Debug, Clone, PartialEq)]
pub struct Qwen35MoeConfig {
    /// Per-expert FFN hidden size (512 for Qwen3.5-MoE).
    pub moe_intermediate_size: u32,
    /// Total number of experts (256 for Qwen3.5-MoE).
    pub num_experts: u32,
    /// Experts activated per token (8 for Qwen3.5-MoE).
    pub num_experts_per_tok: u32,
    /// Shared-expert FFN hidden size (512 for Qwen3.5-MoE). Shared experts
    /// are gated by a separate sigmoid — see ADR-013 Decision 13.
    pub shared_expert_intermediate_size: u32,
}

/// Full architecture config for a Qwen3.5 or Qwen3.5-MoE model.
///
/// Source of truth is the GGUF metadata. See [`Qwen35Config::from_gguf`] for
/// the exact key-to-field mapping — grounded in the apex GGUF dump on
/// 2026-04-23.
#[derive(Debug, Clone, PartialEq)]
pub struct Qwen35Config {
    pub variant: Qwen35Variant,

    // --- Global dims / counts ---
    pub hidden_size: u32,
    pub num_hidden_layers: u32,
    /// Full-attention Q-head count.
    pub num_attention_heads: u32,
    /// Full-attention KV-head count (GQA).
    pub num_key_value_heads: u32,
    /// Full-attention per-head dim (= attention.key_length = attention.value_length; 256 for Qwen3.5).
    pub head_dim: u32,

    // --- Linear-attention (Gated DeltaNet) dims ---
    /// Number of K heads in the linear-attention branch.
    pub linear_num_key_heads: u32,
    /// Number of V heads in the linear-attention branch (>= linear_num_key_heads; GQA).
    pub linear_num_value_heads: u32,
    pub linear_key_head_dim: u32,
    pub linear_value_head_dim: u32,
    /// 1D conv kernel width for the SSM conv1d (4 for Qwen3.5).
    pub linear_conv_kernel_dim: u32,

    // --- Layer stack layout ---
    /// Full-attention layer period (4 for Qwen3.5: layers 3, 7, 11, ... are full attention).
    pub full_attention_interval: u32,
    /// Per-layer kind. Authoritative. Computed from `full_attention_interval`
    /// when `layer_types` is not emitted as an explicit GGUF array (current
    /// llama.cpp/apex convention; kept as a Vec so future metadata can
    /// override on a per-layer basis).
    pub layer_types: Vec<Qwen35LayerKind>,

    // --- RoPE / MROPE ---
    pub partial_rotary_factor: f32, // 0.25 for Qwen3.5 (rotary_dim / head_dim)
    pub rope_theta: f64,            // 1e7 for Qwen3.5
    pub rotary_dim: u32,            // 64 for Qwen3.5 (= partial_rotary_factor * head_dim)
    pub mrope_section: [u32; 4],    // [11, 11, 10, 0] for Qwen3.5
    /// IMROPE interleaved mode. Always `true` for Qwen3.5 (per the model family's
    /// runtime convention — `GGML_ROPE_TYPE_IMROPE == 40`).
    pub mrope_interleaved: bool,

    // --- Norm ---
    pub rms_norm_eps: f32, // 1e-6 for Qwen3.5

    // --- Misc runtime flags ---
    pub max_position_embeddings: u32,
    pub vocab_size: u32,
    pub attn_output_gate: bool,     // true for Qwen3.5
    pub mtp_num_hidden_layers: u32, // 0 if MTP absent (apex GGUF case)
    /// Whether the MTP NextN block carries its own dedicated `embed_tokens` table.
    ///
    /// Qwen3.5 MTP models (e.g. Qwen3.5-Coder) ship `mtp.embed_tokens.weight`; the
    /// HF config flag `mtp_use_dedicated_embeddings == True`.  Qwen3.6 27B + 35B-A3B
    /// instead share the main model's `token_embd.weight` (HF flag `False`); convert
    /// correctly skips emitting `blk.{N}.nextn.embed_tokens.weight` and the loader
    /// must reuse the main verifier embedding table.
    ///
    /// Default `true` matches the historical Qwen3.5 baseline assumed by ADR-013 P14.
    /// Read from GGUF metadata key `{arch}.nextn.use_dedicated_embeddings` when
    /// present, else inferred from tensor-presence at load time.
    pub mtp_use_dedicated_embeddings: bool,

    // --- FFN variant-specific ---
    pub intermediate_size: Option<u32>, // dense: Some(17408); moe: None
    pub moe: Option<Qwen35MoeConfig>,   // dense: None; moe: Some(...)
}

// ---------------------------------------------------------------------
// Parser
// ---------------------------------------------------------------------

fn required_u32(gguf: &GgufFile, key: &str) -> Result<u32> {
    gguf.metadata_u32(key).ok_or_else(|| {
        anyhow!(
            "qwen35 config: required key '{}' missing or wrong type",
            key
        )
    })
}

fn required_f32(gguf: &GgufFile, key: &str) -> Result<f32> {
    gguf.metadata_f32(key).ok_or_else(|| {
        anyhow!(
            "qwen35 config: required key '{}' missing or wrong type",
            key
        )
    })
}

fn required_i32_array_4(gguf: &GgufFile, key: &str) -> Result<[u32; 4]> {
    let mv = gguf
        .metadata(key)
        .ok_or_else(|| anyhow!("qwen35 config: required key '{}' missing", key))?;
    let arr = match mv {
        MetadataValue::Array(a) => a,
        _ => bail!(
            "qwen35 config: key '{}' has type {:?}, expected Array",
            key,
            std::mem::discriminant(mv)
        ),
    };
    if arr.len() != 4 {
        bail!(
            "qwen35 config: key '{}' length {} != 4 (mrope sections)",
            key,
            arr.len()
        );
    }
    let mut out = [0u32; 4];
    for (i, v) in arr.iter().enumerate() {
        out[i] = match v {
            MetadataValue::Int32(x) if *x >= 0 => *x as u32,
            MetadataValue::Uint32(x) => *x,
            MetadataValue::Int8(x) if *x >= 0 => *x as u32,
            MetadataValue::Int16(x) if *x >= 0 => *x as u32,
            MetadataValue::Uint8(x) => *x as u32,
            MetadataValue::Uint16(x) => *x as u32,
            other => bail!(
                "qwen35 config: key '{}' element {} has unexpected variant {:?}",
                key,
                i,
                std::mem::discriminant(other)
            ),
        };
    }
    Ok(out)
}

/// Compute `layer_types` from `num_hidden_layers` and `full_attention_interval`.
///
/// Convention (verified against apex GGUF tensor dump on 2026-04-23 — layers
/// 0, 1, 2 are linear-attention and layer 3 is the first full-attention):
///
/// ```text
/// layer_types[i] = FullAttention   if (i + 1) % interval == 0
///                = LinearAttention otherwise
/// ```
///
/// For the default `full_attention_interval = 4`: layers {3, 7, 11, ...} are
/// FullAttention, all others are LinearAttention.
pub fn default_layer_types(num_hidden_layers: u32, interval: u32) -> Vec<Qwen35LayerKind> {
    if interval == 0 {
        return vec![Qwen35LayerKind::LinearAttention; num_hidden_layers as usize];
    }
    (0..num_hidden_layers)
        .map(|i| {
            if (i + 1) % interval == 0 {
                Qwen35LayerKind::FullAttention
            } else {
                Qwen35LayerKind::LinearAttention
            }
        })
        .collect()
}

impl Qwen35Config {
    /// Parse a Qwen3.5 / Qwen3.5-MoE configuration from loaded GGUF
    /// metadata.
    ///
    /// Required keys (missing → error):
    ///
    /// ```text
    /// {prefix}.block_count                            u32
    /// {prefix}.embedding_length                       u32
    /// {prefix}.attention.head_count                   u32
    /// {prefix}.attention.head_count_kv                u32
    /// {prefix}.attention.key_length                   u32 (= head_dim)
    /// {prefix}.attention.value_length                 u32 (= head_dim)
    /// {prefix}.attention.layer_norm_rms_epsilon       f32
    /// {prefix}.context_length                         u32
    /// {prefix}.rope.freq_base                         f32
    /// {prefix}.rope.dimension_count                   u32 (rotary_dim)
    /// {prefix}.rope.dimension_sections                i32[4]
    /// {prefix}.full_attention_interval                u32
    /// {prefix}.ssm.state_size                         u32 (linear head_dim)
    /// {prefix}.ssm.group_count                        u32 (linear_num_key_heads)
    /// {prefix}.ssm.inner_size                         u32 (= linear_num_value_heads * state_size)
    /// {prefix}.ssm.conv_kernel                        u32 (linear_conv_kernel_dim)
    /// ```
    ///
    /// MoE variant additionally requires:
    ///
    /// ```text
    /// qwen35moe.expert_count                          u32
    /// qwen35moe.expert_used_count                     u32
    /// qwen35moe.expert_feed_forward_length            u32
    /// qwen35moe.expert_shared_feed_forward_length     u32
    /// ```
    ///
    /// Dense variant requires instead:
    ///
    /// ```text
    /// qwen35.feed_forward_length                      u32  (intermediate_size)
    /// ```
    ///
    /// Optional keys (fallback to Qwen3.5 documented defaults if absent):
    ///
    /// ```text
    /// {prefix}.nextn_predict_layers                  u32   default 0
    /// {prefix}.attention.output_gate                  bool  default true
    /// ```
    pub fn from_gguf(gguf: &GgufFile) -> Result<Self> {
        let arch = gguf
            .metadata_string("general.architecture")
            .ok_or_else(|| anyhow!("GGUF missing required key 'general.architecture'"))?;

        let variant = Qwen35Variant::from_arch(arch).ok_or_else(|| {
            anyhow!(
                "general.architecture = {:?} is not a Qwen3.5 variant (expected {:?} or {:?})",
                arch,
                ARCH_QWEN35,
                ARCH_QWEN35MOE
            )
        })?;

        let p = variant.key_prefix();

        let total_block_count = required_u32(gguf, &format!("{p}.block_count"))?;
        let mtp_num_hidden_layers = gguf
            .metadata_u32(&format!("{p}.nextn_predict_layers"))
            .or_else(|| gguf.metadata_u32(&format!("{p}.mtp.num_hidden_layers")))
            .unwrap_or(0);
        // ADR-013 P14 follow-up (2026-04-30): read mtp_use_dedicated_embeddings.
        //
        // Resolution order:
        //   1) Explicit metadata `{arch}.nextn.use_dedicated_embeddings` (forward-compat;
        //      our convert pipeline writes this when MTP is present).
        //   2) Tensor-presence inference: if `blk.{N}.nextn.embed_tokens.weight`
        //      exists in the GGUF, the MTP block has dedicated embeddings; else it
        //      shares the main model's `token_embd.weight`.
        //   3) Fallback (no MTP at all): default `true` to preserve historical Qwen3.5
        //      semantics; the loader is a no-op when mtp_num_hidden_layers == 0.
        //
        // llama.cpp itself has no canonical key for this; it implicitly relies on
        // tensor presence (llama-arch.cpp:759 "NextN/MTP tensors are currently
        // ignored"). Our key is namespaced under the existing `{arch}.nextn.*`
        // family for cleanliness.
        let mtp_use_dedicated_embeddings = gguf
            .metadata(&format!("{p}.nextn.use_dedicated_embeddings"))
            .and_then(|v| match v {
                MetadataValue::Bool(b) => Some(*b),
                _ => None,
            })
            .unwrap_or_else(|| {
                if mtp_num_hidden_layers == 0 {
                    true
                } else {
                    let layer_index = total_block_count.saturating_sub(mtp_num_hidden_layers);
                    let tname = format!("blk.{layer_index}.nextn.embed_tokens.weight");
                    gguf.tensor_info(&tname).is_some()
                }
            });
        if mtp_num_hidden_layers > total_block_count {
            bail!(
                "qwen35 config: {p}.nextn_predict_layers ({}) exceeds {p}.block_count ({})",
                mtp_num_hidden_layers,
                total_block_count
            );
        }
        let num_hidden_layers = total_block_count - mtp_num_hidden_layers;
        let hidden_size = required_u32(gguf, &format!("{p}.embedding_length"))?;
        let num_attention_heads = required_u32(gguf, &format!("{p}.attention.head_count"))?;
        let num_key_value_heads = required_u32(gguf, &format!("{p}.attention.head_count_kv"))?;

        // Qwen3.5 has key_length == value_length == head_dim.
        let key_length = required_u32(gguf, &format!("{p}.attention.key_length"))?;
        let value_length = required_u32(gguf, &format!("{p}.attention.value_length"))?;
        if key_length != value_length {
            bail!(
                "qwen35 config: attention.key_length ({}) != attention.value_length ({}); \
                 Qwen3.5 requires them equal",
                key_length,
                value_length
            );
        }
        let head_dim = key_length;

        let rms_norm_eps = required_f32(gguf, &format!("{p}.attention.layer_norm_rms_epsilon"))?;
        let max_position_embeddings = required_u32(gguf, &format!("{p}.context_length"))?;
        let rope_theta = required_f32(gguf, &format!("{p}.rope.freq_base"))? as f64;
        let rotary_dim = required_u32(gguf, &format!("{p}.rope.dimension_count"))?;
        let mrope_section = required_i32_array_4(gguf, &format!("{p}.rope.dimension_sections"))?;
        let full_attention_interval = required_u32(gguf, &format!("{p}.full_attention_interval"))?;

        // Linear attention (SSM keys).
        let ssm_state_size = required_u32(gguf, &format!("{p}.ssm.state_size"))?;
        let ssm_group_count = required_u32(gguf, &format!("{p}.ssm.group_count"))?;
        let ssm_inner_size = required_u32(gguf, &format!("{p}.ssm.inner_size"))?;
        let ssm_conv_kernel = required_u32(gguf, &format!("{p}.ssm.conv_kernel"))?;
        if ssm_state_size == 0 {
            bail!("qwen35 config: {p}.ssm.state_size must be > 0");
        }
        if ssm_inner_size % ssm_state_size != 0 {
            bail!(
                "qwen35 config: {p}.ssm.inner_size ({}) must be a multiple of \
                 {p}.ssm.state_size ({})",
                ssm_inner_size,
                ssm_state_size
            );
        }
        let linear_num_value_heads = ssm_inner_size / ssm_state_size;

        // Optional keys.
        let attn_output_gate = gguf
            .metadata(&format!("{p}.attention.output_gate"))
            .and_then(|v| match v {
                MetadataValue::Bool(b) => Some(*b),
                _ => None,
            })
            .unwrap_or(true);

        // Partial rotary factor is derived from rotary_dim / head_dim.
        if head_dim == 0 {
            bail!("qwen35 config: head_dim is 0");
        }
        let partial_rotary_factor = (rotary_dim as f32) / (head_dim as f32);

        // Vocab size: prefer explicit metadata, else use tokenizer tokens array length.
        let vocab_size = gguf
            .metadata_u32(&format!("{p}.vocab_size"))
            .or_else(|| {
                gguf.metadata("tokenizer.ggml.tokens")
                    .and_then(|v| match v {
                        MetadataValue::Array(a) => Some(a.len() as u32),
                        _ => None,
                    })
            })
            .ok_or_else(|| {
                anyhow!(
                    "qwen35 config: can't determine vocab_size \
                     (neither {p}.vocab_size nor tokenizer.ggml.tokens present)"
                )
            })?;

        // FFN variant-specific.
        let (intermediate_size, moe) = match variant {
            Qwen35Variant::Dense => {
                let fl = required_u32(gguf, &format!("{p}.feed_forward_length"))?;
                (Some(fl), None)
            }
            Qwen35Variant::Moe => {
                let num_experts = required_u32(gguf, &format!("{p}.expert_count"))?;
                let num_experts_per_tok = required_u32(gguf, &format!("{p}.expert_used_count"))?;
                let moe_intermediate_size =
                    required_u32(gguf, &format!("{p}.expert_feed_forward_length"))?;
                let shared_expert_intermediate_size =
                    required_u32(gguf, &format!("{p}.expert_shared_feed_forward_length"))?;
                (
                    None,
                    Some(Qwen35MoeConfig {
                        moe_intermediate_size,
                        num_experts,
                        num_experts_per_tok,
                        shared_expert_intermediate_size,
                    }),
                )
            }
        };

        let layer_types = default_layer_types(num_hidden_layers, full_attention_interval);

        Ok(Qwen35Config {
            variant,
            hidden_size,
            num_hidden_layers,
            num_attention_heads,
            num_key_value_heads,
            head_dim,
            linear_num_key_heads: ssm_group_count,
            linear_num_value_heads,
            linear_key_head_dim: ssm_state_size,
            linear_value_head_dim: ssm_state_size,
            linear_conv_kernel_dim: ssm_conv_kernel,
            full_attention_interval,
            layer_types,
            partial_rotary_factor,
            rope_theta,
            rotary_dim,
            mrope_section,
            mrope_interleaved: true, // Qwen3.5 convention (see ADR-013 Decision 10).
            rms_norm_eps,
            max_position_embeddings,
            vocab_size,
            attn_output_gate,
            mtp_num_hidden_layers,
            mtp_use_dedicated_embeddings,
            intermediate_size,
            moe,
        })
    }
}

// ---------------------------------------------------------------------
// Wave 5a — Qwen3.6 detection (separate from arch-string dispatch)
// ---------------------------------------------------------------------

/// True if a loaded GGUF's `general.name` metadata identifies it as a
/// Qwen3.6-family model (substring match on `qwen3.6`, case-insensitive).
///
/// **Why a separate helper, not the arch string?** Qwen3.6 reuses the
/// `general.architecture = "qwen35" | "qwen35moe"` metadata values from
/// its Qwen3.5 ancestor (per `project_qwen36_architecture.md` and
/// `arch/entries/qwen35{,moe}.rs`). The two families share ~all of the
/// forward path (linear-attn + gated full-attn + MROPE + tokenizer); the
/// only practical differentiator at GGUF level is the `general.name`
/// string.
///
/// This is the gating signal for the Wave 5a `HF2Q_QWEN36_AUTOREG=1`
/// opt-in: when this returns `true` and the env var is unset,
/// `cmd_generate` errors out with an operator-actionable message rather
/// than silently routing through the autoregressive path. Wave 5b will
/// land a chunk-scan kernel for long-prefill SOTA perf and remove the
/// gate.
///
/// Returns `false` if `general.name` is missing — the conservative
/// default keeps known-good Qwen3.5 GGUFs on the existing path without
/// requiring the new env var.
pub fn is_qwen36_gguf(gguf: &mlx_native::gguf::GgufFile) -> bool {
    gguf.metadata_string("general.name")
        .map(|name| name.to_lowercase().contains("qwen3.6"))
        .unwrap_or(false)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn variant_from_arch() {
        assert_eq!(
            Qwen35Variant::from_arch("qwen35"),
            Some(Qwen35Variant::Dense)
        );
        assert_eq!(
            Qwen35Variant::from_arch("qwen35moe"),
            Some(Qwen35Variant::Moe)
        );
        assert_eq!(Qwen35Variant::from_arch("gemma4"), None);
        assert_eq!(Qwen35Variant::from_arch("qwen3"), None);
        assert_eq!(Qwen35Variant::from_arch(""), None);
    }

    #[test]
    fn layer_types_interval_4() {
        // 40 layers, every 4th is full-attention (layers 3, 7, ..., 39).
        let lt = default_layer_types(40, 4);
        assert_eq!(lt.len(), 40);
        for (i, kind) in lt.iter().enumerate() {
            let want = if (i + 1) % 4 == 0 {
                Qwen35LayerKind::FullAttention
            } else {
                Qwen35LayerKind::LinearAttention
            };
            assert_eq!(*kind, want, "layer {} kind mismatch", i);
        }
        // Expected pattern for first 8 layers: [L, L, L, F, L, L, L, F].
        use Qwen35LayerKind::*;
        assert_eq!(
            &lt[..8],
            &[
                LinearAttention,
                LinearAttention,
                LinearAttention,
                FullAttention,
                LinearAttention,
                LinearAttention,
                LinearAttention,
                FullAttention,
            ]
        );
    }

    #[test]
    fn layer_types_dense_27b_64layer() {
        let lt = default_layer_types(64, 4);
        let full_count = lt
            .iter()
            .filter(|k| **k == Qwen35LayerKind::FullAttention)
            .count();
        assert_eq!(full_count, 16); // 64 / 4 = 16 full-attention layers.
    }

    #[test]
    fn layer_types_interval_zero_all_linear() {
        let lt = default_layer_types(8, 0);
        assert!(lt.iter().all(|k| *k == Qwen35LayerKind::LinearAttention));
    }

    // ------------------------------------------------------------------
    // is_qwen36_gguf — Wave 5a Qwen3.6 detection (general.name substring)
    // ------------------------------------------------------------------

    /// Write a minimal GGUF file with a single string-typed metadata kv
    /// pair and zero tensors. Used by the `is_qwen36_gguf` tests; the
    /// GGUF parser only needs to walk the header + metadata to satisfy
    /// `metadata_string("general.name")`. This mirrors the sibling
    /// helper in `mtp_tests.rs::write_gguf` but adds metadata support
    /// (the existing helper writes zero metadata kvs).
    ///
    /// On-disk layout:
    /// ```text
    /// magic   "GGUF"                 4 bytes
    /// version 3                      u32
    /// n_tensors 0                    u64
    /// n_kv    1                      u64
    /// key_len, key_bytes             u64 + bytes
    /// value_type GGUF_TYPE_STRING(8) u32
    /// value_len, value_bytes         u64 + bytes
    /// (padding to 32-byte alignment) (none required for n_tensors=0)
    /// ```
    fn write_meta_only_gguf(path: &std::path::Path, name_value: Option<&str>) {
        const GGUF_TYPE_STRING: u32 = 8;
        let mut buf: Vec<u8> = Vec::new();
        buf.extend_from_slice(b"GGUF");
        buf.extend_from_slice(&3u32.to_le_bytes());
        buf.extend_from_slice(&0u64.to_le_bytes()); // n_tensors

        let n_kv: u64 = if name_value.is_some() { 1 } else { 0 };
        buf.extend_from_slice(&n_kv.to_le_bytes());

        if let Some(value) = name_value {
            let key = "general.name";
            buf.extend_from_slice(&(key.len() as u64).to_le_bytes());
            buf.extend_from_slice(key.as_bytes());
            buf.extend_from_slice(&GGUF_TYPE_STRING.to_le_bytes());
            buf.extend_from_slice(&(value.len() as u64).to_le_bytes());
            buf.extend_from_slice(value.as_bytes());
        }
        // Align to 32 bytes (alignment field is unset; default is 32 per
        // the GGUF spec — `tensor_data_offset` parsing aligns the same).
        while buf.len() % 32 != 0 {
            buf.push(0);
        }

        std::fs::write(path, &buf).expect("write meta-only gguf");
    }

    #[test]
    fn is_qwen36_gguf_matches_canonical_name() {
        // The real apex GGUF advertises:
        //   general.name = "Qwen3.6 35B A3B Abliterix EGA Abliterated Apex"
        // (per `general.name` from the post-fc85681 convert pipeline).
        // The 27B sibling carries `Qwen3.6-27B`. Both must trigger the gate.
        let tmp = std::env::temp_dir().join(format!("qwen36_pos_{}.gguf", std::process::id()));
        write_meta_only_gguf(&tmp, Some("Qwen3.6 35B A3B Abliterix EGA Abliterated Apex"));
        let gguf = mlx_native::gguf::GgufFile::open(&tmp).expect("open");
        assert!(is_qwen36_gguf(&gguf));
        std::fs::remove_file(&tmp).ok();
    }

    #[test]
    fn is_qwen36_gguf_case_insensitive() {
        // Lowercase, uppercase, and mixed-case all trip the substring check.
        for name in &["qwen3.6-27b", "QWEN3.6-27B", "Some Qwen3.6 model"] {
            let tmp = std::env::temp_dir().join(format!(
                "qwen36_case_{}_{}.gguf",
                std::process::id(),
                name.len()
            ));
            write_meta_only_gguf(&tmp, Some(name));
            let gguf = mlx_native::gguf::GgufFile::open(&tmp).expect("open");
            assert!(is_qwen36_gguf(&gguf), "name {:?} should match", name);
            std::fs::remove_file(&tmp).ok();
        }
    }

    #[test]
    fn is_qwen36_gguf_rejects_qwen35_canonical_name() {
        // Qwen3.5 (no "qwen3.6" substring) must NOT match — it stays on the
        // existing path without the new env var.
        let tmp = std::env::temp_dir().join(format!("qwen35_neg_{}.gguf", std::process::id()));
        write_meta_only_gguf(&tmp, Some("Qwen3.5 27B"));
        let gguf = mlx_native::gguf::GgufFile::open(&tmp).expect("open");
        assert!(!is_qwen36_gguf(&gguf));
        std::fs::remove_file(&tmp).ok();
    }

    #[test]
    fn is_qwen36_gguf_rejects_other_families() {
        for name in &["Gemma 4 26B", "Llama 3 70B", "DeepSeek V3", ""] {
            let tmp = std::env::temp_dir().join(format!(
                "other_family_{}_{}.gguf",
                std::process::id(),
                name.len()
            ));
            write_meta_only_gguf(&tmp, Some(name));
            let gguf = mlx_native::gguf::GgufFile::open(&tmp).expect("open");
            assert!(!is_qwen36_gguf(&gguf), "name {:?} should not match", name);
            std::fs::remove_file(&tmp).ok();
        }
    }

    #[test]
    fn is_qwen36_gguf_returns_false_when_name_missing() {
        // No general.name metadata at all — the conservative default. Real
        // Qwen3.5 GGUFs ship `general.name`, so a missing value is either a
        // truncated test fixture or a non-Qwen file; either way the gate
        // does NOT fire (the existing dispatch handles non-Qwen GGUFs).
        let tmp = std::env::temp_dir().join(format!("no_name_{}.gguf", std::process::id()));
        write_meta_only_gguf(&tmp, None);
        let gguf = mlx_native::gguf::GgufFile::open(&tmp).expect("open");
        assert!(!is_qwen36_gguf(&gguf));
        std::fs::remove_file(&tmp).ok();
    }

    #[test]
    fn key_prefix_roundtrip() {
        assert_eq!(Qwen35Variant::Dense.key_prefix(), ARCH_QWEN35);
        assert_eq!(Qwen35Variant::Moe.key_prefix(), ARCH_QWEN35MOE);
    }

    // ------------------------------------------------------------------
    // Wedge-4 / iter-227 — Qwen3-VL arch dispatch predicates
    // ------------------------------------------------------------------

    /// hf2q's convert pipeline (Wedge-4f) emits `general.architecture =
    /// "qwen3_vl"` (with underscore). The real GGUF inspected on
    /// 2026-05-02 at `.cfa-archive/wedge4f-out/qwen3-vl-2b-q4_0.gguf`
    /// carries this exact string. Pin the predicate against it so a
    /// future convert-pipeline drift does not silently break the
    /// `cmd_generate` / `cmd_serve` dispatch arm.
    #[test]
    fn iter227_recognizes_underscored_qwen3_vl_arch_string() {
        assert!(is_qwen3_vl_arch(ARCH_QWEN3_VL));
        assert!(is_qwen3_vl_arch("qwen3_vl"));
        assert!(!is_qwen3_vl_moe_arch("qwen3_vl"));
    }

    /// llama.cpp upstream emits the no-underscore variant. We recognize
    /// it so a future hf2q convert-pipeline alignment to the upstream
    /// string does not re-break dispatch.
    #[test]
    fn iter227_recognizes_upstream_no_underscore_qwen3vl_arch_string() {
        assert!(is_qwen3_vl_arch(ARCH_QWEN3VL_UPSTREAM));
        assert!(is_qwen3_vl_arch("qwen3vl"));
        assert!(!is_qwen3_vl_moe_arch("qwen3vl"));
    }

    /// llama.cpp upstream's MoE variant (Qwen3-VL-30B-A3B). hf2q does
    /// not yet emit or serve this, but the predicate must distinguish it
    /// from the dense variant so the dispatch error message can be
    /// MoE-specific.
    #[test]
    fn iter227_recognizes_upstream_qwen3vlmoe_arch_string() {
        assert!(is_qwen3_vl_arch(ARCH_QWEN3VLMOE_UPSTREAM));
        assert!(is_qwen3_vl_moe_arch("qwen3vlmoe"));
    }

    /// Regression guard — `is_qwen3_vl_arch` must NOT widen onto the
    /// existing Qwen3.5 / Qwen3.5-MoE arches (which have working
    /// dispatch through `Qwen35LoadedModel::load`) or onto unrelated
    /// families (Gemma 4, BERT, etc.). Iter-227 must be additive.
    #[test]
    fn iter227_does_not_widen_onto_existing_arches() {
        for arch in &[
            "qwen35",
            "qwen35moe",
            "gemma4",
            "gemma3",
            "bert",
            "nomic-bert",
            "llama",
            "qwen3", // base Qwen3 (no VL)
            "qwen2",
            "",
            "totally-fake-arch-name",
        ] {
            assert!(
                !is_qwen3_vl_arch(arch),
                "is_qwen3_vl_arch({arch:?}) must be false (regression guard for iter-227 dispatch)"
            );
            assert!(
                !is_qwen3_vl_moe_arch(arch),
                "is_qwen3_vl_moe_arch({arch:?}) must be false (regression guard for iter-227 dispatch)"
            );
        }
    }

    /// Sanity: `qwen3_vl` is NOT a `Qwen35Variant` — it routes through a
    /// separate dispatch path. This test ensures Qwen35Variant::from_arch
    /// stays narrow to the Qwen3.5 family even after the iter-227 const
    /// additions.
    #[test]
    fn iter227_qwen3_vl_is_not_a_qwen35_variant() {
        assert_eq!(Qwen35Variant::from_arch(ARCH_QWEN3_VL), None);
        assert_eq!(Qwen35Variant::from_arch(ARCH_QWEN3VL_UPSTREAM), None);
        assert_eq!(Qwen35Variant::from_arch(ARCH_QWEN3VLMOE_UPSTREAM), None);
    }

    /// Integration test against the real apex GGUF on disk. Runtime-
    /// skips when artefact absent (existing path-exists check). Path
    /// fixed to `APEX-Q5_K_M.gguf` — the canonical fixture name.
    ///
    /// Verified values (dumped via `llama-gguf` + python parser on 2026-04-23):
    /// - num_hidden_layers = 40
    /// - hidden_size       = 2048
    /// - head_dim          = 256
    /// - rotary_dim        = 64
    /// - mrope_section     = [11, 11, 10, 0]
    /// - rope_theta        ≈ 1e7
    /// - num_experts       = 256
    /// - num_experts_per_tok = 8
    #[test]
    fn parses_real_apex_gguf() {
        let path = std::path::PathBuf::from(
            "/opt/hf2q/models/qwen3.6-35b-a3b-abliterix-ega-abliterated-apex/\
             APEX-Q5_K_M.gguf",
        );
        if !path.exists() {
            eprintln!("skipping: apex GGUF not at expected path");
            return;
        }
        let gguf = match GgufFile::open(&path) {
            Ok(g) => g,
            Err(e) => {
                // mlx-native's GGUF loader supports F32/F16/Q4_0/Q8_0/Q4_K/Q5_K/
                // Q6_K/I16 (verified at /opt/mlx-native/src/gguf/mod.rs:300-313
                // and dequant table :755-765, as of 2026-04-25).  Q5_K mv_id +
                // dequant are wired (see weight_loader.rs:405-411 and
                // /opt/mlx-native/src/ops/quantized_matmul_id_ggml.rs:65-70);
                // only the Q5_K mm_id kernel (large-batch prefill > 8 tokens)
                // is not yet ported.  If GGUF open fails here it's almost
                // certainly an unrelated parse/IO issue, not a Q5_K limitation.
                eprintln!("skipping: apex GGUF open failed ({e})");
                return;
            }
        };
        let cfg = Qwen35Config::from_gguf(&gguf).expect("parse qwen35 config");

        assert_eq!(cfg.variant, Qwen35Variant::Moe);
        assert_eq!(cfg.num_hidden_layers, 40);
        assert_eq!(cfg.hidden_size, 2048);
        assert_eq!(cfg.num_attention_heads, 16);
        assert_eq!(cfg.num_key_value_heads, 2);
        assert_eq!(cfg.head_dim, 256);
        assert_eq!(cfg.linear_num_key_heads, 16);
        assert_eq!(cfg.linear_num_value_heads, 32); // inner_size / state_size = 4096 / 128
        assert_eq!(cfg.linear_key_head_dim, 128);
        assert_eq!(cfg.linear_value_head_dim, 128);
        assert_eq!(cfg.linear_conv_kernel_dim, 4);
        assert_eq!(cfg.full_attention_interval, 4);
        assert_eq!(cfg.rotary_dim, 64);
        assert_eq!(cfg.mrope_section, [11, 11, 10, 0]);
        assert!(cfg.mrope_interleaved);
        assert!((cfg.partial_rotary_factor - 0.25).abs() < 1e-6);
        assert!((cfg.rope_theta - 1e7).abs() < 1.0);
        assert!(cfg.rms_norm_eps > 0.0 && cfg.rms_norm_eps < 1e-5);
        assert_eq!(cfg.layer_types.len(), 40);
        assert_eq!(cfg.layer_types[3], Qwen35LayerKind::FullAttention);
        assert_eq!(cfg.layer_types[0], Qwen35LayerKind::LinearAttention);

        let moe = cfg.moe.as_ref().expect("moe fields");
        assert_eq!(moe.num_experts, 256);
        assert_eq!(moe.num_experts_per_tok, 8);
        assert_eq!(moe.moe_intermediate_size, 512);
        assert_eq!(moe.shared_expert_intermediate_size, 512);
        assert!(cfg.intermediate_size.is_none());

        // MTP stripped from apex; must be 0.
        assert_eq!(cfg.mtp_num_hidden_layers, 0);
    }

    /// End-to-end Q5_K dequant against the real apex GGUF. Picks
    /// `blk.0.attn_gate.weight` (a Q5_K tensor of shape [2048, 4096] =
    /// 8,388,608 values = 32,768 Q5_K blocks) and verifies:
    /// - load_tensor_f32 returns the expected count.
    /// - All values are finite (no NaN / Inf from broken super-block arithmetic).
    /// - The value distribution is non-degenerate (std > 0).
    ///
    /// Opens the 25 GB file and dequantizes ~32K super-blocks; runs
    /// in ~100ms with mmap. Runtime-skips when artefact absent.
    #[test]
    fn dequantizes_real_apex_q5k_tensor() {
        let path = std::path::PathBuf::from(
            "/opt/hf2q/models/qwen3.6-35b-a3b-abliterix-ega-abliterated-apex/\
             APEX-Q5_K_M.gguf",
        );
        if !path.exists() {
            eprintln!("skipping: apex GGUF not at expected path");
            return;
        }
        let device = match mlx_native::MlxDevice::new() {
            Ok(d) => d,
            Err(e) => {
                eprintln!("skipping: no Metal device: {e}");
                return;
            }
        };
        let gguf = GgufFile::open(&path).expect("open apex gguf");
        let buf = gguf
            .load_tensor_f32("blk.0.attn_gate.weight", &device)
            .expect("load Q5_K tensor");

        let got: &[f32] = buf.as_slice().expect("slice");
        assert_eq!(got.len(), 2048 * 4096, "element count");

        // All finite.
        let mut n_nan = 0usize;
        let mut n_inf = 0usize;
        let mut sum = 0.0_f64;
        let mut sum_sq = 0.0_f64;
        for v in got {
            if v.is_nan() {
                n_nan += 1;
            } else if !v.is_finite() {
                n_inf += 1;
            } else {
                sum += *v as f64;
                sum_sq += (*v as f64) * (*v as f64);
            }
        }
        assert_eq!(n_nan, 0, "Q5_K dequant produced NaN values");
        assert_eq!(n_inf, 0, "Q5_K dequant produced Inf values");

        let n = got.len() as f64;
        let mean = sum / n;
        let variance = (sum_sq / n) - mean * mean;
        let stddev = variance.max(0.0).sqrt();
        assert!(
            stddev > 1e-6,
            "Q5_K dequant produced degenerate (all-equal) tensor; stddev = {}",
            stddev
        );
        // Typical attention-gate weights have small magnitudes; sanity bound.
        assert!(
            stddev < 10.0,
            "Q5_K dequant stddev absurdly large: {}",
            stddev
        );

        eprintln!(
            "blk.0.attn_gate.weight (Q5_K → f32): count={}, mean={:.6}, stddev={:.6}",
            got.len(),
            mean,
            stddev
        );
    }
}