qualia_core_db/inference/
tensor_roles.rs

1//! Task #12 / STELLAR §A — **tensor-name → engine GEMM-role** mapping + the ternary policy.
2//!
3//! The transcoder (`p64_weight`) needs to know *what* each tensor is to (a) populate the P64
4//! manifest with real engine roles (so a container boots without a GGUF re-parse) and (b) apply the
5//! §A compression *policy*: **ternary the FFN projections, keep attention + norms + embeddings at
6//! higher fidelity** (ternary norms/attention would wreck coherence).
7//!
8//! Two naming conventions are recognised:
9//! * **GGUF / llama.cpp** — `blk.{L}.attn_q.weight`, `blk.{L}.ffn_gate.weight`, `token_embd.weight`, …
10//! * **Hugging Face safetensor** — `model.layers.{L}.self_attn.q_proj.weight`,
11//!   `model.layers.{L}.mlp.gate_proj.weight`, `model.embed_tokens.weight`, `lm_head.weight`, …
12
13use crate::p64_weight::{
14    P64_LAYER_GLOBAL, P64_ROLE_ATTN_K, P64_ROLE_ATTN_NORM, P64_ROLE_ATTN_OUTPUT, P64_ROLE_ATTN_Q,
15    P64_ROLE_ATTN_SUBLN, P64_ROLE_ATTN_V, P64_ROLE_FFN_DOWN, P64_ROLE_FFN_GATE, P64_ROLE_FFN_NORM,
16    P64_ROLE_FFN_SUBLN, P64_ROLE_FFN_UP, P64_ROLE_OUTPUT, P64_ROLE_OUTPUT_NORM,
17    P64_ROLE_TOKEN_EMBD,
18};
19
20/// A resolved tensor identity: an engine role + its layer (`P64_LAYER_GLOBAL` for non-layer tensors).
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct TensorRole {
23    pub role: u16,
24    pub layer: u16,
25}
26
27/// Extract the layer index from a `blk.{N}.` (GGUF) or `…layers.{N}.` (HF) tensor name.
28fn extract_layer(name: &str) -> Option<u16> {
29    for marker in ["blk.", "layers."] {
30        if let Some(pos) = name.find(marker) {
31            let rest = &name[pos + marker.len()..];
32            let digits: String = rest.chars().take_while(|c| c.is_ascii_digit()).collect();
33            if let Ok(l) = digits.parse::<u16>() {
34                return Some(l);
35            }
36        }
37    }
38    None
39}
40
41/// Map a tensor name (GGUF or HF convention) to its engine role + layer, or `None` if unrecognised.
42pub fn name_to_role(name: &str) -> Option<TensorRole> {
43    // Global (non-layer) tensors first — and `output_norm` BEFORE the bare `output` (it contains it).
44    if name.contains("output_norm") || name == "model.norm.weight" {
45        return Some(TensorRole {
46            role: P64_ROLE_OUTPUT_NORM,
47            layer: P64_LAYER_GLOBAL,
48        });
49    }
50    if name.contains("token_embd") || name.contains("embed_tokens") {
51        return Some(TensorRole {
52            role: P64_ROLE_TOKEN_EMBD,
53            layer: P64_LAYER_GLOBAL,
54        });
55    }
56    if name.contains("lm_head") || name == "output.weight" {
57        return Some(TensorRole {
58            role: P64_ROLE_OUTPUT,
59            layer: P64_LAYER_GLOBAL,
60        });
61    }
62
63    // Per-layer tensors require a layer index.
64    let layer = extract_layer(name)?;
65    let role = if name.contains("attn_q") || name.contains("q_proj") {
66        P64_ROLE_ATTN_Q
67    } else if name.contains("attn_k") || name.contains("k_proj") {
68        P64_ROLE_ATTN_K
69    } else if name.contains("attn_v") || name.contains("v_proj") {
70        P64_ROLE_ATTN_V
71    } else if name.contains("attn_output") || name.contains("o_proj") {
72        P64_ROLE_ATTN_OUTPUT
73    } else if name.contains("ffn_gate") || name.contains("gate_proj") {
74        P64_ROLE_FFN_GATE
75    } else if name.contains("ffn_up") || name.contains("up_proj") {
76        P64_ROLE_FFN_UP
77    } else if name.contains("ffn_down") || name.contains("down_proj") {
78        P64_ROLE_FFN_DOWN
79    } else if name.contains("attn_sub_norm") || name.contains("self_attn_layer_norm") {
80        P64_ROLE_ATTN_SUBLN
81    } else if name.contains("ffn_sub_norm") || name.contains("mlp_layer_norm") {
82        P64_ROLE_FFN_SUBLN
83    } else if name.contains("attn_norm") || name.contains("input_layernorm") {
84        P64_ROLE_ATTN_NORM
85    } else if name.contains("ffn_norm") || name.contains("post_attention_layernorm") {
86        P64_ROLE_FFN_NORM
87    } else {
88        return None;
89    };
90    Some(TensorRole { role, layer })
91}
92
93/// The §A ternary **policy**: only the FFN projection weights (`gate` / `up` / `down`) are
94/// eligible for BitNet-1.58b ternary packing. Attention projections, norms, and embeddings stay at
95/// higher fidelity — ternarising them destroys coherence.
96pub fn ternary_eligible(role: u16) -> bool {
97    matches!(
98        role,
99        P64_ROLE_FFN_GATE | P64_ROLE_FFN_UP | P64_ROLE_FFN_DOWN
100    )
101}
102
103/// True iff this tensor name resolves to an FFN projection (so it is ternary-eligible). A name we
104/// cannot classify is **not** ternarised (fail safe to high fidelity).
105pub fn name_is_ternary_eligible(name: &str) -> bool {
106    name_to_role(name)
107        .map(|r| ternary_eligible(r.role))
108        .unwrap_or(false)
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn gguf_names_map_to_roles() {
117        assert_eq!(
118            name_to_role("blk.3.attn_q.weight"),
119            Some(TensorRole {
120                role: P64_ROLE_ATTN_Q,
121                layer: 3
122            })
123        );
124        assert_eq!(
125            name_to_role("blk.0.ffn_gate.weight"),
126            Some(TensorRole {
127                role: P64_ROLE_FFN_GATE,
128                layer: 0
129            })
130        );
131        assert_eq!(
132            name_to_role("blk.11.ffn_down.weight"),
133            Some(TensorRole {
134                role: P64_ROLE_FFN_DOWN,
135                layer: 11
136            })
137        );
138        assert_eq!(
139            name_to_role("blk.2.attn_norm.weight"),
140            Some(TensorRole {
141                role: P64_ROLE_ATTN_NORM,
142                layer: 2
143            })
144        );
145        assert_eq!(
146            name_to_role("token_embd.weight"),
147            Some(TensorRole {
148                role: P64_ROLE_TOKEN_EMBD,
149                layer: P64_LAYER_GLOBAL
150            })
151        );
152        assert_eq!(
153            name_to_role("output.weight"),
154            Some(TensorRole {
155                role: P64_ROLE_OUTPUT,
156                layer: P64_LAYER_GLOBAL
157            })
158        );
159        assert_eq!(
160            name_to_role("output_norm.weight"),
161            Some(TensorRole {
162                role: P64_ROLE_OUTPUT_NORM,
163                layer: P64_LAYER_GLOBAL
164            })
165        );
166    }
167
168    #[test]
169    fn hf_names_map_to_roles() {
170        assert_eq!(
171            name_to_role("model.layers.5.self_attn.q_proj.weight"),
172            Some(TensorRole {
173                role: P64_ROLE_ATTN_Q,
174                layer: 5
175            })
176        );
177        assert_eq!(
178            name_to_role("model.layers.5.self_attn.o_proj.weight"),
179            Some(TensorRole {
180                role: P64_ROLE_ATTN_OUTPUT,
181                layer: 5
182            })
183        );
184        assert_eq!(
185            name_to_role("model.layers.7.mlp.gate_proj.weight"),
186            Some(TensorRole {
187                role: P64_ROLE_FFN_GATE,
188                layer: 7
189            })
190        );
191        assert_eq!(
192            name_to_role("model.layers.7.mlp.down_proj.weight"),
193            Some(TensorRole {
194                role: P64_ROLE_FFN_DOWN,
195                layer: 7
196            })
197        );
198        assert_eq!(
199            name_to_role("model.layers.7.input_layernorm.weight"),
200            Some(TensorRole {
201                role: P64_ROLE_ATTN_NORM,
202                layer: 7
203            })
204        );
205        assert_eq!(
206            name_to_role("model.layers.7.post_attention_layernorm.weight"),
207            Some(TensorRole {
208                role: P64_ROLE_FFN_NORM,
209                layer: 7
210            })
211        );
212        assert_eq!(
213            name_to_role("model.embed_tokens.weight"),
214            Some(TensorRole {
215                role: P64_ROLE_TOKEN_EMBD,
216                layer: P64_LAYER_GLOBAL
217            })
218        );
219        assert_eq!(
220            name_to_role("lm_head.weight"),
221            Some(TensorRole {
222                role: P64_ROLE_OUTPUT,
223                layer: P64_LAYER_GLOBAL
224            })
225        );
226        assert_eq!(
227            name_to_role("model.norm.weight"),
228            Some(TensorRole {
229                role: P64_ROLE_OUTPUT_NORM,
230                layer: P64_LAYER_GLOBAL
231            })
232        );
233    }
234
235    #[test]
236    fn unknown_names_are_none() {
237        assert_eq!(name_to_role("some.random.tensor"), None);
238        assert_eq!(name_to_role("blk.0.rotary_emb.inv_freq"), None);
239    }
240
241    #[test]
242    fn ternary_policy_is_ffn_only() {
243        // FFN → eligible
244        assert!(name_is_ternary_eligible("blk.0.ffn_gate.weight"));
245        assert!(name_is_ternary_eligible(
246            "model.layers.3.mlp.up_proj.weight"
247        ));
248        assert!(name_is_ternary_eligible("blk.9.ffn_down.weight"));
249        // attention / norms / embeddings → NOT eligible (kept high-fidelity)
250        assert!(!name_is_ternary_eligible("blk.0.attn_q.weight"));
251        assert!(!name_is_ternary_eligible(
252            "model.layers.3.self_attn.o_proj.weight"
253        ));
254        assert!(!name_is_ternary_eligible("blk.0.attn_norm.weight"));
255        assert!(!name_is_ternary_eligible("token_embd.weight"));
256        assert!(!name_is_ternary_eligible("output.weight"));
257        // unknown → fail safe to high fidelity
258        assert!(!name_is_ternary_eligible("mystery.weight"));
259    }
260}
qualia_core_db/inference/tensor_roles.rs

qualia_core_db/inference/
tensor_roles.rs
