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