ferrox_models/
tensor_role.rs

1//! Semantic tensor roles mirroring llama.cpp `llm_tensor` / `LLM_TN`.
2//!
3//! Architecture loaders resolve roles to concrete GGUF names once at
4//! load time. Hot-path kernels never see string names.
5
6/// Logical weight / activation tensor identity inside a decoder block.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
8pub enum TensorRole {
9    TokenEmbd,
10    Output,
11    OutputNorm,
12    AttnNorm,
13    AttnQ,
14    AttnK,
15    AttnV,
16    AttnQkv,
17    AttnOut,
18    AttnQNorm,
19    AttnKNorm,
20    AttnPostNorm,
21    FfnNorm,
22    FfnGate,
23    FfnUp,
24    FfnDown,
25    FfnPostNorm,
26    FfnGateInp,
27    FfnGateExps,
28    FfnUpExps,
29    FfnDownExps,
30    RopeFreqs,
31}
32
33impl TensorRole {
34    /// Default GGUF tensor name for this role (layer-scoped when `layer`
35    /// is `Some`). Matches llama.cpp `LLM_TENSOR_NAMES` conventions.
36    pub fn gguf_name(self, layer: Option<usize>) -> String {
37        match (self, layer) {
38            (Self::TokenEmbd, _) => "token_embd.weight".into(),
39            (Self::Output, _) => "output.weight".into(),
40            (Self::OutputNorm, _) => "output_norm.weight".into(),
41            (Self::RopeFreqs, _) => "rope_freqs.weight".into(),
42            (Self::AttnNorm, Some(i)) => format!("blk.{i}.attn_norm.weight"),
43            (Self::AttnQ, Some(i)) => format!("blk.{i}.attn_q.weight"),
44            (Self::AttnK, Some(i)) => format!("blk.{i}.attn_k.weight"),
45            (Self::AttnV, Some(i)) => format!("blk.{i}.attn_v.weight"),
46            (Self::AttnQkv, Some(i)) => format!("blk.{i}.attn_qkv.weight"),
47            (Self::AttnOut, Some(i)) => format!("blk.{i}.attn_output.weight"),
48            (Self::AttnQNorm, Some(i)) => format!("blk.{i}.attn_q_norm.weight"),
49            (Self::AttnKNorm, Some(i)) => format!("blk.{i}.attn_k_norm.weight"),
50            (Self::AttnPostNorm, Some(i)) => format!("blk.{i}.post_attention_norm.weight"),
51            (Self::FfnNorm, Some(i)) => format!("blk.{i}.ffn_norm.weight"),
52            (Self::FfnGate, Some(i)) => format!("blk.{i}.ffn_gate.weight"),
53            (Self::FfnUp, Some(i)) => format!("blk.{i}.ffn_up.weight"),
54            (Self::FfnDown, Some(i)) => format!("blk.{i}.ffn_down.weight"),
55            (Self::FfnPostNorm, Some(i)) => format!("blk.{i}.post_ffw_norm.weight"),
56            (Self::FfnGateInp, Some(i)) => format!("blk.{i}.ffn_gate_inp.weight"),
57            (Self::FfnGateExps, Some(i)) => format!("blk.{i}.ffn_gate_exps.weight"),
58            (Self::FfnUpExps, Some(i)) => format!("blk.{i}.ffn_up_exps.weight"),
59            (Self::FfnDownExps, Some(i)) => format!("blk.{i}.ffn_down_exps.weight"),
60            (role, None) => panic!("{role:?} requires a layer index"),
61        }
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn layer_scoped_names_match_llama_convention() {
71        assert_eq!(TensorRole::AttnQ.gguf_name(Some(3)), "blk.3.attn_q.weight");
72        assert_eq!(
73            TensorRole::AttnQkv.gguf_name(Some(0)),
74            "blk.0.attn_qkv.weight"
75        );
76        assert_eq!(TensorRole::TokenEmbd.gguf_name(None), "token_embd.weight");
77    }
78}
ferrox_models/tensor_role.rs

ferrox_models/
tensor_role.rs
