rlx_models_core/
gguf_resolve.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Pluggable GGUF tensor-name resolution per `general.architecture`.
17
18use rlx_gguf::GgufFile;
19use std::sync::{Mutex, OnceLock};
20
21use crate::weight_loader::{gguf_to_hf_name, hf_to_gguf_name};
22
23/// Resolve a builder-requested tensor name to the name stored in a GGUF file.
24pub trait GgufTensorNameResolver: Send + Sync {
25    fn matches_arch(&self, arch: &str) -> bool;
26    fn resolve(&self, file: &GgufFile, requested_key: &str) -> Option<String>;
27}
28
29/// HF `model.layers.N.*` ↔ GGUF `blk.N.*` (Llama, Qwen3, Qwen35, …).
30pub struct LlamaFamilyGgufResolver;
31
32impl GgufTensorNameResolver for LlamaFamilyGgufResolver {
33    fn matches_arch(&self, arch: &str) -> bool {
34        matches!(
35            arch,
36            "llama"
37                | "llama4"
38                | "qwen3"
39                | "qwen2"
40                | "qwen35"
41                | "qwen35moe"
42                | "qwen36"
43                | "gemma"
44                | "gemma2"
45                | "mistral"
46        )
47    }
48
49    fn resolve(&self, file: &GgufFile, key: &str) -> Option<String> {
50        if file.tensors.contains_key(key) {
51            return Some(key.to_string());
52        }
53        if let Some(g) = hf_to_gguf_name(key) {
54            if file.tensors.contains_key(&g) {
55                return Some(g);
56            }
57        }
58        if let Some(h) = gguf_to_hf_name(key) {
59            if file.tensors.contains_key(&h) {
60                return Some(h);
61            }
62        }
63        None
64    }
65}
66
67/// Strip common HF prefixes and match verbatim tensor names (architecture-agnostic fallback).
68pub struct PrefixStripGgufResolver;
69
70/// Alias for [`PrefixStripGgufResolver`] (older name).
71pub type PassThroughGgufResolver = PrefixStripGgufResolver;
72
73impl GgufTensorNameResolver for PrefixStripGgufResolver {
74    fn matches_arch(&self, _arch: &str) -> bool {
75        true
76    }
77
78    fn resolve(&self, file: &GgufFile, key: &str) -> Option<String> {
79        let mut k = key.to_string();
80        for prefix in [
81            "model.diffusion_model.",
82            "diffusion_model.",
83            "transformer.",
84            "model.",
85        ] {
86            if let Some(rest) = k.strip_prefix(prefix) {
87                k = rest.to_string();
88                break;
89            }
90        }
91        if file.tensors.contains_key(&k) {
92            return Some(k);
93        }
94        if file.tensors.contains_key(key) {
95            return Some(key.to_string());
96        }
97        None
98    }
99}
100
101/// Gemma 2/3/4: 4 RMSNorms per layer disagree with the Llama 2-norm convention.
102///
103/// Llama treats `post_attention_layernorm` as the pre-FFN norm and aliases it
104/// to `ffn_norm`. Gemma 2/3/4 (V2/V3/V4 layer styles) have a dedicated
105/// `post_attention_norm` between the attention output and the residual add,
106/// *and* a separate `ffn_norm` / `post_ffw_norm` pair around the MLP. Without
107/// this resolver, the Llama mapper would alias `post_attention_layernorm` to
108/// `ffn_norm`, collide with `pre_feedforward_layernorm`, and silently load
109/// the wrong tensor. The tail map is identical across these arches — only
110/// the GGUF arch tag (and runtime details like sliding-window stride) differ.
111pub struct Gemma2GgufResolver;
112
113impl GgufTensorNameResolver for Gemma2GgufResolver {
114    fn matches_arch(&self, arch: &str) -> bool {
115        matches!(
116            arch,
117            "gemma2" | "gemma3" | "gemma3n" | "gemma4" | "gemma4moe"
118        )
119    }
120
121    fn resolve(&self, file: &GgufFile, key: &str) -> Option<String> {
122        // Identity hit first — accept native GGUF names verbatim.
123        if file.tensors.contains_key(key) {
124            return Some(key.to_string());
125        }
126        // Handle the four-norm-per-layer scheme explicitly. The Llama mapper
127        // is wrong for `post_attention_layernorm` (it aliases to `ffn_norm`,
128        // which Gemma 2 reserves for the pre-FFN norm) and has no entry at
129        // all for the `pre_feedforward_layernorm`/`post_feedforward_layernorm`
130        // pair.
131        if let Some(rest) = key.strip_prefix("model.layers.") {
132            if let Some((idx, tail)) = rest.split_once('.') {
133                let gguf_tail = match tail {
134                    "post_attention_layernorm.weight" => Some("post_attention_norm.weight"),
135                    "pre_feedforward_layernorm.weight" => Some("ffn_norm.weight"),
136                    "post_feedforward_layernorm.weight" => Some("post_ffw_norm.weight"),
137                    // Inverse of `gguf_to_hf_name_for_arch` — Gemma 4 GGUF stores
138                    // per-layer output scalars as `layer_output_scale`, not HF
139                    // `self_attn.output_scale`.
140                    "self_attn.output_scale.weight" => Some("layer_output_scale.weight"),
141                    _ => None,
142                };
143                if let Some(t) = gguf_tail {
144                    let g = format!("blk.{idx}.{t}");
145                    if file.tensors.contains_key(&g) {
146                        return Some(g);
147                    }
148                }
149            }
150        }
151        // Fall through to Llama-family mapping for everything else
152        // (input_layernorm, attn/mlp weights, lm_head, embeddings, …).
153        LlamaFamilyGgufResolver.resolve(file, key)
154    }
155}
156
157/// Qwen3.5 native `blk.N.*` names; also accept HF aliases via the Llama mapper.
158pub struct Qwen35NativeGgufResolver;
159
160impl GgufTensorNameResolver for Qwen35NativeGgufResolver {
161    fn matches_arch(&self, arch: &str) -> bool {
162        matches!(arch, "qwen35" | "qwen35moe" | "qwen36")
163    }
164
165    fn resolve(&self, file: &GgufFile, key: &str) -> Option<String> {
166        if file.tensors.contains_key(key) {
167            return Some(key.to_string());
168        }
169        LlamaFamilyGgufResolver.resolve(file, key)
170    }
171}
172
173static CUSTOM_RESOLVERS: Mutex<Vec<Box<dyn GgufTensorNameResolver>>> = Mutex::new(Vec::new());
174static BUILTIN_RESOLVERS: OnceLock<Vec<Box<dyn GgufTensorNameResolver>>> = OnceLock::new();
175
176fn builtin_resolvers() -> &'static Vec<Box<dyn GgufTensorNameResolver>> {
177    BUILTIN_RESOLVERS.get_or_init(|| {
178        vec![
179            Box::new(Qwen35NativeGgufResolver),
180            Box::new(Gemma2GgufResolver),
181            Box::new(LlamaFamilyGgufResolver),
182            Box::new(PrefixStripGgufResolver),
183        ]
184    })
185}
186
187/// Register built-in GGUF resolvers (idempotent). Called automatically on first resolve; \
188/// call from `main` if you register custom resolvers and need ordering guarantees.
189pub fn ensure_builtin_resolvers() {
190    let _ = builtin_resolvers();
191}
192
193/// Register a custom resolver (call before first GGUF load). Later registrations win
194/// among resolvers that match the same architecture.
195pub fn register_gguf_tensor_resolver(resolver: Box<dyn GgufTensorNameResolver>) {
196    CUSTOM_RESOLVERS
197        .lock()
198        .expect("gguf resolver registry lock")
199        .push(resolver);
200}
201
202/// Resolve `requested_key` against tensors in `file` using registered resolvers.
203pub fn resolve_gguf_tensor_name(
204    file: &GgufFile,
205    arch: &str,
206    requested_key: &str,
207) -> Option<String> {
208    for r in builtin_resolvers().iter() {
209        if r.matches_arch(arch) {
210            if let Some(name) = r.resolve(file, requested_key) {
211                return Some(name);
212            }
213        }
214    }
215    let custom = CUSTOM_RESOLVERS
216        .lock()
217        .expect("gguf resolver registry lock");
218    for r in custom.iter() {
219        if r.matches_arch(arch) {
220            if let Some(name) = r.resolve(file, requested_key) {
221                return Some(name);
222            }
223        }
224    }
225    if file.tensors.contains_key(requested_key) {
226        return Some(requested_key.to_string());
227    }
228    if let Some(g) = hf_to_gguf_name(requested_key) {
229        if file.tensors.contains_key(&g) {
230            return Some(g);
231        }
232    }
233    if let Some(h) = gguf_to_hf_name(requested_key) {
234        if file.tensors.contains_key(&h) {
235            return Some(h);
236        }
237    }
238    None
239}
rlx_models_core/gguf_resolve.rs

rlx_models_core/
gguf_resolve.rs
