atlas_core/config/

methods.rs

// SPDX-License-Identifier: AGPL-3.0-only

//! [`ModelConfig`] inherent helper methods. Split out of `config.rs` for
//! file-size budget. Pure derived getters + small predicates over the
//! struct fields.

#![allow(unused_imports)]

use super::{LayerType, ModelConfig};

impl ModelConfig {
    /// Every configured stop-token id, primary first.
    ///
    /// Falls back to `vec![eos_token_id]` when `eos_token_ids` was never populated, so a
    /// hand-built `ModelConfig` and a scalar-EOS checkpoint both behave exactly as before.
    pub fn eos_ids(&self) -> Vec<u32> {
        if self.eos_token_ids.is_empty() {
            vec![self.eos_token_id]
        } else {
            self.eos_token_ids.clone()
        }
    }

    /// Does this token id terminate generation?
    pub fn is_eos(&self, id: u32) -> bool {
        if self.eos_token_ids.is_empty() {
            id == self.eos_token_id
        } else {
            self.eos_token_ids.contains(&id)
        }
    }

    /// GQA ratio: number of Q heads per KV head.
    pub fn gqa_ratio(&self) -> usize {
        self.num_attention_heads
            .checked_div(self.num_key_value_heads)
            .unwrap_or(1)
    }

    /// Layer type for a given layer index.
    /// Falls back to full_attention_interval if layer_types is empty.
    /// Layer kind for ANY index in the checkpoint, including layers past the text stack.
    ///
    /// `layer_type` covers the text stack only. Indices `>= num_hidden_layers` are
    /// MTP/NextN layers and resolve through `mtp_layer_types`; that is what lets GLM-5.3's
    /// layer 45 be represented as the sparse-attention block it actually is, rather than
    /// being appended to the text stack and silently swept into every text-layer loop.
    pub fn layer_type_at(&self, layer_idx: usize) -> Option<LayerType> {
        if layer_idx < self.num_hidden_layers {
            return Some(self.layer_type(layer_idx));
        }
        self.mtp_layer_types
            .get(layer_idx - self.num_hidden_layers)
            .copied()
    }

    /// Layers (text stack only) whose mixer is `deepseek_sparse_attention`.
    pub fn sparse_attention_layers(&self) -> Vec<usize> {
        self.layer_types
            .iter()
            .enumerate()
            .filter(|(_, t)| **t == LayerType::SparseAttention)
            .map(|(i, _)| i)
            .collect()
    }

    /// True when any layer — text stack **or** MTP — needs the sparse-attention
    /// indexer. Scheduling and cache sizing both key off this, so it must not be
    /// answered from `layer_types` alone.
    pub fn has_sparse_attention(&self) -> bool {
        self.layer_types.contains(&LayerType::SparseAttention)
            || self.mtp_layer_types.contains(&LayerType::SparseAttention)
    }

    pub fn layer_type(&self, layer_idx: usize) -> LayerType {
        if !self.layer_types.is_empty() {
            self.layer_types
                .get(layer_idx)
                .cloned()
                .unwrap_or(LayerType::FullAttention)
        } else if self.full_attention_interval > 0
            && (layer_idx + 1).is_multiple_of(self.full_attention_interval)
        {
            LayerType::FullAttention
        } else {
            LayerType::LinearAttention
        }
    }

