metrale_config/

dispatch.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! 2026-09-26: `parse_config`: route a `config.json` to its family parser by the top-level
//! `model_type`, then fill the full stop-token set.
//!
//! Owner: config.
//! Invariants:
//! - On success, `eos_token_ids` is non-empty and its first element is `eos_token_id`.

#![allow(unused_imports)]

use anyhow::{Context, Result};

use super::{
    AttnPositionEncoding, LayerType, ModelConfig, default_conv_kernel, default_partial_rotary,
    default_rms_eps, default_rope_theta, finalize_config, parse_deepseek_v4, parse_gemma4_params,
    parse_glm5_next, parse_kimi_k3, parse_laguna, parse_longcat_ngram, parse_minimax_m2,
    parse_mistral_params, parse_quantization_config, parse_qwen4_exp, parse_step3p7,
    parse_vision_config, resolve_attn_position_encoding, sanitize_kimi_k3_eos, validate_config,
};

fn required_u64(raw: &serde_json::Value, key: &str, model_type: &str) -> Result<u64> {
    let value = raw
        .get(key)
        .with_context(|| format!("{model_type} config missing required field `{key}`"))?;
    value
        .as_u64()
        .with_context(|| format!("{model_type} config field `{key}` must be an unsigned integer"))
}

fn required_nonzero_usize(raw: &serde_json::Value, key: &str, model_type: &str) -> Result<usize> {
    let value = required_u64(raw, key, model_type)? as usize;
    if value == 0 {
        anyhow::bail!("{model_type} config field `{key}` must be greater than zero");
    }
    Ok(value)
}

fn required_u32(raw: &serde_json::Value, key: &str, model_type: &str) -> Result<u32> {
    let value = required_u64(raw, key, model_type)?;
    u32::try_from(value)
        .with_context(|| format!("{model_type} config field `{key}` does not fit in u32"))
}

/// 2026-09-26: Parse a checkpoint `config.json` into a [`ModelConfig`]: the family parse, then
/// `eos_token_ids` from every declared stop id, then the Kimi-K3 stop-set clean-up.
pub fn parse_config(json: &str) -> Result<ModelConfig> {
    let mut config = parse_config_dispatch(json)?;
    resolve_attn_position_encoding(&mut config)?;
    populate_eos_token_ids(&mut config, json);
    sanitize_kimi_k3_eos(&mut config);
    Ok(config)
}

#[cfg(test)]
pub(super) fn populate_eos_token_ids_for_test(config: &mut ModelConfig, json: &str) {
    populate_eos_token_ids(config, json);
}

/// 2026-09-26: Collect every declared stop-token id into `eos_token_ids`, primary first.
///
/// `eos_token_id` may be a scalar or an array, at the top level or inside `text_config`; the
/// family parsers keep one id in the `u32` `eos_token_id`, and this recovers the rest.
/// `config.eos_token_id` is always element 0, whatever the family parser chose (laguna takes
/// the first array entry, step3p7 the last). The other ids follow in config order,
/// de-duplicated. Invalid JSON leaves the config unchanged.
fn populate_eos_token_ids(config: &mut ModelConfig, json: &str) {
    let Ok(raw) = serde_json::from_str::<serde_json::Value>(json) else {
        return;
    };
    let mut ids: Vec<u32> = vec![config.eos_token_id];
    let mut push = |v: &serde_json::Value| match v {
        serde_json::Value::Number(n) => {
            if let Some(x) = n.as_u64() {
                ids.push(x as u32);
            }
        }
        serde_json::Value::Array(a) => {
            for e in a {
                if let Some(x) = e.as_u64() {
                    ids.push(x as u32);
                }
            }
        }
        _ => {}
    };
    if let Some(v) = raw.get("eos_token_id") {
        push(v);
    }
    if let Some(v) = raw.get("text_config").and_then(|t| t.get("eos_token_id")) {
        push(v);
    }
    let mut seen = std::collections::BTreeSet::new();
    ids.retain(|id| seen.insert(*id));
    config.eos_token_ids = ids;
}

