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    pub rope_freqs: Option<Vec<f32>>,
210    /// LongRoPE's two candidate factor sets, kept so the choice between
211    /// them can be made when the *run's* context size is known rather
212    /// than at parse time. llama.cpp picks per request
213    /// (`llama_model::get_rope_factors` reads `cparams.n_ctx_seq`), and
214    /// the two sets are not interchangeable: Phi-4-mini's short set is
215    /// all ones (no correction at all) while its long set reaches 47.
216    /// Choosing from the checkpoint's advertised 131072 when the user
217    /// runs at 4096 is a different model.
218    pub rope_freqs_long: Option<Vec<f32>>,
219    pub rope_freqs_short: Option<Vec<f32>>,
220    /// `<arch>.rope.scaling.original_context_length` — the threshold the
221    /// choice above is made against.
222    pub rope_orig_ctx: Option<usize>,
223    /// Rotary width when it is narrower than `head_dim`
224    /// (`<arch>.rope.dimension_count`, llama.cpp `hparams.n_rot`).
225    /// `None` means the whole head rotates, which is the common case.
226    /// Phi-3/Phi-4 rotate 96 of 128.
227    pub rope_dim: Option<usize>,
228    /// LongRoPE/YaRN magnitude scaling (`<arch>.rope.scaling.attn_factor`,
229    /// llama.cpp `hparams.rope_attn_factor` folded into
230    /// `cparams.yarn_attn_factor` at `llama-context.cpp:231`, then applied
231    /// as ggml `rope_yarn`'s `mscale`, which multiplies *both* `cos` and
232    /// `sin` — so it scales the RoPE'd vector, at every position, whether
233    /// or not any frequency correction is active.
234    ///
235    /// Phi-4-mini ships `1.1902381`. Ignoring it does not merely change
236    /// long-context behaviour: q and k are both scaled, so every attention
237    /// logit is off by `attn_factor²` and the softmax is sharper than the
238    /// model's. Measured symptom: ferrox and llama.cpp diverge from the
239    /// eighth token of a greedy completion on the same GGUF.
240    ///
241    /// `1.0` for every architecture that does not set the key.
242    pub rope_attn_factor: f32,
243    /// RoPE pairing convention for this architecture -- see
244    /// `RopeLayout`. Independently of `rope_freqs`: a Llama checkpoint
245    /// needs both `Norm` pairing *and* the per-band frequency factors.
246    pub rope_layout: RopeLayout,
247    /// How Q/K RMSNorm weights are applied when present (see
248    /// [`crate::capability::QkNormStyle`]).
249    pub qk_norm_style: crate::capability::QkNormStyle,
250    /// Gemma 2+/3 alternating SWA period. When `Some(p)`, layer `il`
251    /// uses sliding-window attention iff `(il + 1) % p != 0` (llama.cpp
252    /// `set_swa_pattern` convention); every `p`-th layer is full-attn.
253    pub swa_pattern: Option<usize>,
254    /// Attention logit soft-capping (Gemma 2+). Applied as
255    /// `softcap * tanh(score / softcap)` before softmax.
256    pub attn_logit_softcap: Option<f32>,
257    /// Final logit soft-capping (Gemma 2+). Applied to lm_head output.
258    pub final_logit_softcap: Option<f32>,
259    /// Input embedding scale (Gemma: `sqrt(hidden_dim)`).
260    pub embedding_scale: Option<f32>,
261    /// Optional override for the attention score scale baked into Q
262    /// *instead of* the kernel's default `1/sqrt(head_dim)`. When set,
263    /// callers must pass `score_scale = 1.0` into the attention kernel
264    /// (llama.cpp Gemma: scale Q then `build_attn(..., 1.0f)`). Prefer
265    /// leaving this `None` when the override equals `1/sqrt(head_dim)`.
266    pub attention_scale: Option<f32>,
267    /// RoPE base used on SWA layers (Gemma 3: defaults to `10000` when
268    /// the GGUF omits `rope.freq_base_swa`; full-attn layers keep
269    /// [`Self::rope_theta`]).
270    pub rope_theta_swa: Option<f32>,
271    /// Dense/MoE FFN activation pairing.
272    pub ffn_activation: FfnActivation,
273    /// Every field on this config that is a best-effort estimate rather
274    /// than a confirmed value from an official config.json / GGUF file.
275    pub best_effort_fields: &'static [&'static str],
276}
277
278/// Dense / expert FFN non-linearity used by the generic decoder.
279#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
280pub enum FfnActivation {
281    /// `silu(gate) * up` with separate gate/up matrices (Llama / Qwen).
282    #[default]
283    Swiglu,
284    /// Phi-3 fused gate+up in one `ffn_up` matrix (`2 * n_ff` rows).
285    SwigluFused,
286    /// Gemma GeGLU: `gelu(gate) * up`.
287    Gelu,
288}
289
290impl ModelConfig {
291    /// Re-picks the LongRoPE factor set now that the run's context size
292    /// is known, matching llama.cpp `llama_model::get_rope_factors`:
293    /// `rope_freqs.weight` (Llama 3) always wins; otherwise the long set
294    /// applies only when the context exceeds
295    /// `rope.scaling.original_context_length`, and the short set
296    /// otherwise.
297    ///
298    /// A no-op for every checkpoint that ships neither set, which is all
299    /// of them except the Phi-3/Phi-4 family today.
300    pub fn apply_runtime_context(&mut self, ctx: usize) {
301        let (Some(orig), true) = (
302            self.rope_orig_ctx,
303            self.rope_freqs_long.is_some() || self.rope_freqs_short.is_some(),
304        ) else {
305            return;
306        };
307        let picked = if ctx > orig {
308            self.rope_freqs_long.as_ref()
309        } else {
310            self.rope_freqs_short.as_ref()
311        };
312        if let Some(f) = picked
313            .or(self.rope_freqs_long.as_ref())
314            .or(self.rope_freqs_short.as_ref())
315        {
316            self.rope_freqs = Some(f.clone());
317        }
318    }
319
320    /// True if layer `layer_idx` (0-indexed) should be built as an
321    /// ordinary dense FFN rather than this model's MoE topology.
322    pub fn layer_is_dense(&self, layer_idx: usize) -> bool {
323        layer_idx < self.n_dense_leading_layers
324    }
325
326    /// Sliding-window size for layer `il`, honouring Gemma-style
327    /// alternating SWA patterns. `None` means full causal attention.
328    pub fn layer_sliding_window(&self, layer_idx: usize) -> Option<usize> {
329        let window = self.sliding_window?;
330        match self.swa_pattern {
331            None => Some(window),
332            Some(period) if period > 1 => {
333                // llama.cpp `set_swa_pattern` (dense_first=false):
334                // `is_swa = (il % period) < (period - 1)` — equivalent to
335                // full attention when `(il + 1) % period == 0`.
336                if (layer_idx + 1).is_multiple_of(period) {
337                    None
338                } else {
339                    Some(window)
340                }
341            }
342            Some(_) => Some(window),
343        }
344    }
345
346    /// RoPE frequency base for layer `il` (SWA layers may differ).
347    pub fn layer_rope_theta(&self, layer_idx: usize) -> f32 {
348        match (self.layer_sliding_window(layer_idx), self.rope_theta_swa) {
349            (Some(_), Some(theta)) => theta,
350            _ => self.rope_theta,
351        }
352    }
353
354    /// Which attention mechanism layer `layer_idx` (0-indexed, ferrox's
355    /// usual convention) uses. For `AttentionKind::Gqa` every layer is
356    /// `LayerAttentionKind::Gqa`; for `AttentionKind::KimiHybrid`, looks
357    /// up `layer_idx + 1` (the real `kda_layers`/`full_attn_layers`
358    /// lists are 1-indexed -- see `KimiHybridAttention`'s doc comment)
359    /// in those real per-layer lists.
360    ///
361    /// # Panics
362    /// If `layer_idx` isn't covered by either list of a `KimiHybrid`
363    /// config -- can't happen for `kimi_k3()`, whose lists are tested
364    /// (`kimi_k3_hybrid_attention_layers_partition_every_layer_exactly_once`)
365    /// to partition every layer with no gaps, but a caller building a
366    /// custom `KimiHybridAttention` must uphold the same invariant.
367    pub fn layer_attention_kind(&self, layer_idx: usize) -> LayerAttentionKind {
368        match &self.attention {
369            AttentionKind::Gqa => LayerAttentionKind::Gqa,
370            AttentionKind::KimiHybrid(hybrid) => {
371                let one_indexed = layer_idx + 1;
372                if hybrid.kda_layers.contains(&one_indexed) {
373                    LayerAttentionKind::KimiKda
374                } else if hybrid.full_attn_layers.contains(&one_indexed) {
375                    LayerAttentionKind::KimiMla
376                } else {
377                    panic!(
378                        "layer {layer_idx} (1-indexed {one_indexed}) is in neither \
379                         kda_layers nor full_attn_layers"
380                    )
381                }
382            }
383        }
384    }
385
386    /// Total parameter count implied by the MoE config, as a sanity
387    /// check against the publicly reported total (this is an order of
388    /// magnitude check, not an exact parameter-count reproduction).
389    pub fn approx_active_params_per_token(&self) -> usize {
390        let attn_params_per_layer = 4 * self.hidden_dim * self.hidden_dim; // q,k,v,o (rough)
391        let active_experts = self.moe.n_experts_active + self.moe.n_shared_experts;
392        let expert_params = active_experts * 3 * self.moe.hidden_dim * self.moe.expert_ffn_dim; // gate,up,down
393        self.n_layers * (attn_params_per_layer + expert_params)
394    }
395}
396
397/// GLM-5.2 (Z.ai) **structural sketch only** — not a supported real
398/// inference path. Real DSA lives in `glm_dsa` / `glm52_decoder` and is
399/// not wired into `Decoder` / `ferrox-server`. This preset drives
400/// smoke/bench with synthetic GQA weights only (~744B / ~40B active
401/// hparams as published placeholders).
402pub fn glm_5_2() -> ModelConfig {
403    ModelConfig {
404        sliding_window: None,
405        name: "glm-5.2",
406        attention: AttentionKind::Gqa,
407        n_layers: 92,
408        hidden_dim: 6144,
409        n_heads: 48,
410        n_kv_heads: 8,
411        head_dim: 128,
412        vocab_size: 151552,
413        rope_theta: 1_000_000.0,
414        rms_norm_eps: 1e-5,
415        moe: MoeLayerConfig {
416            n_experts: 256,
417            n_experts_active: 8,
418            n_shared_experts: 1,
419            hidden_dim: 6144,
420            expert_ffn_dim: 2048,
421            // Sigmoid, not softmax: reading ik_llama.cpp's real GGUF
422            // hparams-loading source (llama-hparams.cpp,
423            // LLM_ARCH_GLM4_MOE case) directly showed GLM4-MoE-family
424            // models default to sigmoid gating with post-selection
425            // score renormalization. GLM-5.2 is presumed to continue
426            // this lineage; not confirmed against GLM-5.2's own
427            // config.json (unavailable in this environment).
428            gating: GatingFunction::Sigmoid,
429            norm_topk_prob: true,
430         expert_group_count: None, expert_group_used_count: None,},
431        // No evidence found (via ik_llama.cpp source or public
432        // reporting) that GLM-5.2 skips MoE on any leading layers;
433        // defaulting to 0 (every layer uses this model's MoE
434        // topology) rather than assuming DeepSeek's convention
435        // applies here too.
436        n_dense_leading_layers: 0,
437        rope_freqs: None,
438        rope_attn_factor: 1.0,
439        rope_dim: None,
440        rope_freqs_long: None,
441        rope_freqs_short: None,
442        rope_orig_ctx: None,
443        // Placeholder GQA path; real GLM-5.2 DSA uses interleaved RoPE
444        // via `glm_dsa`/`mla`, not this preset's Decoder path.
445        rope_layout: RopeLayout::Neox,
446        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
447        swa_pattern: None,
448        attn_logit_softcap: None,
449        final_logit_softcap: None,
450        embedding_scale: None,
451        attention_scale: None,
452        rope_theta_swa: None,
453        ffn_activation: FfnActivation::Swiglu,
454        best_effort_fields: &[
455            "n_layers",
456            "hidden_dim",
457            "n_heads",
458            "n_kv_heads",
459            "head_dim",
460            "rope_theta",
461            "moe.expert_ffn_dim",
462            "moe.n_shared_experts",
463            "moe.gating (sigmoid assumed from GLM4-MoE-family convention found in ik_llama.cpp source, not confirmed for GLM-5.2 specifically)",
464        ],
465    }
466}
467
468/// DeepSeek V4 Pro **structural sketch only** — CSA/HCA is not on this
469/// GQA `Decoder` path. Real primitives live under
470/// `deepseek_v4_attention` / `hyper_connections` and are not assembled
471/// into a served decoder yet. Hparams (~1.6T / ~49B active) are
472/// placeholders for smoke/bench.
473pub fn deepseek_v4_pro() -> ModelConfig {
474    ModelConfig {
475        sliding_window: None,
476        name: "deepseek-v4-pro",
477        attention: AttentionKind::Gqa,
478        n_layers: 96,
479        hidden_dim: 7168,
480        n_heads: 56,
481        n_kv_heads: 8,
482        head_dim: 128,
483        vocab_size: 129280,
484        rope_theta: 1_000_000.0,
485        rms_norm_eps: 1e-6,
486        moe: MoeLayerConfig {
487            n_experts: 385,
488            n_experts_active: 6,
489            n_shared_experts: 1,
490            hidden_dim: 7168,
491            expert_ffn_dim: 2048,
492            // Sigmoid, not softmax: this is the stronger-confidence of
493            // the two sigmoid-gating corrections in this file.
494            // DeepSeek-V3's own published technical report explicitly
495            // documents computing per-expert affinity via sigmoid and
496            // renormalizing only the selected experts' scores to sum
497            // to one; reading ik_llama.cpp's real GGUF hparams-loading
498            // source (llama-hparams.cpp, LLM_ARCH_DEEPSEEK2 case)
499            // confirmed this is exactly what that code path defaults
500            // to for the DeepSeek-2/3 lineage. DeepSeek V4 Pro is
501            // presumed to continue using sigmoid gating for the same
502            // reason; not confirmed against V4 Pro's own config.json.
503            gating: GatingFunction::Sigmoid,
504            norm_topk_prob: true,
505         expert_group_count: None, expert_group_used_count: None,},
506        // DeepSeek-V3's own published technical report documents the
507        // first 3 transformer layers as dense (ordinary FFN, no
508        // expert routing), with MoE starting from layer 4 onward;
509        // ik_llama.cpp's real hparams-loading source
510        // (LLM_KV_LEADING_DENSE_BLOCK_COUNT) confirms this is a real,
511        // loaded GGUF metadata field for the DeepSeek-2/3 lineage.
512        // DeepSeek V4 Pro is presumed to continue this convention;
513        // not confirmed against V4 Pro's own config.json.
514        n_dense_leading_layers: 3,
515        rope_freqs: None,
516        rope_attn_factor: 1.0,
517        rope_dim: None,
518        rope_freqs_long: None,
519        rope_freqs_short: None,
520        rope_orig_ctx: None,
521        // llama.cpp maps LLM_ARCH_DEEPSEEK4 -> LLAMA_ROPE_TYPE_NORM.
522        rope_layout: RopeLayout::Norm,
523        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
524        swa_pattern: None,
525        attn_logit_softcap: None,
526        final_logit_softcap: None,
527        embedding_scale: None,
528        attention_scale: None,
529        rope_theta_swa: None,
530        ffn_activation: FfnActivation::Swiglu,
531        best_effort_fields: &[
532            "n_layers",
533            "hidden_dim",
534            "n_heads",
535            "n_kv_heads",
536            "head_dim",
537            "moe.expert_ffn_dim",
538            "attention_variant (CSA/HCA hybrid NOT implemented, GQA fallback in use)",
539            "moe.gating (sqrtsoftplus: confirmed for real V4 in llama.cpp PR #24162; this preset still uses Sigmoid on the wrong GQA sketch path)",
540            "n_dense_leading_layers (3: same confidence basis as gating above, DeepSeek-V3 technical report + ik_llama.cpp source, not confirmed for V4 Pro)",
541        ],
542    }
543}
544
545/// Kimi K3 **structural sketch only** for the generic GQA `Decoder`.
546/// Real checkpoint work uses the dedicated Kimi stack (`kimi_loader` /
547/// `KimiEngine`); slice-verified, not a full end-to-end run. Do not
548/// treat this preset as a runnable Kimi substitute.
549pub fn kimi_k3() -> ModelConfig {
550    ModelConfig {
551        sliding_window: None,
552        name: "kimi-k3",
553        n_layers: 93,
554        hidden_dim: 7168,
555        // n_heads/n_kv_heads/head_dim describe the Gqa fallback
556        // Decoder actually runs today, not Kimi K3's real attention
557        // (see `attention` below) -- kept at reasonable stand-in
558        // values (matching MLA's num_heads=96 and combined
559        // qk_nope+qk_rope head dim) rather than deleted, so the
560        // placeholder path stays runnable.
561        n_heads: 96,
562        n_kv_heads: 96,
563        head_dim: 192,
564        vocab_size: 163840,
565        // Not present in the published text_config; RoPE only ever
566        // applies to Gated MLA's 64-dim qk_rope_head_dim slice in the
567        // real architecture, and Decoder doesn't implement that slicing
568        // yet, so this remains an unconfirmed placeholder.
569        rope_theta: 1_000_000.0,
570        rms_norm_eps: 1e-5,
571        moe: MoeLayerConfig {
572            n_experts: 896,
573            n_experts_active: 16,
574            n_shared_experts: 2,
575            hidden_dim: 7168,
576            expert_ffn_dim: 3072,
577            // Confirmed directly from the real config.json:
578            // "moe_router_activation_func": "sigmoid".
579            gating: GatingFunction::Sigmoid,
580            norm_topk_prob: true,
581         expert_group_count: None, expert_group_used_count: None,},
582        // Confirmed directly from the real config.json:
583        // "first_k_dense_replace": 1.
584        n_dense_leading_layers: 1,
585        // Kimi K3's real, published attention topology (verified
586        // against huggingface.co/moonshotai/Kimi-K3/config.json's
587        // linear_attn_config block and the real KimiDeltaAttention /
588        // KimiMLAAttention reference implementations in
589        // modeling_kimi_linear.py) -- not yet wired into Decoder's
590        // forward pass, which still runs the Gqa placeholder above for
591        // every layer regardless of this field.
592        attention: AttentionKind::KimiHybrid(KimiHybridAttention {
593            kda_layers: vec![
594                1, 2, 3, 5, 6, 7, 9, 10, 11, 13, 14, 15, 17, 18, 19, 21, 22, 23, 25, 26, 27, 29,
595                30, 31, 33, 34, 35, 37, 38, 39, 41, 42, 43, 45, 46, 47, 49, 50, 51, 53, 54, 55,
596                57, 58, 59, 61, 62, 63, 65, 66, 67, 69, 70, 71, 73, 74, 75, 77, 78, 79, 81, 82,
597                83, 85, 86, 87, 89, 90, 91,
598            ],
599            full_attn_layers: vec![
600                4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84,
601                88, 92, 93,
602            ],
603            mla: MlaConfig {
604                num_heads: 96,
605                q_lora_rank: 1536,
606                kv_lora_rank: 512,
607                qk_nope_head_dim: 128,
608                qk_rope_head_dim: 64,
609                v_head_dim: 128,
610                use_output_gate: true,
611                // Real, confirmed: Kimi K3's `KimiMLAAttention.forward`
612                // never rotates -- see `MlaConfig::rope`'s doc comment.
613                rope: None,
614            },
615            kda: KdaConfig {
616                num_heads: 96,
617                head_dim: 128,
618                short_conv_kernel_size: 4,
619                gate_lower_bound: -5.0,
620                use_full_rank_gate: true,
621            },
622        }),
623        rope_freqs: None,
624        rope_attn_factor: 1.0,
625        rope_dim: None,
626        rope_freqs_long: None,
627        rope_freqs_short: None,
628        rope_orig_ctx: None,
629        // GQA placeholder path only; real Kimi attention is rope-less MLA
630        // or KDA and never reaches Decoder::apply_rope_head.
631        rope_layout: RopeLayout::Neox,
632        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
633        swa_pattern: None,
634        attn_logit_softcap: None,
635        final_logit_softcap: None,
636        embedding_scale: None,
637        attention_scale: None,
638        rope_theta_swa: None,
639        ffn_activation: FfnActivation::Swiglu,
640        best_effort_fields: &[
641            "n_heads/n_kv_heads/head_dim (describe the unimplemented Gqa placeholder, not Kimi K3's real MLA/KDA attention -- see `attention` field)",
642            "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)",
643            "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)",
644        ],
645    }
646}
647
648/// Matches the generated on-disk fixture exactly (hidden_dim, head
649/// counts, ffn_dim, vocab, rope_theta, eps).
650/// Used by `ferrox inspect-run` and the cross-validation test in
651/// `crates/ferrox-models/tests/gguf_roundtrip.rs` to prove the real
652/// GGUF loader + forward pass produce the same numbers as an
653/// independent NumPy reference implementation reading the same file.
654pub fn test_dense_fixture() -> ModelConfig {
655    ModelConfig {
656        sliding_window: None,
657        name: "ferrox-test-dense",
658        attention: AttentionKind::Gqa,
659        n_layers: 2,
660        hidden_dim: 32,
661        n_heads: 4,
662        n_kv_heads: 2,
663        head_dim: 8,
664        vocab_size: 32,
665        rope_theta: 10000.0,
666        rms_norm_eps: 1e-5,
667        moe: MoeLayerConfig {
668            n_experts: 1,
669            n_experts_active: 1,
670            n_shared_experts: 0,
671            hidden_dim: 32,
672            expert_ffn_dim: 32,
673            gating: GatingFunction::Softmax,
674            norm_topk_prob: true,
675            expert_group_count: None,
676            expert_group_used_count: None,
677        },
678        n_dense_leading_layers: 0,
679        rope_freqs: None,
680        rope_attn_factor: 1.0,
681        rope_dim: None,
682        rope_freqs_long: None,
683        rope_freqs_short: None,
684        rope_orig_ctx: None,
685        // Matches the independent reference's split-half apply_rope.
686        rope_layout: RopeLayout::Neox,
687        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
688        swa_pattern: None,
689        attn_logit_softcap: None,
690        final_logit_softcap: None,
691        embedding_scale: None,
692        attention_scale: None,
693        rope_theta_swa: None,
694        ffn_activation: FfnActivation::Swiglu,
695        best_effort_fields: &["this is a synthetic test fixture, not a real model"],
696    }
697}
698
699/// Matches the generated on-disk multi-expert MoE fixture: 4 experts,
700/// top-2 routing, 1 shared
701/// expert, packed 3D expert tensors. Used to verify the previously-
702/// untested multi-expert loading path (`split_expert_tensor` in
703/// `ferrox-models::loader`) against a real file, the same way
704/// `test_dense_fixture` verifies the single-expert path.
705pub fn test_moe_fixture() -> ModelConfig {
706    ModelConfig {
707        sliding_window: None,
708        name: "ferrox-test-moe",
709        attention: AttentionKind::Gqa,
710        n_layers: 2,
711        hidden_dim: 32,
712        n_heads: 4,
713        n_kv_heads: 2,
714        head_dim: 8,
715        vocab_size: 32,
716        rope_theta: 10000.0,
717        rms_norm_eps: 1e-5,
718        moe: MoeLayerConfig {
719            n_experts: 4,
720            n_experts_active: 2,
721            n_shared_experts: 1,
722            hidden_dim: 32,
723            expert_ffn_dim: 32,
724            gating: GatingFunction::Softmax,
725            norm_topk_prob: true,
726            expert_group_count: None,
727            expert_group_used_count: None,
728        },
729        n_dense_leading_layers: 0,
730        rope_freqs: None,
731        rope_attn_factor: 1.0,
732        rope_dim: None,
733        rope_freqs_long: None,
734        rope_freqs_short: None,
735        rope_orig_ctx: None,
736        rope_layout: RopeLayout::Neox,
737        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
738        swa_pattern: None,
739        attn_logit_softcap: None,
740        final_logit_softcap: None,
741        embedding_scale: None,
742        attention_scale: None,
743        rope_theta_swa: None,
744        ffn_activation: FfnActivation::Swiglu,
745        best_effort_fields: &["this is a synthetic multi-expert test fixture, not a real model"],
746    }
747}
748
749/// Matches the generated on-disk mixed-topology fixture: 3 layers, the
750/// first of which is
751/// an ordinary dense FFN and the remaining two are genuine MoE (3
752/// experts, top-1 routing, 1 shared expert each). Used to verify the
753/// "leading dense layers" loading path
754/// (`ModelConfig::layer_is_dense`) against a real file -- the pattern
755/// found in DeepSeek-2/3-family models via ik_llama.cpp's source
756/// (`LLM_KV_LEADING_DENSE_BLOCK_COUNT`), which was previously only
757/// documented, not implemented or tested.
758pub fn test_mixed_fixture() -> ModelConfig {
759    ModelConfig {
760        sliding_window: None,
761        name: "ferrox-test-mixed",
762        attention: AttentionKind::Gqa,
763        n_layers: 3,
764        hidden_dim: 32,
765        n_heads: 4,
766        n_kv_heads: 2,
767        head_dim: 8,
768        vocab_size: 32,
769        rope_theta: 10000.0,
770        rms_norm_eps: 1e-5,
771        moe: MoeLayerConfig {
772            n_experts: 3,
773            n_experts_active: 1,
774            n_shared_experts: 1,
775            hidden_dim: 32,
776            expert_ffn_dim: 32,
777            gating: GatingFunction::Softmax,
778            norm_topk_prob: true,
779            expert_group_count: None,
780            expert_group_used_count: None,
781        },
782        n_dense_leading_layers: 1,
783        rope_freqs: None,
784        rope_attn_factor: 1.0,
785        rope_dim: None,
786        rope_freqs_long: None,
787        rope_freqs_short: None,
788        rope_orig_ctx: None,
789        rope_layout: RopeLayout::Neox,
790        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
791        swa_pattern: None,
792        attn_logit_softcap: None,
793        final_logit_softcap: None,
794        embedding_scale: None,
795        attention_scale: None,
796        rope_theta_swa: None,
797        ffn_activation: FfnActivation::Swiglu,
798        best_effort_fields: &["this is a synthetic mixed dense/MoE test fixture, not a real model"],
799    }
800}
801
802#[cfg(test)]
803mod tests {
804    use super::*;
805
806    #[test]
807    fn rope_layout_for_gguf_architecture_matches_llama_cpp() {
808        // Confirmed against llama.cpp's llama_model_rope_type
809        // (src/llama-model.cpp): llama -> NORM, olmoe/qwen2/phi3/gemma -> NEOX.
810        assert_eq!(RopeLayout::for_gguf_architecture("llama"), RopeLayout::Norm);
811        assert_eq!(
812            RopeLayout::for_gguf_architecture("llama4"),
813            RopeLayout::Norm
814        );
815        assert_eq!(
816            RopeLayout::for_gguf_architecture("deepseek2"),
817            RopeLayout::Norm
818        );
819        assert_eq!(RopeLayout::for_gguf_architecture("olmoe"), RopeLayout::Neox);
820        assert_eq!(RopeLayout::for_gguf_architecture("qwen2"), RopeLayout::Neox);
821        assert_eq!(
822            RopeLayout::for_gguf_architecture("qwen2moe"),
823            RopeLayout::Neox
824        );
825        assert_eq!(RopeLayout::for_gguf_architecture("qwen3"), RopeLayout::Neox);
826        assert_eq!(RopeLayout::for_gguf_architecture("phi3"), RopeLayout::Neox);
827        assert_eq!(
828            RopeLayout::for_gguf_architecture("gemma3"),
829            RopeLayout::Neox
830        );
831        // Unknown architectures keep the historical Neox default at this
832        // helper only; load-time uses capability::resolve_architecture and
833        // fails closed instead of guessing.
834        assert_eq!(
835            RopeLayout::for_gguf_architecture("totally-unknown-arch"),
836            RopeLayout::Neox
837        );
838    }
839
840    #[test]
841    fn all_presets_have_consistent_moe_hidden_dim() {
842        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
843            assert_eq!(
844                cfg.hidden_dim, cfg.moe.hidden_dim,
845                "{}: attention hidden_dim and MoE hidden_dim must match",
846                cfg.name
847            );
848        }
849    }
850
851    #[test]
852    fn all_presets_route_fewer_experts_than_total() {
853        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
854            assert!(
855                cfg.moe.n_experts_active < cfg.moe.n_experts,
856                "{}: active experts must be a sparse subset of total experts",
857                cfg.name
858            );
859        }
860    }
861
862    #[test]
863    fn all_presets_have_divisible_heads() {
864        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
865            assert_eq!(
866                cfg.n_heads % cfg.n_kv_heads,
867                0,
868                "{}: n_heads must be a multiple of n_kv_heads for GQA grouping",
869                cfg.name
870            );
871        }
872    }
873
874    #[test]
875    fn every_preset_declares_its_uncertain_fields() {
876        // This is a documentation-honesty test: any preset with zero
877        // best_effort_fields would be silently overclaiming precision
878        // we don't have. Fail loudly if that ever happens.
879        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
880            assert!(
881                !cfg.best_effort_fields.is_empty(),
882                "{}: must disclose which fields are unconfirmed estimates",
883                cfg.name
884            );
885        }
886    }
887
888    /// Kimi K3's `kda_layers`/`full_attn_layers` were transcribed by
889    /// hand from the real published config.json; this test guards
890    /// against a transcription slip (duplicate, out-of-range, or
891    /// missing layer index) rather than trusting the transcription.
892    #[test]
893    fn kimi_k3_hybrid_attention_layers_partition_every_layer_exactly_once() {
894        let cfg = kimi_k3();
895        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
896            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
897        };
898
899        let mut seen = std::collections::HashSet::new();
900        for &l in hybrid
901            .kda_layers
902            .iter()
903            .chain(hybrid.full_attn_layers.iter())
904        {
905            assert!(
906                (1..=cfg.n_layers).contains(&l),
907                "layer {l} is out of the published 1..={} range",
908                cfg.n_layers
909            );
910            assert!(
911                seen.insert(l),
912                "layer {l} appears in both/either list twice"
913            );
914        }
915        // Dense-vs-MoE (n_dense_leading_layers) and attention-type
916        // (KDA vs Gated MLA) are independent per-layer properties in
917        // the real config -- e.g. layer 1 is both the sole dense
918        // leading layer *and* a KDA layer -- so every one of the 93
919        // layers, dense or not, is covered by exactly one of these two
920        // lists (confirmed: 69 + 24 == 93, not 93 - 1).
921        assert_eq!(
922            hybrid.kda_layers.len() + hybrid.full_attn_layers.len(),
923            cfg.n_layers,
924            "every layer must be assigned exactly one of KDA or Gated MLA"
925        );
926        assert_eq!(
927            hybrid.kda_layers.len(),
928            69,
929            "expected 69 KDA layers per the published config"
930        );
931        assert_eq!(
932            hybrid.full_attn_layers.len(),
933            24,
934            "expected 24 Gated MLA layers per the published config"
935        );
936    }
937
938    #[test]
939    fn layer_attention_kind_is_gqa_for_every_layer_of_a_gqa_model() {
940        let cfg = glm_5_2();
941        for l in 0..cfg.n_layers {
942            assert_eq!(cfg.layer_attention_kind(l), LayerAttentionKind::Gqa);
943        }
944    }
945
946    #[test]
947    fn layer_attention_kind_classifies_every_kimi_k3_layer_without_panicking() {
948        let cfg = kimi_k3();
949        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
950            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
951        };
952        for l in 0..cfg.n_layers {
953            let kind = cfg.layer_attention_kind(l);
954            let one_indexed = l + 1;
955            if hybrid.kda_layers.contains(&one_indexed) {
956                assert_eq!(kind, LayerAttentionKind::KimiKda);
957            } else {
958                assert_eq!(kind, LayerAttentionKind::KimiMla);
959            }
960        }
961    }
962
963    #[test]
964    fn layer_attention_kind_matches_the_real_published_layer_1_and_4() {
965        // Layer 1 (1-indexed, so index 0 here) is published as KDA;
966        // layer 4 (index 3) is published as the first Gated MLA layer.
967        let cfg = kimi_k3();
968        assert_eq!(cfg.layer_attention_kind(0), LayerAttentionKind::KimiKda);
969        assert_eq!(cfg.layer_attention_kind(3), LayerAttentionKind::KimiMla);
970    }
971
972    #[test]
973    fn kimi_k3_mla_q_head_dim_matches_gqa_placeholder_head_dim() {
974        // The Gqa-placeholder head_dim above is deliberately set to
975        // Gated MLA's combined q_head_dim (qk_nope + qk_rope) so the
976        // placeholder path at least reflects a real dimension from the
977        // published config rather than an arbitrary guess.
978        let cfg = kimi_k3();
979        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
980            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
981        };
982        assert_eq!(
983            cfg.head_dim,
984            hybrid.mla.qk_nope_head_dim + hybrid.mla.qk_rope_head_dim
985        );
986    }
987
988    #[test]
989    fn approx_active_params_is_nonzero_and_finite_order_of_magnitude() {
990        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
991            let approx = cfg.approx_active_params_per_token();
992            // Sanity band: active params/token for these models is
993            // reported in the tens of billions; this is a loose
994            // order-of-magnitude check (1e9 to 1e12), not a precise
995            // parameter-count reproduction.
996            assert!(
997                approx > 1_000_000_000 && approx < 1_000_000_000_000,
998                "{}: approx_active_params_per_token={approx} is outside a plausible range",
999                cfg.name
1000            );
1001        }
1002    }
1003}
1004
1005#[cfg(test)]
1006mod longrope_tests {
1007    use super::*;
1008
1009    fn cfg_with_factors() -> ModelConfig {
1010        let mut c = test_dense_fixture();
1011        c.rope_orig_ctx = Some(4096);
1012        c.rope_freqs_short = Some(vec![1.0; 48]);
1013        c.rope_freqs_long = Some((0..48).map(|i| 1.0 + i as f32).collect());
1014        c.rope_freqs = None;
1015        c
1016    }
1017
1018    /// llama.cpp `llama_model::get_rope_factors`: long only when the
1019    /// run's context exceeds `original_context_length`. Phi-4-mini's
1020    /// short set is all ones, so picking long at 4096 would apply a
1021    /// correction the model never asked for at that length.
1022    #[test]
1023    fn long_set_only_above_the_original_context() {
1024        let mut c = cfg_with_factors();
1025        c.apply_runtime_context(4096);
1026        assert_eq!(
1027            c.rope_freqs.as_ref().unwrap()[1],
1028            1.0,
1029            "at the threshold, short"
1030        );
1031
1032        let mut c = cfg_with_factors();
1033        c.apply_runtime_context(4097);
1034        assert_eq!(c.rope_freqs.as_ref().unwrap()[1], 2.0, "above it, long");
1035
1036        let mut c = cfg_with_factors();
1037        c.apply_runtime_context(1024);
1038        assert_eq!(c.rope_freqs.as_ref().unwrap()[1], 1.0, "below it, short");
1039    }
1040
1041    /// `rope_freqs.weight` (Llama 3) is not a LongRoPE set and outranks
1042    /// one, the same precedence llama.cpp gives it. The loader encodes
1043    /// that by leaving the long/short pair empty whenever the explicit
1044    /// tensor is present, so the runtime re-pick has nothing to apply.
1045    #[test]
1046    fn an_explicit_rope_freqs_tensor_is_never_overridden() {
1047        let mut c = test_dense_fixture();
1048        c.rope_freqs = Some(vec![7.0; 48]);
1049        c.rope_orig_ctx = Some(4096);
1050        c.rope_freqs_long = None;
1051        c.rope_freqs_short = None;
1052        c.apply_runtime_context(131072);
1053        assert_eq!(c.rope_freqs.as_ref().unwrap()[0], 7.0);
1054    }
1055
1056    /// A checkpoint with neither set must come back untouched, so the
1057    /// call is free to sit on every load path.
1058    #[test]
1059    fn models_without_longrope_are_untouched() {
1060        let mut c = test_dense_fixture();
1061        c.rope_freqs = None;
1062        c.apply_runtime_context(8192);
1063        assert!(c.rope_freqs.is_none());
1064        assert!(c.rope_orig_ctx.is_none());
1065    }
1066}