    /// Number of attention (KV-cache-consuming) layers: full, sliding, and
    /// sparse. All three write to the paged KV cache — only *which* keys they
    /// read differs (all / a window / a runtime-selected top-k) — so every
    /// consumer sized from this count — KV pool `num_layers`,
    /// `attn_layer_dtypes`, loader `layer_kv_dtypes` indexing — must see them
    /// all. Step 3.7 is the only model emitting `SlidingAttention` layer types
    /// (12 full + 33 sliding); counting full-only there undersized the dtype
    /// vec and panicked the loader at layer 13.
    ///
    /// 🪤 The same omission recurred for `SparseAttention`: GLM-5.3-Flash is
    /// 34 `linear_attention` + 11 `deepseek_sparse_attention`, so a full/sliding
    /// filter returned **0** and the KV pool came out zero-sized ("KV cache block
    /// size is zero", measured 2026-08-28). Delegating to
    /// [`LayerType::is_attention`] is what keeps this honest: the predicate lives
    /// next to the enum, so a new variant is answered in one place.
    pub fn num_attention_layers(&self) -> usize {
        if !self.layer_types.is_empty() {
            self.layer_types.iter().filter(|t| t.is_attention()).count()
        } else {
            self.num_hidden_layers
                .checked_div(self.full_attention_interval)
                .unwrap_or(self.num_hidden_layers)
        }
    }

    /// Number of SSM (linear attention) layers.
    pub fn num_ssm_layers(&self) -> usize {
        if !self.layer_types.is_empty() {
            self.layer_types
                .iter()
                .filter(|t| **t == LayerType::LinearAttention)
                .count()
        } else {
            self.num_hidden_layers - self.num_attention_layers()
        }
    }

    /// Whether this model carries recurrent (SSM / linear-attention) state —
    /// the honest capability signal for the SSM snapshot tiers. Derived from
    /// [`Self::num_ssm_layers`] so the config-level predicate and the runtime
    /// pool predicate (`ssm_pool.num_ssm_layers > 0`) agree by construction
    /// (SSOT). A pure-attention model (dense or MoE) returns `false`:
    /// requesting an SSM tier for it must fail fast, never silently no-op.
    pub fn has_recurrent_state(&self) -> bool {
        self.num_ssm_layers() > 0
    }

    /// Whether this model has MoE routed experts — the capability signal for
    /// the expert-streaming tier. Keyed on config, never on observed expert
    /// tensors (EP ranks legitimately own zero local expert tensors).
    pub fn has_experts(&self) -> bool {
        self.num_experts > 0
    }

    /// Rotary embedding dimension.
    ///
    /// Priority:
    /// 1. Explicit `rotary_dim` field (MiniMax M2 — integer in config.json).
    /// 2. `partial_rotary_factor * head_dim` (Qwen3/Gemma-4 convention — float).
    pub fn rotary_dim(&self) -> usize {
        if self.rotary_dim > 0 {
            self.rotary_dim
        } else {
            (self.partial_rotary_factor * self.head_dim as f64) as usize
        }
    }

    /// SSM projection output size: Q + K + V + Z concatenated.
    pub fn ssm_qkvz_size(&self) -> usize {
        let q = self.linear_num_key_heads * self.linear_key_head_dim;
        let k = self.linear_num_key_heads * self.linear_key_head_dim;
        let v = self.linear_num_value_heads * self.linear_value_head_dim;
        let z = self.linear_num_value_heads * self.linear_value_head_dim;
        q + k + v + z
    }

    /// SSM QKV projection output size (without Z): Q + K + V.
    pub fn ssm_qkv_size(&self) -> usize {
        let q = self.linear_num_key_heads * self.linear_key_head_dim;
        let k = self.linear_num_key_heads * self.linear_key_head_dim;
        let v = self.linear_num_value_heads * self.linear_value_head_dim;
        q + k + v
    }

    /// SSM Z gate projection output size.
    pub fn ssm_z_size(&self) -> usize {
        self.linear_num_value_heads * self.linear_value_head_dim
    }

    /// SSM beta+alpha projection output size.
    pub fn ssm_ba_size(&self) -> usize {
        // beta: num_value_heads, alpha: num_value_heads
        self.linear_num_value_heads * 2
    }

    /// Range of expert indices local to this EP rank.
    /// Returns (start, end) where start is inclusive and end is exclusive.
    pub fn local_expert_range(&self) -> (usize, usize) {
        if self.ep_world_size <= 1 {
            return (0, self.num_experts);
        }
        let per_rank = self.num_experts / self.ep_world_size;
        let start = self.ep_rank * per_rank;
        let end = if self.ep_rank == self.ep_world_size - 1 {
            self.num_experts // last rank gets remainder
        } else {
            start + per_rank
        };
        (start, end)
    }

