spark_runtime/weights/gguf/
dequant_cpu.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! CPU reference dequantization for GGUF/GGML quantized weight blocks.
3//!
4//! This is the correctness oracle and portable fallback for the Atlas GGUF
5//! loader. Every routine reproduces llama.cpp's `dequantize_row_*` index
6//! arithmetic exactly (little-endian throughout). Clarity is preferred over
7//! speed: the GPU dequant kernels are the fast path; these functions are what
8//! those kernels are validated against, and what runs when a ggml type has no
9//! GPU kernel or under `MockGpuBackend` unit tests (which do not execute
10//! kernels).
11
12use anyhow::{Context, Result, bail};
13
14mod blocks;
15
16#[cfg(test)]
17mod tests;
18
19/// A GGML block-quantization type.
20///
21/// The PrismML-private `Q2_0` (ggml id 42) is polymorphic in its group size
22/// (128 in the shipped Ternary-Bonsai file, 64 in the fork master), so the
23/// group is carried in the variant rather than hardcoded. Standard `Tq2_0`
24/// (id 35, group-256, scale-at-end) is a *separate* entry and must not be
25/// conflated with it.
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum GgmlType {
28    F32,
29    F16,
30    Bf16,
31    Q4_0,
32    Q4_1,
33    Q5_0,
34    Q5_1,
35    Q8_0,
36    Q8_1,
37    Q2K,
38    Q3K,
39    Q4K,
40    Q5K,
41    Q6K,
42    Q8K,
43    Tq1_0,
44    Tq2_0,
45    /// PrismML id 42: 2-bit, symmetric, scale-at-front, contiguous codes.
46    /// `group` ∈ {64, 128}; block = 2 + group/4 bytes.
47    Q2_0 {
48        group: usize,
49    },
50}
51
52impl GgmlType {
53    /// Map a raw ggml type id to a `GgmlType`. `q2_0_group` supplies the group
54    /// size for the id-42 PrismML type (ignored for every other id); pass the
55    /// value read from the fork's GGUF metadata (128 default, 64 fork-master).
56    pub fn from_id(id: u32, q2_0_group: usize) -> Result<Self> {
57        Ok(match id {
58            0 => Self::F32,
59            1 => Self::F16,
60            30 => Self::Bf16,
61            2 => Self::Q4_0,
62            3 => Self::Q4_1,
63            6 => Self::Q5_0,
64            7 => Self::Q5_1,
65            8 => Self::Q8_0,
66            9 => Self::Q8_1,
67            10 => Self::Q2K,
68            11 => Self::Q3K,
69            12 => Self::Q4K,
70            13 => Self::Q5K,
71            14 => Self::Q6K,
72            15 => Self::Q8K,
73            34 => Self::Tq1_0,
74            35 => Self::Tq2_0,
75            42 => Self::Q2_0 { group: q2_0_group },
76            other => bail!("unsupported / unknown ggml type id {other}"),
77        })
78    }
79
80    /// Number of dequantized elements produced by one block (`QK` / `QK_K`).
81    pub fn block_size(self) -> usize {
82        match self {
83            Self::F32 | Self::F16 | Self::Bf16 => 1,
84            Self::Q4_0 | Self::Q4_1 | Self::Q5_0 | Self::Q5_1 | Self::Q8_0 | Self::Q8_1 => 32,
85            Self::Q2K
86            | Self::Q3K
87            | Self::Q4K
88            | Self::Q5K
89            | Self::Q6K
90            | Self::Q8K
91            | Self::Tq1_0
92            | Self::Tq2_0 => 256,
93            Self::Q2_0 { group } => group,
94        }
95    }
96
97    /// On-disk byte size of one block.
98    pub fn block_bytes(self) -> Result<usize> {
99        Ok(match self {
100            Self::F32 => 4,
101            Self::F16 | Self::Bf16 => 2,
102            Self::Q4_0 => 18,
103            Self::Q4_1 => 20,
104            Self::Q5_0 => 22,
105            Self::Q5_1 => 24,
106            Self::Q8_0 => 34,
107            Self::Q8_1 => 36,
108            Self::Q2K => 84,
109            Self::Q3K => 110,
110            Self::Q4K => 144,
111            Self::Q5K => 176,
112            Self::Q6K => 210,
113            Self::Q8K => 292,
114            Self::Tq1_0 => 54,
115            Self::Tq2_0 => 66,
116            Self::Q2_0 { group } => {
117                if group == 0 || !group.is_multiple_of(4) {
118                    bail!("Q2_0 group size must be a positive multiple of 4, got {group}");
119                }
120                2 + group / 4
121            }
122        })
123    }
124}
125
126/// True if the CPU reference dequant handles this ggml type id.
127pub(crate) fn supports(id: u32) -> bool {
128    // The group value is irrelevant to the support decision (id 42 is always
129    // supported once a group is known); probe with the default group.
130    GgmlType::from_id(id, 128).is_ok()
131}
132
133/// Dequantize a tensor's raw block bytes into little-endian BF16 host bytes.
134/// `q2_group` is the id-42 PrismML group size (128 default / 64 fork-master);
135/// ignored for every other type. This is the loader's CPU-fallback entry point.
136pub(crate) fn to_bf16_bytes(
137    id: u32,
138    q2_group: usize,
139    raw: &[u8],
140    n_elements: usize,
141) -> Result<Vec<u8>> {
142    let t = GgmlType::from_id(id, q2_group)?;
143    let mut out = vec![0u16; n_elements];
144    dequant_to_bf16(t, raw, n_elements, &mut out)?;
145    let mut bytes = Vec::with_capacity(n_elements * 2);
146    for v in &out {
147        bytes.extend_from_slice(&v.to_le_bytes());
148    }
149    Ok(bytes)
150}
151
152/// Dequantize `n_elements` from packed `block_bytes` into `out` as `f32`.
153///
154/// `block_bytes` must contain at least `ceil(n_elements / block_size)` whole
155/// blocks; `n_elements` must be a whole multiple of the type's block size
156/// (GGUF tensors always are — the last dim is padded to the block size).
157pub fn dequant_to_f32(
158    t: GgmlType,
159    block_bytes: &[u8],
160    n_elements: usize,
161    out: &mut [f32],
162) -> Result<()> {
163    if out.len() < n_elements {
164        bail!("out buffer too small: {} < {}", out.len(), n_elements);
165    }
166
167    // Passthrough / widen types operate on the whole run at once.
168    match t {
169        GgmlType::F32 => {
170            let need = n_elements.checked_mul(4).context("f32 size overflow")?;
171            require_len(block_bytes, need)?;
172            blocks::widen_f32(&block_bytes[..need], &mut out[..n_elements]);
173            return Ok(());
174        }
175        GgmlType::F16 => {
176            let need = n_elements.checked_mul(2).context("f16 size overflow")?;
177            require_len(block_bytes, need)?;
178            blocks::widen_f16(&block_bytes[..need], &mut out[..n_elements]);
179            return Ok(());
180        }
181        GgmlType::Bf16 => {
182            let need = n_elements.checked_mul(2).context("bf16 size overflow")?;
183            require_len(block_bytes, need)?;
184            blocks::widen_bf16(&block_bytes[..need], &mut out[..n_elements]);
185            return Ok(());
186        }
187        _ => {}
188    }
189
190    let qk = t.block_size();
191    if !n_elements.is_multiple_of(qk) {
192        bail!("n_elements {n_elements} is not a multiple of block size {qk} for {t:?}");
193    }
194    let bb = t.block_bytes()?;
195    let n_blocks = n_elements / qk;
196    let need = n_blocks
197        .checked_mul(bb)
198        .context("block byte count overflow")?;
199    require_len(block_bytes, need)?;
200
201    for b in 0..n_blocks {
202        let blk = &block_bytes[b * bb..b * bb + bb];
203        let o = &mut out[b * qk..b * qk + qk];
204        dequant_block(t, blk, o);
205    }
206    Ok(())
207}
208
209/// Dequantize into BF16 bit patterns (round-to-nearest-even), the form the
210/// GGUF loader hands to `WeightTensor { dtype: BF16 }`. Goes via `f32` for a
211/// single, obviously-correct code path.
212pub fn dequant_to_bf16(
213    t: GgmlType,
214    block_bytes: &[u8],
215    n_elements: usize,
216    out: &mut [u16],
217) -> Result<()> {
218    if out.len() < n_elements {
219        bail!("out buffer too small: {} < {}", out.len(), n_elements);
220    }
221    let mut tmp = vec![0f32; n_elements];
222    dequant_to_f32(t, block_bytes, n_elements, &mut tmp)?;
223    for (o, &f) in out[..n_elements].iter_mut().zip(tmp.iter()) {
224        *o = f32_to_bf16_bits(f);
225    }
226    Ok(())
227}
228
229#[inline]
230fn require_len(buf: &[u8], need: usize) -> Result<()> {
231    if buf.len() < need {
232        bail!("input block bytes too small: {} < {}", buf.len(), need);
233    }
234    Ok(())
235}
236
237/// Dispatch one block to its per-type routine. Passthrough arms are reachable
238/// only via callers other than `dequant_to_f32` (which short-circuits them);
239/// they remain correct for a single-element slice.
240fn dequant_block(t: GgmlType, blk: &[u8], out: &mut [f32]) {
241    match t {
242        GgmlType::F32 => blocks::widen_f32(blk, out),
243        GgmlType::F16 => blocks::widen_f16(blk, out),
244        GgmlType::Bf16 => blocks::widen_bf16(blk, out),
245        GgmlType::Q4_0 => blocks::dequant_q4_0(blk, out),
246        GgmlType::Q4_1 => blocks::dequant_q4_1(blk, out),
247        GgmlType::Q5_0 => blocks::dequant_q5_0(blk, out),
248        GgmlType::Q5_1 => blocks::dequant_q5_1(blk, out),
249        GgmlType::Q8_0 => blocks::dequant_q8_0(blk, out),
250        GgmlType::Q8_1 => blocks::dequant_q8_1(blk, out),
251        GgmlType::Q2K => blocks::dequant_q2_k(blk, out),
252        GgmlType::Q3K => blocks::dequant_q3_k(blk, out),
253        GgmlType::Q4K => blocks::dequant_q4_k(blk, out),
254        GgmlType::Q5K => blocks::dequant_q5_k(blk, out),
255        GgmlType::Q6K => blocks::dequant_q6_k(blk, out),
256        GgmlType::Q8K => blocks::dequant_q8_k(blk, out),
257        GgmlType::Tq1_0 => blocks::dequant_tq1_0(blk, out),
258        GgmlType::Tq2_0 => blocks::dequant_tq2_0(blk, out),
259        GgmlType::Q2_0 { group } => blocks::dequant_q2_0_gn(blk, out, group),
260    }
261}
262
263// ---- little-endian scalar helpers (no external deps) ----------------------
264
265/// IEEE binary16 → f32, handling zero/subnormal/inf/nan.
266#[inline]
267pub(super) fn f16_to_f32(h: u16) -> f32 {
268    let sign = if (h >> 15) & 1 == 1 { -1.0f32 } else { 1.0f32 };
269    let exp = (h >> 10) & 0x1f;
270    let mant = (h & 0x3ff) as f32;
271    match exp {
272        0 => sign * mant * 2f32.powi(-24), // subnormal / zero
273        0x1f => {
274            if mant == 0.0 {
275                sign * f32::INFINITY
276            } else {
277                f32::NAN
278            }
279        }
280        _ => sign * (1.0 + mant / 1024.0) * 2f32.powi(exp as i32 - 15),
281    }
282}
283
284#[inline]
285pub(super) fn bf16_to_f32(bits: u16) -> f32 {
286    f32::from_bits(u32::from(bits) << 16)
287}
288
289#[inline]
290pub(super) fn rd_f16(b: &[u8], off: usize) -> f32 {
291    f16_to_f32(u16::from_le_bytes([b[off], b[off + 1]]))
292}
293
294/// f32 → bf16 bits, round-to-nearest-even, NaN preserved (quieted).
295#[inline]
296pub(super) fn f32_to_bf16_bits(f: f32) -> u16 {
297    let bits = f.to_bits();
298    if f.is_nan() {
299        return ((bits >> 16) as u16) | 0x0040;
300    }
301    let rounding_bias = 0x0000_7fff + ((bits >> 16) & 1);
302    (bits.wrapping_add(rounding_bias) >> 16) as u16
303}
spark_runtime/weights/gguf/dequant_cpu.rs

spark_runtime/weights/gguf/
dequant_cpu.rs
