rlx_llama32/
config.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// LLaMA-3.2 configuration — HF `config.json` and GGUF `llama.*` metadata.
5
6use rlx_gguf::{GgufFile, MetaValue};
7use serde::Deserialize;
8use std::path::Path;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Default)]
11#[serde(rename_all = "lowercase")]
12pub enum Llama32RopeType {
13    #[default]
14    Default,
15    #[serde(rename = "llama3")]
16    Llama3,
17}
18
19#[derive(Debug, Clone, Deserialize)]
20pub struct Llama32RopeScaling {
21    pub factor: f32,
22    #[serde(default = "default_low_freq_factor")]
23    pub low_freq_factor: f32,
24    #[serde(default = "default_high_freq_factor")]
25    pub high_freq_factor: f32,
26    pub original_max_position_embeddings: usize,
27    #[serde(default)]
28    pub rope_type: Llama32RopeType,
29}
30
31fn default_low_freq_factor() -> f32 {
32    1.0
33}
34fn default_high_freq_factor() -> f32 {
35    4.0
36}
37
38#[derive(Debug, Clone, Deserialize)]
39pub struct Llama32Config {
40    pub vocab_size: usize,
41    pub hidden_size: usize,
42    pub intermediate_size: usize,
43    pub num_hidden_layers: usize,
44    pub num_attention_heads: usize,
45    pub num_key_value_heads: usize,
46    pub max_position_embeddings: usize,
47
48    #[serde(default = "default_rms_norm_eps")]
49    pub rms_norm_eps: f64,
50    #[serde(default = "default_rope_theta")]
51    pub rope_theta: f64,
52    #[serde(default = "default_hidden_act")]
53    pub hidden_act: String,
54    #[serde(default)]
55    pub tie_word_embeddings: bool,
56    #[serde(default)]
57    pub attention_bias: bool,
58    /// Explicit head dim (Llama 3.x); when absent, derived from hidden/heads.
59    #[serde(default)]
60    pub head_dim: Option<usize>,
61    #[serde(default)]
62    pub rope_scaling: Option<Llama32RopeScaling>,
63    /// RoPE pairing flavor. GGUF Llama weights are permuted by the HF→GGUF
64    /// converter for llama.cpp's interleaved (`NORM`) RoPE, so GGUF-backed
65    /// inference must rotate with [`rlx_ir::RopeStyle::GptJ`]; HF-safetensors
66    /// checkpoints use [`rlx_ir::RopeStyle::NeoX`] (default). Not present in
67    /// HF `config.json`, so skipped during deserialization.
68    #[serde(skip)]
69    pub rope_style: rlx_ir::RopeStyle,
70    /// GGUF `general.architecture` tag when loaded from GGUF (`llama`, `phi3`, …).
71    #[serde(skip)]
72    pub gguf_arch: Option<String>,
73    /// Rotary dimension when it differs from [`head_dim`] (Phi-3 partial RoPE).
74    #[serde(skip)]
75    pub rope_dim: Option<usize>,
76}
77
78fn default_rms_norm_eps() -> f64 {
79    1e-5
80}
81fn default_rope_theta() -> f64 {
82    500_000.0
83}
84fn default_hidden_act() -> String {
85    "silu".into()
86}
87
88impl Llama32Config {
89    pub fn from_file(path: &Path) -> anyhow::Result<Self> {
90        let data = std::fs::read_to_string(path)?;
91        Ok(serde_json::from_str(&data)?)
92    }
93
94    pub fn from_gguf(raw: &GgufFile) -> anyhow::Result<Self> {
95        llama32_cfg_from_gguf(raw)
96    }
97
98    pub fn head_dim(&self) -> usize {
99        self.head_dim
100            .unwrap_or(self.hidden_size / self.num_attention_heads)
101    }
102
103    pub fn kv_group_size(&self) -> usize {
104        self.num_attention_heads / self.num_key_value_heads
105    }
106
107    pub fn q_proj_dim(&self) -> usize {
108        self.num_attention_heads * self.head_dim()
109    }
110
111    pub fn kv_proj_dim(&self) -> usize {
112        self.num_key_value_heads * self.head_dim()
113    }
114
115    /// Leading per-head dims that receive RoPE (equals [`head_dim`] for Llama;
116    /// may be smaller for Phi-3 partial RoPE).
117    pub fn n_rot(&self) -> usize {
118        self.rope_dim
119            .filter(|&r| r > 0 && r <= self.head_dim())
120            .unwrap_or_else(|| self.head_dim())
121    }
122
123    pub fn uses_partial_rope(&self) -> bool {
124        self.n_rot() < self.head_dim()
125    }
126
127    pub fn is_phi_arch(&self) -> bool {
128        matches!(self.gguf_arch.as_deref(), Some("phi3") | Some("phi4"))
129    }
130
131    #[cfg(test)]
132    pub(crate) fn tiny_test() -> Self {
133        Self {
134            vocab_size: 32,
135            hidden_size: 16,
136            intermediate_size: 32,
137            num_hidden_layers: 2,
138            num_attention_heads: 4,
139            num_key_value_heads: 2,
140            max_position_embeddings: 16,
141            rms_norm_eps: 1e-5,
142            rope_theta: 500_000.0,
143            hidden_act: "silu".into(),
144            tie_word_embeddings: false,
145            attention_bias: false,
146            head_dim: None,
147            rope_scaling: None,
148            rope_style: rlx_ir::RopeStyle::NeoX,
149            gguf_arch: None,
150            rope_dim: None,
151        }
152    }
153}
154
155pub fn llama32_cfg_from_gguf(raw: &GgufFile) -> anyhow::Result<Llama32Config> {
156    let arch_prefix = raw
157        .metadata
158        .get("general.architecture")
159        .and_then(MetaValue::as_str)
160        .unwrap_or("llama");
161    let get_meta = |k: &str| -> Option<&MetaValue> {
162        raw.metadata.get(k).or_else(|| {
163            let suffix = k.strip_prefix("llama.")?;
164            if arch_prefix == "llama" {
165                None
166            } else {
167                let arch_key = format!("{arch_prefix}.{suffix}");
168                raw.metadata.get(&arch_key)
169            }
170        })
171    };
172    let get_u32 = |k: &str| -> anyhow::Result<u32> {
173        get_meta(k)
174            .and_then(MetaValue::as_u32)
175            .ok_or_else(|| anyhow::anyhow!("missing GGUF metadata key: {k}"))
176    };
177    let get_f32 = |k: &str| -> Option<f32> {
178        get_meta(k).and_then(|v| match v {
179            MetaValue::F32(x) => Some(*x),
180            _ => None,
181        })
182    };
183    let get_bool = |k: &str| -> Option<bool> {
184        get_meta(k).and_then(|v| match v {
185            MetaValue::Bool(b) => Some(*b),
186            _ => None,
187        })
188    };
189
190    let hidden_size = get_u32("llama.embedding_length")? as usize;
191    let num_attention_heads = get_u32("llama.attention.head_count")? as usize;
192    let head_dim_key = get_u32("llama.attention.key_length")
193        .ok()
194        .map(|v| v as usize);
195    let rope_dim = get_u32("llama.rope.dimension_count")
196        .ok()
197        .map(|v| v as usize);
198    let head_dim = head_dim_key.or(rope_dim);
199
200    let rope_scaling = match get_meta("llama.rope.scaling.type").and_then(MetaValue::as_str) {
201        Some("none") | None => {
202            // Llama 3.x often bakes scaling into rope_freqs.weight; HF fields may be absent.
203            None
204        }
205        Some("linear") | Some("yarn") | Some("longrope") => {
206            let factor = get_f32("llama.rope.scaling.factor")
207                .or_else(|| get_f32("llama.rope.scale_linear"))
208                .unwrap_or(1.0);
209            let original = get_u32("llama.rope.scaling.original_context_length")
210                .map(|v| v as usize)
211                .unwrap_or(8192);
212            Some(Llama32RopeScaling {
213                factor,
214                low_freq_factor: 1.0,
215                high_freq_factor: 4.0,
216                original_max_position_embeddings: original,
217                rope_type: Llama32RopeType::Llama3,
218            })
219        }
220        other => {
221            return Err(anyhow::anyhow!(
222                "unsupported llama.rope.scaling.type: {other:?}"
223            ));
224        }
225    };
226
227    Ok(Llama32Config {
228        vocab_size: infer_vocab_size_from_gguf(raw),
229        hidden_size,
230        intermediate_size: get_u32("llama.feed_forward_length")? as usize,
231        num_hidden_layers: get_u32("llama.block_count")? as usize,
232        num_attention_heads,
233        num_key_value_heads: get_u32("llama.attention.head_count_kv")? as usize,
234        max_position_embeddings: get_u32("llama.context_length").unwrap_or(8192) as usize,
235        rms_norm_eps: get_f32("llama.attention.layer_norm_rms_epsilon").unwrap_or(1e-5) as f64,
236        rope_theta: get_f32("llama.rope.freq_base").unwrap_or(500_000.0) as f64,
237        hidden_act: "silu".into(),
238        tie_word_embeddings: get_bool("llama.tie_word_embeddings").unwrap_or_else(|| {
239            // Llama-2 / TinyLlama GGUF often omits the flag; untied checkpoints
240            // carry a separate `output.weight` tensor.
241            !raw.tensors.contains_key("output.weight")
242        }),
243        attention_bias: false,
244        head_dim,
245        rope_scaling,
246        // Phi-3/4 GGUF uses HF NeoX rotate-half; plain Llama GGUF is GPT-J.
247        rope_style: if matches!(arch_prefix, "phi3" | "phi4") {
248            rlx_ir::RopeStyle::NeoX
249        } else {
250            rlx_ir::RopeStyle::GptJ
251        },
252        gguf_arch: Some(arch_prefix.to_string()),
253        rope_dim: rope_dim.filter(|r| head_dim_key.is_some() && *r <= head_dim_key.unwrap()),
254    })
255}
256
257/// Resolve vocab size from GGUF metadata / tensors. Llama-3 GGUF carries
258/// `llama.vocab_size`; older llama-tagged files (TinyLlama, SmolLM2, …) often
259/// only expose `tokenizer.ggml.tokens` or an embed row count.
260fn infer_vocab_size_from_gguf(raw: &GgufFile) -> usize {
261    if let Some(v) = raw
262        .metadata
263        .get("llama.vocab_size")
264        .and_then(MetaValue::as_u32)
265    {
266        return v as usize;
267    }
268    if let Some(MetaValue::Array(tokens)) = raw.metadata.get("tokenizer.ggml.tokens") {
269        if !tokens.is_empty() {
270            return tokens.len();
271        }
272    }
273    for name in ["token_embd.weight", "model.embed_tokens.weight"] {
274        if let Some(t) = raw.tensors.get(name) {
275            if !t.shape.is_empty() {
276                return t.shape[0];
277            }
278        }
279    }
280    128_256
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286
287    #[test]
288    fn parse_llama32_1b_like() {
289        let json = r#"{
290            "vocab_size": 128256,
291            "hidden_size": 2048,
292            "intermediate_size": 8192,
293            "num_hidden_layers": 16,
294            "num_attention_heads": 32,
295            "num_key_value_heads": 8,
296            "max_position_embeddings": 131072,
297            "rope_theta": 500000.0,
298            "rms_norm_eps": 1e-05,
299            "tie_word_embeddings": true,
300            "rope_scaling": {
301                "factor": 32.0,
302                "high_freq_factor": 4.0,
303                "low_freq_factor": 1.0,
304                "original_max_position_embeddings": 8192,
305                "rope_type": "llama3"
306            }
307        }"#;
308        let cfg: Llama32Config = serde_json::from_str(json).unwrap();
309        assert_eq!(cfg.head_dim(), 64);
310        assert_eq!(cfg.kv_group_size(), 4);
311        assert!(cfg.rope_scaling.is_some());
312    }
313
314    #[test]
315    fn gguf_vocab_inferred_from_tokenizer_tokens() {
316        use rlx_gguf::GgmlType;
317        use std::sync::atomic::{AtomicU64, Ordering};
318
319        static SEQ: AtomicU64 = AtomicU64::new(0);
320        let path = std::env::temp_dir().join(format!(
321            "rlx_llama32_vocab_{}_{}_{}.gguf",
322            std::process::id(),
323            SEQ.fetch_add(1, Ordering::Relaxed),
324            std::time::SystemTime::now()
325                .duration_since(std::time::UNIX_EPOCH)
326                .unwrap()
327                .as_nanos()
328        ));
329
330        let mut buf: Vec<u8> = Vec::new();
331        buf.extend_from_slice(&rlx_gguf::GGUF_MAGIC.to_le_bytes());
332        buf.extend_from_slice(&3u32.to_le_bytes());
333        buf.extend_from_slice(&2u64.to_le_bytes()); // 2 tensors
334        buf.extend_from_slice(&9u64.to_le_bytes()); // metadata keys
335
336        let write_str = |buf: &mut Vec<u8>, k: &str, v: &str| {
337            buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
338            buf.extend_from_slice(k.as_bytes());
339            buf.extend_from_slice(&8u32.to_le_bytes());
340            buf.extend_from_slice(&(v.len() as u64).to_le_bytes());
341            buf.extend_from_slice(v.as_bytes());
342        };
343        let write_u32 = |buf: &mut Vec<u8>, k: &str, v: u32| {
344            buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
345            buf.extend_from_slice(k.as_bytes());
346            buf.extend_from_slice(&4u32.to_le_bytes());
347            buf.extend_from_slice(&v.to_le_bytes());
348        };
349        let write_string_array = |buf: &mut Vec<u8>, k: &str, items: &[String]| {
350            buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
351            buf.extend_from_slice(k.as_bytes());
352            buf.extend_from_slice(&9u32.to_le_bytes());
353            buf.extend_from_slice(&8u32.to_le_bytes());
354            buf.extend_from_slice(&(items.len() as u64).to_le_bytes());
355            for s in items {
356                buf.extend_from_slice(&(s.len() as u64).to_le_bytes());
357                buf.extend_from_slice(s.as_bytes());
358            }
359        };
360
361        write_str(&mut buf, "general.architecture", "llama");
362        write_u32(&mut buf, "llama.embedding_length", 2048);
363        write_u32(&mut buf, "llama.feed_forward_length", 5632);
364        write_u32(&mut buf, "llama.block_count", 22);
365        write_u32(&mut buf, "llama.attention.head_count", 32);
366        write_u32(&mut buf, "llama.attention.head_count_kv", 4);
367        write_u32(&mut buf, "llama.context_length", 2048);
368        write_u32(&mut buf, "llama.rope.freq_base", 10_000);
369        let vocab = 128u32;
370        let tokens: Vec<String> = (0..vocab).map(|i| format!("t{i}")).collect();
371        write_string_array(&mut buf, "tokenizer.ggml.tokens", &tokens);
372
373        let embed_bytes = vocab as u64 * 2048 * 4;
374        for (name, rows, cols, offset) in [
375            ("token_embd.weight", vocab as u64, 2048u64, 0u64),
376            ("output.weight", 2048u64, vocab as u64, embed_bytes),
377        ] {
378            buf.extend_from_slice(&(name.len() as u64).to_le_bytes());
379            buf.extend_from_slice(name.as_bytes());
380            buf.extend_from_slice(&2u32.to_le_bytes());
381            buf.extend_from_slice(&rows.to_le_bytes());
382            buf.extend_from_slice(&cols.to_le_bytes());
383            buf.extend_from_slice(&(GgmlType::F32 as u32).to_le_bytes());
384            buf.extend_from_slice(&offset.to_le_bytes());
385        }
386        while !buf
387            .len()
388            .is_multiple_of(rlx_gguf::DEFAULT_ALIGNMENT as usize)
389        {
390            buf.push(0);
391        }
392        let n_floats = (vocab as usize * 2048) * 2;
393        for _ in 0..n_floats {
394            buf.extend_from_slice(&0f32.to_le_bytes());
395        }
396        std::fs::write(&path, &buf).unwrap();
397
398        let raw = rlx_gguf::GgufFile::from_path(&path).expect("parse tinyllama-like gguf");
399        let cfg = llama32_cfg_from_gguf(&raw).expect("llama32 config");
400        assert_eq!(cfg.vocab_size, vocab as usize);
401        assert!(!cfg.tie_word_embeddings);
402        std::fs::remove_file(path).ok();
403    }
404}
rlx_llama32/config.rs

rlx_llama32/
config.rs