    /// Whether the given expert ID is local to this EP rank.
    pub fn is_local_expert(&self, expert_id: usize) -> bool {
        let (start, end) = self.local_expert_range();
        expert_id >= start && expert_id < end
    }

    /// Range `[start, end)` of a `total`-sized dimension owned by this TP rank.
    /// `total` must be divisible by `tp_world_size`. Returns `(0, total)` when
    /// TP is disabled.
    pub fn tp_shard_range(&self, total: usize) -> (usize, usize) {
        if self.tp_world_size <= 1 {
            return (0, total);
        }
        debug_assert!(
            total.is_multiple_of(self.tp_world_size),
            "tp_shard_range: total={} not divisible by tp_world_size={}",
            total,
            self.tp_world_size,
        );
        let per_rank = total / self.tp_world_size;
        let start = self.tp_rank * per_rank;
        (start, start + per_rank)
    }

    /// Per-rank shard size for a `total`-sized dimension under TP.
    pub fn tp_shard_dim(&self, total: usize) -> usize {
        if self.tp_world_size <= 1 {
            return total;
        }
        total / self.tp_world_size
    }

    /// Weight key prefix for layer-level weights.
    /// Returns `"model.layers"` for flat models (qwen3_next),
    /// or `"model.language_model.layers"` for conditional generation models (qwen3_5_moe).
    pub fn layer_prefix(&self, layer_idx: usize) -> String {
        if self.weight_prefix.is_empty() {
            format!("model.layers.{layer_idx}")
        } else {
            format!("{}.layers.{layer_idx}", self.weight_prefix)
        }
    }

    /// Derive model-agnostic capabilities from this config.
    pub fn capabilities(&self) -> crate::capabilities::ModelCapabilities {
        crate::capabilities::ModelCapabilities::from_config(self)
    }

    // ── Factory sub-dispatch predicates ──
    // Used only by loader_for_config() to select the right weight loader
    // within the qwen3_5_moe model_type family. Not for general use —
    // prefer config fields (attn_gated, nested_config) or capabilities.

    /// Factory use only. Prefer `config.attn_gated` or `config.capabilities()`.
    pub fn is_qwen35(&self) -> bool {
        self.model_type == "qwen3_5_moe"
    }

    /// Factory use only.
    pub fn is_qwen35_dense(&self) -> bool {
        self.model_type == "qwen3_5" && self.num_experts == 0
    }

    /// Factory use only.
    ///
    /// Recognises the upstream `qwen3_vl_moe` model_type (Qwen3-VL MoE)
    /// and Qwen3.5-VL — which ships with `model_type = "qwen3_5"` plus
    /// `architectures = ["Qwen3_5ForConditionalGeneration"]` and a
    /// populated `vision_config` block. The vision_config presence is
    /// the durable signal: the trunk model_type stays `qwen3_5` whether
    /// the checkpoint is text-only or VL, but VL ships an extra
    /// vision encoder which the parser exposes as `config.vision`.
    pub fn is_qwen3_vl(&self) -> bool {
        if self.model_type == "qwen3_vl_moe" {
            return true;
        }
        // Qwen3.5-VL: trunk model_type is `qwen3_5`; the vision tower
        // is detected by the parsed `vision_config` block.
        if self.model_type == "qwen3_5" && self.vision.is_some() {
            return true;
        }
        false
    }