fn parse_config_dispatch(json: &str) -> Result<ModelConfig> {
    let raw: serde_json::Value =
        serde_json::from_str(json).context("Invalid JSON in config.json")?;

    // 2026-09-26: A remote-code checkpoint may declare only `architectures` and no `model_type`.
    // The known architecture names map onto their `model_type`; the map is read only when
    // `model_type` is absent or empty.
    let top_model_type = raw
        .get("model_type")
        .and_then(serde_json::Value::as_str)
        .filter(|s| !s.is_empty())
        .or_else(|| {
            raw.get("architectures")
                .and_then(serde_json::Value::as_array)
                .and_then(|a| a.first())
                .and_then(serde_json::Value::as_str)
                .and_then(|arch| match arch {
                    "LongcatFlashNgramForCausalLM" => Some("longcat_flash_ngram"),
                    "LongcatFlashForCausalLM" => Some("longcat_flash"),
                    "Qwen4ExpForConditionalGeneration"
                    | "Qwen3_8FlashNextForConditionalGeneration" => Some("qwen4_exp"),
                    _ => None,
                })
        })
        .unwrap_or("");

    match top_model_type {
        "qwen3_vl_moe" | "qwen3_5_moe" | "qwen3_5" => {
            let text_config = raw
                .get("text_config")
                .context("qwen3_5_moe config missing text_config")?;
            let mut config: ModelConfig = serde_json::from_value(text_config.clone())
                .context("Failed to parse text_config")?;
            config.model_type = top_model_type.to_string();
            if config.eos_token_id == 0 {
                config.eos_token_id = text_config
                    .get("eos_token_id")
                    .and_then(serde_json::Value::as_u64)
                    .unwrap_or(0) as u32;
            }
            if config.vocab_size == 0 {
                config.vocab_size = raw
                    .get("vocab_size")
                    .and_then(serde_json::Value::as_u64)
                    .unwrap_or(0) as usize;
            }
            if let Some(rope_params) = text_config.get("rope_parameters") {
                if config.rope_theta == default_rope_theta()
                    && let Some(theta) = rope_params
                        .get("rope_theta")
                        .and_then(serde_json::Value::as_f64)
                {
                    config.rope_theta = theta;
                }
                if config.partial_rotary_factor == default_partial_rotary()
                    && let Some(prf) = rope_params
                        .get("partial_rotary_factor")
                        .and_then(serde_json::Value::as_f64)
                {
                    config.partial_rotary_factor = prf;
                }
            }
            config.norm_topk_prob = true;
            config.nested_config = true;
            config.attn_gated = top_model_type != "qwen3_vl_moe";
            if top_model_type == "qwen3_vl_moe" || raw.get("vision_config").is_some() {
                config.vision = parse_vision_config(&raw);
            }
            // 2026-09-26: A MoE checkpoint with interleaved MRoPE becomes `qwen3_6_moe`, the
            // `model_type` the Qwen3.6 MoE kernel targets declare. Dense `qwen3_5` keeps its
            // name and only records the MRoPE fields.
            if let Some(rope_params) = text_config.get("rope_parameters") {
                if let Some(ms) = rope_params.get("mrope_section").and_then(|v| v.as_array())
                    && ms.len() == 3
                {
                    config.mrope_section = [
                        ms[0].as_u64().unwrap_or(0) as usize,
                        ms[1].as_u64().unwrap_or(0) as usize,
                        ms[2].as_u64().unwrap_or(0) as usize,
                    ];
                }
                config.mrope_interleaved = rope_params
                    .get("mrope_interleaved")
                    .and_then(|v| v.as_bool())
                    .unwrap_or(false);
                let is_moe = top_model_type == "qwen3_5_moe" || top_model_type == "qwen3_vl_moe";
                if is_moe
                    && config.mrope_interleaved
                    && config.mrope_section.iter().sum::<usize>() > 0
                {
                    config.model_type = "qwen3_6_moe".to_string();
                }
            }
            // 2026-09-26: Holo-3.1: a `qwen3_5_moe` config with a vision tower, no MTP
            // layers and `image_token_id` 248056 becomes `holo3_1_moe`. The MTP check keeps
            // Qwen3.6-35B-A3B, whose config declares `mtp_num_hidden_layers` 1 and the same
            // image token, off this path.
            if top_model_type == "qwen3_5_moe"
                && config.vision.is_some()
                && config.mtp_num_hidden_layers == 0
                && raw
                    .get("image_token_id")
                    .and_then(serde_json::Value::as_u64)
                    == Some(248_056)
            {
                config.model_type = "holo3_1_moe".to_string();
            }
            finalize_config(&mut config, &raw)?;
            Ok(config)
        }
        "nemotron_h" | "nemotron_h_puzzle" => {
            // 2026-09-26: Puzzle configs put the hybrid schedule in `layers_block_type` /
            // `block_configs` and may leave `num_hidden_layers` null. Rewrite the JSON so serde
            // can deserialize it, then derive the per-layer schedule below.
            let mut raw_mut = raw.clone();
            if top_model_type == "nemotron_h_puzzle" {
                apply_nemotron_puzzle_json(&mut raw_mut)?;
            }
            let mut config: ModelConfig = serde_json::from_value(raw_mut.clone())
                .context("Failed to parse nemotron_h config.json")?;
            if config.num_experts == 0 && config.n_routed_experts > 0 {
                config.num_experts = config.n_routed_experts;
            }
            if config.rms_norm_eps == default_rms_eps() && config.norm_eps > 0.0 {
                config.rms_norm_eps = config.norm_eps;
            }
            if config.linear_conv_kernel_dim == default_conv_kernel() && config.conv_kernel > 0 {
                config.linear_conv_kernel_dim = config.conv_kernel;
            }
            if config.shared_expert_intermediate_size == 0
                && config.moe_shared_expert_intermediate_size > 0
            {
                config.shared_expert_intermediate_size = config.moe_shared_expert_intermediate_size;
            }
            config.attn_gated = false;
            // 2026-09-29: Nemotron-H attention applies no rotary embedding (HF
            // `NemotronHAttention.forward`, the checkpoint's own modelling code and vLLM's
            // `nemotron_h.py` all feed the projected Q/K straight into attention). The
            // config's `rope_theta` / `partial_rotary_factor` are inert HF defaults.
            config.attn_position_encoding = Some(AttnPositionEncoding::None);
            config.weight_prefix = "backbone".to_string();
            if !config.hybrid_override_pattern.is_empty() && config.layer_types.is_empty() {
                config.layer_types = config
                    .hybrid_override_pattern
                    .chars()
                    .map(|c| match c {
                        'M' => LayerType::LinearAttention,
                        'E' => LayerType::Moe,
                        '*' => LayerType::FullAttention,
                        other => panic!("Unknown hybrid_override_pattern char: '{other}'"),
                    })
                    .collect();
            }
            if top_model_type == "nemotron_h_puzzle" {
                apply_nemotron_puzzle_config(&mut config, &raw_mut)?;
            }
            finalize_config(&mut config, &raw_mut)?;
            Ok(config)
        }
        "gemma4" => parse_gemma4_params(&raw),
        "laguna" => parse_laguna(&raw),
        "longcat_flash_ngram" | "longcat_flash" => parse_longcat_ngram(&raw),
        // 2026-09-26: `qwen3_8_flash_next` is a second name for `qwen4_exp`; both use the same
        // parser, which also reads the hyper-connection, indexer and n-gram PLE fields.
        "qwen4_exp" | "qwen3_8_flash_next" => parse_qwen4_exp(&raw),
        "m2m_100" | "nllb" => {
            let mut config = ModelConfig::qwen3_next_80b_nvfp4();
            config.model_type = "m2m_100".to_string();
            config.hidden_size = required_nonzero_usize(&raw, "d_model", top_model_type)?;
            config.num_hidden_layers =
                required_nonzero_usize(&raw, "decoder_layers", top_model_type)?;
            config.intermediate_size =
                required_nonzero_usize(&raw, "decoder_ffn_dim", top_model_type)?;
            config.vocab_size = required_nonzero_usize(&raw, "vocab_size", top_model_type)?;
            config.num_attention_heads =
                required_nonzero_usize(&raw, "decoder_attention_heads", top_model_type)?;
            config.num_key_value_heads = config.num_attention_heads;
            if !config
                .hidden_size
                .is_multiple_of(config.num_attention_heads)
            {
                anyhow::bail!(
                    "{} config has d_model ({}) not divisible by decoder_attention_heads ({})",
                    top_model_type,
                    config.hidden_size,
                    config.num_attention_heads,
                );
            }
            config.head_dim = config.hidden_size / config.num_attention_heads;
            config.max_position_embeddings =
                required_nonzero_usize(&raw, "max_position_embeddings", top_model_type)?;
            config.bos_token_id = required_u32(&raw, "bos_token_id", top_model_type)?;
            config.eos_token_id = required_u32(&raw, "eos_token_id", top_model_type)?;
            config.tie_word_embeddings = true;
            config.attn_gated = false;
            config.weight_prefix = "model.decoder".to_string();
            config.num_experts = 0;
            config.num_experts_per_tok = 1;
            config.moe_intermediate_size = 0;
            config.shared_expert_intermediate_size = 0;
            config.layer_types.clear();
            config.full_attention_interval = 1;
            config.linear_num_key_heads = 0;
            config.linear_key_head_dim = 0;
            config.linear_num_value_heads = 0;
            config.linear_value_head_dim = 0;
            config.mtp_num_hidden_layers = 0;
            config.vision = None;
            config.quantization_config = parse_quantization_config(&raw)?;
            validate_config(&config)?;
            Ok(config)
        }
        "minimax_m2" => parse_minimax_m2(&raw),
        "step3p7" => parse_step3p7(&raw),
        "deepseek_v4" => parse_deepseek_v4(json),
        // 2026-09-26: GLM-5.3-Flash; its fields sit under `text_config`.
        "glm5_next" | "glm5_next_text" => parse_glm5_next(json),
        // 2026-09-26: Kimi K3: the parser unwraps `text_config` when present (a flat
        // `kimi_linear` config is read as is) and renames the family to `kimi_k3`.
        "kimi_k3" | "kimi_linear" => parse_kimi_k3(json),
        _ => {
            let mut config: ModelConfig =
                serde_json::from_str(json).context("Failed to parse config.json")?;
            config.attn_gated = true;
            finalize_config(&mut config, &raw)?;
            Ok(config)
        }
    }
}

