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