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