/// 2026-09-26: Rewrite Puzzle HF JSON so serde can load it as `ModelConfig`.
///
/// - `num_hidden_layers` null, 0 or absent: set to the length of `layers_block_type`, or
///   of `block_configs` when that is missing.
/// - `moe_intermediate_size` / `num_experts_per_tok` absent or 0: set to the maximum over
///   `block_configs`, when some block declares one.
fn apply_nemotron_puzzle_json(raw: &mut serde_json::Value) -> Result<()> {
    let obj = raw
        .as_object_mut()
        .context("nemotron_h_puzzle config.json is not an object")?;
    let n_layers = obj
        .get("layers_block_type")
        .and_then(|v| v.as_array())
        .map(|a| a.len())
        .or_else(|| {
            obj.get("block_configs")
                .and_then(|v| v.as_array())
                .map(|a| a.len())
        })
        .context("nemotron_h_puzzle missing layers_block_type / block_configs")?;
    if obj
        .get("num_hidden_layers")
        .map(|v| v.is_null() || v.as_u64() == Some(0))
        .unwrap_or(true)
    {
        obj.insert("num_hidden_layers".into(), serde_json::json!(n_layers));
    }
    let mut max_inter = 0usize;
    let mut max_topk = 0usize;
    if let Some(blocks) = obj.get("block_configs").and_then(|v| v.as_array()) {
        for b in blocks {
            if let Some(mi) = b.get("moe_intermediate_size").and_then(|v| v.as_u64()) {
                max_inter = max_inter.max(mi as usize);
            }
            if let Some(tk) = b.get("num_experts_per_tok").and_then(|v| v.as_u64()) {
                max_topk = max_topk.max(tk as usize);
            }
        }
    }
    if obj
        .get("moe_intermediate_size")
        .and_then(|v| v.as_u64())
        .unwrap_or(0)
        == 0
        && max_inter > 0
    {
        obj.insert("moe_intermediate_size".into(), serde_json::json!(max_inter));
    }
    if obj
        .get("num_experts_per_tok")
        .and_then(|v| v.as_u64())
        .unwrap_or(0)
        == 0
        && max_topk > 0
    {
        obj.insert("num_experts_per_tok".into(), serde_json::json!(max_topk));
    }
    Ok(())
}

