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    /// Alternating SWA period, llama.cpp's `set_swa_pattern` argument.
262    ///
263    /// `Some(0)` windows every layer and `Some(1)` windows none, which
264    /// are llama.cpp's two degenerate spellings and are NOT the same as
265    /// `None` (no period known, so every layer windows). Any larger `p`
266    /// alternates, with the phase in [`Self::swa_dense_first`].
267    pub swa_pattern: Option<usize>,
268    /// llama.cpp's `dense_first` argument to `set_swa_pattern`, which
269    /// decides WHICH layer of each period is the full-attention one.
270    ///
271    /// `false` puts it last (`il % p == p - 1`), `true` puts it first
272    /// (`il % p == 0`). Getting this wrong is not a near miss: on a
273    /// 32-layer period-4 model the two phases disagree about SIXTEEN
274    /// layers, each of which then attends over the wrong span at full
275    /// speed. `capability::default_swa_layout` carries the per-arch
276    /// value, transcribed from llama.cpp.
277    pub swa_dense_first: bool,
278    /// Attention logit soft-capping (Gemma 2+). Applied as
279    /// `softcap * tanh(score / softcap)` before softmax.
280    pub attn_logit_softcap: Option<f32>,
281    /// Final logit soft-capping (Gemma 2+). Applied to lm_head output.
282    pub final_logit_softcap: Option<f32>,
283    /// Input embedding scale (Gemma: `sqrt(hidden_dim)`).
284    pub embedding_scale: Option<f32>,
285    /// Optional override for the attention score scale baked into Q
286    /// *instead of* the kernel's default `1/sqrt(head_dim)`. When set,
287    /// callers must pass `score_scale = 1.0` into the attention kernel
288    /// (llama.cpp Gemma: scale Q then `build_attn(..., 1.0f)`). Prefer
289    /// leaving this `None` when the override equals `1/sqrt(head_dim)`.
290    pub attention_scale: Option<f32>,
291    /// RoPE base used on SWA layers (Gemma 3: defaults to `10000` when
292    /// the GGUF omits `rope.freq_base_swa`; full-attn layers keep
293    /// [`Self::rope_theta`]).
294    pub rope_theta_swa: Option<f32>,
295    /// Dense/MoE FFN activation pairing.
296    pub ffn_activation: FfnActivation,
297    /// Every field on this config that is a best-effort estimate rather
298    /// than a confirmed value from an official config.json / GGUF file.
299    pub best_effort_fields: &'static [&'static str],
300}
301
302/// Dense / expert FFN non-linearity used by the generic decoder.
303#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
304pub enum FfnActivation {
305    /// `silu(gate) * up` with separate gate/up matrices (Llama / Qwen).
306    #[default]
307    Swiglu,
308    /// Phi-3 fused gate+up in one `ffn_up` matrix (`2 * n_ff` rows).
309    SwigluFused,
310    /// Gemma GeGLU: `gelu(gate) * up`.
311    Gelu,
312}
313
314impl ModelConfig {
315    /// Re-picks the LongRoPE factor set now that the run's context size
316    /// is known, matching llama.cpp `llama_model::get_rope_factors`:
317    /// `rope_freqs.weight` (Llama 3) always wins; otherwise the long set
318    /// applies only when the context exceeds
319    /// `rope.scaling.original_context_length`, and the short set
320    /// otherwise.
321    ///
322    /// A no-op for every checkpoint that ships neither set, which is all
323    /// of them except the Phi-3/Phi-4 family today.
324    ///
325    /// Because it re-picks `rope_freqs` wholesale it would also discard
326    /// a YaRN rewrite folded into that field at parse time (see
327    /// [`Self::rope_freqs`]). No real checkpoint hits that: LongRoPE
328    /// files declare `rope.scaling.type = "longrope"`, which the loader's
329    /// YaRN arm deliberately does not claim, so the two never populate
330    /// the field on the same file.
331    pub fn apply_runtime_context(&mut self, ctx: usize) {
332        let (Some(orig), true) = (
333            self.rope_orig_ctx,
334            self.rope_freqs_long.is_some() || self.rope_freqs_short.is_some(),
335        ) else {
336            return;
337        };
338        let picked = if ctx > orig {
339            self.rope_freqs_long.as_ref()
340        } else {
341            self.rope_freqs_short.as_ref()
342        };
343        if let Some(f) = picked
344            .or(self.rope_freqs_long.as_ref())
345            .or(self.rope_freqs_short.as_ref())
346        {
347            self.rope_freqs = Some(f.clone());
348        }
349    }
350
351    /// True if layer `layer_idx` (0-indexed) should be built as an
352    /// ordinary dense FFN rather than this model's MoE topology.
353    pub fn layer_is_dense(&self, layer_idx: usize) -> bool {
354        layer_idx < self.n_dense_leading_layers
355    }
356
357    /// Sliding-window size for layer `il`, honouring Gemma-style
358    /// alternating SWA patterns. `None` means full causal attention.
359    pub fn layer_sliding_window(&self, layer_idx: usize) -> Option<usize> {
360        let window = self.sliding_window?;
361        // llama.cpp `llama_hparams::set_swa_pattern`
362        // (`src/llama-hparams.cpp:8-22`), both phases:
363        //
364        //   dense_first: is_swa = n_pattern == 0 || (il % n_pattern != 0)
365        //   otherwise:   is_swa = n_pattern == 0 || (il % n_pattern < n_pattern - 1)
366        //
367        // `period == 1` therefore windows NOTHING under either phase,
368        // which is the opposite of `None`. It used to be filtered out
369        // before it reached here and fell back to "every layer", which
370        // is exactly inverted.
371        let sliding = match self.swa_pattern {
372            None | Some(0) => true,
373            Some(period) if self.swa_dense_first => !layer_idx.is_multiple_of(period),
374            Some(period) => layer_idx % period < period - 1,
375        };
376        sliding.then_some(window)
377    }
378
379    /// The narrowest sliding window any layer of this model uses, or
380    /// `None` if every layer is full-causal.
381    ///
382    /// For an alternating-SWA model (gpt-oss, Gemma-3) the
383    /// full-attention layers impose no constraint on the KV block
384    /// layout and the sliding ones impose the window -- so the model's
385    /// constraint is simply the window, present as soon as *any* layer
386    /// slides. A model that is 5/6 full-attention is not 5/6 exempt:
387    /// one mis-aligned sliding layer corrupts the answer.
388    pub fn kv_block_window(&self) -> Option<usize> {
389        (0..self.n_layers).find_map(|il| self.layer_sliding_window(il))
390    }
391
392    /// The window EVERY layer slides by, or `None` if any layer attends
393    /// over the whole history.
394    ///
395    /// This is the opposite question to [`Self::kv_block_window`], and
396    /// the difference is the whole reason both exist. That one asks
397    /// "does any layer constrain the block layout", so one sliding layer
398    /// is enough. This one asks "may a page that has fallen behind the
399    /// window be taken away", and there one *full-attention* layer is
400    /// enough to say no.
401    ///
402    /// A page group holds one block in every layer and is freed as a
403    /// unit, so on an alternating-SWA model (gpt-oss, Gemma-3) freeing
404    /// the group behind the window would take the full-attention layers'
405    /// block with it -- and those layers still read position 0 at every
406    /// step. The result is not a crash: the block is reused by another
407    /// request and the full layers attend over its bytes. So this
408    /// returns `None` for the alternating case, and a mixed-window model
409    /// (were one to appear) gets `None` too rather than the narrowest
410    /// window, because the widest is the one that must still be readable.
411    pub fn uniform_sliding_window(&self) -> Option<usize> {
412        let first = self.layer_sliding_window(0)?;
413        (1..self.n_layers)
414            .all(|il| self.layer_sliding_window(il) == Some(first))
415            .then_some(first)
416    }
417
418    /// The KV cache block layout to use for this model, given the block
419    /// size an operator asked for.
420    ///
421    /// The requested size is rounded *down* to something that divides
422    /// the window (see [`ferrox_core::kv_swa`]), so a config that would
423    /// straddle the window boundary becomes a smaller block rather than
424    /// a startup failure or -- much worse -- a silently wrong mask.
425    pub fn kv_block_layout(&self, desired_block_size: usize) -> ferrox_core::BlockLayout {
426        let window = self.kv_block_window();
427        let block_size = ferrox_core::aligned_block_size(desired_block_size, window);
428        ferrox_core::BlockLayout::new(block_size, window)
429            .expect("aligned_block_size returns a size BlockLayout accepts")
430    }
431
432    /// RoPE frequency base for layer `il` (SWA layers may differ).
433    pub fn layer_rope_theta(&self, layer_idx: usize) -> f32 {
434        match (self.layer_sliding_window(layer_idx), self.rope_theta_swa) {
435            (Some(_), Some(theta)) => theta,
436            _ => self.rope_theta,
437        }
438    }
439
440    /// Which attention mechanism layer `layer_idx` (0-indexed, ferrox's
441    /// usual convention) uses. For `AttentionKind::Gqa` every layer is
442    /// `LayerAttentionKind::Gqa`; for `AttentionKind::KimiHybrid`, looks
443    /// up `layer_idx + 1` (the real `kda_layers`/`full_attn_layers`
444    /// lists are 1-indexed -- see `KimiHybridAttention`'s doc comment)
445    /// in those real per-layer lists.
446    ///
447    /// # Panics
448    /// If `layer_idx` isn't covered by either list of a `KimiHybrid`
449    /// config -- can't happen for `kimi_k3()`, whose lists are tested
450    /// (`kimi_k3_hybrid_attention_layers_partition_every_layer_exactly_once`)
451    /// to partition every layer with no gaps, but a caller building a
452    /// custom `KimiHybridAttention` must uphold the same invariant.
453    pub fn layer_attention_kind(&self, layer_idx: usize) -> LayerAttentionKind {
454        match &self.attention {
455            AttentionKind::Gqa => LayerAttentionKind::Gqa,
456            AttentionKind::KimiHybrid(hybrid) => {
457                let one_indexed = layer_idx + 1;
458                if hybrid.kda_layers.contains(&one_indexed) {
459                    LayerAttentionKind::KimiKda
460                } else if hybrid.full_attn_layers.contains(&one_indexed) {
461                    LayerAttentionKind::KimiMla
462                } else {
463                    panic!(
464                        "layer {layer_idx} (1-indexed {one_indexed}) is in neither \
465                         kda_layers nor full_attn_layers"
466                    )
467                }
468            }
469        }
470    }
471
472    /// Total parameter count implied by the MoE config, as a sanity
473    /// check against the publicly reported total (this is an order of
474    /// magnitude check, not an exact parameter-count reproduction).
475    pub fn approx_active_params_per_token(&self) -> usize {
476        let attn_params_per_layer = 4 * self.hidden_dim * self.hidden_dim; // q,k,v,o (rough)
477        let active_experts = self.moe.n_experts_active + self.moe.n_shared_experts;
478        let expert_params = active_experts * 3 * self.moe.hidden_dim * self.moe.expert_ffn_dim; // gate,up,down
479        self.n_layers * (attn_params_per_layer + expert_params)
480    }
481}
482
483/// GLM-5.2 (Z.ai) **structural sketch only** — not a supported real
484/// inference path. Real DSA lives in `glm_dsa` / `glm52_decoder` and is
485/// not wired into `Decoder` / `ferrox-server`. This preset drives
486/// smoke/bench with synthetic GQA weights only (~744B / ~40B active
487/// hparams as published placeholders).
488pub fn glm_5_2() -> ModelConfig {
489    ModelConfig {
490        sliding_window: None,
491        name: "glm-5.2",
492        attention: AttentionKind::Gqa,
493        n_layers: 92,
494        hidden_dim: 6144,
495        n_heads: 48,
496        n_kv_heads: 8,
497        head_dim: 128,
498        vocab_size: 151552,
499        rope_theta: 1_000_000.0,
500        rms_norm_eps: 1e-5,
501        moe: MoeLayerConfig {
502            expert_weights_scale: 1.0,
503            n_experts: 256,
504            n_experts_active: 8,
505            n_shared_experts: 1,
506            hidden_dim: 6144,
507            expert_ffn_dim: 2048,
508            // Sigmoid, not softmax: reading ik_llama.cpp's real GGUF
509            // hparams-loading source (llama-hparams.cpp,
510            // LLM_ARCH_GLM4_MOE case) directly showed GLM4-MoE-family
511            // models default to sigmoid gating with post-selection
512            // score renormalization. GLM-5.2 is presumed to continue
513            // this lineage; not confirmed against GLM-5.2's own
514            // config.json (unavailable in this environment).
515            gating: GatingFunction::Sigmoid,
516            norm_topk_prob: true,
517         expert_group_count: None, expert_group_used_count: None,},
518        // No evidence found (via ik_llama.cpp source or public
519        // reporting) that GLM-5.2 skips MoE on any leading layers;
520        // defaulting to 0 (every layer uses this model's MoE
521        // topology) rather than assuming DeepSeek's convention
522        // applies here too.
523        n_dense_leading_layers: 0,
524        rope_freqs: None,
525        rope_attn_factor: 1.0,
526        rope_dim: None,
527        rope_freqs_long: None,
528        rope_freqs_short: None,
529        rope_orig_ctx: None,
530        // Placeholder GQA path; real GLM-5.2 DSA uses interleaved RoPE
531        // via `glm_dsa`/`mla`, not this preset's Decoder path.
532        rope_layout: RopeLayout::Neox,
533        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
534        swa_pattern: None,
535        swa_dense_first: false,
536        attn_logit_softcap: None,
537        final_logit_softcap: None,
538        embedding_scale: None,
539        attention_scale: None,
540        rope_theta_swa: None,
541        ffn_activation: FfnActivation::Swiglu,
542        best_effort_fields: &[
543            "n_layers",
544            "hidden_dim",
545            "n_heads",
546            "n_kv_heads",
547            "head_dim",
548            "rope_theta",
549            "moe.expert_ffn_dim",
550            "moe.n_shared_experts",
551            "moe.gating (sigmoid assumed from GLM4-MoE-family convention found in ik_llama.cpp source, not confirmed for GLM-5.2 specifically)",
552        ],
553    }
554}
555
556/// DeepSeek V4 Pro **structural sketch only** — CSA/HCA is not on this
557/// GQA `Decoder` path. Real primitives live under
558/// `deepseek_v4_attention` / `hyper_connections` and are not assembled
559/// into a served decoder yet. Hparams (~1.6T / ~49B active) are
560/// placeholders for smoke/bench.
561pub fn deepseek_v4_pro() -> ModelConfig {
562    ModelConfig {
563        sliding_window: None,
564        name: "deepseek-v4-pro",
565        attention: AttentionKind::Gqa,
566        n_layers: 96,
567        hidden_dim: 7168,
568        n_heads: 56,
569        n_kv_heads: 8,
570        head_dim: 128,
571        vocab_size: 129280,
572        rope_theta: 1_000_000.0,
573        rms_norm_eps: 1e-6,
574        moe: MoeLayerConfig {
575            expert_weights_scale: 1.0,
576            n_experts: 385,
577            n_experts_active: 6,
578            n_shared_experts: 1,
579            hidden_dim: 7168,
580            expert_ffn_dim: 2048,
581            // Sigmoid, not softmax: this is the stronger-confidence of
582            // the two sigmoid-gating corrections in this file.
583            // DeepSeek-V3's own published technical report explicitly
584            // documents computing per-expert affinity via sigmoid and
585            // renormalizing only the selected experts' scores to sum
586            // to one; reading ik_llama.cpp's real GGUF hparams-loading
587            // source (llama-hparams.cpp, LLM_ARCH_DEEPSEEK2 case)
588            // confirmed this is exactly what that code path defaults
589            // to for the DeepSeek-2/3 lineage. DeepSeek V4 Pro is
590            // presumed to continue using sigmoid gating for the same
591            // reason; not confirmed against V4 Pro's own config.json.
592            gating: GatingFunction::Sigmoid,
593            norm_topk_prob: true,
594         expert_group_count: None, expert_group_used_count: None,},
595        // DeepSeek-V3's own published technical report documents the
596        // first 3 transformer layers as dense (ordinary FFN, no
597        // expert routing), with MoE starting from layer 4 onward;
598        // ik_llama.cpp's real hparams-loading source
599        // (LLM_KV_LEADING_DENSE_BLOCK_COUNT) confirms this is a real,
600        // loaded GGUF metadata field for the DeepSeek-2/3 lineage.
601        // DeepSeek V4 Pro is presumed to continue this convention;
602        // not confirmed against V4 Pro's own config.json.
603        n_dense_leading_layers: 3,
604        rope_freqs: None,
605        rope_attn_factor: 1.0,
606        rope_dim: None,
607        rope_freqs_long: None,
608        rope_freqs_short: None,
609        rope_orig_ctx: None,
610        // llama.cpp maps LLM_ARCH_DEEPSEEK4 -> LLAMA_ROPE_TYPE_NORM.
611        rope_layout: RopeLayout::Norm,
612        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
613        swa_pattern: None,
614        swa_dense_first: false,
615        attn_logit_softcap: None,
616        final_logit_softcap: None,
617        embedding_scale: None,
618        attention_scale: None,
619        rope_theta_swa: None,
620        ffn_activation: FfnActivation::Swiglu,
621        best_effort_fields: &[
622            "n_layers",
623            "hidden_dim",
624            "n_heads",
625            "n_kv_heads",
626            "head_dim",
627            "moe.expert_ffn_dim",
628            "attention_variant (CSA/HCA hybrid NOT implemented, GQA fallback in use)",
629            "moe.gating (sqrtsoftplus: confirmed for real V4 in llama.cpp PR #24162; this preset still uses Sigmoid on the wrong GQA sketch path)",
630            "n_dense_leading_layers (3: same confidence basis as gating above, DeepSeek-V3 technical report + ik_llama.cpp source, not confirmed for V4 Pro)",
631        ],
632    }
633}
634
635/// Kimi K3 **structural sketch only** for the generic GQA `Decoder`.
636/// Real checkpoint work uses the dedicated Kimi stack (`kimi_loader` /
637/// `KimiEngine`); slice-verified, not a full end-to-end run. Do not
638/// treat this preset as a runnable Kimi substitute.
639pub fn kimi_k3() -> ModelConfig {
640    ModelConfig {
641        sliding_window: None,
642        name: "kimi-k3",
643        n_layers: 93,
644        hidden_dim: 7168,
645        // n_heads/n_kv_heads/head_dim describe the Gqa fallback
646        // Decoder actually runs today, not Kimi K3's real attention
647        // (see `attention` below) -- kept at reasonable stand-in
648        // values (matching MLA's num_heads=96 and combined
649        // qk_nope+qk_rope head dim) rather than deleted, so the
650        // placeholder path stays runnable.
651        n_heads: 96,
652        n_kv_heads: 96,
653        head_dim: 192,
654        vocab_size: 163840,
655        // Not present in the published text_config; RoPE only ever
656        // applies to Gated MLA's 64-dim qk_rope_head_dim slice in the
657        // real architecture, and Decoder doesn't implement that slicing
658        // yet, so this remains an unconfirmed placeholder.
659        rope_theta: 1_000_000.0,
660        rms_norm_eps: 1e-5,
661        moe: MoeLayerConfig {
662            expert_weights_scale: 1.0,
663            n_experts: 896,
664            n_experts_active: 16,
665            n_shared_experts: 2,
666            hidden_dim: 7168,
667            expert_ffn_dim: 3072,
668            // Confirmed directly from the real config.json:
669            // "moe_router_activation_func": "sigmoid".
670            gating: GatingFunction::Sigmoid,
671            norm_topk_prob: true,
672         expert_group_count: None, expert_group_used_count: None,},
673        // Confirmed directly from the real config.json:
674        // "first_k_dense_replace": 1.
675        n_dense_leading_layers: 1,
676        // Kimi K3's real, published attention topology (verified
677        // against huggingface.co/moonshotai/Kimi-K3/config.json's
678        // linear_attn_config block and the real KimiDeltaAttention /
679        // KimiMLAAttention reference implementations in
680        // modeling_kimi_linear.py) -- not yet wired into Decoder's
681        // forward pass, which still runs the Gqa placeholder above for
682        // every layer regardless of this field.
683        attention: AttentionKind::KimiHybrid(KimiHybridAttention {
684            kda_layers: vec![
685                1, 2, 3, 5, 6, 7, 9, 10, 11, 13, 14, 15, 17, 18, 19, 21, 22, 23, 25, 26, 27, 29,
686                30, 31, 33, 34, 35, 37, 38, 39, 41, 42, 43, 45, 46, 47, 49, 50, 51, 53, 54, 55,
687                57, 58, 59, 61, 62, 63, 65, 66, 67, 69, 70, 71, 73, 74, 75, 77, 78, 79, 81, 82,
688                83, 85, 86, 87, 89, 90, 91,
689            ],
690            full_attn_layers: vec![
691                4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, 84,
692                88, 92, 93,
693            ],
694            mla: MlaConfig {
695                num_heads: 96,
696                q_lora_rank: 1536,
697                kv_lora_rank: 512,
698                qk_nope_head_dim: 128,
699                qk_rope_head_dim: 64,
700                v_head_dim: 128,
701                use_output_gate: true,
702                // Real, confirmed: Kimi K3's `KimiMLAAttention.forward`
703                // never rotates -- see `MlaConfig::rope`'s doc comment.
704                rope: None,
705            },
706            kda: KdaConfig {
707                num_heads: 96,
708                head_dim: 128,
709                short_conv_kernel_size: 4,
710                gate_lower_bound: -5.0,
711                use_full_rank_gate: true,
712            },
713        }),
714        rope_freqs: None,
715        rope_attn_factor: 1.0,
716        rope_dim: None,
717        rope_freqs_long: None,
718        rope_freqs_short: None,
719        rope_orig_ctx: None,
720        // GQA placeholder path only; real Kimi attention is rope-less MLA
721        // or KDA and never reaches Decoder::apply_rope_head.
722        rope_layout: RopeLayout::Neox,
723        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
724        swa_pattern: None,
725        swa_dense_first: false,
726        attn_logit_softcap: None,
727        final_logit_softcap: None,
728        embedding_scale: None,
729        attention_scale: None,
730        rope_theta_swa: None,
731        ffn_activation: FfnActivation::Swiglu,
732        best_effort_fields: &[
733            "n_heads/n_kv_heads/head_dim (describe the unimplemented Gqa placeholder, not Kimi K3's real MLA/KDA attention -- see `attention` field)",
734            "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)",
735            "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)",
736        ],
737    }
738}
739
740/// Matches the generated on-disk fixture exactly (hidden_dim, head
741/// counts, ffn_dim, vocab, rope_theta, eps).
742/// Used by `ferrox inspect-run` and the cross-validation test in
743/// `crates/ferrox-models/tests/gguf_roundtrip.rs` to prove the real
744/// GGUF loader + forward pass produce the same numbers as an
745/// independent NumPy reference implementation reading the same file.
746pub fn test_dense_fixture() -> ModelConfig {
747    ModelConfig {
748        sliding_window: None,
749        name: "ferrox-test-dense",
750        attention: AttentionKind::Gqa,
751        n_layers: 2,
752        hidden_dim: 32,
753        n_heads: 4,
754        n_kv_heads: 2,
755        head_dim: 8,
756        vocab_size: 32,
757        rope_theta: 10000.0,
758        rms_norm_eps: 1e-5,
759        moe: MoeLayerConfig {
760            expert_weights_scale: 1.0,
761            n_experts: 1,
762            n_experts_active: 1,
763            n_shared_experts: 0,
764            hidden_dim: 32,
765            expert_ffn_dim: 32,
766            gating: GatingFunction::Softmax,
767            norm_topk_prob: true,
768            expert_group_count: None,
769            expert_group_used_count: None,
770        },
771        n_dense_leading_layers: 0,
772        rope_freqs: None,
773        rope_attn_factor: 1.0,
774        rope_dim: None,
775        rope_freqs_long: None,
776        rope_freqs_short: None,
777        rope_orig_ctx: None,
778        // Matches the independent reference's split-half apply_rope.
779        rope_layout: RopeLayout::Neox,
780        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
781        swa_pattern: None,
782        swa_dense_first: false,
783        attn_logit_softcap: None,
784        final_logit_softcap: None,
785        embedding_scale: None,
786        attention_scale: None,
787        rope_theta_swa: None,
788        ffn_activation: FfnActivation::Swiglu,
789        best_effort_fields: &["this is a synthetic test fixture, not a real model"],
790    }
791}
792
793/// Matches the generated on-disk multi-expert MoE fixture: 4 experts,
794/// top-2 routing, 1 shared
795/// expert, packed 3D expert tensors. Used to verify the previously-
796/// untested multi-expert loading path (`split_expert_tensor` in
797/// `ferrox-models::loader`) against a real file, the same way
798/// `test_dense_fixture` verifies the single-expert path.
799pub fn test_moe_fixture() -> ModelConfig {
800    ModelConfig {
801        sliding_window: None,
802        name: "ferrox-test-moe",
803        attention: AttentionKind::Gqa,
804        n_layers: 2,
805        hidden_dim: 32,
806        n_heads: 4,
807        n_kv_heads: 2,
808        head_dim: 8,
809        vocab_size: 32,
810        rope_theta: 10000.0,
811        rms_norm_eps: 1e-5,
812        moe: MoeLayerConfig {
813            expert_weights_scale: 1.0,
814            n_experts: 4,
815            n_experts_active: 2,
816            n_shared_experts: 1,
817            hidden_dim: 32,
818            expert_ffn_dim: 32,
819            gating: GatingFunction::Softmax,
820            norm_topk_prob: true,
821            expert_group_count: None,
822            expert_group_used_count: None,
823        },
824        n_dense_leading_layers: 0,
825        rope_freqs: None,
826        rope_attn_factor: 1.0,
827        rope_dim: None,
828        rope_freqs_long: None,
829        rope_freqs_short: None,
830        rope_orig_ctx: None,
831        rope_layout: RopeLayout::Neox,
832        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
833        swa_pattern: None,
834        swa_dense_first: false,
835        attn_logit_softcap: None,
836        final_logit_softcap: None,
837        embedding_scale: None,
838        attention_scale: None,
839        rope_theta_swa: None,
840        ffn_activation: FfnActivation::Swiglu,
841        best_effort_fields: &["this is a synthetic multi-expert test fixture, not a real model"],
842    }
843}
844
845/// Matches the generated on-disk mixed-topology fixture: 3 layers, the
846/// first of which is
847/// an ordinary dense FFN and the remaining two are genuine MoE (3
848/// experts, top-1 routing, 1 shared expert each). Used to verify the
849/// "leading dense layers" loading path
850/// (`ModelConfig::layer_is_dense`) against a real file -- the pattern
851/// found in DeepSeek-2/3-family models via ik_llama.cpp's source
852/// (`LLM_KV_LEADING_DENSE_BLOCK_COUNT`), which was previously only
853/// documented, not implemented or tested.
854pub fn test_mixed_fixture() -> ModelConfig {
855    ModelConfig {
856        sliding_window: None,
857        name: "ferrox-test-mixed",
858        attention: AttentionKind::Gqa,
859        n_layers: 3,
860        hidden_dim: 32,
861        n_heads: 4,
862        n_kv_heads: 2,
863        head_dim: 8,
864        vocab_size: 32,
865        rope_theta: 10000.0,
866        rms_norm_eps: 1e-5,
867        moe: MoeLayerConfig {
868            expert_weights_scale: 1.0,
869            n_experts: 3,
870            n_experts_active: 1,
871            n_shared_experts: 1,
872            hidden_dim: 32,
873            expert_ffn_dim: 32,
874            gating: GatingFunction::Softmax,
875            norm_topk_prob: true,
876            expert_group_count: None,
877            expert_group_used_count: None,
878        },
879        n_dense_leading_layers: 1,
880        rope_freqs: None,
881        rope_attn_factor: 1.0,
882        rope_dim: None,
883        rope_freqs_long: None,
884        rope_freqs_short: None,
885        rope_orig_ctx: None,
886        rope_layout: RopeLayout::Neox,
887        qk_norm_style: crate::capability::QkNormStyle::WholeVector,
888        swa_pattern: None,
889        swa_dense_first: false,
890        attn_logit_softcap: None,
891        final_logit_softcap: None,
892        embedding_scale: None,
893        attention_scale: None,
894        rope_theta_swa: None,
895        ffn_activation: FfnActivation::Swiglu,
896        best_effort_fields: &["this is a synthetic mixed dense/MoE test fixture, not a real model"],
897    }
898}
899
900#[cfg(test)]
901mod tests {
902    use super::*;
903
904    #[test]
905    fn rope_layout_for_gguf_architecture_matches_llama_cpp() {
906        // Confirmed against llama.cpp's llama_model_rope_type
907        // (src/llama-model.cpp): llama -> NORM, olmoe/qwen2/phi3/gemma -> NEOX.
908        assert_eq!(RopeLayout::for_gguf_architecture("llama"), RopeLayout::Norm);
909        assert_eq!(
910            RopeLayout::for_gguf_architecture("llama4"),
911            RopeLayout::Norm
912        );
913        assert_eq!(
914            RopeLayout::for_gguf_architecture("deepseek2"),
915            RopeLayout::Norm
916        );
917        assert_eq!(RopeLayout::for_gguf_architecture("olmoe"), RopeLayout::Neox);
918        assert_eq!(RopeLayout::for_gguf_architecture("qwen2"), RopeLayout::Neox);
919        assert_eq!(
920            RopeLayout::for_gguf_architecture("qwen2moe"),
921            RopeLayout::Neox
922        );
923        assert_eq!(RopeLayout::for_gguf_architecture("qwen3"), RopeLayout::Neox);
924        assert_eq!(RopeLayout::for_gguf_architecture("phi3"), RopeLayout::Neox);
925        assert_eq!(
926            RopeLayout::for_gguf_architecture("gemma3"),
927            RopeLayout::Neox
928        );
929        // Unknown architectures keep the historical Neox default at this
930        // helper only; load-time uses capability::resolve_architecture and
931        // fails closed instead of guessing.
932        assert_eq!(
933            RopeLayout::for_gguf_architecture("totally-unknown-arch"),
934            RopeLayout::Neox
935        );
936    }
937
938    /// gpt-oss's real shape: a 128-token window on every other layer.
939    /// A KV block size of 128 or any divisor of it is fine; 48 or 256
940    /// are not, and the config layer must round down rather than hand
941    /// the cache something it will refuse (or, worse, accept).
942    #[test]
943    fn an_alternating_swa_model_constrains_the_block_layout() {
944        let mut cfg = test_dense_fixture();
945        cfg.n_layers = 24;
946        cfg.sliding_window = Some(128);
947        cfg.swa_pattern = Some(2);
948
949        // Half the layers are full-attention, but the model is still
950        // constrained: one mis-aligned sliding layer is enough.
951        assert!(cfg.layer_sliding_window(1).is_none() || cfg.layer_sliding_window(0).is_none());
952        assert_eq!(cfg.kv_block_window(), Some(128));
953
954        let layout = cfg.kv_block_layout(256);
955        assert_eq!(layout.block_size(), 128, "256 must round down, not up");
956        assert_eq!(layout.sliding_window(), Some(128));
957        assert_eq!(layout.blocks_per_window(), Some(1));
958
959        assert_eq!(cfg.kv_block_layout(48).block_size(), 32);
960        assert_eq!(cfg.kv_block_layout(32).block_size(), 32);
961    }
962
963    /// Gemma-3: window 512, every 6th layer full-attention.
964    #[test]
965    fn a_gemma3_shaped_model_takes_its_window_from_the_sliding_layers() {
966        let mut cfg = test_dense_fixture();
967        cfg.n_layers = 30;
968        cfg.sliding_window = Some(512);
969        cfg.swa_pattern = Some(6);
970        assert!(
971            cfg.layer_sliding_window(5).is_none(),
972            "every 6th layer is full-attention"
973        );
974        assert_eq!(cfg.kv_block_window(), Some(512));
975        assert_eq!(cfg.kv_block_layout(100).block_size(), 64);
976        assert_eq!(cfg.kv_block_layout(64).blocks_per_window(), Some(8));
977    }
978
979    /// The two window questions give OPPOSITE answers on an alternating
980    /// model, and that is the point of having both.
981    ///
982    /// "Does any layer constrain the block layout" is yes, so the block
983    /// size rounds down to the window. "May a page behind the window be
984    /// taken away" is no, because the group holds the full-attention
985    /// layers' blocks too and those layers still read position 0. A
986    /// serving path that read `kv_block_window` for the second question
987    /// would free pages half the layers are still attending over -- not
988    /// a crash, just another request's bytes in this one's answer.
989    #[test]
990    fn only_a_uniformly_windowed_model_may_give_a_page_back() {
991        let mut alternating = test_dense_fixture();
992        alternating.n_layers = 24;
993        alternating.sliding_window = Some(128);
994        alternating.swa_pattern = Some(2);
995        assert_eq!(alternating.kv_block_window(), Some(128));
996        assert_eq!(
997            alternating.uniform_sliding_window(),
998            None,
999            "a full-attention layer forbids the slide"
1000        );
1001
1002        let mut uniform = test_dense_fixture();
1003        uniform.n_layers = 24;
1004        uniform.sliding_window = Some(128);
1005        uniform.swa_pattern = None;
1006        assert_eq!(uniform.uniform_sliding_window(), Some(128));
1007
1008        // `Some(0)` is llama.cpp's spelling of "every layer slides"
1009        // (`set_swa_pattern(0)`), and it is the one that may give a page
1010        // back. `Some(1)` is the OPPOSITE -- no layer slides -- and this
1011        // used to assert the two were the same, which is how the
1012        // inversion stayed invisible.
1013        let mut period_zero = uniform.clone();
1014        period_zero.swa_pattern = Some(0);
1015        assert_eq!(period_zero.uniform_sliding_window(), Some(128));
1016
1017        let mut period_one = uniform.clone();
1018        period_one.swa_pattern = Some(1);
1019        assert_eq!(
1020            period_one.uniform_sliding_window(),
1021            None,
1022            "period 1 windows no layer, so there is no window to slide"
1023        );
1024        assert_eq!(period_one.kv_block_window(), None);
1025
1026        let mut full = test_dense_fixture();
1027        full.sliding_window = None;
1028        assert_eq!(full.uniform_sliding_window(), None);
1029    }
1030
1031    #[test]
1032    fn a_full_causal_model_keeps_the_block_size_it_was_given() {
1033        let mut cfg = test_dense_fixture();
1034        cfg.sliding_window = None;
1035        cfg.swa_pattern = None;
1036        assert_eq!(cfg.kv_block_window(), None);
1037        let layout = cfg.kv_block_layout(48);
1038        assert_eq!(layout.block_size(), 48);
1039        assert_eq!(layout.sliding_window(), None);
1040    }
1041
1042    #[test]
1043    fn all_presets_have_consistent_moe_hidden_dim() {
1044        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
1045            assert_eq!(
1046                cfg.hidden_dim, cfg.moe.hidden_dim,
1047                "{}: attention hidden_dim and MoE hidden_dim must match",
1048                cfg.name
1049            );
1050        }
1051    }
1052
1053    #[test]
1054    fn all_presets_route_fewer_experts_than_total() {
1055        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
1056            assert!(
1057                cfg.moe.n_experts_active < cfg.moe.n_experts,
1058                "{}: active experts must be a sparse subset of total experts",
1059                cfg.name
1060            );
1061        }
1062    }
1063
1064    #[test]
1065    fn all_presets_have_divisible_heads() {
1066        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
1067            assert_eq!(
1068                cfg.n_heads % cfg.n_kv_heads,
1069                0,
1070                "{}: n_heads must be a multiple of n_kv_heads for GQA grouping",
1071                cfg.name
1072            );
1073        }
1074    }
1075
1076    #[test]
1077    fn every_preset_declares_its_uncertain_fields() {
1078        // This is a documentation-honesty test: any preset with zero
1079        // best_effort_fields would be silently overclaiming precision
1080        // we don't have. Fail loudly if that ever happens.
1081        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
1082            assert!(
1083                !cfg.best_effort_fields.is_empty(),
1084                "{}: must disclose which fields are unconfirmed estimates",
1085                cfg.name
1086            );
1087        }
1088    }
1089
1090    /// Kimi K3's `kda_layers`/`full_attn_layers` were transcribed by
1091    /// hand from the real published config.json; this test guards
1092    /// against a transcription slip (duplicate, out-of-range, or
1093    /// missing layer index) rather than trusting the transcription.
1094    #[test]
1095    fn kimi_k3_hybrid_attention_layers_partition_every_layer_exactly_once() {
1096        let cfg = kimi_k3();
1097        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
1098            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
1099        };
1100
1101        let mut seen = std::collections::HashSet::new();
1102        for &l in hybrid
1103            .kda_layers
1104            .iter()
1105            .chain(hybrid.full_attn_layers.iter())
1106        {
1107            assert!(
1108                (1..=cfg.n_layers).contains(&l),
1109                "layer {l} is out of the published 1..={} range",
1110                cfg.n_layers
1111            );
1112            assert!(
1113                seen.insert(l),
1114                "layer {l} appears in both/either list twice"
1115            );
1116        }
1117        // Dense-vs-MoE (n_dense_leading_layers) and attention-type
1118        // (KDA vs Gated MLA) are independent per-layer properties in
1119        // the real config -- e.g. layer 1 is both the sole dense
1120        // leading layer *and* a KDA layer -- so every one of the 93
1121        // layers, dense or not, is covered by exactly one of these two
1122        // lists (confirmed: 69 + 24 == 93, not 93 - 1).
1123        assert_eq!(
1124            hybrid.kda_layers.len() + hybrid.full_attn_layers.len(),
1125            cfg.n_layers,
1126            "every layer must be assigned exactly one of KDA or Gated MLA"
1127        );
1128        assert_eq!(
1129            hybrid.kda_layers.len(),
1130            69,
1131            "expected 69 KDA layers per the published config"
1132        );
1133        assert_eq!(
1134            hybrid.full_attn_layers.len(),
1135            24,
1136            "expected 24 Gated MLA layers per the published config"
1137        );
1138    }
1139
1140    #[test]
1141    fn layer_attention_kind_is_gqa_for_every_layer_of_a_gqa_model() {
1142        let cfg = glm_5_2();
1143        for l in 0..cfg.n_layers {
1144            assert_eq!(cfg.layer_attention_kind(l), LayerAttentionKind::Gqa);
1145        }
1146    }
1147
1148    #[test]
1149    fn layer_attention_kind_classifies_every_kimi_k3_layer_without_panicking() {
1150        let cfg = kimi_k3();
1151        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
1152            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
1153        };
1154        for l in 0..cfg.n_layers {
1155            let kind = cfg.layer_attention_kind(l);
1156            let one_indexed = l + 1;
1157            if hybrid.kda_layers.contains(&one_indexed) {
1158                assert_eq!(kind, LayerAttentionKind::KimiKda);
1159            } else {
1160                assert_eq!(kind, LayerAttentionKind::KimiMla);
1161            }
1162        }
1163    }
1164
1165    #[test]
1166    fn layer_attention_kind_matches_the_real_published_layer_1_and_4() {
1167        // Layer 1 (1-indexed, so index 0 here) is published as KDA;
1168        // layer 4 (index 3) is published as the first Gated MLA layer.
1169        let cfg = kimi_k3();
1170        assert_eq!(cfg.layer_attention_kind(0), LayerAttentionKind::KimiKda);
1171        assert_eq!(cfg.layer_attention_kind(3), LayerAttentionKind::KimiMla);
1172    }
1173
1174    #[test]
1175    fn kimi_k3_mla_q_head_dim_matches_gqa_placeholder_head_dim() {
1176        // The Gqa-placeholder head_dim above is deliberately set to
1177        // Gated MLA's combined q_head_dim (qk_nope + qk_rope) so the
1178        // placeholder path at least reflects a real dimension from the
1179        // published config rather than an arbitrary guess.
1180        let cfg = kimi_k3();
1181        let AttentionKind::KimiHybrid(hybrid) = &cfg.attention else {
1182            panic!("kimi_k3() must use AttentionKind::KimiHybrid");
1183        };
1184        assert_eq!(
1185            cfg.head_dim,
1186            hybrid.mla.qk_nope_head_dim + hybrid.mla.qk_rope_head_dim
1187        );
1188    }
1189
1190    #[test]
1191    fn approx_active_params_is_nonzero_and_finite_order_of_magnitude() {
1192        for cfg in [glm_5_2(), deepseek_v4_pro(), kimi_k3()] {
1193            let approx = cfg.approx_active_params_per_token();
1194            // Sanity band: active params/token for these models is
1195            // reported in the tens of billions; this is a loose
1196            // order-of-magnitude check (1e9 to 1e12), not a precise
1197            // parameter-count reproduction.
1198            assert!(
1199                approx > 1_000_000_000 && approx < 1_000_000_000_000,
1200                "{}: approx_active_params_per_token={approx} is outside a plausible range",
1201                cfg.name
1202            );
1203        }
1204    }
1205}
1206
1207#[cfg(test)]
1208mod longrope_tests {
1209    use super::*;
1210
1211    fn cfg_with_factors() -> ModelConfig {
1212        let mut c = test_dense_fixture();
1213        c.rope_orig_ctx = Some(4096);
1214        c.rope_freqs_short = Some(vec![1.0; 48]);
1215        c.rope_freqs_long = Some((0..48).map(|i| 1.0 + i as f32).collect());
1216        c.rope_freqs = None;
1217        c
1218    }
1219
1220    /// llama.cpp `llama_model::get_rope_factors`: long only when the
1221    /// run's context exceeds `original_context_length`. Phi-4-mini's
1222    /// short set is all ones, so picking long at 4096 would apply a
1223    /// correction the model never asked for at that length.
1224    #[test]
1225    fn long_set_only_above_the_original_context() {
1226        let mut c = cfg_with_factors();
1227        c.apply_runtime_context(4096);
1228        assert_eq!(
1229            c.rope_freqs.as_ref().unwrap()[1],
1230            1.0,
1231            "at the threshold, short"
1232        );
1233
1234        let mut c = cfg_with_factors();
1235        c.apply_runtime_context(4097);
1236        assert_eq!(c.rope_freqs.as_ref().unwrap()[1], 2.0, "above it, long");
1237
1238        let mut c = cfg_with_factors();
1239        c.apply_runtime_context(1024);
1240        assert_eq!(c.rope_freqs.as_ref().unwrap()[1], 1.0, "below it, short");
1241    }
1242
1243    /// `rope_freqs.weight` (Llama 3) is not a LongRoPE set and outranks
1244    /// one, the same precedence llama.cpp gives it. The loader encodes
1245    /// that by leaving the long/short pair empty whenever the explicit
1246    /// tensor is present, so the runtime re-pick has nothing to apply.
1247    #[test]
1248    fn an_explicit_rope_freqs_tensor_is_never_overridden() {
1249        let mut c = test_dense_fixture();
1250        c.rope_freqs = Some(vec![7.0; 48]);
1251        c.rope_orig_ctx = Some(4096);
1252        c.rope_freqs_long = None;
1253        c.rope_freqs_short = None;
1254        c.apply_runtime_context(131072);
1255        assert_eq!(c.rope_freqs.as_ref().unwrap()[0], 7.0);
1256    }
1257
1258    /// A checkpoint with neither set must come back untouched, so the
1259    /// call is free to sit on every load path.
1260    #[test]
1261    fn models_without_longrope_are_untouched() {
1262        let mut c = test_dense_fixture();
1263        c.rope_freqs = None;
1264        c.apply_runtime_context(8192);
1265        assert!(c.rope_freqs.is_none());
1266        assert!(c.rope_orig_ctx.is_none());
1267    }
1268}