    /// Whether to skip NVFP4 quantization of the LM head.
    /// MLA models (kv_lora_rank > 0) lose logit precision under NVFP4.
    /// Gemma-4 dense (31B): the LM head ties to BF16 embed_tokens whose
    /// rows have heavy outliers (final_norm.weight max=510, several
    /// embedding rows in similar range). The runtime BF16→NVFP4 path
    /// uses a single per-tensor absmax for `scale2`, which forces a
    /// coarse scale that loses ~7 bits in normal-magnitude rows. For a
    /// 262 144-row vocab matrix that compounds into the 0.14-margin
    /// argmax flip on creative prompts (verified 2026-05-01 via FP32
    /// lm_head bisection: NVFP4 output had top1=` a` 21.85 vs FP32 BF16
    /// view top1=` a` 21.85 — quantization noise was visible in the
    /// SAME logit channel that flipped the tiebreak). Skipping the
    /// runtime quantization keeps the LM head as plain BF16 dense; the
    /// FP32 lm_head path (gated by `ATLAS_GEMMA4_FP32_LMHEAD=1`) can
    /// then act on full-precision weights without the NVFP4 floor.
    pub fn skip_lm_head_quantization(&self) -> bool {
        // CLI override (`--lm-head-dtype`, set into `lm_head_bf16_override` at serve
        // time) wins. `bf16` keeps the LM head in BF16 instead of runtime-quantizing it
        // to NVFP4 — the 4-bit floor on the final vocab projection is a prime suspect for
        // argmax flips in long structured generation; vLLM keeps lm_head at checkpoint
        // precision. (Replaces the former ATLAS_LMHEAD_BF16 env var; PCND: explicit arg.)
        if let Some(force_bf16) = self.lm_head_bf16_override {
            return force_bf16;
        }
        if self.kv_lora_rank > 0 {
            return true;
        }
        if self.model_type == "laguna" {
            return true;
        }
        if self.model_type == "gemma4" && self.num_experts == 0 {
            // Allow rollback via env for A/B testing.
            return std::env::var("ATLAS_GEMMA4_LMHEAD_NVFP4").ok().as_deref() != Some("1");
        }
        false
    }

    /// Mamba-2 d_inner = mamba_num_heads * mamba_head_dim.
    pub fn mamba2_d_inner(&self) -> usize {
        self.mamba_num_heads * self.mamba_head_dim
    }

    /// Mamba-2 d_xBC = d_inner + 2 * n_groups * ssm_state_size.
    /// This is the dimension that goes through conv1d (x + B + C concatenated).
    pub fn mamba2_d_xbc(&self) -> usize {
        self.mamba2_d_inner() + 2 * self.n_groups * self.ssm_state_size
    }

    /// Mamba-2 in_proj output size = z + xBC + dt.
    pub fn mamba2_in_proj_size(&self) -> usize {
        self.mamba2_d_inner() + self.mamba2_d_xbc() + self.mamba_num_heads
    }

    /// Per-layer SSM hidden state size in bytes (FP32).
    /// Dispatches on SSM architecture: Mamba-2 vs GDN, using config fields.
    pub fn ssm_h_state_bytes(&self) -> usize {
        if self.mamba_num_heads > 0 && self.mamba_head_dim > 0 {
            // Mamba-2: h[num_heads, head_dim, state_size] FP32
            self.mamba_num_heads * self.mamba_head_dim * self.ssm_state_size * 4
        } else {
            // GDN: h[nv, vd, kd] FP32
            self.linear_num_value_heads * self.linear_value_head_dim * self.linear_key_head_dim * 4
        }
    }

    /// Per-layer SSM conv state size in bytes (FP32).
    pub fn ssm_conv_state_bytes(&self) -> usize {
        let d_conv = self.linear_conv_kernel_dim;
        if self.mamba_num_heads > 0 && self.mamba_head_dim > 0 {
            // Mamba-2: conv: [d_xBC, d_conv] FP32
            self.mamba2_d_xbc() * d_conv * 4
        } else {
            // GDN: conv: [conv_dim, d_conv] FP32
            let conv_dim = self.linear_num_key_heads * self.linear_key_head_dim * 2
                + self.linear_num_value_heads * self.linear_value_head_dim;
            conv_dim * d_conv * 4
        }
    }

