ferrox-models 0.12.0

Model loaders and decoder stacks for the Ferrox inference engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

//! Semantic tensor roles mirroring llama.cpp `llm_tensor` / `LLM_TN`.
//!
//! Architecture loaders resolve roles to concrete GGUF names once at
//! load time. Hot-path kernels never see string names.

/// Logical weight / activation tensor identity inside a decoder block.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TensorRole {
    TokenEmbd,
    Output,
    OutputNorm,
    AttnNorm,
    AttnQ,
    AttnK,
    AttnV,
    AttnQkv,
    AttnOut,
    AttnQNorm,
    AttnKNorm,
    AttnPostNorm,
    FfnNorm,
    FfnGate,
    FfnUp,
    FfnDown,
    FfnPostNorm,
    FfnGateInp,
    FfnGateExps,
    FfnUpExps,
    FfnDownExps,
    RopeFreqs,
}

impl TensorRole {
    /// Default GGUF tensor name for this role (layer-scoped when `layer`
    /// is `Some`). Matches llama.cpp `LLM_TENSOR_NAMES` conventions.
    pub fn gguf_name(self, layer: Option<usize>) -> String {
        match (self, layer) {
            (Self::TokenEmbd, _) => "token_embd.weight".into(),
            (Self::Output, _) => "output.weight".into(),
            (Self::OutputNorm, _) => "output_norm.weight".into(),
            (Self::RopeFreqs, _) => "rope_freqs.weight".into(),
            (Self::AttnNorm, Some(i)) => format!("blk.{i}.attn_norm.weight"),
            (Self::AttnQ, Some(i)) => format!("blk.{i}.attn_q.weight"),
            (Self::AttnK, Some(i)) => format!("blk.{i}.attn_k.weight"),
            (Self::AttnV, Some(i)) => format!("blk.{i}.attn_v.weight"),
            (Self::AttnQkv, Some(i)) => format!("blk.{i}.attn_qkv.weight"),
            (Self::AttnOut, Some(i)) => format!("blk.{i}.attn_output.weight"),
            (Self::AttnQNorm, Some(i)) => format!("blk.{i}.attn_q_norm.weight"),
            (Self::AttnKNorm, Some(i)) => format!("blk.{i}.attn_k_norm.weight"),
            (Self::AttnPostNorm, Some(i)) => format!("blk.{i}.post_attention_norm.weight"),
            (Self::FfnNorm, Some(i)) => format!("blk.{i}.ffn_norm.weight"),
            (Self::FfnGate, Some(i)) => format!("blk.{i}.ffn_gate.weight"),
            (Self::FfnUp, Some(i)) => format!("blk.{i}.ffn_up.weight"),
            (Self::FfnDown, Some(i)) => format!("blk.{i}.ffn_down.weight"),
            (Self::FfnPostNorm, Some(i)) => format!("blk.{i}.post_ffw_norm.weight"),
            (Self::FfnGateInp, Some(i)) => format!("blk.{i}.ffn_gate_inp.weight"),
            (Self::FfnGateExps, Some(i)) => format!("blk.{i}.ffn_gate_exps.weight"),
            (Self::FfnUpExps, Some(i)) => format!("blk.{i}.ffn_up_exps.weight"),
            (Self::FfnDownExps, Some(i)) => format!("blk.{i}.ffn_down_exps.weight"),
            (role, None) => panic!("{role:?} requires a layer index"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn layer_scoped_names_match_llama_convention() {
        assert_eq!(TensorRole::AttnQ.gguf_name(Some(3)), "blk.3.attn_q.weight");
        assert_eq!(
            TensorRole::AttnQkv.gguf_name(Some(0)),
            "blk.0.attn_qkv.weight"
        );
        assert_eq!(TensorRole::TokenEmbd.gguf_name(None), "token_embd.weight");
    }
}