rlx_llama32/
runner.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
16use crate::prefill_mode::MetalGgufPrefillMode;
17use crate::{Llama32Config, Llama32Generator, llama32_cfg_from_gguf};
18use anyhow::{Context, Result, anyhow, bail};
19use rlx_cli::{LmRunner, WeightFormat};
20use rlx_gguf::{GgufFile, MetaValue};
21use rlx_qwen3::SampleOpts;
22use rlx_runtime::Device;
23use std::path::{Path, PathBuf};
24
25// ────────────────────────────────────────────────────────────────
26// LLaMA-3.2 runner — Meta Llama 3.x small LMs (1B / 3B).
27// ────────────────────────────────────────────────────────────────
28
29/// Where to load the Llama 3.2 config from.
30///
31/// Type alias of the shared `rlx_runtime::ConfigSource<T>` so the
32/// per-family `*ConfigSource` enums no longer duplicate the same
33/// `Embedded | JsonFile | Explicit(T)` shape. The variant constructors
34/// (`Llama32ConfigSource::Embedded`, etc.) keep working because
35/// type-alias resolution expands the path through the generic enum.
36pub type Llama32ConfigSource = rlx_runtime::ConfigSource<Llama32Config>;
37
38#[derive(Debug, Clone)]
39pub struct Llama32RunnerBuilder {
40    weights: Option<PathBuf>,
41    config: Option<Llama32ConfigSource>,
42    device: Option<Device>,
43    max_seq: Option<usize>,
44    max_memory_gb: Option<f32>,
45    stream: bool,
46    sample: Option<SampleOpts>,
47    format: Option<WeightFormat>,
48    /// `None` = auto (packed when GGUF ≥ 256 MB). `Some(_)` is an explicit override.
49    packed_weights: Option<bool>,
50    /// When false, decode uses one-shot graphs (slower compile, but
51    /// avoids bucketed-cache edge cases on some GPU backends).
52    bucketed_decode_cache: bool,
53}
54
55impl Default for Llama32RunnerBuilder {
56    fn default() -> Self {
57        Self {
58            weights: None,
59            config: None,
60            device: None,
61            max_seq: None,
62            max_memory_gb: None,
63            stream: true,
64            sample: None,
65            format: None,
66            packed_weights: None,
67            bucketed_decode_cache: true,
68        }
69    }
70}
71
72impl Llama32RunnerBuilder {
73    pub fn weights<P: Into<PathBuf>>(mut self, path: P) -> Self {
74        self.weights = Some(path.into());
75        self
76    }
77
78    pub fn format(mut self, fmt: WeightFormat) -> Self {
79        self.format = Some(fmt);
80        self
81    }
82
83    pub fn config(mut self, src: Llama32ConfigSource) -> Self {
84        self.config = Some(src);
85        self
86    }
87
88    pub fn config_value(self, cfg: Llama32Config) -> Self {
89        self.config(Llama32ConfigSource::Explicit(cfg))
90    }
91
92    pub fn device(mut self, d: Device) -> Self {
93        self.device = Some(d);
94        self
95    }
96
97    pub fn max_seq(mut self, n: usize) -> Self {
98        self.max_seq = Some(n);
99        self
100    }
101
102    pub fn max_memory_gb(mut self, gb: f32) -> Self {
103        self.max_memory_gb = Some(gb);
104        self
105    }
106
107    pub fn stream(mut self, on: bool) -> Self {
108        self.stream = on;
109        self
110    }
111
112    pub fn sample(mut self, opts: SampleOpts) -> Self {
113        self.sample = Some(opts);
114        self
115    }
116
117    /// Keep K-quant weights packed in the arena (`Op::DequantMatMul`).
118    /// GGUF only. Supported on CPU, Metal, and MLX.
119    ///
120    /// When unset, large GGUF files (≥ 256 MB on disk) auto-enable packed
121    /// prefill to avoid F32-dequant host memory blowups.
122    pub fn packed_weights(mut self, on: bool) -> Self {
123        self.packed_weights = Some(on);
124        self
125    }
126
127    /// Enable bucketed decode compile cache (default: true).
128    pub fn bucketed_decode_cache(mut self, on: bool) -> Self {
129        self.bucketed_decode_cache = on;
130        self
131    }
132
133    pub fn build(self) -> Result<Llama32Runner> {
134        let weights_path = self
135            .weights
136            .ok_or_else(|| anyhow!("weights path required (call .weights(...))"))?;
137        let format = match self.format {
138            Some(f) => f,
139            None => WeightFormat::from_path(&weights_path)?,
140        };
141        let device = self.device.unwrap_or(Device::Cpu);
142        let max_seq = self.max_seq.unwrap_or(128);
143        let stream = self.stream;
144        let sample = self.sample.unwrap_or_else(SampleOpts::greedy);
145
146        let (cfg, total_bytes_estimate) = match format {
147            WeightFormat::Gguf => load_llama32_gguf_config(&weights_path, self.config.as_ref())?,
148            WeightFormat::Safetensors => {
149                load_llama32_safetensors_config(&weights_path, self.config.as_ref())?
150            }
151        };
152
153        if let Some(cap_gb) = self.max_memory_gb {
154            let est_gb = total_bytes_estimate as f32 / (1024.0 * 1024.0 * 1024.0);
155            if est_gb > cap_gb {
156                bail!(
157                    "weights would dequant to ~{est_gb:.1} GB at F32, exceeds cap {cap_gb:.1} GB"
158                );
159            }
160        }
161
162        let use_packed = self.packed_weights.unwrap_or_else(|| {
163            matches!(format, WeightFormat::Gguf)
164                && std::fs::metadata(&weights_path)
165                    .map(|m| m.len() >= 256 * 1024 * 1024)
166                    .unwrap_or(false)
167        });
168
169        crate::validate_device(&cfg, device, use_packed)?;
170
171        if use_packed && !matches!(format, WeightFormat::Gguf) {
172            bail!(
173                "packed_weights(true) requires a .gguf file; got {:?} for {:?}",
174                format,
175                weights_path
176            );
177        }
178
179        let prefill_mode = if use_packed {
180            if matches!(
181                device,
182                Device::Metal | Device::Cuda | Device::Rocm | Device::Mlx
183            ) {
184                MetalGgufPrefillMode::PackedGguf
185            } else {
186                MetalGgufPrefillMode::CpuF32
187            }
188        } else {
189            MetalGgufPrefillMode::Auto
190        };
191
192        if use_packed {
193            eprintln!(
194                "[llama32-runner] packed_weights=true — Q4 prefill + bucketed decode on {device:?}"
195            );
196        }
197
198        let path_str = weights_path
199            .to_str()
200            .ok_or_else(|| anyhow!("non-utf8 weights path"))?;
201        let mut loader = rlx_core::weight_loader::load_from_path(path_str)?;
202        let mut generator = Llama32Generator::from_loader_at_mode(
203            cfg.clone(),
204            loader.as_mut(),
205            device,
206            &weights_path,
207            prefill_mode,
208        )?
209        .with_compile_seq_cap(max_seq)
210        .with_prefill_cache(8);
211        if self.bucketed_decode_cache {
212            generator = generator.with_decode_cache(max_seq.saturating_add(16).max(64));
213        }
214
215        Ok(Llama32Runner {
216            generator,
217            cfg,
218            sample,
219            stream,
220            device,
221            packed_weights: use_packed,
222        })
223    }
224}
225
226pub struct Llama32Runner {
227    generator: Llama32Generator,
228    cfg: Llama32Config,
229    sample: SampleOpts,
230    stream: bool,
231    device: Device,
232    packed_weights: bool,
233}
234
235impl Llama32Runner {
236    pub fn builder() -> Llama32RunnerBuilder {
237        Llama32RunnerBuilder::default()
238    }
239
240    pub fn config(&self) -> &Llama32Config {
241        &self.cfg
242    }
243
244    pub fn device(&self) -> Device {
245        self.device
246    }
247
248    /// Current sampling options used by [`generate`](Self::generate) and
249    /// [`generate_until`](Self::generate_until).
250    pub fn sample_opts(&self) -> &SampleOpts {
251        &self.sample
252    }
253
254    /// Override the sampling options for subsequent generations. Lets a
255    /// long-lived runner switch between greedy / temperature / top-p on a
256    /// per-request basis without rebuilding (weights + compile caches stay
257    /// warm).
258    pub fn set_sample(&mut self, opts: SampleOpts) {
259        self.sample = opts;
260    }
261
262    /// Single prefill forward; returns last-position logits `[vocab]`.
263    pub fn predict_logits(&mut self, prompt_ids: &[u32]) -> Result<Vec<f32>> {
264        self.generator.prefill_get_last_logits(prompt_ids)
265    }
266
267    pub fn generate_packed(
268        &mut self,
269        prompt_ids: &[u32],
270        n_new: usize,
271        on_token: impl FnMut(u32),
272    ) -> Result<Vec<u32>> {
273        if !self.packed_weights {
274            bail!("generate_packed() only works in packed_weights(true) mode");
275        }
276        self.generate(prompt_ids, n_new, on_token)
277    }
278
279    pub fn generate(
280        &mut self,
281        prompt_ids: &[u32],
282        n_new: usize,
283        mut on_token: impl FnMut(u32),
284    ) -> Result<Vec<u32>> {
285        self.generator.prefill(prompt_ids);
286        let tokens = if self.stream {
287            self.generator
288                .generate_cached_with(n_new, self.sample, &mut on_token)?
289        } else {
290            let toks = self.generator.generate_cached(n_new, self.sample)?;
291            for &t in &toks {
292                on_token(t);
293            }
294            toks
295        };
296        Ok(tokens)
297    }
298
299    /// KV-cached generation with an early stop. `keep_going` is called with
300    /// each freshly sampled id and returns whether to continue; returning
301    /// `false` halts generation *after* that token (which is included in the
302    /// result). Use it to stop on an end-of-sequence id instead of always
303    /// decoding the full `n_new` budget.
304    pub fn generate_until(
305        &mut self,
306        prompt_ids: &[u32],
307        n_new: usize,
308        keep_going: impl FnMut(u32) -> bool,
309    ) -> Result<Vec<u32>> {
310        self.generator.prefill(prompt_ids);
311        self.generator
312            .generate_cached_until(n_new, self.sample, keep_going)
313    }
314}
315
316impl LmRunner for Llama32Runner {
317    fn family(&self) -> &'static str {
318        "llama32"
319    }
320    fn vocab_size(&self) -> usize {
321        self.config().vocab_size
322    }
323    fn predict_logits(&mut self, prompt_ids: &[u32]) -> Result<Vec<f32>> {
324        Llama32Runner::predict_logits(self, prompt_ids)
325    }
326    fn generate(
327        &mut self,
328        prompt_ids: &[u32],
329        n_new: usize,
330        on_token: &mut dyn FnMut(u32) -> bool,
331    ) -> Result<Vec<u32>> {
332        Llama32Runner::generate(self, prompt_ids, n_new, |tok| {
333            let _ = on_token(tok);
334        })
335    }
336}
337
338fn load_llama32_gguf_config(
339    path: &Path,
340    override_src: Option<&Llama32ConfigSource>,
341) -> Result<(Llama32Config, u64)> {
342    let raw = GgufFile::from_path(path).with_context(|| format!("opening {path:?}"))?;
343    let arch = raw
344        .metadata
345        .get("general.architecture")
346        .and_then(MetaValue::as_str)
347        .unwrap_or("llama");
348    const LLAMA_SHAPED_GGUF_ARCHES: &[&str] = &["llama", "phi3", "phi4"];
349    if !LLAMA_SHAPED_GGUF_ARCHES.contains(&arch) {
350        bail!(
351            "{path:?} has architecture {arch:?}; Llama32Runner expects general.architecture ∈ {LLAMA_SHAPED_GGUF_ARCHES:?}"
352        );
353    }
354    let cfg = match override_src {
355        Some(Llama32ConfigSource::Explicit(c)) => c.clone(),
356        Some(Llama32ConfigSource::JsonFile(p)) => {
357            Llama32Config::from_file(p).with_context(|| format!("reading override config {p:?}"))?
358        }
359        Some(Llama32ConfigSource::Embedded) | None => llama32_cfg_from_gguf(&raw)?,
360    };
361    let bytes_est: u64 = raw
362        .tensors
363        .values()
364        .map(|t| (t.n_elements() as u64) * 4)
365        .sum();
366    Ok((cfg, bytes_est))
367}
368
369fn load_llama32_safetensors_config(
370    path: &Path,
371    override_src: Option<&Llama32ConfigSource>,
372) -> Result<(Llama32Config, u64)> {
373    let cfg_path = match override_src {
374        Some(Llama32ConfigSource::Explicit(c)) => {
375            return Ok((c.clone(), default_st_size_estimate(path)));
376        }
377        Some(Llama32ConfigSource::JsonFile(p)) => p.clone(),
378        Some(Llama32ConfigSource::Embedded) => {
379            bail!("ConfigSource::Embedded only valid for GGUF; pass JsonFile for safetensors")
380        }
381        None => path
382            .parent()
383            .ok_or_else(|| anyhow!("weights path has no parent dir"))?
384            .join("config.json"),
385    };
386    let cfg = Llama32Config::from_file(&cfg_path)
387        .with_context(|| format!("reading config {cfg_path:?}"))?;
388    Ok((cfg, default_st_size_estimate(path)))
389}
390
391fn default_st_size_estimate(path: &Path) -> u64 {
392    std::fs::metadata(path).map(|m| m.len()).unwrap_or(0)
393}
rlx_llama32/runner.rs

rlx_llama32/
runner.rs