    /// SSM state normalization dimensions: (num_heads, k_dim, v_dim).
    /// Used by the state normalization kernel to prevent drift.
    pub fn ssm_state_norm_dims(&self) -> (usize, usize, usize) {
        if self.mamba_num_heads > 0 && self.mamba_head_dim > 0 {
            (
                self.mamba_num_heads,
                self.mamba_head_dim,
                self.ssm_state_size,
            )
        } else {
            (
                self.linear_num_value_heads,
                self.linear_key_head_dim,
                self.linear_value_head_dim,
            )
        }
    }

    /// MoE expert input dimension: latent size if LatentMoE, else hidden_size.
    pub fn moe_input_size(&self) -> usize {
        if self.moe_latent_size > 0 {
            self.moe_latent_size
        } else {
            self.hidden_size
        }
    }

    /// Routed expert intermediate size for layer `i`.
    ///
    /// Puzzle checkpoints prune channels non-uniformly across MoE layers;
    /// look up `moe_intermediate_sizes[i]` when populated, else the scalar.
    pub fn moe_intermediate_size_for(&self, layer: usize) -> usize {
        self.moe_intermediate_sizes
            .get(layer)
            .copied()
            .filter(|&s| s > 0)
            .unwrap_or(self.moe_intermediate_size)
    }

    /// Top-K experts per token for layer `i` (Puzzle per-block schedule).
    pub fn num_experts_per_tok_for(&self, layer: usize) -> usize {
        self.num_experts_per_toks
            .get(layer)
            .copied()
            .filter(|&k| k > 0)
            .unwrap_or(self.num_experts_per_tok)
    }

    /// Max routed intermediate across all layers (buffer / scratch sizing).
    pub fn max_moe_intermediate_size(&self) -> usize {
        self.moe_intermediate_sizes
            .iter()
            .copied()
            .max()
            .unwrap_or(0)
            .max(self.moe_intermediate_size)
    }

    /// Number of MoE-only layers (Nemotron-H).
    pub fn num_moe_layers(&self) -> usize {
        self.layer_types
            .iter()
            .filter(|t| **t == LayerType::Moe)
            .count()
    }

    /// Whether every byte of a sequence's per-layer state is represented by
    /// its KV blocks.
    ///
    /// False for models whose PREFILL builds per-sequence state that KV pages
    /// do not carry: GLM-5.3's DSA indexer rows (`Glm5NextDsaState`) and
    /// compressed DeepSeek V4's compressor pool/ring. Every KV-only mechanism
    /// — radix prefix reuse and the `--swap-space-gb` spill image alike — is
    /// unsafe for those models, and this is the single fact both gates below
    /// are asking about.
    fn per_sequence_state_is_kv_complete(&self) -> bool {
        match self.model_type.as_str() {
            "glm5_next" | "glm5_next_text" => false,
            "deepseek_v4" => self.compress_ratios.iter().all(|&ratio| ratio == 0),
            _ => true,
        }
    }

    /// Whether the radix prefix cache captures every state needed to resume
    /// this model exactly. Preflight SSOT for `build_prefix_cache`.
    pub fn kv_only_prefix_cache_is_safe(&self) -> bool {
        self.per_sequence_state_is_kv_complete()
    }

    /// Whether a sequence may be swapped out to the `--swap-space-gb` pool and
    /// restored from it. Preflight SSOT for `resolve_swap_space_gb`.
    ///
    /// `save_sequence_state_dispatch` writes KV blocks plus the `SsmLayerState`
    /// of each `LayerType::LinearAttention` layer, and nothing else; the
    /// swap-out then calls `free_sequence`, which hands every remaining
    /// per-layer state to #821's `release_state`. A model that is not
    /// KV-complete therefore resumes with a freshly ZEROED pool behind a KV
    /// image that assumes a populated one — a silently wrong answer, not a
    /// crash. Distinct from the prefix-cache predicate because they are
    /// distinct guarantees; they happen to have the same answer today.
    pub fn kv_only_swap_out_is_safe(&self) -> bool {
        self.per_sequence_state_is_kv_complete()
    }
}