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