/// 2026-09-26: Map Puzzle `layers_block_type` onto `layer_types`, and the `moe` entries of
/// `block_configs` onto the per-layer MoE sizes (0 for other layers).
fn apply_nemotron_puzzle_config(config: &mut ModelConfig, raw: &serde_json::Value) -> Result<()> {
    let block_types = raw
        .get("layers_block_type")
        .and_then(|v| v.as_array())
        .context("nemotron_h_puzzle missing layers_block_type")?;
    config.layer_types = block_types
        .iter()
        .map(|v| {
            let s = v.as_str().unwrap_or("");
            Ok(match s {
                "mamba" => LayerType::LinearAttention,
                "moe" => LayerType::Moe,
                "attention" => LayerType::FullAttention,
                other => anyhow::bail!("unknown layers_block_type entry: '{other}'"),
            })
        })
        .collect::<Result<Vec<_>>>()?;
    if config.num_hidden_layers == 0 {
        config.num_hidden_layers = config.layer_types.len();
    }
    let n = config.num_hidden_layers;
    let mut inters = vec![0usize; n];
    let mut topks = vec![0usize; n];
    if let Some(blocks) = raw.get("block_configs").and_then(|v| v.as_array()) {
        for (i, b) in blocks.iter().enumerate().take(n) {
            if b.get("block_type").and_then(|v| v.as_str()) != Some("moe") {
                continue;
            }
            inters[i] = b
                .get("moe_intermediate_size")
                .and_then(|v| v.as_u64())
                .unwrap_or(0) as usize;
            topks[i] = b
                .get("num_experts_per_tok")
                .and_then(|v| v.as_u64())
                .unwrap_or(0) as usize;
        }
    }
    config.moe_intermediate_sizes = inters;
    config.num_experts_per_toks = topks;
    if let Some(m) = config
        .moe_intermediate_sizes
        .iter()
        .copied()
        .filter(|&s| s > 0)
        .max()
    {
        config.moe_intermediate_size = m;
    }
    if let Some(m) = config
        .num_experts_per_toks
        .iter()
        .copied()
        .filter(|&k| k > 0)
        .max()
    {
        config.num_experts_per_tok = m;
    }
    Ok(())
}