rlx_gemma/
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::{GemmaConfig, GemmaGenerator, gemma_cfg_from_gguf};
17use anyhow::{Context, Result, anyhow, bail};
18use rlx_cli::{LmRunner, WeightFormat};
19use rlx_core::gguf_support::{
20    GgufModelFamily, ResolveWeightsOptions, assert_gguf_family, gguf_f32_bytes_estimate,
21    resolve_weights_file_with_options,
22};
23use rlx_qwen3::SampleOpts;
24use rlx_runtime::Device;
25use std::path::{Path, PathBuf};
26
27// ────────────────────────────────────────────────────────────────
28// Gemma runner — Meta Llama 3.x small LMs (1B / 3B).
29// ────────────────────────────────────────────────────────────────
30
31/// Where to load the Gemma config from. Alias of the shared
32/// `rlx_runtime::ConfigSource<T>` — same `Embedded | JsonFile | Explicit(T)`
33/// shape as the other family `*ConfigSource` enums.
34pub type GemmaConfigSource = rlx_runtime::ConfigSource<GemmaConfig>;
35
36#[derive(Debug, Clone, Default)]
37pub struct GemmaRunnerBuilder {
38    weights: Option<PathBuf>,
39    config: Option<GemmaConfigSource>,
40    device: Option<Device>,
41    max_seq: Option<usize>,
42    max_memory_gb: Option<f32>,
43    stream: bool,
44    sample: Option<SampleOpts>,
45    format: Option<WeightFormat>,
46    packed_weights: bool,
47}
48
49impl GemmaRunnerBuilder {
50    pub fn weights<P: Into<PathBuf>>(mut self, path: P) -> Self {
51        self.weights = Some(path.into());
52        self
53    }
54
55    pub fn format(mut self, fmt: WeightFormat) -> Self {
56        self.format = Some(fmt);
57        self
58    }
59
60    pub fn config(mut self, src: GemmaConfigSource) -> Self {
61        self.config = Some(src);
62        self
63    }
64
65    pub fn config_value(self, cfg: GemmaConfig) -> Self {
66        self.config(GemmaConfigSource::Explicit(cfg))
67    }
68
69    pub fn device(mut self, d: Device) -> Self {
70        self.device = Some(d);
71        self
72    }
73
74    pub fn max_seq(mut self, n: usize) -> Self {
75        self.max_seq = Some(n);
76        self
77    }
78
79    pub fn max_memory_gb(mut self, gb: f32) -> Self {
80        self.max_memory_gb = Some(gb);
81        self
82    }
83
84    pub fn stream(mut self, on: bool) -> Self {
85        self.stream = on;
86        self
87    }
88
89    pub fn sample(mut self, opts: SampleOpts) -> Self {
90        self.sample = Some(opts);
91        self
92    }
93
94    /// Keep K-quant weights packed in the arena (`Op::DequantMatMul`).
95    /// GGUF only. Uses `Op::DequantMatMul` on the selected device.
96    pub fn packed_weights(mut self, on: bool) -> Self {
97        self.packed_weights = on;
98        self
99    }
100
101    pub fn build(self) -> Result<GemmaRunner> {
102        let resolve = ResolveWeightsOptions {
103            prefer_gguf_substring: Some(rlx_core::DEFAULT_GGUF_PREFER_SUBSTR),
104            ..Default::default()
105        };
106        let weights_path = resolve_weights_file_with_options(
107            self.weights
108                .as_ref()
109                .ok_or_else(|| anyhow!("weights path required (call .weights(...))"))?,
110            &resolve,
111        )?;
112        let format = WeightFormat::resolve(&weights_path, self.format)?;
113        let device = self.device.unwrap_or(Device::Cpu);
114        let max_seq = self.max_seq.unwrap_or(128);
115        let stream = self.stream;
116        let sample = self.sample.unwrap_or_else(SampleOpts::greedy);
117
118        let (cfg, total_bytes_estimate) = match format {
119            WeightFormat::Gguf => load_gemma_gguf_config(&weights_path, self.config.as_ref())?,
120            WeightFormat::Safetensors => {
121                load_gemma_safetensors_config(&weights_path, self.config.as_ref())?
122            }
123        };
124
125        if let Some(cap_gb) = self.max_memory_gb {
126            let est_gb = total_bytes_estimate as f32 / (1024.0 * 1024.0 * 1024.0);
127            if est_gb > cap_gb {
128                bail!(
129                    "weights would dequant to ~{est_gb:.1} GB at F32, exceeds cap {cap_gb:.1} GB"
130                );
131            }
132        }
133
134        crate::capabilities::validate_device(&cfg, device, self.packed_weights)?;
135
136        let path_str = weights_path
137            .to_str()
138            .ok_or_else(|| anyhow!("non-utf8 weights path"))?;
139        let generator = if self.packed_weights {
140            None
141        } else {
142            Some(
143                GemmaGenerator::from_path(cfg.clone(), path_str, device)?
144                    .with_inference_caches(max_seq),
145            )
146        };
147
148        let packed = if self.packed_weights {
149            if !matches!(format, WeightFormat::Gguf) {
150                bail!(
151                    "packed_weights(true) requires a .gguf file; got {:?} for {:?}",
152                    format,
153                    weights_path
154                );
155            }
156            eprintln!(
157                "[gemma-runner] packed_weights=true — Q4 prefill + bucketed decode on {device:?}"
158            );
159            Some(crate::packed_session::GemmaPackedSession::build(
160                cfg.clone(),
161                &weights_path,
162                max_seq,
163                device,
164            )?)
165        } else {
166            None
167        };
168
169        Ok(GemmaRunner {
170            generator,
171            cfg,
172            sample,
173            stream,
174            device,
175            packed,
176        })
177    }
178}
179
180pub struct GemmaRunner {
181    generator: Option<GemmaGenerator>,
182    cfg: GemmaConfig,
183    sample: SampleOpts,
184    stream: bool,
185    device: Device,
186    packed: Option<crate::packed_session::GemmaPackedSession>,
187}
188
189impl GemmaRunner {
190    pub fn builder() -> GemmaRunnerBuilder {
191        GemmaRunnerBuilder::default()
192    }
193
194    pub fn config(&self) -> &GemmaConfig {
195        &self.cfg
196    }
197
198    pub fn device(&self) -> Device {
199        self.device
200    }
201
202    /// Single prefill forward; returns last-position logits `[vocab]`.
203    pub fn predict_logits(&mut self, prompt_ids: &[u32]) -> Result<Vec<f32>> {
204        if let Some(p) = self.packed.as_mut() {
205            return p.predict_logits(prompt_ids);
206        }
207        let generator = self
208            .generator
209            .as_mut()
210            .ok_or_else(|| anyhow!("F32 generator unavailable in packed_weights mode"))?;
211        generator.prefill_get_last_logits(prompt_ids)
212    }
213
214    /// Post-`model.norm` hidden for the last prompt token (packed GGUF path only).
215    pub fn predict_last_hidden(&mut self, prompt_ids: &[u32]) -> Result<Vec<f32>> {
216        self.packed
217            .as_mut()
218            .ok_or_else(|| anyhow!("predict_last_hidden requires packed_weights(true)"))?
219            .predict_last_hidden(prompt_ids)
220    }
221
222    pub fn generate_packed(
223        &mut self,
224        prompt_ids: &[u32],
225        n_new: usize,
226        on_token: impl FnMut(u32),
227    ) -> Result<Vec<u32>> {
228        if self.packed.is_none() {
229            bail!("generate_packed() only works in packed_weights(true) mode");
230        }
231        let sample = self.sample;
232        self.packed
233            .as_mut()
234            .unwrap()
235            .generate(prompt_ids, n_new, sample, on_token)
236    }
237
238    pub fn generate(
239        &mut self,
240        prompt_ids: &[u32],
241        n_new: usize,
242        mut on_token: impl FnMut(u32),
243    ) -> Result<Vec<u32>> {
244        if self.packed.is_some() {
245            return self.generate_packed(prompt_ids, n_new, on_token);
246        }
247        let generator = self
248            .generator
249            .as_mut()
250            .ok_or_else(|| anyhow!("F32 generator unavailable in packed_weights mode"))?;
251        generator.prefill(prompt_ids);
252        let tokens = if self.stream {
253            generator.generate_cached_with(n_new, self.sample, &mut on_token)?
254        } else {
255            let toks = generator.generate_cached(n_new, self.sample)?;
256            for &t in &toks {
257                on_token(t);
258            }
259            toks
260        };
261        Ok(tokens)
262    }
263
264    /// EAGLE3-style generate with a per-step **aux hidden state**
265    /// callback. After each decoded token, `on_aux` receives the
266    /// per-layer pre-attention-norm hidden states from the layer
267    /// indices in `aux_hidden_layer_ids` (one `Vec<f32>` per id,
268    /// each of length `hidden_size`).
269    ///
270    /// This is the entry point the rlx-eagle3 verifier bridge
271    /// (`AuxStateBuffer::write`) consumes — wire it like:
272    ///
273    /// ```ignore
274    /// runner.generate_with_aux(prompt, n, vec![2, 30, 57],
275    ///     |tok| { ... },
276    ///     |aux| { aux_buffer.write(aux); });
277    /// ```
278    ///
279    /// **Not supported in `packed_weights(true)` mode** — packed
280    /// builds use a separate compile path that doesn't yet thread
281    /// the aux tap.
282    pub fn generate_with_aux(
283        &mut self,
284        prompt_ids: &[u32],
285        n_new: usize,
286        aux_hidden_layer_ids: Vec<usize>,
287        mut on_token: impl FnMut(u32),
288        mut on_aux: impl FnMut(Vec<Vec<f32>>),
289    ) -> Result<Vec<u32>> {
290        if self.packed.is_some() {
291            bail!("generate_with_aux is not supported with packed_weights(true)");
292        }
293        if aux_hidden_layer_ids.is_empty() {
294            // No aux ids → falls back to plain generate.
295            return self.generate(prompt_ids, n_new, on_token);
296        }
297        let generator = self
298            .generator
299            .as_mut()
300            .ok_or_else(|| anyhow!("F32 generator unavailable in packed_weights mode"))?;
301        generator.set_aux_hidden_layer_ids(aux_hidden_layer_ids);
302        generator.prefill(prompt_ids);
303        let sample = self.sample;
304        let mut emitted: Vec<u32> = Vec::with_capacity(n_new);
305        for _ in 0..n_new {
306            let tok = generator.step_cached(sample)?;
307            emitted.push(tok);
308            on_token(tok);
309            if let Some(aux) = generator.take_last_aux() {
310                on_aux(aux);
311            }
312        }
313        // Clear the tap so subsequent generate() calls go through
314        // the bucketed fast path again.
315        generator.set_aux_hidden_layer_ids(Vec::new());
316        Ok(emitted)
317    }
318
319    /// Generate after splicing vision/audio rows into pre-scaled text embeddings.
320    pub fn generate_from_embeds(
321        &mut self,
322        prompt_ids: &[u32],
323        inputs_embeds: &[f32],
324        n_new: usize,
325        mut on_token: impl FnMut(u32),
326    ) -> Result<Vec<u32>> {
327        if self.packed.is_some() {
328            bail!("generate_from_embeds is not supported with packed_weights(true)");
329        }
330        let generator = self
331            .generator
332            .as_mut()
333            .ok_or_else(|| anyhow!("F32 generator unavailable in packed_weights mode"))?;
334        let tokens = if self.stream {
335            generator.generate_from_embeds_with(
336                prompt_ids,
337                inputs_embeds,
338                n_new,
339                self.sample,
340                &mut on_token,
341            )?
342        } else {
343            let toks =
344                generator.generate_from_embeds(prompt_ids, inputs_embeds, n_new, self.sample)?;
345            for &t in &toks {
346                on_token(t);
347            }
348            toks
349        };
350        Ok(tokens)
351    }
352
353    /// Build fused inputs + run generation (text LM weights must include `model.embed_tokens.weight`).
354    pub fn generate_multimodal(
355        &mut self,
356        mm_cfg: &crate::multimodal::GemmaMultimodalConfig,
357        token_ids: &[u32],
358        image_embeds: &[f32],
359        audio_embeds: &[f32],
360        video_embeds: &[f32],
361        n_new: usize,
362        mut on_token: impl FnMut(u32),
363    ) -> Result<Vec<u32>> {
364        let generator = self
365            .generator
366            .as_ref()
367            .ok_or_else(|| anyhow!("F32 generator unavailable in packed_weights mode"))?;
368        let embeds = crate::multimodal_embed::build_multimodal_inputs_embeds(
369            generator.weights_cache(),
370            &self.cfg,
371            mm_cfg,
372            token_ids,
373            image_embeds,
374            audio_embeds,
375            video_embeds,
376        )?;
377        let attn_bias = crate::multimodal_mask::build_multimodal_prefill_attn_bias(
378            token_ids, &self.cfg, mm_cfg, 1,
379        );
380        let generator = self
381            .generator
382            .as_mut()
383            .ok_or_else(|| anyhow!("F32 generator unavailable in packed_weights mode"))?;
384        let tokens = if self.stream {
385            generator.generate_from_embeds_with_bias_and_callback(
386                token_ids,
387                &embeds,
388                attn_bias,
389                n_new,
390                self.sample,
391                &mut on_token,
392            )?
393        } else {
394            let toks = generator.generate_from_embeds_with_bias(
395                token_ids,
396                &embeds,
397                attn_bias,
398                n_new,
399                self.sample,
400            )?;
401            for &t in &toks {
402                on_token(t);
403            }
404            toks
405        };
406        Ok(tokens)
407    }
408}
409
410impl LmRunner for GemmaRunner {
411    fn family(&self) -> &'static str {
412        "gemma"
413    }
414    fn vocab_size(&self) -> usize {
415        self.config().vocab_size
416    }
417    fn predict_logits(&mut self, prompt_ids: &[u32]) -> Result<Vec<f32>> {
418        GemmaRunner::predict_logits(self, prompt_ids)
419    }
420    fn generate(
421        &mut self,
422        prompt_ids: &[u32],
423        n_new: usize,
424        on_token: &mut dyn FnMut(u32) -> bool,
425    ) -> Result<Vec<u32>> {
426        // Inherent generate ignores stop signal — drop the bool.
427        GemmaRunner::generate(self, prompt_ids, n_new, |tok| {
428            let _ = on_token(tok);
429        })
430    }
431}
432
433fn load_gemma_gguf_config(
434    path: &Path,
435    override_src: Option<&GemmaConfigSource>,
436) -> Result<(GemmaConfig, u64)> {
437    let raw = assert_gguf_family(path, GgufModelFamily::Gemma)?;
438    let cfg = match override_src {
439        Some(GemmaConfigSource::Explicit(c)) => c.clone(),
440        Some(GemmaConfigSource::JsonFile(p)) => {
441            GemmaConfig::from_file(p).with_context(|| format!("reading override config {p:?}"))?
442        }
443        Some(GemmaConfigSource::Embedded) | None => gemma_cfg_from_gguf(&raw)?,
444    };
445    Ok((cfg, gguf_f32_bytes_estimate(&raw)))
446}
447
448fn load_gemma_safetensors_config(
449    path: &Path,
450    override_src: Option<&GemmaConfigSource>,
451) -> Result<(GemmaConfig, u64)> {
452    let cfg_path = match override_src {
453        Some(GemmaConfigSource::Explicit(c)) => {
454            return Ok((c.clone(), default_st_size_estimate(path)));
455        }
456        Some(GemmaConfigSource::JsonFile(p)) => p.clone(),
457        Some(GemmaConfigSource::Embedded) => {
458            bail!("ConfigSource::Embedded only valid for GGUF; pass JsonFile for safetensors")
459        }
460        None => path
461            .parent()
462            .ok_or_else(|| anyhow!("weights path has no parent dir"))?
463            .join("config.json"),
464    };
465    let cfg = GemmaConfig::from_file(&cfg_path)
466        .with_context(|| format!("reading config {cfg_path:?}"))?;
467    Ok((cfg, default_st_size_estimate(path)))
468}
469
470fn default_st_size_estimate(path: &Path) -> u64 {
471    std::fs::metadata(path).map(|m| m.len()).unwrap_or(0)
472}
rlx_gemma/runner.rs

rlx_gemma/
runner.rs
