rlx_pkg/
from_gguf.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Import a GGUF weight file into an `.rlxp` package.
6//!
7//! Defaults to **weight-only** (`include_graph: false`). Pass
8//! `include_graph: true` for a stub graph; for a real executable graph from
9//! ONNX, use `rlx-bake --features onnx -- import-onnx`.
10
11use crate::tier::StorageTier;
12use crate::write::{ContainerKind, PackedWeight, WriteOptions, write_package};
13use anyhow::{Context, Result};
14use rlx_gguf::{GgmlType, GgufFile, MetaValue};
15use rlx_ir::{DType, Graph, Shape};
16use std::path::Path;
17
18/// Options for [`gguf_to_rlxp`].
19#[derive(Debug, Clone)]
20pub struct GgufImportOptions {
21    pub container: ContainerKind,
22    pub include_graph: bool,
23    pub compress_sidecars: bool,
24    pub auto_tier: bool,
25}
26
27impl Default for GgufImportOptions {
28    fn default() -> Self {
29        Self {
30            container: ContainerKind::Flat,
31            include_graph: false,
32            compress_sidecars: true,
33            auto_tier: true,
34        }
35    }
36}
37
38fn scheme_for(dtype: GgmlType) -> String {
39    match dtype {
40        GgmlType::F32 => "f32".into(),
41        GgmlType::F16 => "f16".into(),
42        GgmlType::BF16 => "bf16".into(),
43        GgmlType::Q4K => "gguf_q4_k".into(),
44        GgmlType::Q5K => "gguf_q5_k".into(),
45        GgmlType::Q6K => "gguf_q6_k".into(),
46        GgmlType::Q8K => "gguf_q8_k".into(),
47        GgmlType::Q8_0 => "gguf_q8_0".into(),
48        GgmlType::Q4_0 => "gguf_q4_0".into(),
49        GgmlType::Q4_1 => "gguf_q4_1".into(),
50        GgmlType::Q2K => "gguf_q2_k".into(),
51        GgmlType::Q3K => "gguf_q3_k".into(),
52        GgmlType::TQ1_0 => "gguf_tq1_0".into(),
53        GgmlType::TQ2_0 => "gguf_tq2_0".into(),
54        other => format!("gguf_{}", format!("{other:?}").to_ascii_lowercase()),
55    }
56}
57
58/// Convert `model.gguf` → `.rlxp` (weights + optional stub graph).
59pub fn gguf_to_rlxp(
60    gguf_path: impl AsRef<Path>,
61    out: impl AsRef<Path>,
62    opts: &GgufImportOptions,
63) -> Result<()> {
64    let gguf_path = gguf_path.as_ref();
65    let f = GgufFile::from_path_mmap(gguf_path)
66        .with_context(|| format!("open gguf {}", gguf_path.display()))?;
67
68    let mut weights = Vec::new();
69    let names: Vec<String> = f.keys().map(|s| s.to_string()).collect();
70    for name in names {
71        let t = f.get(&name).with_context(|| format!("tensor {name}"))?;
72        let bytes = f.tensor_bytes(t)?.to_vec();
73        weights.push(PackedWeight {
74            name: name.clone(),
75            shape: t.shape.clone(),
76            scheme: scheme_for(t.dtype),
77            layout: "bt_nk".into(),
78            data: bytes,
79            rank: None,
80            tier: StorageTier::Hot,
81        });
82    }
83
84    let mut wopts = WriteOptions {
85        name: gguf_path
86            .file_stem()
87            .and_then(|s| s.to_str())
88            .unwrap_or("gguf")
89            .into(),
90        producer: Some("rlx-pkg/gguf_import".into()),
91        container: opts.container,
92        compress_sidecars: opts.compress_sidecars,
93        include_graph: opts.include_graph,
94        write_checksums: true,
95        ..WriteOptions::default()
96    };
97
98    if opts.auto_tier {
99        crate::auto_tier::apply_auto_tier(&mut weights, &Default::default());
100    }
101
102    if let Some(MetaValue::String(arch)) = f.metadata.get("general.architecture") {
103        wopts.sidecars.push((
104            "gguf_architecture".into(),
105            "text/plain".into(),
106            arch.as_bytes().to_vec(),
107        ));
108    }
109
110    let graph = if opts.include_graph {
111        let mut g = Graph::new(&wopts.name);
112        let s = Shape::new(&[1], DType::F32);
113        let x = g.input("x", s);
114        g.set_outputs(vec![x]);
115        g
116    } else {
117        Graph::new(&wopts.name)
118    };
119
120    write_package(out, &graph, &weights, &wopts)
121}
rlx_pkg/from_gguf.rs

rlx_pkg/
from_gguf.rs
