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