Skip to main content

ferrox_models/
loader.rs

1//! Loads a real `Decoder` from an on-disk GGUF file, using the
2//! llama.cpp-style tensor naming convention
3//! (`token_embd.weight`, `blk.N.attn_q.weight`, `blk.N.ffn_gate.weight`
4//! or, for MoE, `blk.N.ffn_gate_exps.weight`, `output_norm.weight`,
5//! `output.weight`). Until this module existed, ferrox could only run
6//! correctly-shaped *random* weights.
7//!
8//! Quantized tensors (Q8_0 / Q4_0) are loaded as `WeightMatrix::Quantized`
9//! backed by `WeightBytes::Mapped` -- a zero-copy view into the same
10//! mmap `GgufFile` already holds, with no intermediate heap copy of the
11//! tensor's bytes at all. So a checkpoint's resident memory is the
12//! mmap page cache, not the mmap plus a second in-process copy of every
13//! weight. `WeightMatrix::apply` dispatches to ferrox-quant's fused
14//! dequant+dot kernels directly against those mapped bytes at inference
15//! time. F32 tensors (norms, embeddings, and any weight not natively
16//! quantized) still copy into an owned `Tensor`, since they're small
17//! relative to the quantized weight matrices and need per-element
18//! access patterns a raw byte view doesn't support as cleanly.
19//!
20//! Verified end to end (see `crates/ferrox-models/tests/gguf_roundtrip.rs`)
21//! against a genuinely Q8_0-quantized, generated on-disk GGUF fixture
22//! for the dense (single-expert) case, and against real OLMoE / Qwen2-MoE
23//! checkpoints for the multi-expert 3D-packed-tensor path.
24
25use ferrox_core::expert_store::{ExpertKey, ExpertSource, ExpertStore};
26use ferrox_core::tensor::Tensor;
27use ferrox_core::weight_matrix::quant_kind_for;
28use ferrox_core::weight_matrix::{QuantKind, WeightBytes, WeightMatrix};
29use ferrox_gguf::{GgmlType, GgufError, GgufValue, ShardedGguf, TensorInfo, TensorSource};
30use ferrox_moe::{ExpertWeights, GatingFunction, MoeLayerConfig};
31use std::sync::Arc;
32use thiserror::Error;
33
34use crate::config::ModelConfig;
35#[cfg(feature = "metal")]
36use crate::decoder::MoePackedQ4Planes;
37use crate::decoder::{AttnWeights, Decoder, ExpertBacking, LayerWeights, MoeWeights};
38
39#[derive(Debug, Error)]
40pub enum LoadError {
41    #[error(transparent)]
42    Gguf(#[from] GgufError),
43    #[error(transparent)]
44    Shard(#[from] ferrox_gguf::ShardError),
45    #[error("tensor '{0}' has unsupported dtype {1:?}")]
46    UnsupportedDtype(String, GgmlType),
47    #[error(
48        "MoE tensor '{0}' is not 3D or its expert count {1} does not match config n_experts {2}"
49    )]
50    ExpertCountMismatch(String, usize, usize),
51    #[error("GGUF file is missing required hparam metadata key '{0}'")]
52    MissingHparam(String),
53    /// `general.architecture` is not in the capability registry — refuse
54    /// to guess RoPE/gating rather than emit fluent-but-wrong logits.
55    #[error(
56        "unsupported GGUF architecture '{0}': not in ferrox's capability registry \
57         (unknown required features fail closed; see ferrox_models::capability)"
58    )]
59    UnsupportedArchitecture(String),
60    /// Architecture exists but must not use the generic GQA decoder.
61    #[error("architecture '{0}' cannot use the generic Decoder: {1}")]
62    DedicatedArchitectureRequired(String, &'static str),
63    /// Metadata advertises a feature the generic decoder does not implement.
64    #[error("architecture '{0}' requires unimplemented feature: {1}")]
65    UnsupportedFeature(String, String),
66    /// The checkpoint carries per-block tensors this build never reads,
67    /// i.e. weights that contribute to the real graph and would simply
68    /// be missing from ours. See [`assert_every_tensor_consumed`].
69    #[error(
70        "checkpoint carries {0} tensor(s) this build never reads, so its graph is not the one \
71         ferrox would run: {1}. This is a missing feature, not a corrupt file. Override with \
72         FERROX_ALLOW_UNKNOWN_TENSORS=1 to load anyway and accept wrong output."
73    )]
74    UnconsumedTensors(usize, String),
75    /// `FERROX_STRICT_KERNELS=1` and the model has weights with no
76    /// kernel on the selected accelerator, i.e. it would run, correctly,
77    /// on a silently slower path. Refusing is the point: a benchmark or
78    /// CI run must not be able to publish a number taken off the
79    /// backend it claims. See [`ferrox_core::kernel_registry`].
80    #[error("{0}")]
81    StrictKernels(String),
82}
83
84/// Architecture-family name strings (GGUF's `general.architecture` value)
85/// known, from reading ik_llama.cpp's `llama-hparams.cpp`
86/// (`LLM_ARCH_DEEPSEEK2`, `LLM_ARCH_GLM4_MOE` cases), to default to
87/// sigmoid MoE gating with post-selection renormalization rather than
88/// softmax. See docs/MODELS.md for the citations behind this list.
89const SIGMOID_GATING_ARCHITECTURES: &[&str] = &["deepseek2", "glm4moe"];
90
91/// Architecture-family names whose real reference implementation skips
92/// renormalizing top-k softmax routing weights after selection (GGUF
93/// carries no metadata key for this -- it's hardcoded per-architecture in
94/// both the real HF `transformers` model code and llama.cpp's
95/// `build_moe_ffn` call sites, not read from the file). Confirmed for
96/// `olmoe` against `OlmoeTopKRouter.forward` in
97/// `transformers/models/olmoe/modeling_olmoe.py` (`config.norm_topk_prob`
98/// is `false` in the real published config.json) and llama.cpp's
99/// `src/models/olmoe.cpp` (`build_moe_ffn(..., false, ...,
100/// LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, ...)`). See
101/// `MoeLayerConfig::norm_topk_prob`'s doc comment for why this matters:
102/// getting it wrong silently produces wrong generation output even
103/// though the file loads and shape-validates fine.
104// Architectures whose reference graphs pass `norm_w=false` to
105// `build_moe_ffn` (llama.cpp) / `norm_topk_prob=false` in HF config.
106// Qwen2-MoE: `.scratch/llama.cpp/src/models/qwen2moe.cpp` — Softmax +
107// `false` for the norm_topk slot. Renormalizing top-k weights made
108// Qwen1.5-MoE greedy decode emit garbage despite shared-expert load.
109const NO_TOPK_RENORMALIZE_ARCHITECTURES: &[&str] = &["olmoe", "qwen2moe"];
110
111fn metadata_u64_any(file: &impl TensorSource, keys: &[String]) -> Option<u64> {
112    keys.iter().find_map(|k| file.metadata_u64(k))
113}
114
115fn metadata_f32_any(file: &impl TensorSource, keys: &[String]) -> Option<f32> {
116    keys.iter()
117        .find_map(|k| file.metadata(k).and_then(GgufValue::as_f32))
118}
119
120impl ModelConfig {
121    /// Derives a `ModelConfig` from a real GGUF file's own hyperparameter
122    /// metadata, following llama.cpp's `general.architecture`-prefixed key
123    /// convention (`{arch}.block_count`, `{arch}.embedding_length`,
124    /// `{arch}.attention.head_count`, `{arch}.expert_count`, ...) rather
125    /// than requiring a hand-written preset to already match the file's
126    /// shape exactly. This is what lets `ferrox-server` (and `ferrox
127    /// run-real`) load an arbitrary checkpoint, not just the three
128    /// hand-tuned presets in `config.rs`.
129    ///
130    /// Fields with no corresponding metadata key fall back to widely-used
131    /// llama.cpp defaults (documented inline) and are listed in the
132    /// returned config's `best_effort_fields`, following the same
133    /// confirmed-vs-estimated discipline as the hand-written presets.
134    pub fn from_gguf(file: &impl TensorSource) -> Result<Self, LoadError> {
135        let arch = file
136            .metadata_str("general.architecture")
137            .ok_or_else(|| LoadError::MissingHparam("general.architecture".to_string()))?
138            .to_string();
139        let arch_profile = crate::capability::resolve_profile(&arch)
140            .ok_or_else(|| LoadError::UnsupportedArchitecture(arch.clone()))?;
141        let rope_layout = match arch_profile.path {
142            crate::capability::ArchPath::GenericGqa { rope }
143            | crate::capability::ArchPath::TestFixture { rope } => rope,
144            crate::capability::ArchPath::DedicatedOnly { reason } => {
145                return Err(LoadError::DedicatedArchitectureRequired(
146                    arch.clone(),
147                    reason,
148                ));
149            }
150            crate::capability::ArchPath::Deferred { reason } => {
151                return Err(LoadError::UnsupportedFeature(
152                    arch.clone(),
153                    format!("architecture deferred from Ferrox text-generation scope: {reason}"),
154                ));
155            }
156        };
157        let qk_norm_style = arch_profile.qk_norm;
158        for (meta_key, feature) in crate::capability::unsupported_feature_keys(&arch) {
159            if let Some(v) = metadata_f32_any(file, std::slice::from_ref(&meta_key)) {
160                if v > 0.0 {
161                    return Err(LoadError::UnsupportedFeature(
162                        arch.clone(),
163                        format!("{feature} (metadata {meta_key}={v})"),
164                    ));
165                }
166            }
167            if let Some(v) = metadata_u64_any(file, std::slice::from_ref(&meta_key)) {
168                if v > 0 {
169                    return Err(LoadError::UnsupportedFeature(
170                        arch.clone(),
171                        feature.to_string(),
172                    ));
173                }
174            }
175        }
176        // Metadata-declared multipliers the generic decoder does not
177        // apply. Unlike the tensor-consumption gate, nothing about these
178        // is visible in the weights, so a Granite checkpoint would load
179        // and answer at the wrong scale. See
180        // `capability::unsupported_scaling_keys`.
181        for (meta_key, feature, no_op) in crate::capability::unsupported_scaling_keys(&arch) {
182            if let Some(v) = metadata_f32_any(file, std::slice::from_ref(&meta_key)) {
183                if (v - no_op).abs() > 1e-6 {
184                    return Err(LoadError::UnsupportedFeature(
185                        arch.clone(),
186                        format!("{feature} (metadata {meta_key}={v})"),
187                    ));
188                }
189            }
190        }
191        let key = |suffix: &str| format!("{arch}.{suffix}");
192
193        let name: &'static str = Box::leak(
194            file.metadata_str("general.name")
195                .unwrap_or(&arch)
196                .to_string()
197                .into_boxed_str(),
198        );
199
200        let n_layers =
201            file.metadata_u64(&key("block_count"))
202                .ok_or_else(|| LoadError::MissingHparam(key("block_count")))? as usize;
203        // Baichuan is one architecture string covering two positional
204        // schemes: 7B rotates, 13B uses ALiBi and no RoPE at all
205        // (`src/models/baichuan.cpp:11-14`, `:57-58`, where `inp_pos` is
206        // `nullptr` for 13B, so `ggml_rope_ext` is never reached).
207        // llama.cpp decides that on the layer count and says so in a
208        // comment: "TODO: become GGUF KV parameter". There is therefore
209        // no key for `capability::unsupported_feature_keys` to test and
210        // no tensor for `assert_every_tensor_consumed` to miss. A
211        // Baichuan-13B checkpoint loads clean and is rotated anyway.
212        // Refuse it here, where the layer count is known.
213        if arch == "baichuan" && n_layers == 40 {
214            return Err(LoadError::UnsupportedFeature(
215                arch.clone(),
216                "Baichuan-13B (block_count=40) uses ALiBi and no RoPE, decided by layer \
217                 count with no GGUF key to declare it; the generic decoder would rotate \
218                 every Q/K head instead. Baichuan-7B (block_count=32) is unaffected"
219                    .to_string(),
220            ));
221        }
222        let hidden_dim = file
223            .metadata_u64(&key("embedding_length"))
224            .ok_or_else(|| LoadError::MissingHparam(key("embedding_length")))?
225            as usize;
226        let n_heads = file
227            .metadata_u64(&key("attention.head_count"))
228            .ok_or_else(|| LoadError::MissingHparam(key("attention.head_count")))?
229            as usize;
230
231        let mut best_effort_fields: Vec<&'static str> = Vec::new();
232
233        let n_kv_heads = file
234            .metadata_u64(&key("attention.head_count_kv"))
235            .map(|v| v as usize)
236            .unwrap_or_else(|| {
237                best_effort_fields.push("n_kv_heads (no attention.head_count_kv key; assumed equal to n_heads, i.e. plain MHA)");
238                n_heads
239            });
240        let head_dim = file
241            .metadata_u64(&key("attention.key_length"))
242            .map(|v| v as usize)
243            .unwrap_or_else(|| {
244                best_effort_fields.push(
245                    "head_dim (no attention.key_length key; derived as hidden_dim / n_heads)",
246                );
247                hidden_dim / n_heads
248            });
249        let v_head_dim = file
250            .metadata_u64(&key("attention.value_length"))
251            .map(|v| v as usize)
252            .unwrap_or(head_dim);
253        if v_head_dim != head_dim {
254            return Err(LoadError::UnsupportedFeature(
255                arch.clone(),
256                format!(
257                    "split K/V head dims (key_length={head_dim}, value_length={v_head_dim}); \
258                     generic decoder requires equal head dims"
259                ),
260            ));
261        }
262        let vocab_size = file
263            .metadata("tokenizer.ggml.tokens")
264            .and_then(|v| match v {
265                GgufValue::Array(items) => Some(items.len()),
266                _ => None,
267            })
268            .or_else(|| file.metadata_u64(&key("vocab_size")).map(|v| v as usize))
269            .unwrap_or_else(|| {
270                best_effort_fields.push("vocab_size (no tokenizer.ggml.tokens array or {arch}.vocab_size key; fell back to output.weight's own row count)");
271                // `output.weight`'s real raw shape is `[hidden_dim,
272                // vocab_size]` (ggml's fastest-first `ne[]` order --
273                // see `load_weight_matrix`'s doc comment), so vocab_size
274                // is the *last* element, not the first.
275                file.find_tensor("output.weight")
276                    .and_then(|t| t.shape.last().copied())
277                    .unwrap_or(0) as usize
278            });
279        let rope_theta = metadata_f32_any(file, &[key("rope.freq_base")]).unwrap_or_else(|| {
280            best_effort_fields.push("rope_theta (no rope.freq_base key; defaulted to 10000.0)");
281            10000.0
282        });
283        let rms_norm_eps = metadata_f32_any(
284            file,
285            &[
286                key("attention.layer_norm_rms_epsilon"),
287                key("attention.layer_norm_epsilon"),
288            ],
289        )
290        .unwrap_or_else(|| {
291            best_effort_fields
292                .push("rms_norm_eps (no layer_norm_rms_epsilon key; defaulted to 1e-5)");
293            1e-5
294        });
295
296        let n_experts = metadata_u64_any(file, &[key("expert_count")]).unwrap_or(0) as usize;
297        let is_moe = n_experts > 1;
298
299        let n_experts_active = if is_moe {
300            metadata_u64_any(file, &[key("expert_used_count")]).unwrap_or_else(|| {
301                best_effort_fields
302                    .push("moe.n_experts_active (no expert_used_count key; defaulted to 2)");
303                2
304            }) as usize
305        } else {
306            1
307        };
308        // Prefer the GGUF hparam when present. Qwen2MoE (and some other
309        // HF→GGUF exports) omit `expert_shared_count` but still ship
310        // `blk.N.ffn_{gate,up,down}_shexp.weight` — without a tensor-
311        // presence fallback those weights are silently dropped and the
312        // model runs with a large chunk of active FFN missing.
313        let n_shared_experts = match metadata_u64_any(file, &[key("expert_shared_count")]) {
314            Some(n) => n as usize,
315            None if is_moe && file.find_tensor("blk.0.ffn_gate_shexp.weight").is_some() => {
316                best_effort_fields.push(
317                    "moe.n_shared_experts (no expert_shared_count; inferred 1 from blk.0.ffn_gate_shexp.weight)",
318                );
319                1
320            }
321            None => 0,
322        };
323        // MoE GGUFs often only set `feed_forward_length` (OLMoE=1024,
324        // Qwen2-MoE=5632 for the shared expert). `expert_feed_forward_length`
325        // is optional. llama.cpp `qwen2moe.cpp` uses
326        // `n_ff_exp = n_ff_exp ? n_ff_exp : n_ff / n_expert_used` (1408 for
327        // Qwen1.5-MoE); the shared expert keeps the full `n_ff` (5632).
328        let feed_forward_length = metadata_u64_any(file, &[key("feed_forward_length")]);
329        let expert_ffn_dim = metadata_u64_any(file, &[key("expert_feed_forward_length")])
330            .or_else(|| {
331                feed_forward_length.map(|ff| {
332                    if is_moe && n_experts_active > 0 {
333                        ff / n_experts_active as u64
334                    } else {
335                        ff
336                    }
337                })
338            })
339            .unwrap_or_else(|| {
340                best_effort_fields.push(
341                    "moe.expert_ffn_dim (no expert_feed_forward_length/feed_forward_length; defaulted to 4x hidden_dim)",
342                );
343                (hidden_dim * 4) as u64
344            }) as usize;
345        let n_dense_leading_layers =
346            metadata_u64_any(file, &[key("leading_dense_block_count")]).unwrap_or(0) as usize;
347
348        // ik_llama.cpp's real gating-function hparam
349        // (LLM_KV_EXPERT_GATING_FUNC: 1=softmax, 2=sigmoid) if the file
350        // carries it; otherwise fall back to the same architecture-name
351        // convention the hand-written presets in config.rs use (see
352        // docs/MODELS.md for the citations behind that list).
353        let gating = match metadata_u64_any(file, &[key("expert_gating_func")]) {
354            Some(2) => GatingFunction::Sigmoid,
355            Some(1) => GatingFunction::Softmax,
356            _ => {
357                if SIGMOID_GATING_ARCHITECTURES.contains(&arch.as_str()) {
358                    GatingFunction::Sigmoid
359                } else {
360                    if is_moe {
361                        best_effort_fields.push(
362                            "moe.gating (no expert_gating_func key and architecture not in the known-sigmoid list; defaulted to softmax)",
363                        );
364                    }
365                    GatingFunction::Softmax
366                }
367            }
368        };
369
370        // `{arch}.expert_weights_norm` (llama.cpp
371        // `LLM_KV_EXPERT_WEIGHTS_NORM`) is the real metadata key for
372        // whether the selected experts' weights are renormalised. Most
373        // checkpoints do not carry it, which is why the fallback below
374        // exists at all -- but when one does, the file's own answer wins
375        // over an architecture-name guess.
376        let norm_topk_prob = match file.metadata_bool(&key("expert_weights_norm")) {
377            Some(v) => v,
378            None => {
379                // See `NO_TOPK_RENORMALIZE_ARCHITECTURES`'s doc comment:
380                // an architecture-name lookup, the same convention
381                // `gating`'s fallback above uses.
382                if is_moe && matches!(gating, GatingFunction::Softmax) {
383                    best_effort_fields.push(
384                        "moe.norm_topk_prob (no expert_weights_norm key; defaulted by architecture-name lookup against NO_TOPK_RENORMALIZE_ARCHITECTURES)",
385                    );
386                }
387                !NO_TOPK_RENORMALIZE_ARCHITECTURES.contains(&arch.as_str())
388            }
389        };
390
391        // `{arch}.expert_weights_scale` (`LLM_KV_EXPERT_WEIGHTS_SCALE`).
392        // llama.cpp's `build_moe_ffn` skips the multiply for both 0.0 and
393        // 1.0, so both mean "no scaling" and both land on 1.0 here.
394        let expert_weights_scale = metadata_f32_any(file, &[key("expert_weights_scale")])
395            .filter(|s| *s != 0.0)
396            .unwrap_or(1.0);
397
398        // Real GGUF key (`{arch}.attention.sliding_window`, confirmed
399        // against `gguf-py/gguf/constants.py`'s real
400        // `LLM_KV_ATTENTION_SLIDING_WINDOW`). Some checkpoints
401        // (confirmed for real published Qwen1.5-MoE/Qwen2-MoE GGUFs)
402        // carry a nonzero window value even when the model's own
403        // config disables sliding-window attention entirely
404        // (`use_sliding_window: false`) -- llama.cpp's own convention
405        // is that a window of 0 means "unused," so only a real nonzero
406        // value here is treated as active.
407        let sliding_window = metadata_u64_any(file, &[key("attention.sliding_window")])
408            .map(|v| v as usize)
409            .filter(|&w| w > 0);
410
411        // Gemma alternating SWA period (`attention.sliding_window_pattern`).
412        // llama.cpp: gemma2 defaults period=2, gemma3 defaults period=6 when
413        // the pattern key is absent. A missing key must NOT mean "all SWA".
414        let swa_pattern = metadata_u64_any(file, &[key("attention.sliding_window_pattern")])
415            .map(|v| v as usize)
416            .filter(|&p| p > 1)
417            .or_else(|| {
418                sliding_window?;
419                // llama.cpp hardcodes the period per architecture and
420                // only lets the metadata key override it, so a missing
421                // key is *not* "every layer windowed" — see
422                // `capability::default_swa_pattern`.
423                crate::capability::default_swa_pattern(&arch).or(
424                    // Any Gemma variant not named in the table keeps the
425                    // gemma3+ period rather than going uniform.
426                    match arch_profile.family {
427                        crate::capability::DecoderFamily::GemmaFamily => Some(6),
428                        _ => None,
429                    },
430                )
431            });
432
433        let attn_logit_softcap = metadata_f32_any(
434            file,
435            &[
436                key("attention.logit_softcapping"),
437                key("attn_logit_softcapping"),
438            ],
439        )
440        .filter(|&v| v > 0.0);
441        let final_logit_softcap =
442            metadata_f32_any(file, &[key("final_logit_softcapping")]).filter(|&v| v > 0.0);
443
444        // Gemma: embeddings are scaled by sqrt(hidden_dim) at input.
445        let embedding_scale = if matches!(
446            arch_profile.family,
447            crate::capability::DecoderFamily::GemmaFamily
448        ) {
449            Some((hidden_dim as f32).sqrt())
450        } else {
451            None
452        };
453
454        // Gemma's f_attention_scale equals 1/sqrt(n_embd_head_k) for non-27B,
455        // which is already what `causal_gqa_attention` applies. Do not also
456        // pre-scale Q (that double-scales scores vs llama.cpp's
457        // `build_attn(..., 1.0f)` after an explicit Q scale).
458        let attention_scale = None;
459
460        // SWA-layer RoPE base. `llama_hparams` defaults it to 10000 and
461        // the Gemma-3 lineage relies on that default; the architectures
462        // in `swa_rope_base_follows_model` instead seed it from the
463        // model's own base before the key can override.
464        let rope_theta_swa = if sliding_window.is_some() {
465            let fallback = if crate::capability::swa_rope_base_follows_model(&arch) {
466                rope_theta
467            } else {
468                10_000.0
469            };
470            Some(
471                metadata_f32_any(
472                    file,
473                    &[key("rope.freq_base_swa"), key("rope_freq_base_swa")],
474                )
475                .unwrap_or(fallback),
476            )
477        } else {
478            None
479        };
480
481        let ffn_activation = match arch_profile.family {
482            crate::capability::DecoderFamily::GemmaFamily => crate::config::FfnActivation::Gelu,
483            crate::capability::DecoderFamily::PhiFamily => {
484                crate::config::FfnActivation::SwigluFused
485            }
486            _ => crate::config::FfnActivation::Swiglu,
487        };
488
489        // Llama 3/3.1/3.2's real per-band RoPE frequency correction: one
490        // model-level tensor (`TENSOR_NOT_REQUIRED`, `TENSOR_DUPLICATED`
491        // for every layer but the first in the real llama.cpp source --
492        // i.e. every layer shares this same array), not per-layer. See
493        // `ferrox_core::attention::apply_rope_with_freq_factors`'s doc
494        // comment for why this matters.
495        let rope_freqs = load_f32_vec_optional(file, "rope_freqs.weight")?;
496
497        // Phi-3/Phi-4 LongRoPE: two per-band factor tensors instead of
498        // Llama's single `rope_freqs.weight`, selected by context size
499        // (llama.cpp `llama_model::get_rope_factors`: `rope_freqs` wins if
500        // present, else `rope_long` when the run's context exceeds
501        // `rope.scaling.original_context_length`, else `rope_short`).
502        //
503        // The selection here uses the checkpoint's own advertised context
504        // length, which is what llama.cpp defaults `n_ctx` to. A run that
505        // caps the context below `original_context_length` should use the
506        // short set; ferrox's config is built before the context size is
507        // known, so that case is not yet handled — recorded as a
508        // best-effort field rather than silently assumed correct.
509        let rope_orig_ctx = metadata_u64_any(file, &[key("rope.scaling.original_context_length")])
510            .map(|v| v as usize);
511        // `rope_freqs.weight` outranks the LongRoPE pair (llama.cpp
512        // `get_rope_factors` checks it first), so a checkpoint carrying
513        // it never populates these and the runtime re-pick below cannot
514        // overwrite a Llama-3 correction with a Phi one.
515        let (rope_freqs_long, rope_freqs_short) = if rope_freqs.is_some() {
516            (None, None)
517        } else {
518            (
519                load_f32_vec_optional(file, "rope_factors_long.weight")?,
520                load_f32_vec_optional(file, "rope_factors_short.weight")?,
521            )
522        };
523        // Provisional pick from the checkpoint's own advertised context;
524        // `ModelConfig::apply_runtime_context` re-picks once the run's
525        // `--ctx-size` is known, which is the number llama.cpp decides on.
526        let rope_freqs = match (rope_freqs, rope_orig_ctx) {
527            (Some(f), _) => Some(f),
528            (None, Some(orig)) => {
529                let model_ctx = metadata_u64_any(file, &[key("context_length")])
530                    .unwrap_or(orig as u64) as usize;
531                if model_ctx > orig {
532                    rope_freqs_long.clone().or_else(|| rope_freqs_short.clone())
533                } else {
534                    rope_freqs_short.clone().or_else(|| rope_freqs_long.clone())
535                }
536            }
537            (None, None) => None,
538        };
539
540        // Partial rotary: only when the file says the rotary width is
541        // narrower than a head. Equal values mean "whole head", which is
542        // the same thing as `None` and stays `None` so nothing downstream
543        // has to special-case it.
544        let rope_dim = metadata_u64_any(file, &[key("rope.dimension_count")])
545            .map(|d| d as usize)
546            .filter(|d| *d > 0 && *d < head_dim);
547
548        // See `ModelConfig::rope_attn_factor`.
549        let rope_attn_factor = metadata_f32_any(file, &[key("rope.scaling.attn_factor")])
550            .filter(|f| f.is_finite() && *f > 0.0)
551            .unwrap_or(1.0);
552
553        // YaRN long-context scaling. `rope.scaling.attn_factor` above is
554        // only YaRN's *magnitude* term (ggml `rope_yarn`'s `mscale`); the
555        // frequency half -- which bands get interpolated toward the
556        // trained context and which stay extrapolated -- lives in
557        // `rope.scaling.type` + `rope.scaling.factor`, and ferrox read
558        // neither before this. A YaRN checkpoint was therefore roped as
559        // if it declared no scaling at all: right near position 0 and
560        // progressively wrong further in, i.e. the failure that reads as
561        // long-prompt quality decay rather than as a bug.
562        //
563        // The rewrite is folded into `rope_freqs`, the same per-band
564        // divisor array Llama-3's `rope_freqs.weight` supplies (ggml
565        // divides each band's theta by it), so it rides the existing CPU
566        // and Metal RoPE paths unchanged. When a file carries both, the
567        // two corrections compose by multiplication, as they do in
568        // llama.cpp (`ggml_rope_cache_init` divides by `freq_factors`
569        // *and then* runs `rope_yarn`).
570        // Linear scaling, which was silently DROPPED before this.
571        //
572        // `rope.scaling.type = "linear"` with factor s means rotating
573        // position `p/s` instead of `p`. Since the angle is `p * freq`,
574        // that is exactly `p * (freq / s)`, and `rope_freqs` already
575        // divides each band's frequency. So a uniform vector of `s`
576        // expresses it exactly and rides the existing CPU and Metal RoPE
577        // paths unchanged, the same way YaRN does below.
578        //
579        // Before this, the type was compared against "yarn" and anything
580        // else returned None, so a checkpoint declaring linear scaling
581        // with factor 4 loaded and roped at UNSCALED positions where
582        // llama.cpp divides them by 4. It answered as a different model
583        // with no error. Affects the long-context community rescales
584        // (`*-16k`, `*-32k` Llama-2 derivatives).
585        let rope_freqs = match linear_scaling_from_gguf(file, &arch) {
586            None => rope_freqs,
587            Some(factor) => {
588                let rotary_dim = rope_dim.unwrap_or(head_dim);
589                if rotary_dim == 0 || !rotary_dim.is_multiple_of(2) {
590                    best_effort_fields.push(
591                        "rope_freqs (linear scaling declared but the rotary width is odd; \
592                         scaling not applied)",
593                    );
594                    rope_freqs
595                } else {
596                    let linear = vec![factor; rotary_dim / 2];
597                    match rope_freqs {
598                        None => Some(linear),
599                        // Compose by multiplication, as a file carrying
600                        // its own `rope_freqs.weight` tensor and a
601                        // declared linear factor means both.
602                        Some(own) if own.len() == linear.len() => {
603                            Some(own.iter().zip(linear.iter()).map(|(a, b)| a * b).collect())
604                        }
605                        Some(own) => {
606                            best_effort_fields.push(
607                                "rope_freqs (linear scaling declared but the file's own \
608                                 rope_freqs tensor has a different width; scaling not applied)",
609                            );
610                            Some(own)
611                        }
612                    }
613                }
614            }
615        };
616        let rope_freqs = match yarn_scaling_from_gguf(file, &arch, rope_orig_ctx) {
617            None => rope_freqs,
618            Some(scaling) => {
619                let rotary_dim = rope_dim.unwrap_or(head_dim);
620                if rotary_dim == 0 || !rotary_dim.is_multiple_of(2) {
621                    best_effort_fields.push(
622                        "rope_freqs (YaRN declared but the rotary width is odd; scaling not applied)",
623                    );
624                    rope_freqs
625                } else {
626                    let yarn =
627                        ferrox_core::attention::yarn_freq_factors(scaling, rotary_dim, rope_theta);
628                    match rope_freqs {
629                        None => Some(yarn),
630                        Some(own) if own.len() == yarn.len() => {
631                            Some(own.iter().zip(yarn.iter()).map(|(a, b)| a * b).collect())
632                        }
633                        Some(own) => {
634                            best_effort_fields.push(
635                                "rope_freqs (YaRN declared alongside a per-band factor tensor of a \
636                                 different width; the file's own tensor is used unscaled)",
637                            );
638                            Some(own)
639                        }
640                    }
641                }
642            }
643        };
644
645        // RoPE layout comes from the capability registry above (fail-
646        // closed). Getting this wrong for `llama` (needs Norm) was the
647        // real root cause of the Llama-3.1-8B early-stop/wrong-logits bug.
648
649        if best_effort_fields.is_empty() {
650            best_effort_fields.push(
651                "none -- every field above was read directly from this file's own GGUF metadata",
652            );
653        }
654
655        Ok(ModelConfig {
656            name,
657            n_layers,
658            hidden_dim,
659            n_heads,
660            n_kv_heads,
661            head_dim,
662            vocab_size,
663            rope_theta,
664            rms_norm_eps,
665            // No GGUF file encodes a hybrid KDA/Gated-MLA attention
666            // topology today; every real checkpoint loaded this way
667            // runs the standard Gqa path.
668            attention: crate::config::AttentionKind::Gqa,
669            sliding_window,
670            swa_pattern,
671            moe: MoeLayerConfig {
672                n_experts: n_experts.max(1),
673                n_experts_active,
674                n_shared_experts,
675                hidden_dim,
676                expert_ffn_dim,
677                gating,
678                norm_topk_prob,
679                expert_group_count: metadata_u64_any(file, &[key("expert_group_count")])
680                    .map(|v| v as usize)
681                    .filter(|&c| c > 1),
682                expert_group_used_count: metadata_u64_any(file, &[key("expert_group_used_count")])
683                    .map(|v| v as usize)
684                    .filter(|&c| c > 0),
685                expert_weights_scale,
686            },
687            n_dense_leading_layers,
688            rope_freqs,
689            rope_layout,
690            qk_norm_style,
691            attn_logit_softcap,
692            final_logit_softcap,
693            embedding_scale,
694            attention_scale,
695            rope_attn_factor,
696            rope_dim,
697            rope_freqs_long,
698            rope_freqs_short,
699            rope_orig_ctx,
700            rope_theta_swa,
701            ffn_activation,
702            best_effort_fields: Box::leak(best_effort_fields.into_boxed_slice()),
703        })
704    }
705}
706
707impl crate::sampling::RecommendedSampling {
708    /// The sampling a GGUF recommends for itself, from the
709    /// `general.sampling.*` metadata keys llama.cpp's converter writes
710    /// when the source checkpoint carried a `generation_config.json`.
711    ///
712    /// This is the GGUF half of FreeToken's `load_generation_sampling`
713    /// (`python/freetoken/utils/hf.py:92`), which checks the GGUF
714    /// metadata *first* and only falls back to a `generation_config.json`
715    /// sidecar for non-GGUF checkpoints -- a GGUF is a single file and
716    /// has no sidecar to read.
717    ///
718    /// Key names are llama.cpp's own (`general.sampling.temp`, not
719    /// `temperature`). Each key is independent: a file that names only
720    /// `top_k` recommends only `top_k`, and the two fields it did not
721    /// mention stay `None` so the server's own defaults keep speaking
722    /// for them.
723    ///
724    /// `temp` / `top_p` are read as float *or* integer, because a
725    /// converter that wrote `temp = 1` stores a GGUF integer and
726    /// dropping that value would silently serve the checkpoint greedy --
727    /// the exact repetition-loop failure the recommendation exists to
728    /// prevent.
729    pub fn from_gguf(file: &impl TensorSource) -> Self {
730        let number = |k: &str| -> Option<f32> {
731            file.metadata(k)
732                .and_then(|v| v.as_f32().or_else(|| v.as_u64().map(|u| u as f32)))
733        };
734        crate::sampling::RecommendedSampling {
735            temperature: number("general.sampling.temp"),
736            top_p: number("general.sampling.top_p"),
737            top_k: file
738                .metadata("general.sampling.top_k")
739                .and_then(|v| v.as_u64())
740                .map(|v| v as usize),
741        }
742    }
743}
744
745/// The `linear` RoPE scaling factor, if this file declares one.
746///
747/// Deliberately separate from [`yarn_scaling_from_gguf`]: YaRN needs an
748/// original context length and per-band betas, and linear needs neither.
749/// Any factor at or below one is not a correction, and is treated as
750/// absent rather than applied as a no-op.
751fn linear_scaling_from_gguf(file: &impl TensorSource, arch: &str) -> Option<f32> {
752    let key = |suffix: &str| format!("{arch}.{suffix}");
753    let scaling_type = file.metadata_str(&key("rope.scaling.type"))?;
754    if !scaling_type.eq_ignore_ascii_case("linear") {
755        return None;
756    }
757    metadata_f32_any(file, &[key("rope.scaling.factor")]).filter(|f| f.is_finite() && *f > 1.0)
758}
759
760/// The YaRN RoPE scaling a GGUF declares, or `None` when this file
761/// declares none that changes the rotation.
762///
763/// llama.cpp's key names (`llama-arch.cpp`
764/// `LLM_KV_ROPE_SCALING_TYPE` / `_FACTOR`): `<arch>.rope.scaling.type`
765/// is a string (`"none"`, `"linear"`, `"yarn"`, `"longrope"`) and
766/// `<arch>.rope.scaling.factor` the ratio of served to trained context.
767/// `beta_fast` / `beta_slow` are read from both the plain and the
768/// `yarn_`-prefixed spelling and otherwise fall back to the reference's
769/// own defaults (32.0 / 1.0), which is what a real checkpoint relies on
770/// -- almost none of them write those two keys.
771///
772/// `None` is returned for every case where applying YaRN would be a
773/// guess or a no-op rather than a correction, so that no checkpoint's
774/// rotation moves without the file having asked for it:
775///
776/// * a scaling type other than `yarn` (`linear` divides positions,
777///   `longrope` rides the `rope_factors_long`/`_short` tensors this
778///   loader already reads -- neither is this rewrite, and treating them
779///   as YaRN would rope them wrong in a *new* way instead of leaving
780///   them as they are),
781/// * a missing, non-finite or `<= 1.0` factor (the reference's own
782///   `get_mscale` treats `scale <= 1` as unscaled, and a factor of 1.0
783///   makes every band's divisor exactly 1.0 anyway),
784/// * a missing `rope.scaling.original_context_length` -- the trained
785///   context is what the correction range is measured against, and
786///   inventing one (say, from `context_length`, which on a YaRN file is
787///   the *extended* length) would put the ramp in the wrong place and
788///   quietly rope the checkpoint at frequencies nobody trained.
789fn yarn_scaling_from_gguf(
790    file: &impl TensorSource,
791    arch: &str,
792    orig_ctx: Option<usize>,
793) -> Option<ferrox_core::attention::YarnScaling> {
794    let key = |suffix: &str| format!("{arch}.{suffix}");
795    let scaling_type = file.metadata_str(&key("rope.scaling.type"))?;
796    if !scaling_type.eq_ignore_ascii_case("yarn") {
797        return None;
798    }
799    let factor = metadata_f32_any(file, &[key("rope.scaling.factor")])
800        .filter(|f| f.is_finite() && *f > 1.0)?;
801    let orig_max_pos = orig_ctx?;
802    let beta = |suffix: &str, default: f32| -> f32 {
803        metadata_f32_any(
804            file,
805            &[
806                key(&format!("rope.scaling.{suffix}")),
807                key(&format!("rope.scaling.yarn_{suffix}")),
808            ],
809        )
810        .filter(|v| v.is_finite() && *v > 0.0)
811        .unwrap_or(default)
812    };
813    Some(ferrox_core::attention::YarnScaling {
814        factor,
815        beta_fast: beta("beta_fast", 32.0),
816        beta_slow: beta("beta_slow", 1.0),
817        orig_max_pos,
818        // No GGUF key carries the reference's `truncate` flag, and its
819        // default is `true`; a file that wanted the fractional range
820        // would have no way to say so here.
821        truncate: true,
822    })
823}
824
825pub(crate) fn find_info<'a>(
826    file: &'a impl TensorSource,
827    name: &str,
828) -> Result<&'a TensorInfo, LoadError> {
829    file.find_tensor(name)
830        .ok_or_else(|| LoadError::Gguf(GgufError::TensorNotFound(name.to_string())))
831}
832
833/// Like `load_f32_vec`, but for tensors that only exist on some
834/// checkpoints (e.g. `attn_q_norm`/`attn_k_norm` -- OLMoE-style
835/// per-projection QK-RMSNorm applied to the full q_proj/k_proj output
836/// before RoPE, confirmed against `OlmoeAttention.forward` in
837/// `transformers/models/olmoe/modeling_olmoe.py`: `q_norm(q_proj(x))`,
838/// `k_norm(k_proj(x))`, both plain RMSNorm over the whole projected
839/// width, not per-head). Absent for every other preset/fixture this
840/// loader already handles -- `None` there is correct, not a missing
841/// feature.
842/// Loads the five gpt-oss-only tensors for one layer.
843///
844/// Every one of them is **required**: a gpt-oss checkpoint that is
845/// missing any of these is not a gpt-oss checkpoint ferrox can run, and
846/// quietly substituting zeros would reintroduce exactly the
847/// silently-wrong-graph failure this path exists to remove. The lengths
848/// are asserted against the config for the same reason — a bias of the
849/// wrong width would otherwise be applied to a `zip`-truncated prefix
850/// and produce a plausible, wrong answer.
851///
852/// Shapes follow `src/models/openai-moe.cpp::load_arch_tensors`:
853/// `attn_sinks {n_head}`, `attn_output.bias {n_embd}`,
854/// `ffn_gate_inp.bias {n_expert}`, `ffn_{gate,up}_exps.bias
855/// {n_ff_exp, n_expert}`, `ffn_down_exps.bias {n_embd, n_expert}`.
856/// GGUF stores the fastest dimension first, so the 2-D bias tensors
857/// arrive expert-major and split by simple chunking.
858fn load_gpt_oss_layer(
859    file: &impl TensorSource,
860    l: usize,
861    config: &ModelConfig,
862) -> Result<crate::decoder::GptOssLayer, LoadError> {
863    let n_experts = config.moe.n_experts;
864    let ff = config.moe.expert_ffn_dim;
865
866    let want = |name: &str, got: usize, expect: usize| -> Result<(), LoadError> {
867        if got == expect {
868            Ok(())
869        } else {
870            Err(LoadError::UnsupportedFeature(
871                config.name.to_string(),
872                format!("{name} has {got} elements, expected {expect}"),
873            ))
874        }
875    };
876
877    let attn_sinks = load_f32_vec(file, &format!("blk.{l}.attn_sinks.weight"))?;
878    want(
879        &format!("blk.{l}.attn_sinks.weight"),
880        attn_sinks.len(),
881        config.n_heads,
882    )?;
883    let o_bias = load_f32_vec(file, &format!("blk.{l}.attn_output.bias"))?;
884    want(
885        &format!("blk.{l}.attn_output.bias"),
886        o_bias.len(),
887        config.hidden_dim,
888    )?;
889    let router_bias = load_f32_vec(file, &format!("blk.{l}.ffn_gate_inp.bias"))?;
890    want(
891        &format!("blk.{l}.ffn_gate_inp.bias"),
892        router_bias.len(),
893        n_experts,
894    )?;
895
896    let gate_b = load_f32_vec(file, &format!("blk.{l}.ffn_gate_exps.bias"))?;
897    want(
898        &format!("blk.{l}.ffn_gate_exps.bias"),
899        gate_b.len(),
900        n_experts * ff,
901    )?;
902    let up_b = load_f32_vec(file, &format!("blk.{l}.ffn_up_exps.bias"))?;
903    want(
904        &format!("blk.{l}.ffn_up_exps.bias"),
905        up_b.len(),
906        n_experts * ff,
907    )?;
908    let down_b = load_f32_vec(file, &format!("blk.{l}.ffn_down_exps.bias"))?;
909    want(
910        &format!("blk.{l}.ffn_down_exps.bias"),
911        down_b.len(),
912        n_experts * config.hidden_dim,
913    )?;
914
915    let expert_bias = (0..n_experts)
916        .map(|e| ferrox_moe::ExpertBias {
917            gate: gate_b[e * ff..(e + 1) * ff].to_vec(),
918            up: up_b[e * ff..(e + 1) * ff].to_vec(),
919            down: down_b[e * config.hidden_dim..(e + 1) * config.hidden_dim].to_vec(),
920        })
921        .collect();
922
923    Ok(crate::decoder::GptOssLayer {
924        attn_sinks,
925        o_bias,
926        router_bias,
927        expert_bias,
928    })
929}
930
931pub(crate) fn load_f32_vec_optional(
932    file: &impl TensorSource,
933    name: &str,
934) -> Result<Option<Vec<f32>>, LoadError> {
935    if file.find_tensor(name).is_none() {
936        return Ok(None);
937    }
938    Ok(Some(load_f32_vec(file, name)?))
939}
940
941/// Slice `n` rows starting at `start` out of a quantized matrix without
942/// dequantizing: every `Quantized` kind stores one interleaved block
943/// buffer per row (fixed `row_bytes`), so a row range is a contiguous
944/// byte range. Mapped sources stay zero-copy (sub-range of the same
945/// mmap); other backings get an owned copy. Returns `None` for non-
946/// quantized matrices (F32 / MXFP4) — callers fall back to dequant.
947fn slice_quantized_rows(m: &WeightMatrix, start: usize, n: usize) -> Option<WeightMatrix> {
948    let WeightMatrix::Quantized {
949        data,
950        rows,
951        cols,
952        kind,
953    } = m
954    else {
955        return None;
956    };
957    let total = data.len();
958    if *rows == 0 || total % *rows != 0 || start + n > *rows {
959        return None;
960    }
961    let row_bytes = total / *rows;
962    let (b0, b1) = (start * row_bytes, (start + n) * row_bytes);
963    let bytes = match data {
964        WeightBytes::Mapped { mmap, range } => WeightBytes::Mapped {
965            mmap: mmap.clone(),
966            range: range.start + b0..range.start + b1,
967        },
968        other => WeightBytes::Owned(other.as_slice()[b0..b1].to_vec()),
969    };
970    Some(WeightMatrix::Quantized {
971        data: bytes,
972        rows: n,
973        cols: *cols,
974        kind: *kind,
975    })
976}
977
978/// Loads Q/K/V projections: prefers split `attn_{q,k,v}.weight`, falls
979/// back to fused `attn_qkv.weight` (Phi-3 / some Qwen GGUFs) by
980/// slicing quantized rows (zero-copy for mmapped GGUFs; dequant only
981/// for non-quantized storage). Mirrors llama.cpp `create_tensor_qkv`.
982fn load_qkv_projections(
983    file: &impl TensorSource,
984    layer: usize,
985    config: &ModelConfig,
986) -> Result<(WeightMatrix, WeightMatrix, WeightMatrix), LoadError> {
987    let q_name = format!("blk.{layer}.attn_q.weight");
988    let k_name = format!("blk.{layer}.attn_k.weight");
989    let v_name = format!("blk.{layer}.attn_v.weight");
990    let fused_name = format!("blk.{layer}.attn_qkv.weight");
991
992    if file.find_tensor(&q_name).is_some() {
993        return Ok((
994            load_weight_matrix(file, &q_name)?,
995            load_weight_matrix(file, &k_name)?,
996            load_weight_matrix(file, &v_name)?,
997        ));
998    }
999    if file.find_tensor(&fused_name).is_none() {
1000        return Err(LoadError::Gguf(GgufError::TensorNotFound(q_name)));
1001    }
1002
1003    let fused = load_weight_matrix(file, &fused_name)?;
1004    let q_rows = config.n_heads * config.head_dim;
1005    let kv_rows = config.n_kv_heads * config.head_dim;
1006    let expected = q_rows + 2 * kv_rows;
1007    if fused.rows() != expected {
1008        // Phi-3 sometimes stores Q as full n_embd (== q_rows when MHA).
1009        return Err(LoadError::UnsupportedFeature(
1010            config.name.to_string(),
1011            format!(
1012                "{fused_name} has {} rows; expected q+k+v = {} \
1013                 (n_heads*head_dim + 2*n_kv_heads*head_dim)",
1014                fused.rows(),
1015                expected
1016            ),
1017        ));
1018    }
1019    let cols = fused.cols();
1020    // Quantized fused tensor: split by row ranges without dequantizing,
1021    // keeping Q/K/V on the quantized (Metal-capable) matvec path.
1022    if let (Some(q), Some(k), Some(v)) = (
1023        slice_quantized_rows(&fused, 0, q_rows),
1024        slice_quantized_rows(&fused, q_rows, kv_rows),
1025        slice_quantized_rows(&fused, q_rows + kv_rows, kv_rows),
1026    ) {
1027        return Ok((q, k, v));
1028    }
1029    // Non-quantized storage: dequant once and split.
1030    let mut full = Vec::with_capacity(fused.rows() * cols);
1031    for r in 0..fused.rows() {
1032        full.extend_from_slice(&fused.dequant_row(r));
1033    }
1034    let q = WeightMatrix::F32(Tensor::new(
1035        full[..q_rows * cols].to_vec(),
1036        vec![q_rows, cols],
1037    ));
1038    let k = WeightMatrix::F32(Tensor::new(
1039        full[q_rows * cols..(q_rows + kv_rows) * cols].to_vec(),
1040        vec![kv_rows, cols],
1041    ));
1042    let v = WeightMatrix::F32(Tensor::new(
1043        full[(q_rows + kv_rows) * cols..].to_vec(),
1044        vec![kv_rows, cols],
1045    ));
1046    Ok((q, k, v))
1047}
1048
1049/// Dense-layer FFN tensors: standard gate/up/down, or Phi-3 fused
1050/// `ffn_up` with `2 * expert_ffn_dim` rows and no separate gate.
1051fn load_dense_expert(
1052    file: &impl TensorSource,
1053    layer: usize,
1054    config: &ModelConfig,
1055) -> Result<ExpertWeights, LoadError> {
1056    let gate_name = format!("blk.{layer}.ffn_gate.weight");
1057    let up_name = format!("blk.{layer}.ffn_up.weight");
1058    let down_name = format!("blk.{layer}.ffn_down.weight");
1059    if file.find_tensor(&gate_name).is_some() {
1060        return Ok(ExpertWeights {
1061            gate: load_weight_matrix(file, &gate_name)?,
1062            up: load_weight_matrix(file, &up_name)?,
1063            down: load_weight_matrix(file, &down_name)?,
1064        });
1065    }
1066    // Phi-3 fused SwiGLU: up is [hidden, 2*ff], first half gate, second up.
1067    let fused = load_weight_matrix(file, &up_name)?;
1068    let ff = config.moe.expert_ffn_dim;
1069    if fused.rows() != 2 * ff {
1070        return Err(LoadError::UnsupportedFeature(
1071            config.name.to_string(),
1072            format!(
1073                "{up_name} has {} rows without a companion ffn_gate; \
1074                 expected fused SwiGLU with 2*ffn_dim = {} rows",
1075                fused.rows(),
1076                2 * ff
1077            ),
1078        ));
1079    }
1080    let cols = fused.cols();
1081    // Quantized fused gate+up: split by rows, no dequant (Metal-capable).
1082    if let (Some(gate), Some(up)) = (
1083        slice_quantized_rows(&fused, 0, ff),
1084        slice_quantized_rows(&fused, ff, ff),
1085    ) {
1086        return Ok(ExpertWeights {
1087            gate,
1088            up,
1089            down: load_weight_matrix(file, &down_name)?,
1090        });
1091    }
1092    let mut full = Vec::with_capacity(fused.rows() * cols);
1093    for r in 0..fused.rows() {
1094        full.extend_from_slice(&fused.dequant_row(r));
1095    }
1096    let gate = WeightMatrix::F32(Tensor::new(full[..ff * cols].to_vec(), vec![ff, cols]));
1097    let up = WeightMatrix::F32(Tensor::new(full[ff * cols..].to_vec(), vec![ff, cols]));
1098    Ok(ExpertWeights {
1099        gate,
1100        up,
1101        down: load_weight_matrix(file, &down_name)?,
1102    })
1103}
1104
1105/// Widen a raw plain-float tensor (`F32` / `F16` / `BF16`) to `f32`.
1106///
1107/// The three unquantized element types are handled identically at every
1108/// call site (eager widening to an owned buffer -- none of them has a
1109/// block structure a fused dot kernel could exploit), and each of the
1110/// seven GGUF loaders used to inline the same two-way match. F16 had no
1111/// arm in any of them, which made every `*-f16.gguf` a hard
1112/// `UnsupportedDtype` even though the type was parsed and sized.
1113pub(crate) fn widen_plain_float(
1114    dtype: GgmlType,
1115    raw: &[u8],
1116    name: &str,
1117) -> Result<Vec<f32>, LoadError> {
1118    match dtype {
1119        GgmlType::F32 => {
1120            let mut out = Vec::with_capacity(raw.len() / 4);
1121            for chunk in raw.as_chunks::<4>().0 {
1122                out.push(f32::from_le_bytes(*chunk));
1123            }
1124            Ok(out)
1125        }
1126        GgmlType::F16 => ferrox_quant::dequant_f16(raw)
1127            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::F16)),
1128        GgmlType::BF16 => ferrox_quant::dequant_bf16(raw)
1129            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::BF16)),
1130        // MXFP4 is accepted as a weight matrix and as an MoE expert
1131        // tensor, and `WeightMatrix::dequant` already calls this
1132        // dequantizer, so refusing it here made a 1-D MXFP4 norm or
1133        // bias a hard load error on a checkpoint whose 2-D tensors of
1134        // the same type load fine. That contradicted this function's
1135        // own contract, which is to widen whatever the loaders accept.
1136        GgmlType::MXFP4 => ferrox_quant::dequant_mxfp4_gguf(raw)
1137            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::MXFP4)),
1138        other => Err(LoadError::UnsupportedDtype(name.to_string(), other)),
1139    }
1140}
1141
1142pub(crate) fn load_f32_vec(file: &impl TensorSource, name: &str) -> Result<Vec<f32>, LoadError> {
1143    let info = find_info(file, name)?;
1144    let raw = file.tensor_bytes(name)?;
1145    match info.dtype {
1146        GgmlType::F32 | GgmlType::F16 | GgmlType::BF16 => widen_plain_float(info.dtype, raw, name),
1147        GgmlType::Q8_0 => ferrox_quant::dequant_q8_0(raw)
1148            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q8_0)),
1149        GgmlType::Q4_0 => ferrox_quant::dequant_q4_0(raw)
1150            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q4_0)),
1151        GgmlType::Q4K => ferrox_quant::dequant_q4_k(raw)
1152            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q4K)),
1153        GgmlType::Q5K => ferrox_quant::dequant_q5_k(raw)
1154            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q5K)),
1155        GgmlType::Q6K => ferrox_quant::dequant_q6_k(raw)
1156            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q6K)),
1157        GgmlType::Q2K => ferrox_quant::dequant_q2_k(raw)
1158            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q2K)),
1159        GgmlType::Q3K => ferrox_quant::dequant_q3_k(raw)
1160            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q3K)),
1161        GgmlType::Q4_1 => ferrox_quant::dequant_q4_1(raw)
1162            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q4_1)),
1163        GgmlType::Q5_0 => ferrox_quant::dequant_q5_0(raw)
1164            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q5_0)),
1165        GgmlType::Q5_1 => ferrox_quant::dequant_q5_1(raw)
1166            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q5_1)),
1167        GgmlType::Q8_1 => ferrox_quant::dequant_q8_1(raw)
1168            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::Q8_1)),
1169        GgmlType::IQ4NL => ferrox_quant::dequant_iq4_nl(raw)
1170            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ4NL)),
1171        GgmlType::IQ4XS => ferrox_quant::dequant_iq4_xs(raw)
1172            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ4XS)),
1173        // The codebook-grid tiers. Rare on the 1-D tensors this
1174        // function widens (norms and biases are almost always F32),
1175        // but a dtype ferrox can decode should never be rejected here
1176        // just because the *other* dispatch table below knows it --
1177        // that split is how a supported format turns into a load
1178        // failure on the one checkpoint that uses it.
1179        GgmlType::IQ1S => ferrox_quant::dequant_iq1_s(raw)
1180            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ1S)),
1181        GgmlType::IQ1M => ferrox_quant::dequant_iq1_m(raw)
1182            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ1M)),
1183        GgmlType::IQ2XXS => ferrox_quant::dequant_iq2_xxs(raw)
1184            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ2XXS)),
1185        GgmlType::IQ2XS => ferrox_quant::dequant_iq2_xs(raw)
1186            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ2XS)),
1187        GgmlType::IQ2S => ferrox_quant::dequant_iq2_s(raw)
1188            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ2S)),
1189        GgmlType::IQ3XXS => ferrox_quant::dequant_iq3_xxs(raw)
1190            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ3XXS)),
1191        GgmlType::IQ3S => ferrox_quant::dequant_iq3_s(raw)
1192            .map_err(|_| LoadError::UnsupportedDtype(name.to_string(), GgmlType::IQ3S)),
1193        other => Err(LoadError::UnsupportedDtype(name.to_string(), other)),
1194    }
1195}
1196
1197/// Loads a 2D weight matrix, keeping Q8_0/Q4_0 tensors quantized (raw
1198/// bytes copied out, never dequantized) and only expanding truly F32
1199/// tensors. This is the memory- and bandwidth-saving path: for a
1200/// multi-billion-parameter checkpoint the difference between this and
1201/// "dequant everything on load" is the difference between fitting in
1202/// RAM and not.
1203pub(crate) fn load_weight_matrix(
1204    file: &impl TensorSource,
1205    name: &str,
1206) -> Result<WeightMatrix, LoadError> {
1207    let info = find_info(file, name)?;
1208    // GGUF's on-disk `ne[]` shape array is fastest-varying-dimension-first
1209    // (ggml convention), i.e. `[in_features, out_features]` for a 2D
1210    // weight matrix -- the *reverse* of the row-major `[rows, cols]` =
1211    // `[out_features, in_features]` order `WeightMatrix`/`matmul_f32`
1212    // need. Reversed here once so every consumer below gets the correct
1213    // orientation. Before this reversal existed, every 2D tensor in an
1214    // externally-produced GGUF file was silently loaded transposed -- a
1215    // real bug found by running a real downloaded checkpoint
1216    // (TinyLlama-1.1B-Chat, e.g. `attn_k.weight`'s real raw shape is
1217    // `[2048, 256]` = `[hidden_dim, kv_dim]` = `[in, out]`) -- found
1218    // as a real transposition bug affecting every externally-produced
1219    // GGUF file, caught by serving a real downloaded checkpoint.
1220    let shape: Vec<usize> = info.shape.iter().rev().map(|&d| d as usize).collect();
1221    let (rows, cols) = match shape.as_slice() {
1222        [r, c] => (*r, *c),
1223        other => {
1224            return Err(LoadError::UnsupportedDtype(
1225                format!("{name} (expected 2D, got shape {other:?})"),
1226                info.dtype,
1227            ))
1228        }
1229    };
1230
1231    match info.dtype {
1232        // BF16 has no block/scale structure to keep quantized-in-place
1233        // the way Q4_0/Q8_0/K-quants do -- there's no fused dot kernel
1234        // that would make sense for a plain narrowed float, so it's
1235        // eagerly widened to an owned f32 Tensor exactly like F32
1236        // tensors already are.
1237        GgmlType::F32 | GgmlType::F16 | GgmlType::BF16 => {
1238            let data = load_f32_vec(file, name)?;
1239            Ok(WeightMatrix::F32(Tensor::new(data, shape)))
1240        }
1241        other => match quant_kind_for(other) {
1242            Some(kind) => {
1243                let (mmap, range) = file.tensor_mapped_range(name)?;
1244                #[cfg(feature = "metal")]
1245                ferrox_metal::gpu::register_weight_mmap(Arc::clone(&mmap));
1246                Ok(WeightMatrix::Quantized {
1247                    data: WeightBytes::Mapped { mmap, range },
1248                    rows,
1249                    cols,
1250                    kind,
1251                })
1252            }
1253            None => Err(LoadError::UnsupportedDtype(name.to_string(), other)),
1254        },
1255    }
1256}
1257
1258/// Splits a packed 3D MoE expert tensor `blk.N.ffn_{gate,up,down}_exps.weight`
1259/// (shape `[n_experts, out_dim, in_dim]`) into per-expert `WeightMatrix`es,
1260/// slicing raw bytes directly (quantized tensors stay quantized; block
1261/// boundaries never cross expert boundaries since `in_dim` is a whole
1262/// number of quantization blocks). Matches llama.cpp/ik_llama.cpp layout
1263/// confirmed on real OLMoE and Qwen2-MoE GGUF checkpoints.
1264pub(crate) fn split_expert_tensor(
1265    file: &impl TensorSource,
1266    name: &str,
1267    n_experts: usize,
1268) -> Result<Vec<WeightMatrix>, LoadError> {
1269    let info = find_info(file, name)?;
1270    // Real raw shape is `[in_dim, out_dim, n_experts]` (ggml's
1271    // fastest-first `ne[]` order -- see `load_weight_matrix`'s doc
1272    // comment for the confirmed 2D case this generalizes from). `n_experts`
1273    // is the slowest-varying (last, i.e. outermost/most-major) dimension,
1274    // so each expert's `out_dim*in_dim` block is contiguous with experts
1275    // back-to-back in the mmap.
1276    if info.shape.len() != 3 || info.shape[2] as usize != n_experts {
1277        let file_experts = info.shape.last().map(|&d| d as usize).unwrap_or(0);
1278        return Err(LoadError::ExpertCountMismatch(
1279            name.to_string(),
1280            file_experts,
1281            n_experts,
1282        ));
1283    }
1284    let out_dim = info.shape[1] as usize;
1285    let in_dim = info.shape[0] as usize;
1286    let raw = file.tensor_bytes(name)?;
1287
1288    match info.dtype {
1289        GgmlType::F32 | GgmlType::F16 | GgmlType::BF16 => {
1290            let all = crate::loader::widen_plain_float(info.dtype, raw, name)?;
1291            let per_expert = out_dim * in_dim;
1292            Ok((0..n_experts)
1293                .map(|e| {
1294                    WeightMatrix::F32(Tensor::new(
1295                        all[e * per_expert..(e + 1) * per_expert].to_vec(),
1296                        vec![out_dim, in_dim],
1297                    ))
1298                })
1299                .collect())
1300        }
1301        other => match quant_kind_for(other) {
1302            Some(kind) => {
1303                let (mmap, full_range) = file.tensor_mapped_range(name)?;
1304                #[cfg(feature = "metal")]
1305                ferrox_metal::gpu::register_weight_mmap(Arc::clone(&mmap));
1306                let bytes_per_expert = raw.len() / n_experts;
1307                Ok((0..n_experts)
1308                    .map(|e| WeightMatrix::Quantized {
1309                        data: WeightBytes::Mapped {
1310                            mmap: Arc::clone(&mmap),
1311                            range: (full_range.start + e * bytes_per_expert)
1312                                ..(full_range.start + (e + 1) * bytes_per_expert),
1313                        },
1314                        rows: out_dim,
1315                        cols: in_dim,
1316                        kind,
1317                    })
1318                    .collect())
1319            }
1320            None => Err(LoadError::UnsupportedDtype(name.to_string(), other)),
1321        },
1322    }
1323}
1324
1325/// When every routed expert is mmap-backed with a Metal simdgroup-GEMM
1326/// kind and back-to-back slices, record the combined gate/up/down planes
1327/// for Metal packed MoE. Gate/up/down may differ in kind (e.g. Q4_K /
1328/// Q4_K / Q8_0) but must be uniform across experts per role.
1329#[cfg(feature = "metal")]
1330fn try_build_moe_packed_q4_planes(experts: &[ExpertWeights]) -> Option<MoePackedQ4Planes> {
1331    use ferrox_core::weight_matrix::{QuantKind, WeightBytes};
1332    use std::sync::Arc;
1333
1334    if experts.is_empty() {
1335        return None;
1336    }
1337
1338    fn mapped_sg(m: &WeightMatrix) -> Option<(WeightBytes, usize, &'static str)> {
1339        match m {
1340            WeightMatrix::Quantized {
1341                data: WeightBytes::Mapped { mmap, range },
1342                rows,
1343                kind,
1344                ..
1345            } => {
1346                let kind_str = match kind {
1347                    QuantKind::Q4_0 => "Q4_0",
1348                    QuantKind::Q5_0 => "Q5_0",
1349                    QuantKind::Q4K => "Q4_K",
1350                    QuantKind::Q5K => "Q5_K",
1351                    QuantKind::Q6K => "Q6_K",
1352                    QuantKind::Q8_0 => "Q8_0",
1353                    QuantKind::IQ4XS => "IQ4_XS",
1354                    _ => return None,
1355                };
1356                let _ = ferrox_metal::gpu::mul_mm_sg_meta(kind_str)?;
1357                Some((
1358                    WeightBytes::Mapped {
1359                        mmap: Arc::clone(mmap),
1360                        range: range.clone(),
1361                    },
1362                    *rows,
1363                    kind_str,
1364                ))
1365            }
1366            _ => None,
1367        }
1368    }
1369
1370    let (gate0, ffn_rows, gate_kind) = mapped_sg(&experts[0].gate)?;
1371    let (up0, up_rows, up_kind) = mapped_sg(&experts[0].up)?;
1372    let (down0, hidden_rows, down_kind) = mapped_sg(&experts[0].down)?;
1373    if up_rows != ffn_rows {
1374        return None;
1375    }
1376    let WeightBytes::Mapped {
1377        mmap: gate_mmap,
1378        range: gate0_range,
1379    } = &gate0
1380    else {
1381        return None;
1382    };
1383    let WeightBytes::Mapped {
1384        mmap: up_mmap,
1385        range: up0_range,
1386    } = &up0
1387    else {
1388        return None;
1389    };
1390    let WeightBytes::Mapped {
1391        mmap: down_mmap,
1392        range: down0_range,
1393    } = &down0
1394    else {
1395        return None;
1396    };
1397
1398    let gate_stride = gate0_range.len();
1399    let up_stride = up0_range.len();
1400    let down_stride = down0_range.len();
1401    if gate_stride == 0 || up_stride == 0 || down_stride == 0 {
1402        return None;
1403    }
1404
1405    let n = experts.len();
1406    for (i, ex) in experts.iter().enumerate().skip(1) {
1407        let (g, fr, gk) = mapped_sg(&ex.gate)?;
1408        let (u, ur, uk) = mapped_sg(&ex.up)?;
1409        let (d, hr, dk) = mapped_sg(&ex.down)?;
1410        if gk != gate_kind || uk != up_kind || dk != down_kind {
1411            return None;
1412        }
1413        let WeightBytes::Mapped { mmap, range } = &g else {
1414            return None;
1415        };
1416        if fr != ffn_rows {
1417            return None;
1418        }
1419        if !Arc::ptr_eq(mmap, gate_mmap)
1420            || range.len() != gate_stride
1421            || range.start != gate0_range.start + i * gate_stride
1422        {
1423            return None;
1424        }
1425        let WeightBytes::Mapped { mmap, range } = &u else {
1426            return None;
1427        };
1428        if ur != ffn_rows
1429            || !Arc::ptr_eq(mmap, up_mmap)
1430            || range.len() != up_stride
1431            || range.start != up0_range.start + i * up_stride
1432        {
1433            return None;
1434        }
1435        let WeightBytes::Mapped { mmap, range } = &d else {
1436            return None;
1437        };
1438        if hr != hidden_rows
1439            || !Arc::ptr_eq(mmap, down_mmap)
1440            || range.len() != down_stride
1441            || range.start != down0_range.start + i * down_stride
1442        {
1443            return None;
1444        }
1445    }
1446
1447    Some(MoePackedQ4Planes::new(
1448        WeightBytes::Mapped {
1449            mmap: Arc::clone(gate_mmap),
1450            range: gate0_range.start..gate0_range.start + n * gate_stride,
1451        },
1452        WeightBytes::Mapped {
1453            mmap: Arc::clone(up_mmap),
1454            range: up0_range.start..up0_range.start + n * up_stride,
1455        },
1456        WeightBytes::Mapped {
1457            mmap: Arc::clone(down_mmap),
1458            range: down0_range.start..down0_range.start + n * down_stride,
1459        },
1460        gate_stride,
1461        up_stride,
1462        down_stride,
1463        n,
1464        ffn_rows,
1465        hidden_rows,
1466        gate_kind,
1467        up_kind,
1468        down_kind,
1469    ))
1470}
1471
1472/// One matrix's place inside a store-backed expert's combined byte
1473/// buffer (gate bytes, then up, then down, concatenated by
1474/// `GgufExpertSource::read_expert`).
1475#[derive(Debug, Clone, Copy)]
1476pub struct StoredMatrixSpec {
1477    pub offset: usize,
1478    pub len: usize,
1479    pub rows: usize,
1480    pub cols: usize,
1481    pub kind: QuantKind,
1482}
1483
1484/// Byte-range layout of one store-backed routed expert.
1485#[derive(Debug, Clone, Copy)]
1486pub struct StoredExpertLayout {
1487    pub gate: StoredMatrixSpec,
1488    pub up: StoredMatrixSpec,
1489    pub down: StoredMatrixSpec,
1490}
1491
1492impl StoredExpertLayout {
1493    pub fn total_bytes(&self) -> usize {
1494        self.gate.len + self.up.len + self.down.len
1495    }
1496
1497    /// Builds temporary zero-copy `WeightMatrix` views over a leased
1498    /// buffer. Each view's `WeightBytes::Shared` clone of the lease's
1499    /// `Arc` keeps the cache entry pinned for the view's lifetime.
1500    pub fn materialize(&self, lease: &ferrox_core::expert_store::ExpertLease) -> ExpertWeights {
1501        let mk = |spec: &StoredMatrixSpec| WeightMatrix::Quantized {
1502            data: WeightBytes::Shared {
1503                buf: lease.shared_buf(),
1504                range: spec.offset..spec.offset + spec.len,
1505            },
1506            rows: spec.rows,
1507            cols: spec.cols,
1508            kind: spec.kind,
1509        };
1510        ExpertWeights {
1511            gate: mk(&self.gate),
1512            up: mk(&self.up),
1513            down: mk(&self.down),
1514        }
1515    }
1516}
1517
1518/// [`ExpertSource`] over a (possibly sharded) GGUF checkpoint: each
1519/// expert's gate/up/down byte ranges are read positionally from the
1520/// owning shard file and concatenated, so a store miss touches exactly
1521/// that expert's bytes -- no mmap of the expert region, no shared seek
1522/// cursor.
1523pub struct GgufExpertSource {
1524    files: Vec<std::fs::File>,
1525    /// (layer, expert) -> the three (file index, offset, len) segments
1526    /// in gate/up/down order.
1527    segments: std::collections::HashMap<ExpertKey, [(usize, u64, usize); 3]>,
1528}
1529
1530impl ExpertSource for GgufExpertSource {
1531    fn expert_len(&self, key: ExpertKey) -> Option<usize> {
1532        self.segments
1533            .get(&key)
1534            .map(|segs| segs.iter().map(|&(_, _, len)| len).sum())
1535    }
1536
1537    fn read_expert(&self, key: ExpertKey) -> std::io::Result<Vec<u8>> {
1538        let segs = self
1539            .segments
1540            .get(&key)
1541            .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::NotFound, format!("{key:?}")))?;
1542        let total: usize = segs.iter().map(|&(_, _, len)| len).sum();
1543        let mut buf = vec![0u8; total];
1544        let mut written = 0;
1545        for &(fi, offset, len) in segs {
1546            let dst = &mut buf[written..written + len];
1547            #[cfg(unix)]
1548            {
1549                use std::os::unix::fs::FileExt;
1550                self.files[fi].read_exact_at(dst, offset)?;
1551            }
1552            #[cfg(not(unix))]
1553            {
1554                use std::io::{Read, Seek, SeekFrom};
1555                let mut f = &self.files[fi];
1556                f.seek(SeekFrom::Start(offset))?;
1557                f.read_exact(dst)?;
1558            }
1559            written += len;
1560        }
1561        Ok(buf)
1562    }
1563}
1564
1565/// Collects the per-expert `(file, offset, len)` segments and layout
1566/// for one packed 3D expert tensor -- the store-backed counterpart of
1567/// `split_expert_tensor`, sharing its shape/offset math. Only
1568/// quantized dtypes are supported (an F32/BF16 expert tensor keeps the
1569/// resident path; the store exists for the quantized multi-hundred-GB
1570/// case).
1571/// One packed 3D expert tensor's store-backed description: the owning
1572/// shard index, each expert's `(offset, len)` within that shard file,
1573/// and the matrix spec shared by every expert's slice.
1574struct StoredTensorSpecs {
1575    shard: usize,
1576    per_expert: Vec<(u64, usize)>,
1577    spec: StoredMatrixSpec,
1578}
1579
1580fn stored_expert_specs(
1581    file: &ShardedGguf,
1582    name: &str,
1583    n_experts: usize,
1584) -> Result<Option<StoredTensorSpecs>, LoadError> {
1585    let info = find_info(file, name)?;
1586    if info.shape.len() != 3 || info.shape[2] as usize != n_experts {
1587        let file_experts = info.shape.last().map(|&d| d as usize).unwrap_or(0);
1588        return Err(LoadError::ExpertCountMismatch(
1589            name.to_string(),
1590            file_experts,
1591            n_experts,
1592        ));
1593    }
1594    let out_dim = info.shape[1] as usize;
1595    let in_dim = info.shape[0] as usize;
1596    let Some(kind) = quant_kind_for(info.dtype) else {
1597        return Ok(None); // F32/BF16 (or unsupported): resident fallback
1598    };
1599    let shard = file
1600        .tensor_shard_index(name)
1601        .expect("find_info succeeded, shard index must exist");
1602    // The mmap range of a tensor within a GgufFile IS its byte offset
1603    // range within that shard file (the mmap covers the whole file).
1604    let (_, full_range) = file.tensor_mapped_range(name)?;
1605    let total_len = full_range.end - full_range.start;
1606    let bytes_per_expert = total_len / n_experts;
1607    let per_expert: Vec<(u64, usize)> = (0..n_experts)
1608        .map(|e| {
1609            (
1610                (full_range.start + e * bytes_per_expert) as u64,
1611                bytes_per_expert,
1612            )
1613        })
1614        .collect();
1615    let spec = StoredMatrixSpec {
1616        offset: 0, // caller assigns the position within the combined buffer
1617        len: bytes_per_expert,
1618        rows: out_dim,
1619        cols: in_dim,
1620        kind,
1621    };
1622    Ok(Some(StoredTensorSpecs {
1623        shard,
1624        per_expert,
1625        spec,
1626    }))
1627}
1628
1629impl Decoder {
1630    /// Loads real weights from `path` for the given `config`. `config`
1631    /// supplies the architecture shape (layer count, head counts, MoE
1632    /// topology); tensor names are resolved against it using the
1633    /// llama.cpp naming convention described in the module docs.
1634    ///
1635    /// A `config.moe.n_experts <= 1` model is treated as dense: expert
1636    /// weights are read from the plain `blk.N.ffn_{gate,up,down}.weight`
1637    /// tensor names rather than the packed 3D `_exps` variant.
1638    pub fn from_gguf(
1639        path: impl AsRef<std::path::Path>,
1640        config: ModelConfig,
1641    ) -> Result<Self, LoadError> {
1642        Self::from_gguf_with_expert_cache(path, config, None)
1643    }
1644
1645    /// Like `from_gguf`, but with `expert_cache_bytes: Some(budget)`
1646    /// routed experts are NOT loaded resident: each layer holds only
1647    /// byte-range layouts, and expert bytes are read on demand through
1648    /// one bounded, lease-protected `ExpertStore` shared by every
1649    /// layer (a single global byte budget; see
1650    /// `ferrox_core::expert_store`). Dense layers, shared experts,
1651    /// attention, embeddings, and the output head stay resident/mapped
1652    /// exactly as before -- only routed experts stream. Layers whose
1653    /// expert tensors are F32/BF16 fall back to resident loading (the
1654    /// store exists for the quantized case). Output is bit-identical
1655    /// to the resident path -- same bytes, same kernels -- pinned by
1656    /// the roundtrip suite's equivalence test.
1657    pub fn from_gguf_with_expert_cache(
1658        path: impl AsRef<std::path::Path>,
1659        mut config: ModelConfig,
1660        expert_cache_bytes: Option<u64>,
1661    ) -> Result<Self, LoadError> {
1662        let path = path.as_ref();
1663        let file = ShardedGguf::open(path)?;
1664
1665        // gpt-oss carries five per-layer tensors the generic GQA layer
1666        // structs have no home for, and reuses `post_attention_norm` for
1667        // a *different* norm slot than Gemma does. Both are decided by
1668        // the architecture string, so resolve it once here. See
1669        // `crate::decoder::GptOssWeights`.
1670        let arch = file
1671            .metadata_str("general.architecture")
1672            .unwrap_or_default()
1673            .to_string();
1674        let is_gpt_oss = arch == "gpt-oss";
1675        let mut gpt_oss_layers: Vec<crate::decoder::GptOssLayer> = Vec::new();
1676
1677        // One store for the whole model (keys are (layer, expert)),
1678        // built up-front with every stored expert's segments; created
1679        // only when the cache is enabled AND some layer can use it.
1680        let mut store_segments: std::collections::HashMap<ExpertKey, [(usize, u64, usize); 3]> =
1681            std::collections::HashMap::new();
1682        let mut stored_layouts: Vec<Option<Vec<StoredExpertLayout>>> = Vec::new();
1683
1684        // Loaded like any other weight matrix: a quantized embedding
1685        // table stays quantized (zero-copy mmap) and token lookup
1686        // dequantizes one row via `WeightMatrix::dequant_row`, instead
1687        // of the whole vocabulary tensor being widened to f32 up front.
1688        let embedding = load_weight_matrix(&file, "token_embd.weight")?;
1689
1690        let mut layers = Vec::with_capacity(config.n_layers);
1691        let mut refined_qk_norm = config.qk_norm_style;
1692        for l in 0..config.n_layers {
1693            let (q_proj, k_proj, v_proj) = load_qkv_projections(&file, l, &config)?;
1694            let q_norm = load_f32_vec_optional(&file, &format!("blk.{l}.attn_q_norm.weight"))?;
1695            let k_norm = load_f32_vec_optional(&file, &format!("blk.{l}.attn_k_norm.weight"))?;
1696            // Refine WholeVector vs PerHead from the first observed norm length.
1697            if let Some(ref w) = q_norm {
1698                if w.len() == config.head_dim {
1699                    refined_qk_norm = crate::capability::QkNormStyle::PerHead;
1700                } else if w.len() == config.n_heads * config.head_dim {
1701                    refined_qk_norm = crate::capability::QkNormStyle::WholeVector;
1702                } else {
1703                    return Err(LoadError::UnsupportedFeature(
1704                        config.name.to_string(),
1705                        format!(
1706                            "blk.{l}.attn_q_norm.weight length {} matches neither head_dim={} \
1707                             nor n_heads*head_dim={}",
1708                            w.len(),
1709                            config.head_dim,
1710                            config.n_heads * config.head_dim
1711                        ),
1712                    ));
1713                }
1714            }
1715            let attn = AttnWeights {
1716                q_proj,
1717                k_proj,
1718                v_proj,
1719                o_proj: load_weight_matrix(&file, &format!("blk.{l}.attn_output.weight"))?,
1720                norm_weight: load_f32_vec(&file, &format!("blk.{l}.attn_norm.weight"))?,
1721                q_norm,
1722                k_norm,
1723                // Qwen2/Qwen2-MoE-family real QKV bias (`attn_{q,k,v}.bias`,
1724                // real config `qkv_bias`, `o_proj` has none) -- see
1725                // `AttnWeights::q_bias`'s doc comment.
1726                q_bias: load_f32_vec_optional(&file, &format!("blk.{l}.attn_q.bias"))?,
1727                k_bias: load_f32_vec_optional(&file, &format!("blk.{l}.attn_k.bias"))?,
1728                v_bias: load_f32_vec_optional(&file, &format!("blk.{l}.attn_v.bias"))?,
1729                // gpt-oss ships `post_attention_norm` but applies it in
1730                // Gemma's *other* slot: llama.cpp's openai-moe graph
1731                // norms `ffn_inp` with it after the attention residual,
1732                // i.e. it is the pre-FFN norm, not a post-attention one.
1733                // It is read below into `MoeWeights::norm_weight`.
1734                post_attn_norm: if is_gpt_oss {
1735                    None
1736                } else {
1737                    load_f32_vec_optional(&file, &format!("blk.{l}.post_attention_norm.weight"))?
1738                },
1739                post_ffn_norm: load_f32_vec_optional(
1740                    &file,
1741                    &format!("blk.{l}.post_ffw_norm.weight"),
1742                )?,
1743            };
1744
1745            // Leading dense layers (see ModelConfig::layer_is_dense's
1746            // doc comment) load from the plain dense tensor names
1747            // regardless of this model's global MoE topology, matching
1748            // the DeepSeek-2/3-family convention found in
1749            // ik_llama.cpp's source. A model with n_experts<=1
1750            // globally (the dense test fixture) is dense on every
1751            // layer either way.
1752            let is_dense_layer = config.layer_is_dense(l) || config.moe.n_experts <= 1;
1753            let n_experts = if is_dense_layer {
1754                1
1755            } else {
1756                config.moe.n_experts
1757            };
1758            let experts: ExpertBacking = if is_dense_layer {
1759                ExpertBacking::Resident(vec![load_dense_expert(&file, l, &config)?])
1760            } else {
1761                // Try store-backed layouts first when the cache is
1762                // enabled; fall back to resident when any of the three
1763                // tensors isn't a supported quantized dtype.
1764                let stored = if expert_cache_bytes.is_some() {
1765                    let g = stored_expert_specs(
1766                        &file,
1767                        &format!("blk.{l}.ffn_gate_exps.weight"),
1768                        n_experts,
1769                    )?;
1770                    let u = stored_expert_specs(
1771                        &file,
1772                        &format!("blk.{l}.ffn_up_exps.weight"),
1773                        n_experts,
1774                    )?;
1775                    let d = stored_expert_specs(
1776                        &file,
1777                        &format!("blk.{l}.ffn_down_exps.weight"),
1778                        n_experts,
1779                    )?;
1780                    match (g, u, d) {
1781                        (Some(gt), Some(ut), Some(dt)) => {
1782                            let mut layouts = Vec::with_capacity(n_experts);
1783                            for e in 0..n_experts {
1784                                let key = ExpertKey {
1785                                    layer: l as u32,
1786                                    expert: e as u32,
1787                                };
1788                                store_segments.insert(
1789                                    key,
1790                                    [
1791                                        (gt.shard, gt.per_expert[e].0, gt.per_expert[e].1),
1792                                        (ut.shard, ut.per_expert[e].0, ut.per_expert[e].1),
1793                                        (dt.shard, dt.per_expert[e].0, dt.per_expert[e].1),
1794                                    ],
1795                                );
1796                                let mut gate = gt.spec;
1797                                let mut up = ut.spec;
1798                                let mut down = dt.spec;
1799                                gate.offset = 0;
1800                                up.offset = gate.len;
1801                                down.offset = gate.len + up.len;
1802                                layouts.push(StoredExpertLayout { gate, up, down });
1803                            }
1804                            Some(layouts)
1805                        }
1806                        _ => None,
1807                    }
1808                } else {
1809                    None
1810                };
1811                match stored {
1812                    Some(layouts) => {
1813                        // Placeholder; the shared store is attached in a
1814                        // second pass below once every layer's segments
1815                        // are collected.
1816                        stored_layouts.push(Some(layouts));
1817                        ExpertBacking::Resident(Vec::new())
1818                    }
1819                    None => {
1820                        let gates = split_expert_tensor(
1821                            &file,
1822                            &format!("blk.{l}.ffn_gate_exps.weight"),
1823                            n_experts,
1824                        )?;
1825                        let ups = split_expert_tensor(
1826                            &file,
1827                            &format!("blk.{l}.ffn_up_exps.weight"),
1828                            n_experts,
1829                        )?;
1830                        let downs = split_expert_tensor(
1831                            &file,
1832                            &format!("blk.{l}.ffn_down_exps.weight"),
1833                            n_experts,
1834                        )?;
1835                        ExpertBacking::Resident(
1836                            gates
1837                                .into_iter()
1838                                .zip(ups)
1839                                .zip(downs)
1840                                .map(|((gate, up), down)| ExpertWeights { gate, up, down })
1841                                .collect(),
1842                        )
1843                    }
1844                }
1845            };
1846            if stored_layouts.len() < layers.len() + 1 {
1847                stored_layouts.push(None);
1848            }
1849
1850            let shared_experts: Vec<ExpertWeights> =
1851                if config.moe.n_shared_experts > 0 && !is_dense_layer {
1852                    vec![ExpertWeights {
1853                        gate: load_weight_matrix(&file, &format!("blk.{l}.ffn_gate_shexp.weight"))?,
1854                        up: load_weight_matrix(&file, &format!("blk.{l}.ffn_up_shexp.weight"))?,
1855                        down: load_weight_matrix(&file, &format!("blk.{l}.ffn_down_shexp.weight"))?,
1856                    }]
1857                } else {
1858                    Vec::new()
1859                };
1860
1861            let router = if !is_dense_layer {
1862                load_weight_matrix(&file, &format!("blk.{l}.ffn_gate_inp.weight"))?
1863            } else {
1864                // dense layer: no real router; a zero [1, hidden] matrix
1865                // always selects the single expert deterministically.
1866                WeightMatrix::F32(Tensor::zeros(vec![1, config.hidden_dim]))
1867            };
1868
1869            let n_for_counts = match &experts {
1870                ExpertBacking::Resident(v) if v.is_empty() => n_experts,
1871                other => other.n_experts(),
1872            };
1873            let activation_counts = (0..n_for_counts)
1874                .map(|_| std::sync::atomic::AtomicU64::new(0))
1875                .collect();
1876            // Qwen2-MoE-specific real tensor (`blk.N.ffn_gate_inp_shexp.weight`,
1877            // real on-disk shape `[hidden_dim]`, confirmed against
1878            // llama.cpp's real `qwen2moe.cpp`) -- see
1879            // `MoeWeights::shared_expert_gate`'s doc comment. Presence
1880            // of the tensor itself is the real signal (not an
1881            // architecture-name list): every other supported
1882            // architecture's checkpoints simply don't carry this
1883            // tensor, so this naturally stays `None` there.
1884            let shared_expert_gate = if is_dense_layer {
1885                None
1886            } else {
1887                load_f32_vec_optional(&file, &format!("blk.{l}.ffn_gate_inp_shexp.weight"))?
1888            };
1889            #[cfg(feature = "metal")]
1890            let packed_q4 = match &experts {
1891                ExpertBacking::Resident(v) if !v.is_empty() => try_build_moe_packed_q4_planes(v),
1892                _ => None,
1893            };
1894            // DeepSeek-V3's aux-loss-free selection bias. The on-disk
1895            // name carries no `ffn_` prefix -- llama.cpp's
1896            // `LLM_TENSOR_FFN_EXP_PROBS_B` maps to `blk.%d.exp_probs_b`
1897            // (`llama-arch.cpp:416`, `gguf-py/gguf/constants.py:1240`).
1898            // Optional: only the DeepSeek-V3-lineage MoE recipes carry
1899            // it, and this same generic loader serves OLMoE / Qwen2-MoE /
1900            // Mixtral, which do not.
1901            let exp_probs_bias = if is_dense_layer {
1902                None
1903            } else {
1904                load_f32_vec_optional(&file, &format!("blk.{l}.exp_probs_b.bias"))?
1905            };
1906            if let Some(bias) = &exp_probs_bias {
1907                if bias.len() != config.moe.n_experts {
1908                    return Err(LoadError::UnsupportedFeature(
1909                        arch.clone(),
1910                        format!(
1911                            "blk.{l}.exp_probs_b.bias has {} entries but the model has {} experts",
1912                            bias.len(),
1913                            config.moe.n_experts
1914                        ),
1915                    ));
1916                }
1917                // Grouped selection masks the *biased* scores before the
1918                // global top-k (`build_moe_ffn`, the `n_expert_groups > 1`
1919                // block). ferrox's `route_top_k_grouped` takes a fixed
1920                // count from every group instead, which is a different
1921                // algorithm, so combining the two here would be a guess.
1922                // Refuse rather than route wrongly.
1923                if config.moe.expert_group_count.is_some() {
1924                    return Err(LoadError::UnsupportedFeature(
1925                        arch.clone(),
1926                        format!(
1927                            "blk.{l}.exp_probs_b.bias together with expert groups \
1928                             ({:?}): llama.cpp masks the biased scores per group \
1929                             before a global top-k, which is not the per-group \
1930                             top-k ferrox implements",
1931                            config.moe.expert_group_count
1932                        ),
1933                    ));
1934                }
1935            }
1936            let moe = MoeWeights {
1937                router,
1938                experts,
1939                shared_experts,
1940                shared_expert_gate,
1941                exp_probs_bias,
1942                norm_weight: if is_gpt_oss {
1943                    load_f32_vec(&file, &format!("blk.{l}.post_attention_norm.weight"))?
1944                } else {
1945                    load_f32_vec(&file, &format!("blk.{l}.ffn_norm.weight"))?
1946                },
1947                activation_counts,
1948                #[cfg(feature = "metal")]
1949                packed_q4,
1950            };
1951
1952            if is_gpt_oss {
1953                gpt_oss_layers.push(load_gpt_oss_layer(&file, l, &config)?);
1954            }
1955
1956            layers.push(LayerWeights { attn, moe });
1957        }
1958
1959        let final_norm = load_f32_vec(&file, "output_norm.weight")?;
1960        // Many small Llama/Gemma-family GGUFs tie the lm-head to
1961        // `token_embd.weight` and omit `output.weight` (llama.cpp
1962        // `llama_model_loader` falls back the same way). Prefer the
1963        // explicit head when present.
1964        let output_head = match load_weight_matrix(&file, "output.weight") {
1965            Ok(w) => w,
1966            Err(_) => load_weight_matrix(&file, "token_embd.weight")?,
1967        };
1968
1969        // Second pass: attach the one shared store to every
1970        // store-backed layer. Opening the shard files fresh (plain
1971        // `File` handles for positional reads, not mmaps) keeps the
1972        // stored experts' bytes out of the process's mapped footprint
1973        // entirely.
1974        if !store_segments.is_empty() {
1975            let budget = expert_cache_bytes
1976                .expect("store_segments only populated when a cache budget is set")
1977                as usize;
1978            let files: Result<Vec<std::fs::File>, std::io::Error> =
1979                file.shard_paths().iter().map(std::fs::File::open).collect();
1980            let files = files.map_err(GgufError::from)?;
1981            let store = std::sync::Arc::new(ExpertStore::new(
1982                GgufExpertSource {
1983                    files,
1984                    segments: store_segments,
1985                },
1986                budget,
1987            ));
1988            for (l, layer) in layers.iter_mut().enumerate() {
1989                if let Some(layouts) = stored_layouts.get_mut(l).and_then(Option::take) {
1990                    layer.moe.experts = ExpertBacking::Stored {
1991                        store: std::sync::Arc::clone(&store),
1992                        layouts,
1993                        layer: l as u32,
1994                    };
1995                }
1996            }
1997        }
1998
1999        config.qk_norm_style = refined_qk_norm;
2000
2001        let family = crate::capability::resolve_profile(
2002            file.metadata_str("general.architecture").unwrap_or("llama"),
2003        )
2004        .map(|p| p.family)
2005        .unwrap_or(crate::capability::DecoderFamily::StandardGqa);
2006        let memory_kind = crate::capability::resolve_profile(
2007            file.metadata_str("general.architecture").unwrap_or("llama"),
2008        )
2009        .map(|p| p.memory)
2010        .unwrap_or(crate::capability::MemoryKind::KvGqa);
2011        let execution_plan = crate::execution_plan::ExecutionPlan::from_config(
2012            &config,
2013            family,
2014            memory_kind,
2015            crate::execution_plan::ExecutionPlan::probe_metal_caps(),
2016        );
2017
2018        let decoder = Decoder {
2019            config,
2020            embedding,
2021            layers,
2022            final_norm,
2023            output_head,
2024            gpu_vram_budget_bytes: None,
2025            gpt_oss: if is_gpt_oss {
2026                Some(crate::decoder::GptOssWeights {
2027                    layers: gpt_oss_layers,
2028                })
2029            } else {
2030                None
2031            },
2032            #[cfg(feature = "metal")]
2033            metal_attn_kv: std::sync::Mutex::new(None),
2034            execution_plan,
2035            plan_cache: std::sync::Mutex::new(std::collections::HashMap::new()),
2036        };
2037        // Resolve every kernel the model will need while we still have a
2038        // load-time error path to report it on, then seal: from here a
2039        // lookup that misses is an unpredicted slow path and says so.
2040        decoder.probe_kernels();
2041        ferrox_core::kernel_registry::seal_or_error()
2042            .map_err(|e| LoadError::StrictKernels(e.to_string()))?;
2043        // `ModelConfig` is parsed from a *different* handle on the same
2044        // file (the CLI opens its own `GgufFile`, then hands the config
2045        // here), so the model-level tensors it consumed were recorded on
2046        // that handle, not this one. Replay them before the gate, or
2047        // every Llama-3.x checkpoint reads as carrying an unread
2048        // `rope_freqs.weight` it in fact uses on every RoPE call.
2049        for name in crate::config::MODEL_LEVEL_TENSORS_READ_BY_CONFIG {
2050            file.note_consumed(name);
2051        }
2052        assert_every_tensor_consumed(&file)?;
2053        Ok(decoder)
2054    }
2055}
2056
2057/// Tensor-name prefixes a text-generation load legitimately never
2058/// reads. Everything here is consumed by a *different* code path, not by
2059/// nothing: multimodal projector planes belong to `mmproj`, and the
2060/// per-shard split bookkeeping is metadata, not weights.
2061const IGNORED_TENSOR_PREFIXES: &[&str] = &["mm.", "v.", "mmproj.", "resampler.", "audio."];
2062
2063/// Fails the load when the checkpoint carries tensors this build never
2064/// looked at.
2065///
2066/// A tensor nobody reads is not a harmless extra: it is a term of the
2067/// real graph that ours is missing. gpt-oss ships `blk.N.attn_sinks`
2068/// and ferrox has no attention-sink code anywhere, so the file loads,
2069/// runs at full speed, and emits a different distribution than the model
2070/// it claims to be; the newer MoE recipes ship `ffn_exp_probs_b` the
2071/// same way. Both are silent today, and both are exactly what the
2072/// architecture registry cannot catch, because the architecture *string*
2073/// is one ferrox does support — it is the checkpoint that carries more
2074/// than the registry entry promises.
2075///
2076/// This is deliberately the last check in the load: by here every loader
2077/// arm has had its chance to ask for what it needs, so what is left over
2078/// is what nothing in this build knows about.
2079///
2080/// `FERROX_ALLOW_UNKNOWN_TENSORS=1` downgrades it to a warning, for the
2081/// case where a human has decided the missing term does not matter (a
2082/// bias tensor of zeros, an auxiliary head that never runs). The default
2083/// is refusal: a wrong answer is worse than no answer.
2084pub fn assert_every_tensor_consumed(file: &ShardedGguf) -> Result<(), LoadError> {
2085    let mut left: Vec<String> = file
2086        .unconsumed_tensors()
2087        .into_iter()
2088        .filter(|n| !IGNORED_TENSOR_PREFIXES.iter().any(|p| n.starts_with(p)))
2089        .collect();
2090    if left.is_empty() {
2091        return Ok(());
2092    }
2093    left.sort();
2094    let shown = left.iter().take(8).cloned().collect::<Vec<_>>().join(", ");
2095    let listing = if left.len() > 8 {
2096        format!("{shown}, … (+{} more)", left.len() - 8)
2097    } else {
2098        shown
2099    };
2100    if matches!(
2101        std::env::var("FERROX_ALLOW_UNKNOWN_TENSORS")
2102            .ok()
2103            .as_deref(),
2104        Some("1") | Some("true") | Some("on")
2105    ) {
2106        eprintln!(
2107            "ferrox: WARNING — {} tensor(s) in this checkpoint are never read \
2108             ({listing}); output may be wrong (FERROX_ALLOW_UNKNOWN_TENSORS=1)",
2109            left.len()
2110        );
2111        return Ok(());
2112    }
2113    Err(LoadError::UnconsumedTensors(left.len(), listing))
2114}
2115
2116#[cfg(test)]
2117mod tests {
2118
2119    /// A quantized 1-D tensor loads through the shared helper.
2120    ///
2121    /// This used to be six copies of `load_f32_vec`, and they had
2122    /// drifted badly: this one decoded twenty dtypes while the five
2123    /// architecture loaders decoded three (F32/F16/BF16). A quantizer
2124    /// that emits a Q8_0 norm or bias -- ordinary for aggressive
2125    /// quants -- loaded on the generic path and was rejected with
2126    /// `UnsupportedDtype` on GLM-5.2, Kimi, DeepSeek-MLA, Gemma-4 and
2127    /// the hybrid stack.
2128    ///
2129    /// This file's own comment predicted exactly that, about the same
2130    /// split one level down: "a dtype ferrox can decode should never be
2131    /// rejected here just because the *other* dispatch table below
2132    /// knows it -- that split is how a supported format turns into a
2133    /// load failure on the one checkpoint that uses it."
2134    #[test]
2135    fn a_quantized_one_dimensional_tensor_widens_through_the_shared_helper() {
2136        let values: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.25).collect();
2137        let quantized = ferrox_quant::quantize_q8_0(&values);
2138
2139        struct OneTensor {
2140            info: TensorInfo,
2141            bytes: Vec<u8>,
2142        }
2143        impl TensorSource for OneTensor {
2144            fn metadata(&self, _key: &str) -> Option<&ferrox_gguf::GgufValue> {
2145                None
2146            }
2147            fn find_tensor(&self, name: &str) -> Option<&TensorInfo> {
2148                (name == self.info.name).then_some(&self.info)
2149            }
2150            fn tensor_bytes(&self, _name: &str) -> Result<&[u8], GgufError> {
2151                Ok(&self.bytes)
2152            }
2153            fn tensor_mapped_range(
2154                &self,
2155                name: &str,
2156            ) -> Result<
2157                (
2158                    std::sync::Arc<ferrox_gguf::MmapHandle>,
2159                    std::ops::Range<usize>,
2160                ),
2161                GgufError,
2162            > {
2163                // Never reached: `load_f32_vec` widens from bytes.
2164                Err(GgufError::TensorNotFound(name.to_string()))
2165            }
2166        }
2167
2168        let source = OneTensor {
2169            info: TensorInfo {
2170                name: "blk.0.attn_norm.weight".to_string(),
2171                shape: vec![64],
2172                dtype: GgmlType::Q8_0,
2173                offset: 0,
2174            },
2175            bytes: quantized,
2176        };
2177
2178        let widened = load_f32_vec(&source, "blk.0.attn_norm.weight")
2179            .expect("a Q8_0 norm must load, not report an unsupported dtype");
2180        assert_eq!(widened.len(), values.len());
2181        for (got, want) in widened.iter().zip(values.iter()) {
2182            assert!(
2183                (got - want).abs() < 0.05,
2184                "q8_0 round trip: got {got}, want {want}"
2185            );
2186        }
2187    }
2188    use super::*;
2189    use byteorder::{LittleEndian, WriteBytesExt};
2190    use std::io::Write;
2191
2192    fn write_string(buf: &mut Vec<u8>, s: &str) {
2193        buf.write_u64::<LittleEndian>(s.len() as u64).unwrap();
2194        buf.write_all(s.as_bytes()).unwrap();
2195    }
2196
2197    fn write_kv_str(buf: &mut Vec<u8>, key: &str, val: &str) {
2198        write_string(buf, key);
2199        buf.write_u32::<LittleEndian>(8).unwrap(); // type = string
2200        write_string(buf, val);
2201    }
2202
2203    /// A minimal, tensor-free GGUF byte buffer declaring only
2204    /// `general.architecture` (no `{arch}.block_count` or any other
2205    /// hparam key) -- the shape a stripped-down or malformed file might
2206    /// take, and the exact case `ModelConfig::from_gguf` must reject
2207    /// loudly rather than silently default around.
2208    fn build_arch_only_gguf(arch: &str) -> Vec<u8> {
2209        let mut buf = Vec::new();
2210        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
2211            .unwrap();
2212        buf.write_u32::<LittleEndian>(3).unwrap(); // version
2213        buf.write_u64::<LittleEndian>(0).unwrap(); // tensor_count
2214        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
2215        write_kv_str(&mut buf, "general.architecture", arch);
2216        buf
2217    }
2218
2219    #[test]
2220    fn model_config_from_gguf_fails_loudly_when_required_hparams_are_missing() {
2221        let tmp =
2222            std::env::temp_dir().join(format!("ferrox_test_arch_only_{}.gguf", std::process::id()));
2223        // Use a registered architecture so the failure is MissingHparam,
2224        // not UnsupportedArchitecture.
2225        std::fs::write(&tmp, build_arch_only_gguf("llama")).unwrap();
2226        let file = ferrox_gguf::GgufFile::open(&tmp).expect("minimal header must still parse");
2227        std::fs::remove_file(&tmp).ok();
2228
2229        match ModelConfig::from_gguf(&file) {
2230            Err(LoadError::MissingHparam(key)) => {
2231                assert_eq!(key, "llama.block_count");
2232            }
2233            other => panic!(
2234                "expected LoadError::MissingHparam for a file with no hparam keys, got {other:?}"
2235            ),
2236        }
2237    }
2238
2239    #[test]
2240    fn model_config_from_gguf_fails_closed_on_unknown_architecture() {
2241        let tmp = std::env::temp_dir().join(format!(
2242            "ferrox_test_unknown_arch_{}.gguf",
2243            std::process::id()
2244        ));
2245        std::fs::write(&tmp, build_arch_only_gguf("bogus-arch-with-no-hparams")).unwrap();
2246        let file = ferrox_gguf::GgufFile::open(&tmp).expect("minimal header must still parse");
2247        std::fs::remove_file(&tmp).ok();
2248
2249        match ModelConfig::from_gguf(&file) {
2250            Err(LoadError::UnsupportedArchitecture(arch)) => {
2251                assert_eq!(arch, "bogus-arch-with-no-hparams");
2252            }
2253            other => panic!(
2254                "expected LoadError::UnsupportedArchitecture for an unregistered arch, got {other:?}"
2255            ),
2256        }
2257    }
2258
2259    fn write_kv_f32(buf: &mut Vec<u8>, key: &str, val: f32) {
2260        write_string(buf, key);
2261        buf.write_u32::<LittleEndian>(6).unwrap(); // type = float32
2262        buf.write_f32::<LittleEndian>(val).unwrap();
2263    }
2264
2265    /// `arch` plus one f32 hparam, so a metadata-only feature gate can be
2266    /// exercised without building a whole checkpoint.
2267    fn build_arch_plus_f32_gguf(arch: &str, key: &str, val: f32) -> Vec<u8> {
2268        let mut buf = Vec::new();
2269        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
2270            .unwrap();
2271        buf.write_u32::<LittleEndian>(3).unwrap(); // version
2272        buf.write_u64::<LittleEndian>(0).unwrap(); // tensor_count
2273        buf.write_u64::<LittleEndian>(2).unwrap(); // kv_count
2274        write_kv_str(&mut buf, "general.architecture", arch);
2275        write_kv_f32(&mut buf, key, val);
2276        buf
2277    }
2278
2279    fn config_error_for(arch: &str, key: &str, val: f32, tag: &str) -> LoadError {
2280        let tmp = std::env::temp_dir().join(format!("ferrox_test_scale_{tag}.gguf"));
2281        std::fs::write(&tmp, build_arch_plus_f32_gguf(arch, key, val)).unwrap();
2282        let file = ferrox_gguf::GgufFile::open(&tmp).expect("minimal header must still parse");
2283        std::fs::remove_file(&tmp).ok();
2284        ModelConfig::from_gguf(&file).expect_err("must not succeed")
2285    }
2286
2287    /// Granite / MiniCPM / Command-R multipliers are hparams, not
2288    /// tensors, so `assert_every_tensor_consumed` cannot see them: a
2289    /// checkpoint declaring one loads, runs at full speed, and computes
2290    /// a differently-scaled graph than it was trained as. Refuse by name
2291    /// until the math lands.
2292    #[test]
2293    fn a_declared_multiplier_this_decoder_does_not_apply_is_refused_by_name() {
2294        for (key, val) in [
2295            ("granite.logit_scale", 6.0f32),
2296            ("granite.residual_scale", 0.22),
2297            ("granite.embedding_scale", 12.0),
2298            ("granite.attention.scale", 0.015_625),
2299        ] {
2300            let tag = key.replace('.', "_");
2301            match config_error_for("granite", key, val, &tag) {
2302                LoadError::UnsupportedFeature(arch, msg) => {
2303                    assert_eq!(arch, "granite");
2304                    assert!(msg.contains(key), "error must name the key: {msg}");
2305                }
2306                other => panic!("expected UnsupportedFeature for {key}, got {other:?}"),
2307            }
2308        }
2309    }
2310
2311    /// The gate must not fire on a multiplier that is a no-op. A file
2312    /// writing `residual_scale = 1.0` describes the graph ferrox already
2313    /// computes, and refusing it would be a false alarm. llama.cpp's
2314    /// `f_attention_scale` uses `0.0` rather than `1.0` as its "unset"
2315    /// sentinel, so the two are checked against their own no-op values.
2316    #[test]
2317    fn a_multiplier_that_is_a_no_op_is_not_refused() {
2318        for (key, val) in [
2319            ("granite.logit_scale", 1.0f32),
2320            ("granite.residual_scale", 1.0),
2321            ("granite.embedding_scale", 1.0),
2322            ("granite.attention.scale", 0.0),
2323        ] {
2324            let tag = format!("noop_{}", key.replace('.', "_"));
2325            // The file carries no `block_count`, so the load still fails
2326            // -- but on the *missing hparam*, having passed this gate.
2327            match config_error_for("granite", key, val, &tag) {
2328                LoadError::MissingHparam(k) => assert_eq!(k, "granite.block_count"),
2329                other => panic!("no-op {key}={val} must pass the scaling gate, got {other:?}"),
2330            }
2331        }
2332    }
2333
2334    /// One GGUF metadata value, in the three types these header-only
2335    /// fixtures need.
2336    enum Kv<'a> {
2337        Str(&'a str),
2338        U32(u32),
2339        F32(f32),
2340    }
2341
2342    /// A tensor-free GGUF carrying exactly `kvs` -- enough for
2343    /// `ModelConfig::from_gguf` to run without a single weight on disk.
2344    fn build_metadata_gguf(kvs: &[(&str, Kv)]) -> Vec<u8> {
2345        let mut buf = Vec::new();
2346        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
2347            .unwrap();
2348        buf.write_u32::<LittleEndian>(3).unwrap(); // version
2349        buf.write_u64::<LittleEndian>(0).unwrap(); // tensor_count
2350        buf.write_u64::<LittleEndian>(kvs.len() as u64).unwrap();
2351        for (k, v) in kvs {
2352            match v {
2353                Kv::Str(s) => write_kv_str(&mut buf, k, s),
2354                Kv::U32(n) => {
2355                    write_string(&mut buf, k);
2356                    buf.write_u32::<LittleEndian>(4).unwrap(); // type = uint32
2357                    buf.write_u32::<LittleEndian>(*n).unwrap();
2358                }
2359                Kv::F32(f) => write_kv_f32(&mut buf, k, *f),
2360            }
2361        }
2362        buf
2363    }
2364
2365    fn open_metadata_gguf(tag: &str, kvs: &[(&str, Kv)]) -> ferrox_gguf::GgufFile {
2366        let tmp = std::env::temp_dir().join(format!("ferrox_test_meta_{tag}.gguf"));
2367        std::fs::write(&tmp, build_metadata_gguf(kvs)).unwrap();
2368        let file = ferrox_gguf::GgufFile::open(&tmp).expect("header-only file must parse");
2369        std::fs::remove_file(&tmp).ok();
2370        file
2371    }
2372
2373    /// A minimal `llama` hparam set (64-wide single head, base 10000)
2374    /// plus whatever RoPE-scaling keys a test wants to add.
2375    fn llama_config_with(tag: &str, extra: &[(&str, Kv)]) -> ModelConfig {
2376        let mut kvs: Vec<(&str, Kv)> = vec![
2377            ("general.architecture", Kv::Str("llama")),
2378            ("llama.block_count", Kv::U32(1)),
2379            ("llama.embedding_length", Kv::U32(64)),
2380            ("llama.attention.head_count", Kv::U32(1)),
2381            ("llama.attention.head_count_kv", Kv::U32(1)),
2382            ("llama.attention.key_length", Kv::U32(64)),
2383            ("llama.rope.freq_base", Kv::F32(10_000.0)),
2384        ];
2385        for (k, v) in extra {
2386            kvs.push((
2387                k,
2388                match v {
2389                    Kv::Str(s) => Kv::Str(s),
2390                    Kv::U32(n) => Kv::U32(*n),
2391                    Kv::F32(f) => Kv::F32(*f),
2392                },
2393            ));
2394        }
2395        ModelConfig::from_gguf(&open_metadata_gguf(tag, &kvs)).expect("fixture must load")
2396    }
2397
2398    /// A checkpoint that declares YaRN gets the per-band divisors the
2399    /// reference's `"yarn"` arm implies, folded into `rope_freqs` so the
2400    /// existing RoPE kernels apply them. Expected values are hand-derived
2401    /// from `_find_correction_dim` for this fixture (rotary width 64,
2402    /// base 10000, original context 131072): `low = 22`, `high = 35`.
2403    ///
2404    /// Before this, ferrox read neither `rope.scaling.type` nor
2405    /// `rope.scaling.factor`, so this file roped exactly like an
2406    /// unscaled one -- correct near position 0, progressively wrong
2407    /// further in.
2408    #[test]
2409    fn a_gguf_declaring_yarn_gets_its_rope_frequencies_rewritten() {
2410        let cfg = llama_config_with(
2411            "yarn",
2412            &[
2413                ("llama.rope.scaling.type", Kv::Str("yarn")),
2414                ("llama.rope.scaling.factor", Kv::F32(8.0)),
2415                (
2416                    "llama.rope.scaling.original_context_length",
2417                    Kv::U32(131_072),
2418                ),
2419            ],
2420        );
2421        let factors = cfg
2422            .rope_freqs
2423            .expect("a YaRN checkpoint must carry rewritten per-band frequencies");
2424        assert_eq!(factors.len(), 32, "one divisor per rotation band");
2425        assert!(
2426            (factors[0] - 1.0).abs() < 1e-6,
2427            "the fastest band is left extrapolated, got {}",
2428            factors[0]
2429        );
2430        let ramp = (31.0 - 22.0) / (35.0 - 22.0);
2431        let want = 1.0 / (ramp / 8.0 + (1.0 - ramp));
2432        assert!(
2433            (factors[31] - want).abs() < 1e-4,
2434            "slowest band: got {}, reference {want}",
2435            factors[31]
2436        );
2437    }
2438
2439    /// The rewrite must not fire on a file that did not ask for it. A
2440    /// scaling type ferrox does not implement (`linear`, `longrope`) is
2441    /// left exactly as it was rather than being roped as YaRN, which
2442    /// would be a new kind of wrong rather than the current known one.
2443    /// `rope.scaling.type = "linear"` must actually scale.
2444    ///
2445    /// Rotating position `p/s` is the same as rotating `p` with every
2446    /// band's frequency divided by `s`, and `rope_freqs` is exactly a
2447    /// per-band frequency divisor, so a uniform vector of `s` expresses
2448    /// linear scaling with no new code on the RoPE paths.
2449    ///
2450    /// Before this, the scaling type was compared against "yarn" and
2451    /// anything else returned None, so such a file loaded and roped at
2452    /// unscaled positions: a different model, no error.
2453    #[test]
2454    fn linear_scaling_is_applied_as_a_uniform_frequency_divisor() {
2455        let cfg = llama_config_with(
2456            "linear",
2457            &[
2458                ("llama.rope.scaling.type", Kv::Str("linear")),
2459                ("llama.rope.scaling.factor", Kv::F32(4.0)),
2460            ],
2461        );
2462        let freqs = cfg
2463            .rope_freqs
2464            .as_ref()
2465            .expect("linear scaling must produce frequency factors");
2466        assert_eq!(freqs.len(), cfg.head_dim / 2, "one factor per rotated pair");
2467        assert!(
2468            freqs.iter().all(|f| (*f - 4.0).abs() < 1e-6),
2469            "linear scaling is uniform across bands, unlike YaRN: got {freqs:?}"
2470        );
2471    }
2472
2473    /// A factor that corrects nothing is not a correction.
2474    #[test]
2475    fn a_linear_factor_of_one_is_treated_as_absent() {
2476        assert!(llama_config_with(
2477            "linear_one",
2478            &[
2479                ("llama.rope.scaling.type", Kv::Str("linear")),
2480                ("llama.rope.scaling.factor", Kv::F32(1.0)),
2481            ],
2482        )
2483        .rope_freqs
2484        .is_none());
2485    }
2486
2487    #[test]
2488    fn a_gguf_without_yarn_scaling_keeps_its_rope_frequencies_untouched() {
2489        assert!(llama_config_with("noscale", &[]).rope_freqs.is_none());
2490        // Linear scaling is NOT "no scaling". It used to land here,
2491        // asserted as `is_none()`, on the reasoning that leaving
2492        // positions alone beat roping them wrong in a new way. Both are
2493        // wrong output: llama.cpp divides the positions by the factor.
2494        // See `linear_scaling_is_applied_as_a_uniform_frequency_divisor`.
2495        // YaRN with a no-op factor is not a correction either.
2496        assert!(llama_config_with(
2497            "yarn_factor_one",
2498            &[
2499                ("llama.rope.scaling.type", Kv::Str("yarn")),
2500                ("llama.rope.scaling.factor", Kv::F32(1.0)),
2501                (
2502                    "llama.rope.scaling.original_context_length",
2503                    Kv::U32(131_072),
2504                ),
2505            ],
2506        )
2507        .rope_freqs
2508        .is_none());
2509    }
2510
2511    /// The correction range is measured against the context the
2512    /// checkpoint was *trained* at, so a file that declares YaRN without
2513    /// `rope.scaling.original_context_length` leaves the rotation alone
2514    /// rather than inventing a trained length (`context_length` on such
2515    /// a file is the *extended* one, which would put the ramp in the
2516    /// wrong place at every band).
2517    #[test]
2518    fn yarn_without_an_original_context_length_is_not_guessed_at() {
2519        let cfg = llama_config_with(
2520            "yarn_noctx",
2521            &[
2522                ("llama.rope.scaling.type", Kv::Str("yarn")),
2523                ("llama.rope.scaling.factor", Kv::F32(8.0)),
2524            ],
2525        );
2526        assert!(cfg.rope_freqs.is_none());
2527    }
2528
2529    /// `general.sampling.*` is the checkpoint's own recommendation, and
2530    /// only the keys the file carries become one: a file naming just
2531    /// `top_k` must leave temperature and top_p to the server's
2532    /// defaults.
2533    #[test]
2534    fn gguf_sampling_metadata_is_read_as_the_checkpoints_recommendation() {
2535        use crate::sampling::RecommendedSampling;
2536        let full = RecommendedSampling::from_gguf(&open_metadata_gguf(
2537            "sampling_full",
2538            &[
2539                ("general.architecture", Kv::Str("llama")),
2540                ("general.sampling.temp", Kv::F32(1.0)),
2541                ("general.sampling.top_k", Kv::U32(20)),
2542                ("general.sampling.top_p", Kv::F32(0.95)),
2543            ],
2544        ));
2545        assert_eq!(
2546            full,
2547            RecommendedSampling {
2548                temperature: Some(1.0),
2549                top_p: Some(0.95),
2550                top_k: Some(20),
2551            }
2552        );
2553
2554        let partial = RecommendedSampling::from_gguf(&open_metadata_gguf(
2555            "sampling_partial",
2556            &[
2557                ("general.architecture", Kv::Str("llama")),
2558                ("general.sampling.top_k", Kv::U32(40)),
2559            ],
2560        ));
2561        assert_eq!(partial.top_k, Some(40));
2562        assert_eq!(partial.temperature, None);
2563        assert_eq!(partial.top_p, None);
2564    }
2565
2566    /// A converter that wrote `temp = 1` stores a GGUF integer, not a
2567    /// float. Dropping it would serve a checkpoint that asked for
2568    /// temperature 1.0 at the framework's greedy default -- the
2569    /// repetition-loop failure the recommendation exists to prevent.
2570    #[test]
2571    fn an_integer_valued_sampling_temp_is_still_a_recommendation() {
2572        let recommended = crate::sampling::RecommendedSampling::from_gguf(&open_metadata_gguf(
2573            "sampling_int_temp",
2574            &[
2575                ("general.architecture", Kv::Str("llama")),
2576                ("general.sampling.temp", Kv::U32(1)),
2577            ],
2578        ));
2579        assert_eq!(recommended.temperature, Some(1.0));
2580    }
2581
2582    /// The overwhelming majority of checkpoints recommend nothing, and
2583    /// those must keep ferrox's existing defaults exactly.
2584    #[test]
2585    fn a_gguf_without_sampling_metadata_recommends_nothing() {
2586        let recommended = crate::sampling::RecommendedSampling::from_gguf(&open_metadata_gguf(
2587            "sampling_absent",
2588            &[("general.architecture", Kv::Str("llama"))],
2589        ));
2590        assert!(recommended.is_empty());
2591    }
2592
2593    #[test]
2594    fn model_config_from_gguf_rejects_dedicated_architectures() {
2595        let tmp = std::env::temp_dir().join(format!(
2596            "ferrox_test_dedicated_arch_{}.gguf",
2597            std::process::id()
2598        ));
2599        std::fs::write(&tmp, build_arch_only_gguf("deepseek4")).unwrap();
2600        let file = ferrox_gguf::GgufFile::open(&tmp).expect("minimal header must still parse");
2601        std::fs::remove_file(&tmp).ok();
2602
2603        match ModelConfig::from_gguf(&file) {
2604            Err(LoadError::DedicatedArchitectureRequired(arch, _)) => {
2605                assert_eq!(arch, "deepseek4");
2606            }
2607            other => panic!(
2608                "expected LoadError::DedicatedArchitectureRequired for deepseek4, got {other:?}"
2609            ),
2610        }
2611    }
2612
2613    /// The same Q5_K block bytes cross-validated against an independent
2614    /// Python reference in `ferrox-quant`'s own tests, reused here for
2615    /// the same full-path proof as the Q6_K test below.
2616    #[rustfmt::skip]
2617    const Q5_K_TEST_BLOCK: [u8; 176] = [
2618        0x66, 0x2a, 0x66, 0x2a, 0x01, 0x01, 0x01, 0x01, 0x4f, 0x4b, 0x10, 0x12, 0x41, 0xe2, 0xc1,
2619        0xb1, 0x72, 0x2f, 0x20, 0x07, 0x31, 0x0c, 0x38, 0xb3, 0x9c, 0xb8, 0xad, 0x2f, 0x9a, 0xea,
2620        0x17, 0xd0, 0xee, 0x93, 0x9e, 0x3e, 0x74, 0xbb, 0x28, 0x18, 0x39, 0x25, 0xb6, 0x09, 0x18,
2621        0x29, 0x1c, 0x1d, 0x29, 0x41, 0x40, 0x0a, 0x74, 0x7d, 0xfd, 0x21, 0xdd, 0x6d, 0x45, 0x73,
2622        0x0e, 0x1e, 0xc0, 0x4a, 0xfc, 0xf3, 0x8e, 0x24, 0x6b, 0x34, 0x7d, 0xbe, 0x94, 0xde, 0x59,
2623        0x7a, 0x35, 0x30, 0x36, 0x0a, 0xf9, 0x4a, 0x9b, 0xa2, 0x26, 0x21, 0xa2, 0xfa, 0xdf, 0x4b,
2624        0x29, 0x64, 0x6f, 0xbb, 0xca, 0x0f, 0x3c, 0xda, 0x20, 0xf4, 0x93, 0x86, 0xab, 0x6e, 0xb9,
2625        0xe5, 0xd5, 0xa0, 0x82, 0xd6, 0x41, 0xff, 0x12, 0xbc, 0x34, 0xbb, 0xab, 0xb8, 0x20, 0x2f,
2626        0xbb, 0x5f, 0x0c, 0x10, 0xcf, 0x49, 0xc5, 0x86, 0x5c, 0xdf, 0xff, 0x78, 0x44, 0x26, 0x3b,
2627        0xc2, 0x23, 0x3d, 0x2b, 0xe9, 0x00, 0x12, 0xf8, 0xea, 0xe2, 0x9e, 0x5e, 0x50, 0x20, 0x9f,
2628        0x9d, 0x8d, 0x7d, 0x7f, 0xcc, 0x1d, 0x0e, 0x13, 0xf8, 0xc2, 0xf1, 0x3d, 0x08, 0x2f, 0x23,
2629        0x13, 0xac, 0x0d, 0xa7, 0xe7, 0x20, 0xa3, 0x90, 0xb7, 0xc8, 0x28,
2630    ];
2631
2632    fn build_single_q5_k_tensor_gguf() -> Vec<u8> {
2633        let mut buf = Vec::new();
2634        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
2635            .unwrap();
2636        buf.write_u32::<LittleEndian>(3).unwrap(); // version
2637        buf.write_u64::<LittleEndian>(1).unwrap(); // tensor_count
2638        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
2639
2640        write_kv_str(&mut buf, "general.architecture", "ferrox-q5k-test");
2641
2642        write_string(&mut buf, "test.weight");
2643        buf.write_u32::<LittleEndian>(2).unwrap(); // n_dims
2644                                                   // Real GGUF ne[] order is fastest-varying-first, i.e. [cols,
2645                                                   // rows] -- reversed from the semantic [rows, cols] this tensor
2646                                                   // represents (1 row, 256 cols / 1 Q5_K block).
2647        buf.write_u64::<LittleEndian>(256).unwrap(); // cols (1 Q5_K block)
2648        buf.write_u64::<LittleEndian>(1).unwrap(); // rows
2649        buf.write_u32::<LittleEndian>(13).unwrap(); // dtype tag: Q5_K
2650        buf.write_u64::<LittleEndian>(0).unwrap(); // offset
2651
2652        while buf.len() % 32 != 0 {
2653            buf.push(0);
2654        }
2655        buf.extend_from_slice(&Q5_K_TEST_BLOCK);
2656        buf
2657    }
2658
2659    /// How far a fused dot may sit from an exact dequantized dot.
2660    ///
2661    /// Two regimes, and one fixed number cannot describe both. With
2662    /// `FERROX_CPU_INT_DOT` off the activation stays f32 and only
2663    /// rounding separates the two. With it on, the activation is
2664    /// quantized to int8 at `d = amax / 127`, which is the flag both
2665    /// binaries turn on by default and the reason the Q5_K and Q6_K
2666    /// cases failed against a flat `1e-2`.
2667    ///
2668    /// The bound grows with the L2 norm of the row, NOT the L1. Each
2669    /// element carries an independent rounding of up to `d/2`, so the
2670    /// dot's error is a sum of independent terms whose standard
2671    /// deviation is `d/sqrt(12) * ||w||_2`. Bounding by the worst case
2672    /// `d/2 * ||w||_1` instead assumes every rounding aligns with its
2673    /// weight's sign, which on this fixture gives 0.347 against a dot
2674    /// of 2.77: 12% of the value, loose enough that injecting a 5%
2675    /// error still passed. Measured here, the real error is 1.8 sigma,
2676    /// so four sigma keeps better than 2x headroom while still failing
2677    /// that 5% injection.
2678    fn fused_dot_tolerance(weights: &[f32], x: &[f32], exact_bound: f32) -> f32 {
2679        if !ferrox_core::weight_matrix::cpu_int_dot_enabled() {
2680            return exact_bound;
2681        }
2682        let amax = x.iter().fold(0.0f32, |a, v| a.max(v.abs()));
2683        let l2 = weights.iter().map(|w| w * w).sum::<f32>().sqrt();
2684        4.0 * (amax / 127.0) / 12f32.sqrt() * l2 + exact_bound
2685    }
2686
2687    #[test]
2688    fn load_weight_matrix_handles_a_real_on_disk_q5_k_tensor_end_to_end() {
2689        let tmp = std::env::temp_dir().join(format!(
2690            "ferrox_test_q5k_tensor_{}.gguf",
2691            std::process::id()
2692        ));
2693        std::fs::write(&tmp, build_single_q5_k_tensor_gguf()).unwrap();
2694        let file = ferrox_gguf::GgufFile::open(&tmp).expect("real Q5_K GGUF file must parse");
2695        std::fs::remove_file(&tmp).ok();
2696
2697        let matrix = load_weight_matrix(&file, "test.weight").expect("Q5_K tensor must load");
2698        assert_eq!(matrix.rows(), 1);
2699        assert_eq!(matrix.cols(), 256);
2700        match &matrix {
2701            WeightMatrix::Quantized { kind, data, .. } => {
2702                assert_eq!(*kind, QuantKind::Q5K);
2703                assert!(
2704                    data.is_mapped(),
2705                    "Q5_K tensors should take the zero-copy mmap path, same as Q8_0/Q4_0"
2706                );
2707            }
2708            _ => panic!("expected a Quantized matrix for a Q5_K tensor"),
2709        }
2710
2711        let expected = ferrox_quant::dequant_q5_k(&Q5_K_TEST_BLOCK).unwrap();
2712        let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.013).sin()).collect();
2713        let expected_dot: f32 = expected.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
2714
2715        let got = matrix.apply(&x);
2716        assert_eq!(got.len(), 1);
2717        assert!(
2718            (got[0] - expected_dot).abs() < fused_dot_tolerance(&expected, &x, 1e-2),
2719            "end-to-end loaded+applied Q5_K matrix diverged from direct dequant: got={} expected={}",
2720            got[0],
2721            expected_dot
2722        );
2723    }
2724
2725    /// The same Q6_K block bytes cross-validated against an independent
2726    /// Python reference in `ferrox-quant`'s own tests; reused here to
2727    /// prove the *full*
2728    /// path -- real on-disk GGUF bytes, parsed by `ferrox-gguf`, read
2729    /// through `GgufFile::tensor_mapped_range`, dispatched by
2730    /// `WeightMatrix::apply` to `ferrox_quant::dot_q6_k_f32` -- produces
2731    /// the same result as directly dequantizing those bytes, not just
2732    /// that the isolated kernel is correct in unit-test isolation.
2733    #[rustfmt::skip]
2734    const Q6_K_TEST_BLOCK: [u8; 210] = [
2735        0xe0, 0xa5, 0x40, 0x5c, 0x8d, 0x3a, 0x0a, 0x26, 0xfb, 0x4b, 0x6e, 0x9a, 0xdf, 0x3e, 0xa3,
2736        0xc4, 0xf8, 0x2b, 0x1d, 0x95, 0x76, 0x7d, 0x3b, 0xcd, 0xfd, 0xef, 0xc2, 0x0b, 0x07, 0x63,
2737        0x29, 0xfb, 0x81, 0x57, 0xbe, 0xbe, 0x06, 0xf7, 0x3a, 0x92, 0xc4, 0x43, 0xff, 0xad, 0xac,
2738        0x7e, 0x0f, 0x00, 0x2a, 0x4f, 0xf0, 0xf8, 0xa9, 0xfa, 0x3c, 0x90, 0x6d, 0x73, 0x2d, 0x5a,
2739        0xe6, 0xc6, 0x46, 0xf2, 0x0d, 0x55, 0x4c, 0x25, 0x38, 0x71, 0x2b, 0x35, 0x38, 0x82, 0x16,
2740        0x37, 0x5f, 0x32, 0x61, 0x02, 0xdd, 0x2f, 0x6f, 0x7b, 0x1f, 0xb4, 0x1a, 0x1b, 0x3e, 0x4f,
2741        0x11, 0xa3, 0x17, 0x40, 0x5a, 0x5f, 0x76, 0xcd, 0x19, 0x27, 0x9b, 0xc7, 0xc8, 0xf7, 0xf7,
2742        0xee, 0xf4, 0x86, 0xd9, 0xfd, 0xa7, 0xfe, 0x9e, 0xac, 0x70, 0x53, 0x5b, 0x76, 0xfb, 0x39,
2743        0xf8, 0x4b, 0x98, 0xfe, 0xd0, 0x06, 0x21, 0x4c, 0x4d, 0xbe, 0x10, 0x2b, 0x06, 0x65, 0xc9,
2744        0x5e, 0xf9, 0x95, 0x72, 0xae, 0x99, 0xd9, 0x7e, 0x15, 0xbd, 0x5e, 0x6d, 0xe8, 0x25, 0x8a,
2745        0xd5, 0x99, 0xc6, 0x6b, 0x69, 0xc7, 0x84, 0xc6, 0xa4, 0xf7, 0xb9, 0x6d, 0x68, 0x45, 0x0e,
2746        0x65, 0x69, 0xeb, 0xe6, 0xeb, 0xe9, 0x28, 0xa6, 0xb9, 0x96, 0xf2, 0xe8, 0xa7, 0x9b, 0x6e,
2747        0x79, 0x8a, 0x68, 0x65, 0x59, 0x98, 0x8b, 0x44, 0x41, 0x98, 0x9a, 0x56, 0x01, 0x01, 0x01,
2748        0x02, 0x01, 0x01, 0x01, 0x01, 0x02, 0x01, 0x02, 0x02, 0x01, 0x01, 0x01, 0x02, 0x1f, 0x25,
2749    ];
2750
2751    fn build_single_q6_k_tensor_gguf() -> Vec<u8> {
2752        let mut buf = Vec::new();
2753        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
2754            .unwrap();
2755        buf.write_u32::<LittleEndian>(3).unwrap(); // version
2756        buf.write_u64::<LittleEndian>(1).unwrap(); // tensor_count
2757        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
2758
2759        write_kv_str(&mut buf, "general.architecture", "ferrox-q6k-test");
2760
2761        write_string(&mut buf, "test.weight");
2762        buf.write_u32::<LittleEndian>(2).unwrap(); // n_dims
2763                                                   // Real GGUF ne[] order is fastest-varying-first, i.e. [cols, rows].
2764        buf.write_u64::<LittleEndian>(256).unwrap(); // cols (1 Q6_K block)
2765        buf.write_u64::<LittleEndian>(1).unwrap(); // rows
2766        buf.write_u32::<LittleEndian>(14).unwrap(); // dtype tag: Q6_K
2767        buf.write_u64::<LittleEndian>(0).unwrap(); // offset
2768
2769        while buf.len() % 32 != 0 {
2770            buf.push(0);
2771        }
2772        buf.extend_from_slice(&Q6_K_TEST_BLOCK);
2773        buf
2774    }
2775
2776    #[test]
2777    fn load_weight_matrix_handles_a_real_on_disk_q6_k_tensor_end_to_end() {
2778        let tmp = std::env::temp_dir().join(format!(
2779            "ferrox_test_q6k_tensor_{}.gguf",
2780            std::process::id()
2781        ));
2782        std::fs::write(&tmp, build_single_q6_k_tensor_gguf()).unwrap();
2783        let file = ferrox_gguf::GgufFile::open(&tmp).expect("real Q6_K GGUF file must parse");
2784        std::fs::remove_file(&tmp).ok();
2785
2786        let matrix = load_weight_matrix(&file, "test.weight").expect("Q6_K tensor must load");
2787        assert_eq!(matrix.rows(), 1);
2788        assert_eq!(matrix.cols(), 256);
2789        match &matrix {
2790            WeightMatrix::Quantized { kind, data, .. } => {
2791                assert_eq!(*kind, QuantKind::Q6K);
2792                assert!(
2793                    data.is_mapped(),
2794                    "Q6_K tensors should take the zero-copy mmap path, same as Q8_0/Q4_0"
2795                );
2796            }
2797            _ => panic!("expected a Quantized matrix for a Q6_K tensor"),
2798        }
2799
2800        let expected = ferrox_quant::dequant_q6_k(&Q6_K_TEST_BLOCK).unwrap();
2801        let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.013).sin()).collect();
2802        let expected_dot: f32 = expected.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
2803
2804        let got = matrix.apply(&x);
2805        assert_eq!(got.len(), 1);
2806        assert!(
2807            (got[0] - expected_dot).abs() < fused_dot_tolerance(&expected, &x, 1e-2),
2808            "end-to-end loaded+applied Q6_K matrix diverged from direct dequant: got={} expected={}",
2809            got[0],
2810            expected_dot
2811        );
2812    }
2813
2814    fn build_single_bf16_tensor_gguf(rows: u64, cols: u64, values: &[f32]) -> Vec<u8> {
2815        let mut buf = Vec::new();
2816        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
2817            .unwrap();
2818        buf.write_u32::<LittleEndian>(3).unwrap(); // version
2819        buf.write_u64::<LittleEndian>(1).unwrap(); // tensor_count
2820        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
2821
2822        write_kv_str(&mut buf, "general.architecture", "ferrox-bf16-test");
2823
2824        write_string(&mut buf, "test.weight");
2825        buf.write_u32::<LittleEndian>(2).unwrap(); // n_dims
2826                                                   // Real GGUF ne[] order is fastest-varying-first, i.e. [cols, rows].
2827        buf.write_u64::<LittleEndian>(cols).unwrap();
2828        buf.write_u64::<LittleEndian>(rows).unwrap();
2829        buf.write_u32::<LittleEndian>(30).unwrap(); // dtype tag: BF16
2830        buf.write_u64::<LittleEndian>(0).unwrap(); // offset
2831
2832        while buf.len() % 32 != 0 {
2833            buf.push(0);
2834        }
2835        for &v in values {
2836            // Real bf16 truncation (round-toward-zero, matching a real
2837            // writer closely enough for round-trip test purposes): top
2838            // 16 bits of the f32 bit pattern.
2839            let bf16_bits = (v.to_bits() >> 16) as u16;
2840            buf.extend_from_slice(&bf16_bits.to_le_bytes());
2841        }
2842        buf
2843    }
2844
2845    #[test]
2846    fn load_weight_matrix_handles_a_real_on_disk_bf16_tensor_end_to_end() {
2847        // Values with zero low-mantissa bits, so f32->bf16 truncation
2848        // is lossless and this is an exact-equality check.
2849        let values: Vec<f32> = vec![1.0, -2.5, 0.0, 4.0, -8.0, 16.0];
2850        let tmp = std::env::temp_dir().join(format!(
2851            "ferrox_test_bf16_tensor_{}.gguf",
2852            std::process::id()
2853        ));
2854        std::fs::write(&tmp, build_single_bf16_tensor_gguf(2, 3, &values)).unwrap();
2855        let file = ferrox_gguf::GgufFile::open(&tmp).expect("real BF16 GGUF file must parse");
2856        std::fs::remove_file(&tmp).ok();
2857
2858        let matrix = load_weight_matrix(&file, "test.weight").expect("BF16 tensor must load");
2859        assert_eq!(matrix.rows(), 2);
2860        assert_eq!(matrix.cols(), 3);
2861        match &matrix {
2862            WeightMatrix::F32(tensor) => {
2863                assert_eq!(tensor.data, values, "BF16 must widen to f32 exactly");
2864            }
2865            _ => panic!("expected an F32 matrix for a BF16 tensor (no fused dot kernel for it)"),
2866        }
2867    }
2868
2869    fn build_single_f16_tensor_gguf(rows: u64, cols: u64, values: &[f32]) -> Vec<u8> {
2870        let mut buf = Vec::new();
2871        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
2872            .unwrap();
2873        buf.write_u32::<LittleEndian>(3).unwrap(); // version
2874        buf.write_u64::<LittleEndian>(1).unwrap(); // tensor_count
2875        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
2876
2877        write_kv_str(&mut buf, "general.architecture", "ferrox-f16-test");
2878
2879        write_string(&mut buf, "test.weight");
2880        buf.write_u32::<LittleEndian>(2).unwrap(); // n_dims
2881        buf.write_u64::<LittleEndian>(cols).unwrap();
2882        buf.write_u64::<LittleEndian>(rows).unwrap();
2883        buf.write_u32::<LittleEndian>(1).unwrap(); // dtype tag: F16
2884        buf.write_u64::<LittleEndian>(0).unwrap(); // offset
2885
2886        while buf.len() % 32 != 0 {
2887            buf.push(0);
2888        }
2889        for &v in values {
2890            buf.extend_from_slice(&half::f16::from_f32(v).to_le_bytes());
2891        }
2892        buf
2893    }
2894
2895    /// `GgmlType::F16` was parsed and sized but had no dequant arm in any
2896    /// of the seven loaders, so every `*-f16.gguf` was a hard
2897    /// `UnsupportedDtype`. Values are exactly representable in f16, so
2898    /// this is an exact-equality check.
2899    #[test]
2900    fn load_weight_matrix_handles_a_real_on_disk_f16_tensor_end_to_end() {
2901        let values: Vec<f32> = vec![1.0, -2.5, 0.0, 4.0, -8.0, 16.0];
2902        let tmp = std::env::temp_dir().join(format!(
2903            "ferrox_test_f16_tensor_{}.gguf",
2904            std::process::id()
2905        ));
2906        std::fs::write(&tmp, build_single_f16_tensor_gguf(2, 3, &values)).unwrap();
2907        let file = ferrox_gguf::GgufFile::open(&tmp).expect("real F16 GGUF file must parse");
2908        std::fs::remove_file(&tmp).ok();
2909
2910        let matrix = load_weight_matrix(&file, "test.weight").expect("F16 tensor must load");
2911        assert_eq!(matrix.rows(), 2);
2912        assert_eq!(matrix.cols(), 3);
2913        match &matrix {
2914            WeightMatrix::F32(tensor) => {
2915                assert_eq!(tensor.data, values, "F16 must widen to f32 exactly");
2916            }
2917            _ => panic!("expected an F32 matrix for an F16 tensor (no fused dot kernel for it)"),
2918        }
2919
2920        // The same tensor read as a plain vector (norm weights, biases and
2921        // the router all take this path, not `load_weight_matrix`).
2922        let tmp =
2923            std::env::temp_dir().join(format!("ferrox_test_f16_vec_{}.gguf", std::process::id()));
2924        std::fs::write(&tmp, build_single_f16_tensor_gguf(2, 3, &values)).unwrap();
2925        let file = ferrox_gguf::GgufFile::open(&tmp).expect("real F16 GGUF file must parse");
2926        std::fs::remove_file(&tmp).ok();
2927        assert_eq!(load_f32_vec(&file, "test.weight").unwrap(), values);
2928    }
2929
2930    fn build_single_q5_1_tensor_gguf() -> Vec<u8> {
2931        let mut buf = Vec::new();
2932        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
2933            .unwrap();
2934        buf.write_u32::<LittleEndian>(3).unwrap(); // version
2935        buf.write_u64::<LittleEndian>(1).unwrap(); // tensor_count
2936        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
2937
2938        write_kv_str(&mut buf, "general.architecture", "ferrox-q5-1-test");
2939
2940        write_string(&mut buf, "test.weight");
2941        buf.write_u32::<LittleEndian>(2).unwrap(); // n_dims
2942                                                   // Real GGUF ne[] order is fastest-varying-first, i.e. [cols, rows].
2943        buf.write_u64::<LittleEndian>(32).unwrap(); // cols (1 Q5_1 block)
2944        buf.write_u64::<LittleEndian>(1).unwrap(); // rows
2945        buf.write_u32::<LittleEndian>(7).unwrap(); // dtype tag: Q5_1
2946        buf.write_u64::<LittleEndian>(0).unwrap(); // offset
2947
2948        while buf.len() % 32 != 0 {
2949            buf.push(0);
2950        }
2951        // d=0.25 (f16 0x3400), m=1.5 (f16 0x3E00) -- both exact in f16,
2952        // hand-verified bit patterns to avoid pulling in the `half`
2953        // crate just for two test constants. qh varied, qs a real
2954        // (non-degenerate) pattern.
2955        buf.extend_from_slice(&0x3400u16.to_le_bytes());
2956        buf.extend_from_slice(&0x3E00u16.to_le_bytes());
2957        buf.extend_from_slice(&[0x9au8, 0x3c, 0xf0, 0x0f]);
2958        buf.extend_from_slice(&(0..16u8).map(|i| i | ((15 - i) << 4)).collect::<Vec<u8>>());
2959        buf
2960    }
2961
2962    #[test]
2963    fn load_weight_matrix_handles_a_real_on_disk_q5_1_tensor_end_to_end() {
2964        let tmp = std::env::temp_dir().join(format!(
2965            "ferrox_test_q5_1_tensor_{}.gguf",
2966            std::process::id()
2967        ));
2968        std::fs::write(&tmp, build_single_q5_1_tensor_gguf()).unwrap();
2969        let file = ferrox_gguf::GgufFile::open(&tmp).expect("real Q5_1 GGUF file must parse");
2970        std::fs::remove_file(&tmp).ok();
2971
2972        let matrix = load_weight_matrix(&file, "test.weight").expect("Q5_1 tensor must load");
2973        assert_eq!(matrix.rows(), 1);
2974        assert_eq!(matrix.cols(), 32);
2975        let raw = file.tensor_bytes("test.weight").unwrap();
2976        let expected = ferrox_quant::dequant_q5_1(raw).unwrap();
2977        match &matrix {
2978            WeightMatrix::Quantized { kind, data, .. } => {
2979                assert_eq!(*kind, QuantKind::Q5_1);
2980                assert!(data.is_mapped());
2981            }
2982            _ => panic!("expected a Quantized matrix for a Q5_1 tensor"),
2983        }
2984
2985        let x: Vec<f32> = (0..32).map(|i| ((i as f32) * 0.017).cos()).collect();
2986        let expected_dot: f32 = expected.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
2987        let got = matrix.apply(&x);
2988        assert_eq!(got.len(), 1);
2989        assert!(
2990            (got[0] - expected_dot).abs() < 1e-2,
2991            "end-to-end loaded+applied Q5_1 matrix diverged from direct dequant: got={} expected={}",
2992            got[0],
2993            expected_dot
2994        );
2995    }
2996
2997    // Same bytes as ferrox-quant's own Q3_K_TEST_BLOCK (Python-cross-
2998    // validated there); duplicated here to build a real on-disk GGUF
2999    // file, matching this file's existing per-format test convention
3000    // (see Q6_K_TEST_BLOCK above).
3001    const Q3_K_TEST_BLOCK: [u8; 110] = [
3002        0x56, 0xf2, 0xb4, 0x2b, 0xd5, 0x6f, 0x51, 0x71, 0x3c, 0x0a, 0xb9, 0x1d, 0xd0, 0xb9, 0x3b,
3003        0xb3, 0x0f, 0xff, 0x8c, 0xb2, 0x83, 0x3a, 0x3d, 0x24, 0xb1, 0x12, 0x56, 0xe3, 0x23, 0x54,
3004        0xf2, 0xfa, 0x7f, 0xdf, 0x31, 0xe1, 0x18, 0x26, 0x6e, 0xcd, 0x5b, 0x38, 0xee, 0xbd, 0x9f,
3005        0x8c, 0x57, 0x47, 0x0b, 0x11, 0xcb, 0xfb, 0xb4, 0x83, 0xa0, 0x4e, 0x0b, 0xd4, 0xa7, 0x85,
3006        0xe0, 0x60, 0xf3, 0xb3, 0xe3, 0x95, 0x43, 0xc6, 0x05, 0x05, 0x77, 0x53, 0xed, 0x23, 0xcc,
3007        0x6a, 0x0e, 0x89, 0xa1, 0x79, 0x85, 0xf6, 0x6e, 0x5a, 0x23, 0x63, 0xbe, 0x53, 0xfa, 0xa2,
3008        0x2b, 0xe9, 0xcd, 0xce, 0xf8, 0x3d, 0x6f, 0xd0, 0x42, 0x6e, 0x3b, 0x7f, 0x23, 0x26, 0xd3,
3009        0xb9, 0x18, 0xbf, 0xa4, 0x34,
3010    ];
3011
3012    fn build_single_q3_k_tensor_gguf() -> Vec<u8> {
3013        let mut buf = Vec::new();
3014        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
3015            .unwrap();
3016        buf.write_u32::<LittleEndian>(3).unwrap(); // version
3017        buf.write_u64::<LittleEndian>(1).unwrap(); // tensor_count
3018        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
3019
3020        write_kv_str(&mut buf, "general.architecture", "ferrox-q3k-test");
3021
3022        write_string(&mut buf, "test.weight");
3023        buf.write_u32::<LittleEndian>(2).unwrap(); // n_dims
3024                                                   // Real GGUF ne[] order is fastest-varying-first, i.e. [cols, rows].
3025        buf.write_u64::<LittleEndian>(256).unwrap(); // cols (1 Q3_K block)
3026        buf.write_u64::<LittleEndian>(1).unwrap(); // rows
3027        buf.write_u32::<LittleEndian>(11).unwrap(); // dtype tag: Q3_K
3028        buf.write_u64::<LittleEndian>(0).unwrap(); // offset
3029
3030        while buf.len() % 32 != 0 {
3031            buf.push(0);
3032        }
3033        buf.extend_from_slice(&Q3_K_TEST_BLOCK);
3034        buf
3035    }
3036
3037    #[test]
3038    fn load_weight_matrix_handles_a_real_on_disk_q3_k_tensor_end_to_end() {
3039        let tmp = std::env::temp_dir().join(format!(
3040            "ferrox_test_q3k_tensor_{}.gguf",
3041            std::process::id()
3042        ));
3043        std::fs::write(&tmp, build_single_q3_k_tensor_gguf()).unwrap();
3044        let file = ferrox_gguf::GgufFile::open(&tmp).expect("real Q3_K GGUF file must parse");
3045        std::fs::remove_file(&tmp).ok();
3046
3047        let matrix = load_weight_matrix(&file, "test.weight").expect("Q3_K tensor must load");
3048        assert_eq!(matrix.rows(), 1);
3049        assert_eq!(matrix.cols(), 256);
3050        match &matrix {
3051            WeightMatrix::Quantized { kind, data, .. } => {
3052                assert_eq!(*kind, QuantKind::Q3K);
3053                assert!(data.is_mapped());
3054            }
3055            _ => panic!("expected a Quantized matrix for a Q3_K tensor"),
3056        }
3057
3058        let expected = ferrox_quant::dequant_q3_k(&Q3_K_TEST_BLOCK).unwrap();
3059        let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.013).sin()).collect();
3060        let expected_dot: f32 = expected.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
3061
3062        let got = matrix.apply(&x);
3063        assert_eq!(got.len(), 1);
3064        assert!(
3065            (got[0] - expected_dot).abs() < fused_dot_tolerance(&expected, &x, 1e-1),
3066            "end-to-end loaded+applied Q3_K matrix diverged from direct dequant: got={} expected={}",
3067            got[0],
3068            expected_dot
3069        );
3070    }
3071
3072    // Same bytes as ferrox-quant's own IQ4_XS_TEST_BLOCK (Python-cross-
3073    // validated there); duplicated here to build a real on-disk GGUF
3074    // file, matching this file's existing per-format test convention.
3075    const IQ4_XS_TEST_BLOCK: [u8; 136] = [
3076        0x5c, 0x33, 0xb4, 0x39, 0xd1, 0x64, 0x97, 0x82, 0xcb, 0xbd, 0x88, 0x95, 0xf3, 0x60, 0x2a,
3077        0xb5, 0xe7, 0x24, 0xd3, 0xee, 0xfe, 0x71, 0x13, 0xbe, 0x70, 0x84, 0x48, 0x79, 0x7b, 0x3e,
3078        0xf0, 0x55, 0xdc, 0xb2, 0xb2, 0xde, 0x32, 0xa1, 0x5b, 0x02, 0x01, 0xdc, 0x2a, 0xbb, 0xf7,
3079        0x0b, 0x8a, 0x88, 0xdd, 0x0b, 0x02, 0x7e, 0x5e, 0x76, 0x87, 0x30, 0x1e, 0x1c, 0xcf, 0x48,
3080        0xd7, 0x61, 0xf3, 0x51, 0x52, 0x17, 0x98, 0x0a, 0x87, 0xcf, 0x02, 0x91, 0xc8, 0xee, 0xc0,
3081        0x91, 0x69, 0x2a, 0x4f, 0x64, 0x68, 0xa7, 0xb2, 0xe6, 0x98, 0x21, 0x81, 0x75, 0x53, 0x2a,
3082        0x8d, 0x12, 0xae, 0xe0, 0xea, 0x0c, 0x75, 0xff, 0x22, 0x5e, 0x25, 0x19, 0xda, 0x2e, 0x51,
3083        0x4e, 0x81, 0xdc, 0x0e, 0x78, 0x86, 0xd7, 0x58, 0xb5, 0xb7, 0xf6, 0x45, 0xa9, 0x0a, 0x83,
3084        0xfd, 0x2a, 0x12, 0x7d, 0xf0, 0x12, 0x97, 0xe2, 0xfe, 0xf4, 0xd0, 0xa2, 0x11, 0x14, 0x78,
3085        0xdb,
3086    ];
3087
3088    fn build_single_iq4_xs_tensor_gguf() -> Vec<u8> {
3089        let mut buf = Vec::new();
3090        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
3091            .unwrap();
3092        buf.write_u32::<LittleEndian>(3).unwrap(); // version
3093        buf.write_u64::<LittleEndian>(1).unwrap(); // tensor_count
3094        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
3095
3096        write_kv_str(&mut buf, "general.architecture", "ferrox-iq4xs-test");
3097
3098        write_string(&mut buf, "test.weight");
3099        buf.write_u32::<LittleEndian>(2).unwrap(); // n_dims
3100                                                   // Real GGUF ne[] order is fastest-varying-first, i.e. [cols, rows].
3101        buf.write_u64::<LittleEndian>(256).unwrap(); // cols (1 IQ4_XS block)
3102        buf.write_u64::<LittleEndian>(1).unwrap(); // rows
3103        buf.write_u32::<LittleEndian>(23).unwrap(); // dtype tag: IQ4_XS
3104        buf.write_u64::<LittleEndian>(0).unwrap(); // offset
3105
3106        while buf.len() % 32 != 0 {
3107            buf.push(0);
3108        }
3109        buf.extend_from_slice(&IQ4_XS_TEST_BLOCK);
3110        buf
3111    }
3112
3113    #[test]
3114    fn load_weight_matrix_handles_a_real_on_disk_iq4_xs_tensor_end_to_end() {
3115        let tmp = std::env::temp_dir().join(format!(
3116            "ferrox_test_iq4xs_tensor_{}.gguf",
3117            std::process::id()
3118        ));
3119        std::fs::write(&tmp, build_single_iq4_xs_tensor_gguf()).unwrap();
3120        let file = ferrox_gguf::GgufFile::open(&tmp).expect("real IQ4_XS GGUF file must parse");
3121        std::fs::remove_file(&tmp).ok();
3122
3123        let matrix = load_weight_matrix(&file, "test.weight").expect("IQ4_XS tensor must load");
3124        assert_eq!(matrix.rows(), 1);
3125        assert_eq!(matrix.cols(), 256);
3126        match &matrix {
3127            WeightMatrix::Quantized { kind, data, .. } => {
3128                assert_eq!(*kind, QuantKind::IQ4XS);
3129                assert!(data.is_mapped());
3130            }
3131            _ => panic!("expected a Quantized matrix for an IQ4_XS tensor"),
3132        }
3133
3134        let expected = ferrox_quant::dequant_iq4_xs(&IQ4_XS_TEST_BLOCK).unwrap();
3135        let x: Vec<f32> = (0..256).map(|i| ((i as f32) * 0.013).sin()).collect();
3136        let expected_dot: f32 = expected.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
3137
3138        let got = matrix.apply(&x);
3139        assert_eq!(got.len(), 1);
3140        assert!(
3141            (got[0] - expected_dot).abs() < 1e-1,
3142            "end-to-end loaded+applied IQ4_XS matrix diverged from direct dequant: got={} expected={}",
3143            got[0],
3144            expected_dot
3145        );
3146    }
3147
3148    // Same bytes as ferrox-quant's own IQ low-bit test blocks
3149    // (Python-cross-validated there against the real compiled ggml
3150    // implementation), duplicated as literals for the same reason as
3151    // IQ4_XS_TEST_BLOCK above.
3152    const IQ1_S_TEST_BLOCK: [u8; 50] = [
3153        0x0a, 0x2f, 0xfa, 0x06, 0x1e, 0x37, 0x6f, 0xe3, 0x62, 0xd0, 0xb6, 0xa4, 0x25, 0xae, 0x76,
3154        0x14, 0x72, 0x5b, 0xfa, 0x05, 0xd1, 0xf1, 0x2a, 0x4c, 0xad, 0x29, 0xae, 0xf4, 0xcf, 0x0c,
3155        0x96, 0x51, 0x58, 0x03, 0x6d, 0xd3, 0x10, 0x92, 0x70, 0xff, 0x61, 0x58, 0xc8, 0x30, 0x25,
3156        0x64, 0x49, 0x85, 0xc0, 0x24,
3157    ];
3158    const IQ2_XXS_TEST_BLOCK: [u8; 66] = [
3159        0x29, 0x30, 0xd9, 0x33, 0x95, 0x4c, 0x08, 0x1e, 0xad, 0x79, 0x49, 0xf2, 0x8d, 0x5f, 0x93,
3160        0xea, 0x78, 0x18, 0x98, 0xb9, 0x94, 0x14, 0xad, 0xce, 0xca, 0x1d, 0xab, 0x81, 0x53, 0x4a,
3161        0x68, 0xd0, 0x59, 0x96, 0x36, 0x5d, 0xbe, 0x20, 0xc4, 0xff, 0xe4, 0x2c, 0xcd, 0x2f, 0x4f,
3162        0x4f, 0x67, 0x53, 0xc6, 0xd5, 0xa2, 0xfb, 0xc7, 0xf3, 0xe2, 0x6b, 0xf1, 0x99, 0x23, 0x1e,
3163        0x2d, 0x5e, 0x8c, 0x78, 0xc2, 0x31,
3164    ];
3165    const IQ3_XXS_TEST_BLOCK: [u8; 98] = [
3166        0x71, 0x31, 0x16, 0x0a, 0x79, 0x04, 0x5d, 0x87, 0xae, 0x2a, 0x4a, 0x43, 0xfd, 0x02, 0xba,
3167        0x6c, 0x10, 0x42, 0x80, 0xe5, 0x1d, 0x08, 0x22, 0xcb, 0x21, 0x54, 0xf9, 0xaa, 0x8e, 0xc2,
3168        0xf2, 0x34, 0x66, 0x1e, 0x2a, 0xef, 0x19, 0xae, 0x48, 0x47, 0x29, 0xa0, 0x72, 0xd1, 0x31,
3169        0xc0, 0x65, 0x49, 0xde, 0x79, 0x32, 0xe6, 0x4d, 0xb6, 0x55, 0x3f, 0x4d, 0xf1, 0x18, 0xbb,
3170        0x18, 0x59, 0x4c, 0x31, 0xa3, 0xb2, 0x34, 0xdd, 0xf6, 0x4a, 0x91, 0x51, 0x3f, 0x3e, 0x40,
3171        0x69, 0xad, 0xbf, 0x1a, 0xd0, 0x05, 0xfb, 0xbe, 0x8b, 0x0b, 0xdd, 0xdf, 0x7d, 0x94, 0x74,
3172        0x92, 0x3e, 0xff, 0x04, 0x2a, 0xc4, 0xea, 0xc9,
3173    ];
3174
3175    #[rustfmt::skip]
3176    const MXFP4_GGUF_TEST_BLOCKS: [u8; 68] = [0x79, 0xb4, 0x8d, 0xe2, 0x62, 0x5d, 0xbb, 0x9d, 0x54, 0xe6, 0xdb, 0x94, 0x59, 0x7d, 0x28, 0xf9, 0x79, 0x7a, 0xfc, 0xc1, 0xfa, 0x1e, 0x53, 0x5b, 0x0e, 0xc2, 0x5a, 0x2f, 0x0c, 0x82, 0x4d, 0xcb, 0x11, 0x28, 0x7b, 0x7c, 0xb6, 0x45, 0xe0, 0xb0, 0x52, 0x40, 0x51, 0xec, 0x30, 0x1a, 0xd2, 0x17, 0xf3, 0xbb, 0xfc, 0x7c, 0x8f, 0xf0, 0x67, 0x83, 0x88, 0x9d, 0x79, 0xdb, 0xf4, 0x45, 0x29, 0x78, 0xe6, 0xf4, 0x99, 0xea];
3177
3178    fn build_single_iq_lowbit_tensor_gguf(
3179        arch: &str,
3180        tag: u32,
3181        cols: u64,
3182        block: &[u8],
3183    ) -> Vec<u8> {
3184        let mut buf = Vec::new();
3185        buf.write_u32::<LittleEndian>(ferrox_gguf::GGUF_MAGIC)
3186            .unwrap();
3187        buf.write_u32::<LittleEndian>(3).unwrap(); // version
3188        buf.write_u64::<LittleEndian>(1).unwrap(); // tensor_count
3189        buf.write_u64::<LittleEndian>(1).unwrap(); // kv_count
3190        write_kv_str(&mut buf, "general.architecture", arch);
3191        write_string(&mut buf, "test.weight");
3192        buf.write_u32::<LittleEndian>(2).unwrap(); // n_dims
3193        buf.write_u64::<LittleEndian>(cols).unwrap();
3194        buf.write_u64::<LittleEndian>(1).unwrap(); // rows
3195        buf.write_u32::<LittleEndian>(tag).unwrap();
3196        buf.write_u64::<LittleEndian>(0).unwrap(); // offset
3197        while buf.len() % 32 != 0 {
3198            buf.push(0);
3199        }
3200        buf.extend_from_slice(block);
3201        buf
3202    }
3203
3204    /// A structurally valid block of `len` bytes for any of the
3205    /// codebook-grid formats: every bit pattern is a legal code in all
3206    /// of them (the grid indices are bounded by their own bit widths),
3207    /// so a deterministic byte fill is a real block, not a fixture that
3208    /// happens to avoid the interesting paths. Only the f16 scale needs
3209    /// pinning, and only so the comparison below can't be NaN-vs-NaN.
3210    fn pseudo_iq_block(len: usize, seed: u32) -> Vec<u8> {
3211        let mut s = seed;
3212        let mut out = Vec::with_capacity(len);
3213        for _ in 0..len {
3214            s ^= s << 13;
3215            s ^= s >> 17;
3216            s ^= s << 5;
3217            out.push((s >> 24) as u8);
3218        }
3219        out
3220    }
3221
3222    /// End-to-end load+apply for the codebook-grid low-bit formats the
3223    /// published Dynamic GGUFs are built from: a real on-disk tensor of
3224    /// each type must load zero-copy as the right `QuantKind` and
3225    /// produce the same matvec result as dequantizing the block
3226    /// directly. That is the property this test exists for -- the
3227    /// *values* are pinned against real ggml in `ferrox-quant`; what
3228    /// can only break here is the tag -> kind -> block-stride chain,
3229    /// and a wrong stride silently reads the neighbouring row.
3230    /// Dtype tags (19/29/16/17/22/18/21/39) verified against ggml.h's
3231    /// enum ggml_type.
3232    #[test]
3233    fn load_weight_matrix_handles_real_on_disk_iq_lowbit_tensors_end_to_end() {
3234        type DequantFn = fn(&[u8]) -> Result<Vec<f32>, ferrox_quant::QuantError>;
3235        // IQ1_M carries no f16 scale field; its scale is reassembled
3236        // from the four scale words' top nibbles, and the top nibble of
3237        // the last one supplies the f16 sign + high exponent bits.
3238        // Pinning it to 0x2 keeps the exponent out of the all-ones
3239        // NaN/Inf pattern whatever the rest of the fill does. The other
3240        // three do carry a leading f16 `d`, pinned for the same reason.
3241        let mut iq1m = pseudo_iq_block(ferrox_quant::IQ1_M_BLOCK_BYTES, 0x2907_31A0);
3242        iq1m[55] = (iq1m[55] & 0x0F) | 0x20;
3243        let mut iq2xs = pseudo_iq_block(ferrox_quant::IQ2_XS_BLOCK_BYTES, 0x2107_31A1);
3244        let mut iq2s = pseudo_iq_block(ferrox_quant::IQ2_S_BLOCK_BYTES, 0x2207_31A2);
3245        let mut iq3s = pseudo_iq_block(ferrox_quant::IQ3_S_BLOCK_BYTES, 0x2307_31A3);
3246        for blk in [&mut iq2xs, &mut iq2s, &mut iq3s] {
3247            blk[0..2].copy_from_slice(&half::f16::from_f32(0.115).to_le_bytes());
3248        }
3249        let cases: [(&str, u32, &[u8], QuantKind, DequantFn); 8] = [
3250            (
3251                "iq1s",
3252                19,
3253                &IQ1_S_TEST_BLOCK,
3254                QuantKind::IQ1S,
3255                ferrox_quant::dequant_iq1_s,
3256            ),
3257            (
3258                "iq1m",
3259                29,
3260                &iq1m,
3261                QuantKind::IQ1M,
3262                ferrox_quant::dequant_iq1_m,
3263            ),
3264            (
3265                "iq2xxs",
3266                16,
3267                &IQ2_XXS_TEST_BLOCK,
3268                QuantKind::IQ2XXS,
3269                ferrox_quant::dequant_iq2_xxs,
3270            ),
3271            (
3272                "iq2xs",
3273                17,
3274                &iq2xs,
3275                QuantKind::IQ2XS,
3276                ferrox_quant::dequant_iq2_xs,
3277            ),
3278            (
3279                "iq2s",
3280                22,
3281                &iq2s,
3282                QuantKind::IQ2S,
3283                ferrox_quant::dequant_iq2_s,
3284            ),
3285            (
3286                "iq3xxs",
3287                18,
3288                &IQ3_XXS_TEST_BLOCK,
3289                QuantKind::IQ3XXS,
3290                ferrox_quant::dequant_iq3_xxs,
3291            ),
3292            (
3293                "iq3s",
3294                21,
3295                &iq3s,
3296                QuantKind::IQ3S,
3297                ferrox_quant::dequant_iq3_s,
3298            ),
3299            (
3300                "mxfp4_gguf",
3301                39,
3302                &MXFP4_GGUF_TEST_BLOCKS,
3303                QuantKind::Mxfp4Gguf,
3304                ferrox_quant::dequant_mxfp4_gguf,
3305            ),
3306        ];
3307        for (name, tag, block, kind, dequant) in cases {
3308            let expected = dequant(block).unwrap();
3309            let cols = expected.len();
3310            let tmp = std::env::temp_dir().join(format!("ferrox_test_{name}_tensor.gguf"));
3311            std::fs::write(
3312                &tmp,
3313                build_single_iq_lowbit_tensor_gguf(name, tag, cols as u64, block),
3314            )
3315            .unwrap();
3316            let file = ferrox_gguf::GgufFile::open(&tmp).expect("file must parse");
3317            std::fs::remove_file(&tmp).ok();
3318
3319            let matrix =
3320                load_weight_matrix(&file, "test.weight").expect("low-bit tensor must load");
3321            assert_eq!((matrix.rows(), matrix.cols()), (1, cols), "{name}");
3322            match &matrix {
3323                WeightMatrix::Quantized { kind: k, data, .. } => {
3324                    assert_eq!(*k, kind, "{name}");
3325                    assert!(data.is_mapped(), "{name} must load zero-copy");
3326                }
3327                _ => panic!("expected a Quantized matrix for {name}"),
3328            }
3329
3330            let x: Vec<f32> = (0..cols).map(|i| ((i as f32) * 0.013).sin()).collect();
3331            let expected_dot: f32 = expected.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
3332            let got = matrix.apply(&x);
3333            assert!(
3334                (got[0] - expected_dot).abs() < 1e-1,
3335                "{name}: loaded+applied diverged from direct dequant: got={} expected={}",
3336                got[0],
3337                expected_dot
3338            );
3339        }
3340    }
3341
3342    #[test]
3343    fn qwen2moe_disables_topk_renorm() {
3344        assert!(
3345            NO_TOPK_RENORMALIZE_ARCHITECTURES.contains(&"qwen2moe"),
3346            "qwen2moe must have norm_topk_prob=false (llama.cpp build_moe_ffn norm_w=false)"
3347        );
3348    }
3349
3350    /// An architecture that uses no RoPE must not reach the generic
3351    /// decoder, which rotates unconditionally.
3352    ///
3353    /// All five of these were admitted as `GenericGqa { rope: Neox }`.
3354    /// Nothing downstream could have caught it: `bloom` and `refact`
3355    /// hardcode their ALiBi slope in `load_arch_hparams` with no GGUF
3356    /// key, so the metadata gates above see nothing, and `mpt` carries
3357    /// no tensor the generic loader fails to consume, so
3358    /// `assert_every_tensor_consumed` sees nothing either. It would have
3359    /// loaded, run at full speed, and answered from rotated positions.
3360    #[test]
3361    fn an_architecture_with_no_rope_is_refused_by_name() {
3362        for arch in ["gpt2", "mpt", "refact", "bloom", "jais"] {
3363            let file = open_metadata_gguf(
3364                &format!("norope_{arch}"),
3365                &[("general.architecture", Kv::Str(arch))],
3366            );
3367            match ModelConfig::from_gguf(&file) {
3368                Err(LoadError::DedicatedArchitectureRequired(got, reason)) => {
3369                    assert_eq!(got, arch);
3370                    assert!(
3371                        reason.contains("ALiBi") || reason.contains("position embeddings"),
3372                        "{arch}: the refusal must name what is missing, got {reason:?}"
3373                    );
3374                }
3375                other => panic!("{arch} must be refused, got {other:?}"),
3376            }
3377        }
3378    }
3379
3380    /// Baichuan is one `general.architecture` string covering two
3381    /// positional schemes, and llama.cpp picks between them on the layer
3382    /// count alone (`src/models/baichuan.cpp:11-14`, with its own "TODO:
3383    /// become GGUF KV parameter"). So the 13B is the MiniCPM case: no
3384    /// key to gate on and no tensor to miss.
3385    #[test]
3386    fn baichuan_13b_is_refused_because_it_uses_alibi_and_the_7b_is_not() {
3387        let thirteen_b = open_metadata_gguf(
3388            "baichuan13b",
3389            &[
3390                ("general.architecture", Kv::Str("baichuan")),
3391                ("baichuan.block_count", Kv::U32(40)),
3392            ],
3393        );
3394        match ModelConfig::from_gguf(&thirteen_b) {
3395            Err(LoadError::UnsupportedFeature(arch, msg)) => {
3396                assert_eq!(arch, "baichuan");
3397                assert!(msg.contains("ALiBi"), "{msg}");
3398                assert!(
3399                    msg.contains("40"),
3400                    "the refusal must name the layer count: {msg}"
3401                );
3402            }
3403            other => panic!("Baichuan-13B must be refused, got {other:?}"),
3404        }
3405
3406        // The 7B rotates exactly as the generic decoder does, so it must
3407        // pass this gate. It still fails later, on the next missing
3408        // hparam, which is what proves the gate let it through.
3409        let seven_b = open_metadata_gguf(
3410            "baichuan7b",
3411            &[
3412                ("general.architecture", Kv::Str("baichuan")),
3413                ("baichuan.block_count", Kv::U32(32)),
3414            ],
3415        );
3416        match ModelConfig::from_gguf(&seven_b) {
3417            Err(LoadError::MissingHparam(key)) => assert_eq!(key, "baichuan.embedding_length"),
3418            other => panic!("Baichuan-7B must pass the ALiBi gate, got {other:?}"),
3419        }
3420    }
3421}