spark_runtime/weights/
gguf.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Generic GGUF weight loader for Atlas.
4//!
5//! Loads any GGUF checkpoint the same way [`super::SafetensorsLoader`] loads
6//! safetensors: mmap the file, walk its tensors, land each one GPU-resident in
7//! a [`WeightStore`] keyed by HuggingFace name. Unlike safetensors, GGUF tensors
8//! are quantized in-file, so each tensor is *dequantized to BF16* on the way in.
9//!
10//! Per the project design directive we **prefer GPU dequant**: upload the raw
11//! quantized GGUF block bytes h2d, run a device dequant kernel that writes BF16
12//! into fresh device memory, and hand that BF16 [`WeightTensor`] to the store.
13//! The per-arch model loaders do the downstream NVFP4 requantize — this loader's
14//! job ends at clean BF16. A pure-CPU reference dequant is the fallback for ggml
15//! types lacking a GPU kernel and the correctness oracle under `MockGpuBackend`
16//! (which cannot execute kernels).
17//!
18//! GGUF `dims` are ggml-order (fastest-varying first); Atlas/HF shapes are the
19//! reverse, so each tensor's shape is reversed before it enters the store.
20//!
21//! The PrismML `Q2_0` (id 42) group size is not encoded in the type id. It
22//! defaults to group-128 (the shipped Ternary-Bonsai layout); set
23//! `ATLAS_GGUF_Q2_GROUP=64` for the fork-master group-64 layout.
24
25mod config;
26mod container;
27mod dequant_cpu;
28mod dequant_gpu;
29mod names;
30mod sidecar;
31mod value_transform;
32
33pub use config::config_from_gguf_dir;
34
35use anyhow::{Context, Result, bail};
36use std::collections::HashMap;
37use std::path::{Path, PathBuf};
38
39use super::{WeightDtype, WeightStore, WeightTensor, check_oom_guard, evict_page_cache};
40use crate::gpu::{DevicePtr, GpuBackend};
41
42/// Locate the backbone GGUF weight file in `dir`. Returns the
43/// lexicographically-first non-mmproj `*.gguf` (also the first shard,
44/// `*-00001-of-*`, of a split file). The mmproj vision sidecar is excluded here
45/// and loaded separately (see `sidecar::find_mmproj`); a dir that is *only* an
46/// mmproj falls back to the first file so the caller still gets a path to error
47/// on.
48pub fn find_gguf(dir: &Path) -> Option<PathBuf> {
49    let mut candidates: Vec<PathBuf> = std::fs::read_dir(dir)
50        .ok()?
51        .filter_map(|e| e.ok())
52        .map(|e| e.path())
53        .filter(|p| p.extension().and_then(|x| x.to_str()) == Some("gguf"))
54        .collect();
55    candidates.sort();
56    candidates
57        .iter()
58        .find(|p| !sidecar::is_mmproj(p))
59        .cloned()
60        .or_else(|| candidates.into_iter().next())
61}
62
63/// True when the native keep-packed Q2_0 decode path is enabled
64/// (`ATLAS_GGUF_NATIVE_Q2=1`). Off by default: the loader dequants every id-42
65/// tensor to BF16 exactly as before, so the default path is byte-identical.
66/// When on, the "big" transform-free FFN projections (see
67/// [`names::is_keep_packed_proj`]) are uploaded as raw `block_q2_0` blocks and
68/// tagged [`WeightDtype::PackedQ2_0`] for in-kernel dequant at decode.
69fn native_q2_enabled() -> bool {
70    std::env::var("ATLAS_GGUF_NATIVE_Q2").ok().as_deref() == Some("1")
71}
72
73/// The id-42 PrismML group size, from `ATLAS_GGUF_Q2_GROUP` (default 128).
74fn q2_group_usize() -> usize {
75    match std::env::var("ATLAS_GGUF_Q2_GROUP").ok().as_deref() {
76        Some("64") => 64,
77        _ => 128,
78    }
79}
80
81/// Map a group size to the container's `Q2Group` (for on-disk byte sizing).
82fn q2_group_variant(g: usize) -> container::Q2Group {
83    if g == 64 {
84        container::Q2Group::G64
85    } else {
86        container::Q2Group::G128
87    }
88}
89
90/// Loads weights from a GGUF file, dequantizing every tensor to BF16 on the GPU.
91///
92/// Mirrors [`super::SafetensorsLoader`] so the two are interchangeable behind the
93/// [`super::WeightLoader`] trait and the serve call-site can pick one on file
94/// type.
95pub struct GgufLoader {
96    /// EP rank (0-based). Only used when `ep_world_size > 1`.
97    pub ep_rank: usize,
98    /// EP world size. When > 1, remote expert slices are skipped.
99    pub ep_world_size: usize,
100    /// Total number of MoE experts in the model (for EP partitioning).
101    pub num_experts: usize,
102    /// Override for the peak-memory multiplier in the pre-flight OOM check.
103    pub peak_memory_multiplier: Option<f64>,
104}
105
106impl Default for GgufLoader {
107    fn default() -> Self {
108        Self::new()
109    }
110}
111
112impl GgufLoader {
113    /// Create a loader with no expert parallelism (loads all tensors).
114    pub fn new() -> Self {
115        Self {
116            ep_rank: 0,
117            ep_world_size: 1,
118            num_experts: 0,
119            peak_memory_multiplier: None,
120        }
121    }
122
123    /// Create a loader with EP-aware expert filtering.
124    pub fn with_ep(ep_rank: usize, ep_world_size: usize, num_experts: usize) -> Self {
125        Self {
126            ep_rank,
127            ep_world_size,
128            num_experts,
129            peak_memory_multiplier: None,
130        }
131    }
132
133    /// True if expert `idx` lives on a remote EP rank and should be skipped.
134    fn should_skip_expert(&self, idx: usize) -> bool {
135        if self.ep_world_size <= 1 || self.num_experts == 0 {
136            return false;
137        }
138        let per_rank = self.num_experts / self.ep_world_size;
139        let local_start = self.ep_rank * per_rank;
140        let local_end = if self.ep_rank == self.ep_world_size - 1 {
141            self.num_experts
142        } else {
143            local_start + per_rank
144        };
145        idx < local_start || idx >= local_end
146    }
147
148    /// Split a dequantized, stacked expert buffer into per-expert `WeightTensor`s
149    /// that alias offsets into the single BF16 device allocation. `shape[0]` is
150    /// the expert count; each expert tensor is `shape[1..]`.
151    fn emit_experts(
152        &self,
153        weights: &mut HashMap<String, WeightTensor>,
154        base_ptr: DevicePtr,
155        shape: &[usize],
156        layer: usize,
157        proj: &str,
158        skipped: &mut usize,
159    ) -> Result<()> {
160        let count = *shape
161            .first()
162            .context("stacked expert tensor has no leading expert dimension")?;
163        let per_elems: usize = shape[1..].iter().product();
164        let per_bytes = per_elems * WeightDtype::BF16.byte_size();
165        let expert_shape: Vec<usize> = shape[1..].to_vec();
166        for e in 0..count {
167            if self.should_skip_expert(e) {
168                *skipped += 1;
169                continue;
170            }
171            let ptr = base_ptr.offset(e * per_bytes);
172            let name = names::expert_name(layer, proj, e);
173            weights.insert(
174                name,
175                WeightTensor {
176                    ptr,
177                    shape: expert_shape.clone(),
178                    dtype: WeightDtype::BF16,
179                },
180            );
181        }
182        Ok(())
183    }
184}
185
186/// Dequant one tensor's raw block bytes to a BF16 device buffer. Prefer-GPU /
187/// CPU-fallback: use the GPU kernel when one exists for `id` (and `force_cpu`
188/// is unset), else the CPU reference dequant (host BF16 → single h2d).
189fn dequant_to_device(
190    gpu: &dyn GpuBackend,
191    id: u32,
192    raw: &[u8],
193    num_elements: usize,
194    q2_group: usize,
195    force_cpu: bool,
196) -> Result<DevicePtr> {
197    if !force_cpu && dequant_gpu::supports(id) {
198        let q_ptr = gpu.alloc(raw.len())?;
199        gpu.copy_h2d(raw, q_ptr)?;
200        let bf16_ptr = dequant_gpu::to_bf16(gpu, id, q_ptr, num_elements, q2_group)
201            .with_context(|| format!("GPU dequant failed for ggml type {id}"))?;
202        gpu.free(q_ptr)?;
203        Ok(bf16_ptr)
204    } else if dequant_cpu::supports(id) {
205        let host = dequant_cpu::to_bf16_bytes(id, q2_group, raw, num_elements)
206            .with_context(|| format!("CPU dequant failed for ggml type {id}"))?;
207        debug_assert_eq!(host.len(), num_elements * WeightDtype::BF16.byte_size());
208        let bf16_ptr = gpu.alloc(host.len())?;
209        gpu.copy_h2d(&host, bf16_ptr)?;
210        Ok(bf16_ptr)
211    } else {
212        bail!("No GPU or CPU dequant available for ggml type {id}");
213    }
214}
215
216/// Pre-flight: estimate the total BF16 footprint (dequant expands quantized
217/// blocks) and bail before allocating if it won't fit under the reserve.
218fn preflight_oom(
219    gpu: &dyn GpuBackend,
220    est_bf16_bytes: usize,
221    reserve_bytes: usize,
222    multiplier: Option<f64>,
223) -> Result<()> {
224    // Small transient overhead: the raw quantized scratch buffer coexists with
225    // its BF16 output for one tensor at a time (freed immediately after).
226    let overhead = multiplier.unwrap_or(1.1);
227    let peak = (est_bf16_bytes as f64 * overhead) as usize;
228    let free = gpu.free_memory()?;
229    let gb = |b: usize| b as f64 / (1024.0 * 1024.0 * 1024.0);
230    tracing::info!(
231        "GGUF pre-flight: ~{:.2} GB BF16 after dequant, {:.1}x overhead = {:.2} GB peak, \
232         {:.2} GB free, {:.1} GB reserve",
233        gb(est_bf16_bytes),
234        overhead,
235        gb(peak),
236        gb(free),
237        gb(reserve_bytes),
238    );
239    if peak + reserve_bytes > free {
240        bail!(
241            "Pre-flight OOM: GGUF dequant to BF16 needs ~{:.2} GB peak + {:.1} GB reserve, \
242             only {:.2} GB free. Use a smaller model or lower --oom-guard-mb.",
243            gb(peak),
244            gb(reserve_bytes),
245            gb(free),
246        );
247    }
248    Ok(())
249}
250
251impl super::WeightLoader for GgufLoader {
252    fn load(
253        &self,
254        model_dir: &Path,
255        gpu: &dyn GpuBackend,
256        oom_reserve_bytes: usize,
257    ) -> Result<WeightStore> {
258        let path = find_gguf(model_dir)
259            .with_context(|| format!("No .gguf file found in {}", model_dir.display()))?;
260        tracing::info!("Loading GGUF weights from {}", path.display());
261
262        let force_cpu = std::env::var("ATLAS_GGUF_FORCE_CPU").ok().as_deref() == Some("1");
263        let native_q2 = native_q2_enabled();
264        let q2_group = q2_group_usize();
265        let q2_variant = q2_group_variant(q2_group);
266        if native_q2 {
267            tracing::info!(
268                "ATLAS_GGUF_NATIVE_Q2=1: keeping id-42 FFN projections packed (group {q2_group})"
269            );
270        }
271
272        // ── Backbone (text model) ──
273        let (bb_file, bb_mmap, bb_gguf) = sidecar::open_gguf(&path)?;
274        let arch = bb_gguf
275            .get_str("general.architecture")
276            .unwrap_or("llama")
277            .to_lowercase();
278
279        // Qwen3.5/3.6 GDN-hybrid GGUFs (llama.cpp `qwen35` converter) encode a
280        // handful of GDN / RMSNorm tensor VALUES differently than Atlas's
281        // kernels expect (norm +1 offset, `A_log = ln(-ssm_a)`, and a value-head
282        // reorder). Read the GDN head geometry once so `load_pass` can invert
283        // them per tensor (see `value_transform`).
284        let is_qwen35 = value_transform::is_qwen35(&arch);
285        let gdn = if is_qwen35 {
286            value_transform::gdn_dims(&bb_gguf, &arch)
287        } else {
288            None
289        };
290        if is_qwen35 && gdn.is_none() {
291            bail!(
292                "GGUF arch '{arch}' is qwen35-family but the SSM metadata keys \
293                 ({arch}.ssm.*) are missing; cannot apply GDN value transforms"
294            );
295        }
296
297        // ── Optional mmproj vision-tower sidecar ──
298        // Open it (if present) up front so the pre-flight OOM check covers both
299        // files. mmaps are virtual, so holding two at once costs no RAM.
300        let mmproj_path = sidecar::find_mmproj(model_dir, &path);
301        let mmproj = match &mmproj_path {
302            Some(mp) => {
303                tracing::info!("Found mmproj vision sidecar {}", mp.display());
304                Some(sidecar::open_gguf(mp)?)
305            }
306            None => None,
307        };
308        let mmproj_arch = mmproj.as_ref().map(|(_, _, g)| {
309            g.get_str("general.architecture")
310                .unwrap_or("clip")
311                .to_lowercase()
312        });
313
314        // Pre-flight: combined BF16 footprint of both files.
315        let mut est = sidecar::est_bf16(&bb_gguf, &arch);
316        if let (Some((_, _, mm_gguf)), Some(mm_arch)) = (mmproj.as_ref(), mmproj_arch.as_ref()) {
317            est += sidecar::est_bf16(mm_gguf, mm_arch);
318        }
319        preflight_oom(gpu, est, oom_reserve_bytes, self.peak_memory_multiplier)?;
320
321        let mut weights: HashMap<String, WeightTensor> = HashMap::new();
322        let mut skipped = 0usize;
323
324        // Pass 1: backbone → weights.
325        sidecar::load_pass(
326            self,
327            gpu,
328            &bb_gguf,
329            &bb_mmap,
330            &arch,
331            gdn,
332            force_cpu,
333            native_q2,
334            q2_group,
335            q2_variant,
336            &mut weights,
337            &mut skipped,
338        )?;
339        drop(bb_gguf);
340        drop(bb_mmap);
341        evict_page_cache(&bb_file);
342
343        // Pass 2: mmproj sidecar → SAME weights map (clip names land under
344        // `model.visual.*`, disjoint from the backbone's `model.layers.*`).
345        // No GDN transforms (gdn = None) and no expert fan-out for clip.
346        if let (Some((mm_file, mm_mmap, mm_gguf)), Some(mm_arch)) = (mmproj, mmproj_arch) {
347            let before = weights.len();
348            // mmproj is a `clip` tower — no qwen35 FFN names, so native_q2 is
349            // irrelevant there; pass false to keep it on the plain BF16 path.
350            sidecar::load_pass(
351                self,
352                gpu,
353                &mm_gguf,
354                &mm_mmap,
355                &mm_arch,
356                None,
357                force_cpu,
358                false,
359                q2_group,
360                q2_variant,
361                &mut weights,
362                &mut skipped,
363            )?;
364            tracing::info!(
365                "Merged {} mmproj tensors (arch '{}') into the weight store",
366                weights.len() - before,
367                mm_arch,
368            );
369            drop(mm_gguf);
370            drop(mm_mmap);
371            evict_page_cache(&mm_file);
372        }
373
374        if skipped > 0 {
375            tracing::info!("EP: skipped {} remote expert slices", skipped);
376        }
377        check_oom_guard(gpu, oom_reserve_bytes, "weight loading (GGUF)")?;
378        tracing::info!("Loaded {} weight tensors (GGUF → BF16)", weights.len());
379        Ok(WeightStore::from_map(weights))
380    }
381}
382
383#[cfg(test)]
384mod real_file_test;
385
386#[cfg(all(test, feature = "cuda"))]
387mod gpu_validate_test;
388
389#[cfg(test)]
390mod tests;
spark_runtime/weights/gguf.rs

spark_runtime/weights/
gguf.rs
