Skip to main content

ferrox_models/
config.rs

1//! Architecture configs. Prefer GGUF / config.json over preset defaults.
2//! Unconfirmed preset fields must be listed in `best_effort_fields`.
3//! What actually runs: `docs/MODELS.md`.
4
5use ferrox_moe::{GatingFunction, MoeLayerConfig};
6
7/// Model-level (not per-layer) tensors `ModelConfig::from_gguf` reads.
8///
9/// The config is parsed from its own file handle, so these lookups are
10/// invisible to the handle the weight loader tracks consumption on.
11/// `loader::assert_every_tensor_consumed` replays them; anything added
12/// here must actually be *used*, not merely read, or the gate stops
13/// meaning what it says.
14pub const MODEL_LEVEL_TENSORS_READ_BY_CONFIG: &[&str] = &[
15    "rope_freqs.weight",
16    "rope_factors_long.weight",
17    "rope_factors_short.weight",
18];
19
20/// Which attention mechanism a model uses. `Gqa` (grouped-query
21/// attention + RoPE, uniform across every layer) is the only variant
22/// `ferrox-core`/`ferrox-models::decoder` actually implement today --
23/// it's what every preset runs through, including the two whose real
24/// published attention differs (DeepSeek V4 Pro's CSA/HCA, Kimi K3's
25/// hybrid KDA/Gated-MLA). `KimiHybrid` exists so Kimi K3's real,
26/// cited attention hyperparameters are captured accurately rather than
27/// silently discarded, even though `Decoder` itself still runs the
28/// GQA path for every layer (the dedicated Kimi decoder is the one
29/// consumer of the hybrid variant today).
30#[derive(Debug, Clone)]
31pub enum AttentionKind {
32    Gqa,
33    KimiHybrid(KimiHybridAttention),
34}
35
36/// Which concrete attention mechanism a single 0-indexed layer uses --
37/// the resolved answer `Decoder` needs per layer once it dispatches on
38/// `AttentionKind` instead of always running GQA (see
39/// `ModelConfig::layer_attention_kind`; only the dedicated Kimi
40/// decoder actually dispatches on it today).
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum LayerAttentionKind {
43    Gqa,
44    /// KDA (Kimi Delta Attention) -- see `ferrox_models::kda`.
45    KimiKda,
46    /// Gated MLA -- see `ferrox_models::mla`.
47    KimiMla,
48}
49
50/// Kimi K3's real attention topology, transcribed from the published
51/// `huggingface.co/moonshotai/Kimi-K3/config.json`'s `linear_attn_config`
52/// block. `kda_layers`/`full_attn_layers` are kept exactly as published
53/// -- **1-indexed** (layer 1 is the model's first transformer layer),
54/// not `ferrox`'s usual 0-indexed `layers` slice -- so a caller wiring
55/// this into `Decoder` must subtract 1 before indexing.
56#[derive(Debug, Clone)]
57pub struct KimiHybridAttention {
58    /// 1-indexed layers using KDA (Kimi Delta Attention: gated
59    /// linear/recurrent attention with a short causal conv). 69 of 93
60    /// layers.
61    pub kda_layers: Vec<usize>,
62    /// 1-indexed layers using Gated MLA (DeepSeek-style multi-head
63    /// latent attention with an output gate). 24 of 93 layers.
64    pub full_attn_layers: Vec<usize>,
65    pub mla: MlaConfig,
66    pub kda: KdaConfig,
67}
68
69/// Gated MLA (multi-head latent attention) hyperparameters, verified
70/// against Kimi K3's real `config.json` `text_config` block and the
71/// real `KimiMLAAttention` reference implementation
72/// (`modeling_kimi_linear.py`).
73#[derive(Debug, Clone)]
74pub struct MlaConfig {
75    pub num_heads: usize,
76    pub q_lora_rank: usize,
77    pub kv_lora_rank: usize,
78    pub qk_nope_head_dim: usize,
79    pub qk_rope_head_dim: usize,
80    pub v_head_dim: usize,
81    /// Kimi K3's addition on top of standard DeepSeek-style MLA:
82    /// `attn_output *= sigmoid(g_proj(hidden_states))` before `o_proj`.
83    pub use_output_gate: bool,
84    /// `None` reproduces Kimi K3's real, confirmed behavior: no rotary
85    /// embedding at all (`mla.rs`'s module doc comment; the real
86    /// `KimiMLAAttention.forward` asserts `use_nope` and never calls a
87    /// rotary function). `Some` is for architectures whose decoupled
88    /// `q_rot`/`k_rot` slices genuinely are position-rotated -- e.g.
89    /// GLM-5.2, whose real `config.json` (`zai-org/GLM-5.2`) sets
90    /// `rope_interleave: true` for its main attention (confirmed
91    /// against llama.cpp PR #25407's `LLAMA_ROPE_TYPE_NORM` rope call
92    /// on `q_pe`/`k_pe` in `src/models/glm-dsa.cpp`).
93    pub rope: Option<MlaRopeConfig>,
94}
95
96/// RoPE parameters for the decoupled `q_rot`/`k_rot` slices of an MLA
97/// attention layer that does apply rotation (unlike Kimi K3 -- see
98/// `MlaConfig::rope`'s doc comment). Always the interleaved convention
99/// (`ferrox_core::attention::apply_rope_interleaved`) for every real
100/// architecture confirmed so far to use this (GLM-5.2's
101/// `rope_interleave: true`); a separate split-half variant isn't wired
102/// in here since no confirmed real user of it exists yet.
103#[derive(Debug, Clone, Copy)]
104pub struct MlaRopeConfig {
105    pub theta: f32,
106}
107
108/// KDA (Kimi Delta Attention) hyperparameters, verified against Kimi
109/// K3's real `config.json` `linear_attn_config` block and the real
110/// gated delta-rule reference implementation in
111/// `fla-org/flash-linear-attention`'s `fla/ops/kda/naive.py` (the
112/// exact recurrence: decay state by `exp(g)`, then add a rank-1
113/// `beta * k ⊗ (v - kᵀS)` correction, then read `o = qᵀS`) and
114/// `fla/ops/kda/gate.py` (the lower-bounded gate:
115/// `g = gate_lower_bound * sigmoid(exp(A_log) * (raw_g + dt_bias))`,
116/// and `beta = sigmoid(raw_beta)`).
117#[derive(Debug, Clone)]
118pub struct KdaConfig {
119    pub num_heads: usize,
120    pub head_dim: usize,
121    pub short_conv_kernel_size: usize,
122    pub gate_lower_bound: f32,
123    pub use_full_rank_gate: bool,
124}
125
126/// Which RoPE pairing convention a model uses. Confirmed against
127/// llama.cpp's `llama_model_rope_type` (`src/llama-model.cpp`):
128/// `Norm` is adjacent-pair / GPT-J (`LLAMA_ROPE_TYPE_NORM`); `Neox` is
129/// split-half / GPT-NeoX (`LLAMA_ROPE_TYPE_NEOX`). Getting this wrong
130/// silently produces fluent-but-wrong logits (the real Llama-3.1-8B
131/// early-stop bug: ferrox applied NeoX to a Norm architecture).
132#[derive(Debug, Clone, Copy, PartialEq, Eq)]
133pub enum RopeLayout {
134    /// Adjacent pairs `(2*i, 2*i+1)` -- llama.cpp `LLAMA_ROPE_TYPE_NORM`.
135    /// Used by `llama` (including Llama 3/3.1/3.2), `deepseek2`,
136    /// `mistral3`, and related families. (`llama4` is DedicatedOnly — MoE
137    /// graph — but its RoPE type in the inventory is still Norm.)
138    Norm,
139    /// Split-half pairs `(i, i+half)` -- llama.cpp `LLAMA_ROPE_TYPE_NEOX`.
140    /// Used by `olmoe`, `qwen2`/`qwen2moe`/`qwen3`, `phi3`, `gemma*`, and
141    /// related families. Ferrox's historical default before architecture-
142    /// aware dispatch existed.
143    Neox,
144}
145
146impl RopeLayout {
147    /// Maps a GGUF `general.architecture` string onto the RoPE pairing
148    /// llama.cpp selects for that family. Prefer
149    /// [`crate::capability::resolve_architecture`] for load-time
150    /// decisions — unknown architectures must fail closed there rather
151    /// than guessing. This helper remains for tests and call sites that
152    /// already know the arch is registered; unknowns still return `Neox`
153    /// only as a last-resort historical default.
154    pub fn for_gguf_architecture(arch: &str) -> Self {
155        match crate::capability::resolve_profile(arch) {
156            Some(p) => p.rope,
157            // Unknown: do not invent Norm for a Qwen/Phi/Gemma-shaped
158            // string that happened to miss the registry.
159            None => RopeLayout::Neox,
160        }
161    }
162}
163
164#[derive(Debug, Clone)]
165pub struct ModelConfig {
166    pub name: &'static str,
167    pub n_layers: usize,
168    pub hidden_dim: usize,
169    pub n_heads: usize,
170    pub n_kv_heads: usize,
171    pub head_dim: usize,
172    pub vocab_size: usize,
173    pub rope_theta: f32,
174    pub rms_norm_eps: f32,
175    pub moe: MoeLayerConfig,
176    /// `Gqa` for every preset except Kimi K3. `Decoder`'s forward pass
177    /// does not yet branch on this -- see `AttentionKind`'s doc
178    /// comment.
179    pub attention: AttentionKind,
180    /// Mistral/Mixtral/Qwen2-family sliding-window attention: when
181    /// set, every layer attends only to the most recent `N` cached
182    /// positions instead of the full causal history (see
183    /// `ferrox_core::attention::causal_gqa_attention_windowed`'s doc
184    /// comment for the real source citations). `None` for every
185    /// architecture that doesn't use this (most models, including
186    /// Qwen1.5/Qwen2-MoE's real published config, which sets
187    /// `use_sliding_window: false` despite carrying a `sliding_window`
188    /// value -- so this field being `None`/`Some` must come from that
189    /// enable flag, not just the window-size field's presence).
190    pub sliding_window: Option<usize>,
191    /// How many of the model's *first* layers use an ordinary dense
192    /// FFN (no expert routing at all) rather than the model's MoE
193    /// topology. Found by reading ik_llama.cpp's real GGUF
194    /// hparams-loading source (`LLM_KV_LEADING_DENSE_BLOCK_COUNT`):
195    /// DeepSeek-2/3-family models don't apply MoE uniformly to every
196    /// layer -- the first few layers are always dense. Zero means
197    /// "every layer uses this model's MoE topology," the default for
198    /// architectures that don't do this.
199    pub n_dense_leading_layers: usize,
200    /// Llama 3/3.1/3.2's real per-band RoPE frequency correction (the
201    /// `rope_freqs.weight` GGUF tensor, `head_dim/2` elements,
202    /// `TENSOR_NOT_REQUIRED` so most architectures leave this `None`).
203    /// See `ferrox_core::attention::apply_rope_with_freq_factors`'s doc
204    /// comment for the real source and the real bug this closes: without
205    /// it, every RoPE angle for a Llama-3-family checkpoint is computed
206    /// slightly wrong, an error that compounds with position and
207    /// eventually produces wrong logits (a spurious early EOS was the
208    /// observed real symptom).
209    ///
210    /// Not only a tensor: this is the *resolved* per-band divisor array,
211    /// so a checkpoint declaring `rope.scaling.type = "yarn"` gets its
212    /// YaRN frequency rewrite folded in here too (see
213    /// `ferrox_core::attention::yarn_freq_factors`, and
214    /// `loader::yarn_scaling_from_gguf` for what the file has to declare
215    /// before that happens). A file carrying both a tensor and a YaRN
216    /// declaration composes them by multiplication, as llama.cpp does
217    /// (`ggml_rope_cache_init` divides by `freq_factors` and *then*
218    /// runs `rope_yarn`). Consumers must therefore treat this as "the
219    /// correction to apply", not as "the tensor this file shipped".
220    pub rope_freqs: Option<Vec<f32>>,
221    /// LongRoPE's two candidate factor sets, kept so the choice between
222    /// them can be made when the *run's* context size is known rather
223    /// than at parse time. llama.cpp picks per request
224    /// (`llama_model::get_rope_factors` reads `cparams.n_ctx_seq`), and
225    /// the two sets are not interchangeable: Phi-4-mini's short set is
226    /// all ones (no correction at all) while its long set reaches 47.
227    /// Choosing from the checkpoint's advertised 131072 when the user
228    /// runs at 4096 is a different model.
229    pub rope_freqs_long: Option<Vec<f32>>,
230    pub rope_freqs_short: Option<Vec<f32>>,
231    /// `<arch>.rope.scaling.original_context_length` — the threshold the
232    /// choice above is made against.
233    pub rope_orig_ctx: Option<usize>,
234    /// Rotary width when it is narrower than `head_dim`
235    /// (`<arch>.rope.dimension_count`, llama.cpp `hparams.n_rot`).
236    /// `None` means the whole head rotates, which is the common case.
237    /// Phi-3/Phi-4 rotate 96 of 128.
238    pub rope_dim: Option<usize>,
239    /// LongRoPE/YaRN magnitude scaling (`<arch>.rope.scaling.attn_factor`,
240    /// llama.cpp `hparams.rope_attn_factor` folded into
241    /// `cparams.yarn_attn_factor` at `llama-context.cpp:231`, then applied
242    /// as ggml `rope_yarn`'s `mscale`, which multiplies *both* `cos` and
243    /// `sin` — so it scales the RoPE'd vector, at every position, whether
244    /// or not any frequency correction is active.
245    ///
246    /// Phi-4-mini ships `1.1902381`. Ignoring it does not merely change
247    /// long-context behaviour: q and k are both scaled, so every attention
248    /// logit is off by `attn_factor²` and the softmax is sharper than the
249    /// model's. Measured symptom: ferrox and llama.cpp diverge from the
250    /// eighth token of a greedy completion on the same GGUF.
251    ///
252    /// `1.0` for every architecture that does not set the key.
253    pub rope_attn_factor: f32,
254    /// RoPE pairing convention for this architecture -- see
255    /// `RopeLayout`. Independently of `rope_freqs`: a Llama checkpoint
256    /// needs both `Norm` pairing *and* the per-band frequency factors.
257    pub rope_layout: RopeLayout,
258    /// How Q/K RMSNorm weights are applied when present (see
259    /// [`crate::capability::QkNormStyle`]).
260    pub qk_norm_style: crate::capability::QkNormStyle,
261    /// Gemma 2+/3 alternating SWA period. When `Some(p)`, layer `il`
262    /// uses sliding-window attention iff `(il + 1) % p != 0` (llama.cpp
263    /// `set_swa_pattern` convention); every `p`-th layer is full-attn.
264    pub swa_pattern: Option<usize>,
265    /// Attention logit soft-capping (Gemma 2+). Applied as
266    /// `softcap * tanh(score / softcap)` before softmax.
267    pub attn_logit_softcap: Option<f32>,
268    /// Final logit soft-capping (Gemma 2+). Applied to lm_head output.
269    pub final_logit_softcap: Option<f32>,
270    /// Input embedding scale (Gemma: `sqrt(hidden_dim)`).
271    pub embedding_scale: Option<f32>,
272    /// Optional override for the attention score scale baked into Q
273    /// *instead of* the kernel's default `1/sqrt(head_dim)`. When set,
274    /// callers must pass `score_scale = 1.0` into the attention kernel
275    /// (llama.cpp Gemma: scale Q then `build_attn(..., 1.0f)`). Prefer
276    /// leaving this `None` when the override equals `1/sqrt(head_dim)`.
277    pub attention_scale: Option<f32>,
278    /// RoPE base used on SWA layers (Gemma 3: defaults to `10000` when
279    /// the GGUF omits `rope.freq_base_swa`; full-attn layers keep
280    /// [`Self::rope_theta`]).
281    pub rope_theta_swa: Option<f32>,
282    /// Dense/MoE FFN activation pairing.
283    pub ffn_activation: FfnActivation,
284    /// Every field on this config that is a best-effort estimate rather
285    /// than a confirmed value from an official config.json / GGUF file.
286    pub best_effort_fields: &'static [&'static str],
287}
288
289/// Dense / expert FFN non-linearity used by the generic decoder.
290#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
291pub enum FfnActivation {
292    /// `silu(gate) * up` with separate gate/up matrices (Llama / Qwen).
293    #[default]
294    Swiglu,
295    /// Phi-3 fused gate+up in one `ffn_up` matrix (`2 * n_ff` rows).
296    SwigluFused,
297    /// Gemma GeGLU: `gelu(gate) * up`.
298    Gelu,
299}
300
301impl ModelConfig {
302    /// Re-picks the LongRoPE factor set now that the run's context size
303    /// is known, matching llama.cpp `llama_model::get_rope_factors`:
304    /// `rope_freqs.weight` (Llama 3) always wins; otherwise the long set
305    /// applies only when the context exceeds
306    /// `rope.scaling.original_context_length`, and the short set
307    /// otherwise.
308    ///
309    /// A no-op for every checkpoint that ships neither set, which is all
310    /// of them except the Phi-3/Phi-4 family today.
311    ///
312    /// Because it re-picks `rope_freqs` wholesale it would also discard
313    /// a YaRN rewrite folded into that field at parse time (see
314    /// [`Self::rope_freqs`]). No real checkpoint hits that: LongRoPE
315    /// files declare `rope.scaling.type = "longrope"`, which the loader's
316    /// YaRN arm deliberately does not claim, so the two never populate
317    /// the field on the same file.
318    pub fn apply_runtime_context(&mut self, ctx: usize) {
319        let (Some(orig), true) = (
320            self.rope_orig_ctx,
321            self.rope_freqs_long.is_some() || self.rope_freqs_short.is_some(),
322        ) else {
323            return;
324        };
325        let picked = if ctx > orig {
326            self.rope_freqs_long.as_ref()
327        } else {
328            self.rope_freqs_short.as_ref()
329        };
330        if let Some(f) = picked
331            .or(self.rope_freqs_long.as_ref())
332            .or(self.rope_freqs_short.as_ref())
333        {
334            self.rope_freqs = Some(f.clone());
335        }
336    }
337
338    /// True if layer `layer_idx` (0-indexed) should be built as an
339    /// ordinary dense FFN rather than this model's MoE topology.
340    pub fn layer_is_dense(&self, layer_idx: usize) -> bool {
341        layer_idx < self.n_dense_leading_layers
342    }
343
344    /// Sliding-window size for layer `il`, honouring Gemma-style
345    /// alternating SWA patterns. `None` means full causal attention.
346    pub fn layer_sliding_window(&self, layer_idx: usize) -> Option<usize> {
347        let window = self.sliding_window?;
348        match self.swa_pattern {
349            None => Some(window),
350            Some(period) if period > 1 => {
351                // llama.cpp `set_swa_pattern` (dense_first=false):
352                // `is_swa = (il % period) < (period - 1)` — equivalent to
353                // full attention when `(il + 1) % period == 0`.
354                if (layer_idx + 1).is_multiple_of(period) {
355                    None
356                } else {
357                    Some(window)
358                }
359            }
360            Some(_) => Some(window),
361        }
362    }
363
364    /// The narrowest sliding window any layer of this model uses, or
365    /// `None` if every layer is full-causal.
366    ///
367    /// For an alternating-SWA model (gpt-oss, Gemma-3) the
368    /// full-attention layers impose no constraint on the KV block
369    /// layout and the sliding ones impose the window -- so the model's
370    /// constraint is simply the window, present as soon as *any* layer
371    /// slides. A model that is 5/6 full-attention is not 5/6 exempt:
372    /// one mis-aligned sliding layer corrupts the answer.
373    pub fn kv_block_window(&self) -> Option<usize> {
374        (0..self.n_layers).find_map(|il| self.layer_sliding_window(il))
375    }
376
377    /// The KV cache block layout to use for this model, given the block
378    /// size an operator asked for.
379    ///
380    /// The requested size is rounded *down* to something that divides
381    /// the window (see [`ferrox_core::kv_swa`]), so a config that would
382    /// straddle the window boundary becomes a smaller block rather than
383    /// a startup failure or -- much worse -- a silently wrong mask.
384    pub fn kv_block_layout(&self, desired_block_size: usize) -> ferrox_core::BlockLayout {
385        let window = self.kv_block_window();
386        let block_size = ferrox_core::aligned_block_size(desired_block_size, window);
387        ferrox_core::BlockLayout::new(block_size, window)
388            .expect("aligned_block_size returns a size BlockLayout accepts")
389    }
390
391    /// RoPE frequency base for layer `il` (SWA layers may differ).
392    pub fn layer_rope_theta(&self, layer_idx: usize) -> f32 {
393        match (self.layer_sliding_window(layer_idx), self.rope_theta_swa) {
394            (Some(_), Some(theta)) => theta,
395            _ => self.rope_theta,
396        }
397    }
398
399    /// Which attention mechanism layer `layer_idx` (0-indexed, ferrox's
400    /// usual convention) uses. For `AttentionKind::Gqa` every layer is
401    /// `LayerAttentionKind::Gqa`; for `AttentionKind::KimiHybrid`, looks
402    /// up `layer_idx + 1` (the real `kda_layers`/`full_attn_layers`
403    /// lists are 1-indexed -- see `KimiHybridAttention`'s doc comment)
404    /// in those real per-layer lists.
405    ///
406    /// # Panics
407    /// If `layer_idx` isn't covered by either list of a `KimiHybrid`
408    /// config -- can't happen for `kimi_k3()`, whose lists are tested
409    /// (`kimi_k3_hybrid_attention_layers_partition_every_layer_exactly_once`)
410    /// to partition every layer with no gaps, but a caller building a
411    /// custom `KimiHybridAttention` must uphold the same invariant.
412    pub fn layer_attention_kind(&self, layer_idx: usize) -> LayerAttentionKind {
413        match &self.attention {
414            AttentionKind::Gqa => LayerAttentionKind::Gqa,
415            AttentionKind::KimiHybrid(hybrid) => {
416                let one_indexed = layer_idx + 1;
417                if hybrid.kda_layers.contains(&one_indexed) {
418                    LayerAttentionKind::KimiKda
419                } else if hybrid.full_attn_layers.contains(&one_indexed) {
420                    LayerAttentionKind::KimiMla
421                } else {
422                    panic!(
423                        "layer {layer_idx} (1-indexed {one_indexed}) is in neither \
424                         kda_layers nor full_attn_layers"
425                    )
426                }
427            }
428        }
429    }
430
431    /// Total parameter count implied by the MoE config, as a sanity
432    /// check against the publicly reported total (this is an order of
433    /// magnitude check, not an exact parameter-count reproduction).
434    pub fn approx_active_params_per_token(&self) -> usize {
435        let attn_params_per_layer = 4 * self.hidden_dim * self.hidden_dim; // q,k,v,o (rough)
436        let active_experts = self.moe.n_experts_active + self.moe.n_shared_experts;
437        let expert_params = active_experts * 3 * self.moe.hidden_dim * self.moe.expert_ffn_dim; // gate,up,down
438        self.n_layers * (attn_params_per_layer + expert_params)
439    }
440}
441
442/// GLM-5.2 (Z.ai) **structural sketch only** — not a supported real
443/// inference path. Real DSA lives in `glm_dsa` / `glm52_decoder` and is
444/// not wired into `Decoder` / `ferrox-server`. This preset drives
445/// smoke/bench with synthetic GQA weights only (~744B / ~40B active
446/// hparams as published placeholders).
447pub fn glm_5_2() -> ModelConfig {
448    ModelConfig {
449        sliding_window: None,
450        name: "glm-5.2",
451        attention: AttentionKind::Gqa,
452        n_layers: 92,
453        hidden_dim: 6144,
454        n_heads: 48,
455        n_kv_heads: 8,
456        head_dim: 128,
457        vocab_size: 151552,
458        rope_theta: 1_000_000.0,
459        rms_norm_eps: 1e-5,
460        moe: MoeLayerConfig {
461            expert_weights_scale: 1.0,
462            n_experts: 256,
463            n_experts_active: 8,
464            n_shared_experts: 1,
465            hidden_dim: 6144,
466            expert_ffn_dim: 2048,
467            // Sigmoid, not softmax: reading ik_llama.cpp's real GGUF
468            // hparams-loading source (llama-hparams.cpp,
469            // LLM_ARCH_GLM4_MOE case) directly showed GLM4-MoE-family
470            // models default to sigmoid gating with post-selection
471            // score renormalization. GLM-5.2 is presumed to continue
472            // this lineage; not confirmed against GLM-5.2's own
473            // config.json (unavailable in this environment).
474            gating: GatingFunction::Sigmoid,
475            norm_topk_prob: true,
476         expert_group_count: None, expert_group_used_count: None,},
477        // No evidence found (via ik_llama.cpp source or public
478        // reporting) that GLM-5.2 skips MoE on any leading layers;
479        // defaulting to 0 (every layer uses this model's MoE
480        // topology) rather than assuming DeepSeek's convention
481        // applies here too.
482        n_dense_leading_layers: 0,
483        rope_freqs: None,
484        rope_attn_factor: 1.0,
485        rope_dim: None,
486        rope_freqs_long: None,
487        rope_freqs_short: None,
488        rope_orig_ctx: None,
489        // Placeholder GQA path; real GLM-5.2 DSA uses interleaved RoPE
490        // via `glm_dsa`/`mla`, not this preset's Decoder path.
491        rope_layout: RopeLayout::Neox,
492        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
493        swa_pattern: None,
494        attn_logit_softcap: None,
495        final_logit_softcap: None,
496        embedding_scale: None,
497        attention_scale: None,
498        rope_theta_swa: None,
499        ffn_activation: FfnActivation::Swiglu,
500        best_effort_fields: &[
501            "n_layers",
502            "hidden_dim",
503            "n_heads",
504            "n_kv_heads",
505            "head_dim",
506            "rope_theta",
507            "moe.expert_ffn_dim",
508            "moe.n_shared_experts",
509            "moe.gating (sigmoid assumed from GLM4-MoE-family convention found in ik_llama.cpp source, not confirmed for GLM-5.2 specifically)",
510        ],
511    }
512}
513
514/// DeepSeek V4 Pro **structural sketch only** — CSA/HCA is not on this
515/// GQA `Decoder` path. Real primitives live under
516/// `deepseek_v4_attention` / `hyper_connections` and are not assembled
517/// into a served decoder yet. Hparams (~1.6T / ~49B active) are
518/// placeholders for smoke/bench.
519pub fn deepseek_v4_pro() -> ModelConfig {
520    ModelConfig {
521        sliding_window: None,
522        name: "deepseek-v4-pro",
523        attention: AttentionKind::Gqa,
524        n_layers: 96,
525        hidden_dim: 7168,
526        n_heads: 56,
527        n_kv_heads: 8,
528        head_dim: 128,
529        vocab_size: 129280,
530        rope_theta: 1_000_000.0,
531        rms_norm_eps: 1e-6,
532        moe: MoeLayerConfig {
533            expert_weights_scale: 1.0,
534            n_experts: 385,
535            n_experts_active: 6,
536            n_shared_experts: 1,
537            hidden_dim: 7168,
538            expert_ffn_dim: 2048,
539            // Sigmoid, not softmax: this is the stronger-confidence of
540            // the two sigmoid-gating corrections in this file.
541            // DeepSeek-V3's own published technical report explicitly
542            // documents computing per-expert affinity via sigmoid and
543            // renormalizing only the selected experts' scores to sum
544            // to one; reading ik_llama.cpp's real GGUF hparams-loading
545            // source (llama-hparams.cpp, LLM_ARCH_DEEPSEEK2 case)
546            // confirmed this is exactly what that code path defaults
547            // to for the DeepSeek-2/3 lineage. DeepSeek V4 Pro is
548            // presumed to continue using sigmoid gating for the same
549            // reason; not confirmed against V4 Pro's own config.json.
550            gating: GatingFunction::Sigmoid,
551            norm_topk_prob: true,
552         expert_group_count: None, expert_group_used_count: None,},
553        // DeepSeek-V3's own published technical report documents the
554        // first 3 transformer layers as dense (ordinary FFN, no
555        // expert routing), with MoE starting from layer 4 onward;
556        // ik_llama.cpp's real hparams-loading source
557        // (LLM_KV_LEADING_DENSE_BLOCK_COUNT) confirms this is a real,
558        // loaded GGUF metadata field for the DeepSeek-2/3 lineage.
559        // DeepSeek V4 Pro is presumed to continue this convention;
560        // not confirmed against V4 Pro's own config.json.
561        n_dense_leading_layers: 3,
562        rope_freqs: None,
563        rope_attn_factor: 1.0,
564        rope_dim: None,
565        rope_freqs_long: None,
566        rope_freqs_short: None,
567        rope_orig_ctx: None,
568        // llama.cpp maps LLM_ARCH_DEEPSEEK4 -> LLAMA_ROPE_TYPE_NORM.
569        rope_layout: RopeLayout::Norm,
570        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
571        swa_pattern: None,
572        attn_logit_softcap: None,
573        final_logit_softcap: None,
574        embedding_scale: None,
575        attention_scale: None,
576        rope_theta_swa: None,
577        ffn_activation: FfnActivation::Swiglu,
578        best_effort_fields: &[
579            "n_layers",
580            "hidden_dim",
581            "n_heads",
582            "n_kv_heads",
583            "head_dim",
584            "moe.expert_ffn_dim",
585            "attention_variant (CSA/HCA hybrid NOT implemented, GQA fallback in use)",
586            "moe.gating (sqrtsoftplus: confirmed for real V4 in llama.cpp PR #24162; this preset still uses Sigmoid on the wrong GQA sketch path)",
587            "n_dense_leading_layers (3: same confidence basis as gating above, DeepSeek-V3 technical report + ik_llama.cpp source, not confirmed for V4 Pro)",
588        ],
589    }
590}
591
592/// Kimi K3 **structural sketch only** for the generic GQA `Decoder`.
593/// Real checkpoint work uses the dedicated Kimi stack (`kimi_loader` /
594/// `KimiEngine`); slice-verified, not a full end-to-end run. Do not
595/// treat this preset as a runnable Kimi substitute.
596pub fn kimi_k3() -> ModelConfig {
597    ModelConfig {
598        sliding_window: None,
599        name: "kimi-k3",
600        n_layers: 93,
601        hidden_dim: 7168,
602        // n_heads/n_kv_heads/head_dim describe the Gqa fallback
603        // Decoder actually runs today, not Kimi K3's real attention
604        // (see `attention` below) -- kept at reasonable stand-in
605        // values (matching MLA's num_heads=96 and combined
606        // qk_nope+qk_rope head dim) rather than deleted, so the
607        // placeholder path stays runnable.
608        n_heads: 96,
609        n_kv_heads: 96,
610        head_dim: 192,
611        vocab_size: 163840,
612        // Not present in the published text_config; RoPE only ever
613        // applies to Gated MLA's 64-dim qk_rope_head_dim slice in the
614        // real architecture, and Decoder doesn't implement that slicing
615        // yet, so this remains an unconfirmed placeholder.
616        rope_theta: 1_000_000.0,
617        rms_norm_eps: 1e-5,
618        moe: MoeLayerConfig {
619            expert_weights_scale: 1.0,
620            n_experts: 896,
621            n_experts_active: 16,
622            n_shared_experts: 2,
623            hidden_dim: 7168,
624            expert_ffn_dim: 3072,
625            // Confirmed directly from the real config.json:
626            // "moe_router_activation_func": "sigmoid".
627            gating: GatingFunction::Sigmoid,
628            norm_topk_prob: true,
629         expert_group_count: None, expert_group_used_count: None,},
630        // Confirmed directly from the real config.json:
631        // "first_k_dense_replace": 1.
632        n_dense_leading_layers: 1,
633        // Kimi K3's real, published attention topology (verified
634        // against huggingface.co/moonshotai/Kimi-K3/config.json's
635        // linear_attn_config block and the real KimiDeltaAttention /
636        // KimiMLAAttention reference implementations in
637        // modeling_kimi_linear.py) -- not yet wired into Decoder's
638        // forward pass, which still runs the Gqa placeholder above for
639        // every layer regardless of this field.
640        attention: AttentionKind::KimiHybrid(KimiHybridAttention {
641            kda_layers: vec![
642                1, 2, 3, 5, 6, 7, 9, 10, 11, 13, 14, 15, 17, 18, 19, 21, 22, 23, 25, 26, 27, 29,
643                30, 31, 33, 34, 35, 37, 38, 39, 41, 42, 43, 45, 46, 47, 49, 50, 51, 53, 54, 55,
644                57, 58, 59, 61, 62, 63, 65, 66, 67, 69, 70, 71, 73, 74, 75, 77, 78, 79, 81, 82,
645                83, 85, 86, 87, 89, 90, 91,
646            ],
647            full_attn_layers: vec![
648                4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84,
649                88, 92, 93,
650            ],
651            mla: MlaConfig {
652                num_heads: 96,
653                q_lora_rank: 1536,
654                kv_lora_rank: 512,
655                qk_nope_head_dim: 128,
656                qk_rope_head_dim: 64,
657                v_head_dim: 128,
658                use_output_gate: true,
659                // Real, confirmed: Kimi K3's `KimiMLAAttention.forward`
660                // never rotates -- see `MlaConfig::rope`'s doc comment.
661                rope: None,
662            },
663            kda: KdaConfig {
664                num_heads: 96,
665                head_dim: 128,
666                short_conv_kernel_size: 4,
667                gate_lower_bound: -5.0,
668                use_full_rank_gate: true,
669            },
670        }),
671        rope_freqs: None,
672        rope_attn_factor: 1.0,
673        rope_dim: None,
674        rope_freqs_long: None,
675        rope_freqs_short: None,
676        rope_orig_ctx: None,
677        // GQA placeholder path only; real Kimi attention is rope-less MLA
678        // or KDA and never reaches Decoder::apply_rope_head.
679        rope_layout: RopeLayout::Neox,
680        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
681        swa_pattern: None,
682        attn_logit_softcap: None,
683        final_logit_softcap: None,
684        embedding_scale: None,
685        attention_scale: None,
686        rope_theta_swa: None,
687        ffn_activation: FfnActivation::Swiglu,
688        best_effort_fields: &[
689            "n_heads/n_kv_heads/head_dim (describe the unimplemented Gqa placeholder, not Kimi K3's real MLA/KDA attention -- see `attention` field)",
690            "rope_theta (not present in the published config; real architecture only applies RoPE to Gated MLA's qk_rope_head_dim slice, which Decoder doesn't implement)",
691            "entire preset beyond hyperparameters (the real 2.8T-parameter checkpoint has not been run end to end; only real slices have, via the dedicated kimi_decoder/kimi_loader stack -- see docs/MODELS.md)",
692        ],
693    }
694}
695
696/// Matches the generated on-disk fixture exactly (hidden_dim, head
697/// counts, ffn_dim, vocab, rope_theta, eps).
698/// Used by `ferrox inspect-run` and the cross-validation test in
699/// `crates/ferrox-models/tests/gguf_roundtrip.rs` to prove the real
700/// GGUF loader + forward pass produce the same numbers as an
701/// independent NumPy reference implementation reading the same file.
702pub fn test_dense_fixture() -> ModelConfig {
703    ModelConfig {
704        sliding_window: None,
705        name: "ferrox-test-dense",
706        attention: AttentionKind::Gqa,
707        n_layers: 2,
708        hidden_dim: 32,
709        n_heads: 4,
710        n_kv_heads: 2,
711        head_dim: 8,
712        vocab_size: 32,
713        rope_theta: 10000.0,
714        rms_norm_eps: 1e-5,
715        moe: MoeLayerConfig {
716            expert_weights_scale: 1.0,
717            n_experts: 1,
718            n_experts_active: 1,
719            n_shared_experts: 0,
720            hidden_dim: 32,
721            expert_ffn_dim: 32,
722            gating: GatingFunction::Softmax,
723            norm_topk_prob: true,
724            expert_group_count: None,
725            expert_group_used_count: None,
726        },
727        n_dense_leading_layers: 0,
728        rope_freqs: None,
729        rope_attn_factor: 1.0,
730        rope_dim: None,
731        rope_freqs_long: None,
732        rope_freqs_short: None,
733        rope_orig_ctx: None,
734        // Matches the independent reference's split-half apply_rope.
735        rope_layout: RopeLayout::Neox,
736        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
737        swa_pattern: None,
738        attn_logit_softcap: None,
739        final_logit_softcap: None,
740        embedding_scale: None,
741        attention_scale: None,
742        rope_theta_swa: None,
743        ffn_activation: FfnActivation::Swiglu,
744        best_effort_fields: &["this is a synthetic test fixture, not a real model"],
745    }
746}
747
748/// Matches the generated on-disk multi-expert MoE fixture: 4 experts,
749/// top-2 routing, 1 shared
750/// expert, packed 3D expert tensors. Used to verify the previously-
751/// untested multi-expert loading path (`split_expert_tensor` in
752/// `ferrox-models::loader`) against a real file, the same way
753/// `test_dense_fixture` verifies the single-expert path.
754pub fn test_moe_fixture() -> ModelConfig {
755    ModelConfig {
756        sliding_window: None,
757        name: "ferrox-test-moe",
758        attention: AttentionKind::Gqa,
759        n_layers: 2,
760        hidden_dim: 32,
761        n_heads: 4,
762        n_kv_heads: 2,
763        head_dim: 8,
764        vocab_size: 32,
765        rope_theta: 10000.0,
766        rms_norm_eps: 1e-5,
767        moe: MoeLayerConfig {
768            expert_weights_scale: 1.0,
769            n_experts: 4,
770            n_experts_active: 2,
771            n_shared_experts: 1,
772            hidden_dim: 32,
773            expert_ffn_dim: 32,
774            gating: GatingFunction::Softmax,
775            norm_topk_prob: true,
776            expert_group_count: None,
777            expert_group_used_count: None,
778        },
779        n_dense_leading_layers: 0,
780        rope_freqs: None,
781        rope_attn_factor: 1.0,
782        rope_dim: None,
783        rope_freqs_long: None,
784        rope_freqs_short: None,
785        rope_orig_ctx: None,
786        rope_layout: RopeLayout::Neox,
787        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
788        swa_pattern: None,
789        attn_logit_softcap: None,
790        final_logit_softcap: None,
791        embedding_scale: None,
792        attention_scale: None,
793        rope_theta_swa: None,
794        ffn_activation: FfnActivation::Swiglu,
795        best_effort_fields: &["this is a synthetic multi-expert test fixture, not a real model"],
796    }
797}
798
799/// Matches the generated on-disk mixed-topology fixture: 3 layers, the
800/// first of which is
801/// an ordinary dense FFN and the remaining two are genuine MoE (3
802/// experts, top-1 routing, 1 shared expert each). Used to verify the
803/// "leading dense layers" loading path
804/// (`ModelConfig::layer_is_dense`) against a real file -- the pattern
805/// found in DeepSeek-2/3-family models via ik_llama.cpp's source
806/// (`LLM_KV_LEADING_DENSE_BLOCK_COUNT`), which was previously only
807/// documented, not implemented or tested.
808pub fn test_mixed_fixture() -> ModelConfig {
809    ModelConfig {
810        sliding_window: None,
811        name: "ferrox-test-mixed",
812        attention: AttentionKind::Gqa,
813        n_layers: 3,
814        hidden_dim: 32,
815        n_heads: 4,
816        n_kv_heads: 2,
817        head_dim: 8,
818        vocab_size: 32,
819        rope_theta: 10000.0,
820        rms_norm_eps: 1e-5,
821        moe: MoeLayerConfig {
822            expert_weights_scale: 1.0,
823            n_experts: 3,
824            n_experts_active: 1,
825            n_shared_experts: 1,
826            hidden_dim: 32,
827            expert_ffn_dim: 32,
828            gating: GatingFunction::Softmax,
829            norm_topk_prob: true,
830            expert_group_count: None,
831            expert_group_used_count: None,
832        },
833        n_dense_leading_layers: 1,
834        rope_freqs: None,
835        rope_attn_factor: 1.0,
836        rope_dim: None,
837        rope_freqs_long: None,
838        rope_freqs_short: None,
839        rope_orig_ctx: None,
840        rope_layout: RopeLayout::Neox,
841        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
842        swa_pattern: None,
843        attn_logit_softcap: None,
844        final_logit_softcap: None,
845        embedding_scale: None,
846        attention_scale: None,
847        rope_theta_swa: None,
848        ffn_activation: FfnActivation::Swiglu,
849        best_effort_fields: &["this is a synthetic mixed dense/MoE test fixture, not a real model"],
850    }
851}
852
853#[cfg(test)]
854mod tests {
855    use super::*;
856
857    #[test]
858    fn rope_layout_for_gguf_architecture_matches_llama_cpp() {
859        // Confirmed against llama.cpp's llama_model_rope_type
860        // (src/llama-model.cpp): llama -> NORM, olmoe/qwen2/phi3/gemma -> NEOX.
861        assert_eq!(RopeLayout::for_gguf_architecture("llama"), RopeLayout::Norm);
862        assert_eq!(
863            RopeLayout::for_gguf_architecture("llama4"),
864            RopeLayout::Norm
865        );
866        assert_eq!(
867            RopeLayout::for_gguf_architecture("deepseek2"),
868            RopeLayout::Norm
869        );
870        assert_eq!(RopeLayout::for_gguf_architecture("olmoe"), RopeLayout::Neox);
871        assert_eq!(RopeLayout::for_gguf_architecture("qwen2"), RopeLayout::Neox);
872        assert_eq!(
873            RopeLayout::for_gguf_architecture("qwen2moe"),
874            RopeLayout::Neox
875        );
876        assert_eq!(RopeLayout::for_gguf_architecture("qwen3"), RopeLayout::Neox);
877        assert_eq!(RopeLayout::for_gguf_architecture("phi3"), RopeLayout::Neox);
878        assert_eq!(
879            RopeLayout::for_gguf_architecture("gemma3"),
880            RopeLayout::Neox
881        );
882        // Unknown architectures keep the historical Neox default at this
883        // helper only; load-time uses capability::resolve_architecture and
884        // fails closed instead of guessing.
885        assert_eq!(
886            RopeLayout::for_gguf_architecture("totally-unknown-arch"),
887            RopeLayout::Neox
888        );
889    }
890
891    /// gpt-oss's real shape: a 128-token window on every other layer.
892    /// A KV block size of 128 or any divisor of it is fine; 48 or 256
893    /// are not, and the config layer must round down rather than hand
894    /// the cache something it will refuse (or, worse, accept).
895    #[test]
896    fn an_alternating_swa_model_constrains_the_block_layout() {
897        let mut cfg = test_dense_fixture();
898        cfg.n_layers = 24;
899        cfg.sliding_window = Some(128);
900        cfg.swa_pattern = Some(2);
901
902        // Half the layers are full-attention, but the model is still
903        // constrained: one mis-aligned sliding layer is enough.
904        assert!(cfg.layer_sliding_window(1).is_none() || cfg.layer_sliding_window(0).is_none());
905        assert_eq!(cfg.kv_block_window(), Some(128));
906
907        let layout = cfg.kv_block_layout(256);
908        assert_eq!(layout.block_size(), 128, "256 must round down, not up");
909        assert_eq!(layout.sliding_window(), Some(128));
910        assert_eq!(layout.blocks_per_window(), Some(1));
911
912        assert_eq!(cfg.kv_block_layout(48).block_size(), 32);
913        assert_eq!(cfg.kv_block_layout(32).block_size(), 32);
914    }
915
916    /// Gemma-3: window 512, every 6th layer full-attention.
917    #[test]
918    fn a_gemma3_shaped_model_takes_its_window_from_the_sliding_layers() {
919        let mut cfg = test_dense_fixture();
920        cfg.n_layers = 30;
921        cfg.sliding_window = Some(512);
922        cfg.swa_pattern = Some(6);
923        assert!(
924            cfg.layer_sliding_window(5).is_none(),
925            "every 6th layer is full-attention"
926        );
927        assert_eq!(cfg.kv_block_window(), Some(512));
928        assert_eq!(cfg.kv_block_layout(100).block_size(), 64);
929        assert_eq!(cfg.kv_block_layout(64).blocks_per_window(), Some(8));
930    }
931
932    #[test]
933    fn a_full_causal_model_keeps_the_block_size_it_was_given() {
934        let mut cfg = test_dense_fixture();
935        cfg.sliding_window = None;
936        cfg.swa_pattern = None;
937        assert_eq!(cfg.kv_block_window(), None);
938        let layout = cfg.kv_block_layout(48);
939        assert_eq!(layout.block_size(), 48);
940        assert_eq!(layout.sliding_window(), None);
941    }
942
943    #[test]
944    fn all_presets_have_consistent_moe_hidden_dim() {
945        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
946            assert_eq!(
947                cfg.hidden_dim, cfg.moe.hidden_dim,
948                "{}: attention hidden_dim and MoE hidden_dim must match",
949                cfg.name
950            );
951        }
952    }
953
954    #[test]
955    fn all_presets_route_fewer_experts_than_total() {
956        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
957            assert!(
958                cfg.moe.n_experts_active < cfg.moe.n_experts,
959                "{}: active experts must be a sparse subset of total experts",
960                cfg.name
961            );
962        }
963    }
964
965    #[test]
966    fn all_presets_have_divisible_heads() {
967        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
968            assert_eq!(
969                cfg.n_heads % cfg.n_kv_heads,
970                0,
971                "{}: n_heads must be a multiple of n_kv_heads for GQA grouping",
972                cfg.name
973            );
974        }
975    }
976
977    #[test]
978    fn every_preset_declares_its_uncertain_fields() {
979        // This is a documentation-honesty test: any preset with zero
980        // best_effort_fields would be silently overclaiming precision
981        // we don't have. Fail loudly if that ever happens.
982        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
983            assert!(
984                !cfg.best_effort_fields.is_empty(),
985                "{}: must disclose which fields are unconfirmed estimates",
986                cfg.name
987            );
988        }
989    }
990
991    /// Kimi K3's `kda_layers`/`full_attn_layers` were transcribed by
992    /// hand from the real published config.json; this test guards
993    /// against a transcription slip (duplicate, out-of-range, or
994    /// missing layer index) rather than trusting the transcription.
995    #[test]
996    fn kimi_k3_hybrid_attention_layers_partition_every_layer_exactly_once() {
997        let cfg = kimi_k3();
998        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
999            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
1000        };
1001
1002        let mut seen = std::collections::HashSet::new();
1003        for &l in hybrid
1004            .kda_layers
1005            .iter()
1006            .chain(hybrid.full_attn_layers.iter())
1007        {
1008            assert!(
1009                (1..=cfg.n_layers).contains(&l),
1010                "layer {l} is out of the published 1..={} range",
1011                cfg.n_layers
1012            );
1013            assert!(
1014                seen.insert(l),
1015                "layer {l} appears in both/either list twice"
1016            );
1017        }
1018        // Dense-vs-MoE (n_dense_leading_layers) and attention-type
1019        // (KDA vs Gated MLA) are independent per-layer properties in
1020        // the real config -- e.g. layer 1 is both the sole dense
1021        // leading layer *and* a KDA layer -- so every one of the 93
1022        // layers, dense or not, is covered by exactly one of these two
1023        // lists (confirmed: 69 + 24 == 93, not 93 - 1).
1024        assert_eq!(
1025            hybrid.kda_layers.len() + hybrid.full_attn_layers.len(),
1026            cfg.n_layers,
1027            "every layer must be assigned exactly one of KDA or Gated MLA"
1028        );
1029        assert_eq!(
1030            hybrid.kda_layers.len(),
1031            69,
1032            "expected 69 KDA layers per the published config"
1033        );
1034        assert_eq!(
1035            hybrid.full_attn_layers.len(),
1036            24,
1037            "expected 24 Gated MLA layers per the published config"
1038        );
1039    }
1040
1041    #[test]
1042    fn layer_attention_kind_is_gqa_for_every_layer_of_a_gqa_model() {
1043        let cfg = glm_5_2();
1044        for l in 0..cfg.n_layers {
1045            assert_eq!(cfg.layer_attention_kind(l), LayerAttentionKind::Gqa);
1046        }
1047    }
1048
1049    #[test]
1050    fn layer_attention_kind_classifies_every_kimi_k3_layer_without_panicking() {
1051        let cfg = kimi_k3();
1052        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
1053            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
1054        };
1055        for l in 0..cfg.n_layers {
1056            let kind = cfg.layer_attention_kind(l);
1057            let one_indexed = l + 1;
1058            if hybrid.kda_layers.contains(&one_indexed) {
1059                assert_eq!(kind, LayerAttentionKind::KimiKda);
1060            } else {
1061                assert_eq!(kind, LayerAttentionKind::KimiMla);
1062            }
1063        }
1064    }
1065
1066    #[test]
1067    fn layer_attention_kind_matches_the_real_published_layer_1_and_4() {
1068        // Layer 1 (1-indexed, so index 0 here) is published as KDA;
1069        // layer 4 (index 3) is published as the first Gated MLA layer.
1070        let cfg = kimi_k3();
1071        assert_eq!(cfg.layer_attention_kind(0), LayerAttentionKind::KimiKda);
1072        assert_eq!(cfg.layer_attention_kind(3), LayerAttentionKind::KimiMla);
1073    }
1074
1075    #[test]
1076    fn kimi_k3_mla_q_head_dim_matches_gqa_placeholder_head_dim() {
1077        // The Gqa-placeholder head_dim above is deliberately set to
1078        // Gated MLA's combined q_head_dim (qk_nope + qk_rope) so the
1079        // placeholder path at least reflects a real dimension from the
1080        // published config rather than an arbitrary guess.
1081        let cfg = kimi_k3();
1082        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
1083            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
1084        };
1085        assert_eq!(
1086            cfg.head_dim,
1087            hybrid.mla.qk_nope_head_dim + hybrid.mla.qk_rope_head_dim
1088        );
1089    }
1090
1091    #[test]
1092    fn approx_active_params_is_nonzero_and_finite_order_of_magnitude() {
1093        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
1094            let approx = cfg.approx_active_params_per_token();
1095            // Sanity band: active params/token for these models is
1096            // reported in the tens of billions; this is a loose
1097            // order-of-magnitude check (1e9 to 1e12), not a precise
1098            // parameter-count reproduction.
1099            assert!(
1100                approx > 1_000_000_000 && approx < 1_000_000_000_000,
1101                "{}: approx_active_params_per_token={approx} is outside a plausible range",
1102                cfg.name
1103            );
1104        }
1105    }
1106}
1107
1108#[cfg(test)]
1109mod longrope_tests {
1110    use super::*;
1111
1112    fn cfg_with_factors() -> ModelConfig {
1113        let mut c = test_dense_fixture();
1114        c.rope_orig_ctx = Some(4096);
1115        c.rope_freqs_short = Some(vec![1.0; 48]);
1116        c.rope_freqs_long = Some((0..48).map(|i| 1.0 + i as f32).collect());
1117        c.rope_freqs = None;
1118        c
1119    }
1120
1121    /// llama.cpp `llama_model::get_rope_factors`: long only when the
1122    /// run's context exceeds `original_context_length`. Phi-4-mini's
1123    /// short set is all ones, so picking long at 4096 would apply a
1124    /// correction the model never asked for at that length.
1125    #[test]
1126    fn long_set_only_above_the_original_context() {
1127        let mut c = cfg_with_factors();
1128        c.apply_runtime_context(4096);
1129        assert_eq!(
1130            c.rope_freqs.as_ref().unwrap()[1],
1131            1.0,
1132            "at the threshold, short"
1133        );
1134
1135        let mut c = cfg_with_factors();
1136        c.apply_runtime_context(4097);
1137        assert_eq!(c.rope_freqs.as_ref().unwrap()[1], 2.0, "above it, long");
1138
1139        let mut c = cfg_with_factors();
1140        c.apply_runtime_context(1024);
1141        assert_eq!(c.rope_freqs.as_ref().unwrap()[1], 1.0, "below it, short");
1142    }
1143
1144    /// `rope_freqs.weight` (Llama 3) is not a LongRoPE set and outranks
1145    /// one, the same precedence llama.cpp gives it. The loader encodes
1146    /// that by leaving the long/short pair empty whenever the explicit
1147    /// tensor is present, so the runtime re-pick has nothing to apply.
1148    #[test]
1149    fn an_explicit_rope_freqs_tensor_is_never_overridden() {
1150        let mut c = test_dense_fixture();
1151        c.rope_freqs = Some(vec![7.0; 48]);
1152        c.rope_orig_ctx = Some(4096);
1153        c.rope_freqs_long = None;
1154        c.rope_freqs_short = None;
1155        c.apply_runtime_context(131072);
1156        assert_eq!(c.rope_freqs.as_ref().unwrap()[0], 7.0);
1157    }
1158
1159    /// A checkpoint with neither set must come back untouched, so the
1160    /// call is free to sit on every load path.
1161    #[test]
1162    fn models_without_longrope_are_untouched() {
1163        let mut c = test_dense_fixture();
1164        c.rope_freqs = None;
1165        c.apply_runtime_context(8192);
1166        assert!(c.rope_freqs.is_none());
1167        assert!(c.rope_orig_ctx.is_none());
1168    }
1169}