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

spark_runtime/weights/
gguf.rs
