fuel_quantized/
tq_mx_types.rs

1// Derived in part from llama.cpp / ggml (MIT License, Copyright (c) 2023-2026 The ggml authors).
2// The full licence text is in THIRD-PARTY-NOTICES.md at the repository root. No licence identifier line
3// is stamped here: the file derives from somebody else's copyright.
4//! `TQ1_0`, `TQ2_0` (ternary) and `MXFP4`, `NVFP4` (4-bit float) block formats
5//! (GAP-211 C). DECODE-ONLY.
6//!
7//! # Provenance
8//!
9//! Block layouts and dequantization are llama.cpp's / ggml's (MIT License,
10//! Copyright (c) 2023-2026 The ggml authors), checked against the source at
11//! commit `9d57ce456c` (the pin recorded on [`GgmlDType`]):
12//!
13//! - `ggml/src/ggml-common.h`: `block_tq1_0`, `block_tq2_0`, `block_mxfp4`,
14//!   `block_nvfp4`, `kvalues_fp4`.
15//! - `ggml/src/ggml-quants.c`: `dequantize_row_tq1_0`, `dequantize_row_tq2_0`,
16//!   `dequantize_row_mxfp4`, `dequantize_row_nvfp4`.
17//! - `ggml/src/ggml-impl.h`: `ggml_e8m0_to_fp32_half`, `ggml_ue4m3_to_fp32`.
18//! - `ggml/include/ggml.h`: `GGML_TYPE_TQ1_0 = 34`, `TQ2_0 = 35`, `MXFP4 = 39`,
19//!   `NVFP4 = 40`.
20//!
21//! # Decode-only
22//!
23//! Like the grid-table IQ formats, these load, dequantize and multiply
24//! (dequantize-then-dot) but cannot be produced here:
25//! `GgmlDType::supports_quantize` is `false` and the fallible quantize entry
26//! points refuse them first.
27
28use crate::k_quants::{GgmlType, QK_K};
29use fuel_ir::quantized::GgmlDType;
30use half::f16;
31
32pub const QK_MXFP4: usize = 32;
33pub const QK_NVFP4: usize = 64;
34const QK_NVFP4_SUB: usize = 16;
35
36/// `kvalues_fp4` from `ggml-common.h`: E2M1 values, doubled (`kvalues_mxfp4` is
37/// an alias).
38pub const KVALUES_FP4: [i8; 16] = [0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12];
39
40#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
41#[repr(C)]
42pub struct BlockTq1_0 {
43    pub(crate) qs: [u8; (QK_K - 4 * QK_K / 64) / 5],
44    pub(crate) qh: [u8; QK_K / 64],
45    pub(crate) d: f16,
46}
47const _: () = assert!(std::mem::size_of::<BlockTq1_0>() == 54);
48
49#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
50#[repr(C)]
51pub struct BlockTq2_0 {
52    pub(crate) qs: [u8; QK_K / 4],
53    pub(crate) d: f16,
54}
55const _: () = assert!(std::mem::size_of::<BlockTq2_0>() == 66);
56
57#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
58#[repr(C)]
59pub struct BlockMxfp4 {
60    pub(crate) e: u8,
61    pub(crate) qs: [u8; QK_MXFP4 / 2],
62}
63const _: () = assert!(std::mem::size_of::<BlockMxfp4>() == 17);
64
65#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
66#[repr(C)]
67pub struct BlockNvfp4 {
68    pub(crate) d: [u8; QK_NVFP4 / QK_NVFP4_SUB],
69    pub(crate) qs: [u8; QK_NVFP4 / 2],
70}
71const _: () = assert!(std::mem::size_of::<BlockNvfp4>() == 36);
72
73/// `ggml_e8m0_to_fp32_half`: half of `2^(x-127)`, with the two smallest codes
74/// as denormals. Code 255 (NaN upstream) is not special-cased upstream either.
75fn e8m0_to_f32_half(x: u8) -> f32 {
76    let bits: u32 = if x < 2 {
77        0x0020_0000u32 << x
78    } else {
79        u32::from(x - 1) << 23
80    };
81    f32::from_bits(bits)
82}
83
84/// `x * 2^e` for an exponent in the normal range; exact (no `powi` rounding).
85fn ldexp(x: f32, e: i32) -> f32 {
86    x * f32::from_bits(((e + 127) as u32) << 23)
87}
88
89/// `ggml_ue4m3_to_fp32`: unsigned 4-exponent-bit (bias 7), 3-mantissa-bit float,
90/// times 0.5 to match the doubled `kvalues_fp4`. Codes 0 and 0x7F decode to 0.
91fn ue4m3_to_f32(x: u8) -> f32 {
92    if x == 0 || x == 0x7F {
93        return 0.0;
94    }
95    let exp = i32::from((x >> 3) & 0xF);
96    let man = f32::from(x & 0x7);
97    let raw = if exp == 0 {
98        ldexp(man, -9)
99    } else {
100        ldexp(1.0 + man / 8.0, exp - 7)
101    };
102    raw * 0.5
103}
104
105/// One base-3 digit of a TQ1_0 byte, as upstream extracts it: -1, 0 or +1.
106fn tq1_digit(byte: u8, pow3: u8) -> f32 {
107    let q = byte.wrapping_mul(pow3);
108    let xi = ((u16::from(q) * 3) >> 8) as i16;
109    f32::from(xi - 1)
110}
111
112impl BlockTq1_0 {
113    fn decode(&self, y: &mut [f32]) {
114        const POW3: [u8; 5] = [1, 3, 9, 27, 81];
115        let d = self.d.to_f32();
116        let mut out = 0;
117        // 48 bytes of `qs`: one run of 32, then one of 16 (upstream's two loops).
118        let (head, tail) = self.qs.split_at(self.qs.len() - self.qs.len() % 32);
119        let runs32 = head.as_chunks::<32>().0.iter().map(|c| &c[..]);
120        let runs16 = tail.as_chunks::<16>().0.iter().map(|c| &c[..]);
121        for chunk in runs32.chain(runs16) {
122            for p in POW3 {
123                for q in chunk {
124                    y[out] = tq1_digit(*q, p) * d;
125                    out += 1;
126                }
127            }
128        }
129        for p in &POW3[..4] {
130            for q in &self.qh {
131                y[out] = tq1_digit(*q, *p) * d;
132                out += 1;
133            }
134        }
135        debug_assert_eq!(out, QK_K);
136    }
137}
138
139impl BlockTq2_0 {
140    fn decode(&self, y: &mut [f32]) {
141        let d = self.d.to_f32();
142        let mut out = 0;
143        for chunk in self.qs.as_chunks::<32>().0 {
144            for l in 0..4 {
145                for q in chunk {
146                    let v = ((q >> (l * 2)) & 3) as i8;
147                    y[out] = f32::from(v - 1) * d;
148                    out += 1;
149                }
150            }
151        }
152    }
153}
154
155impl BlockMxfp4 {
156    fn decode(&self, y: &mut [f32]) {
157        let d = e8m0_to_f32_half(self.e);
158        for (j, q) in self.qs.iter().enumerate() {
159            y[j] = f32::from(KVALUES_FP4[usize::from(q & 0x0F)]) * d;
160            y[j + QK_MXFP4 / 2] = f32::from(KVALUES_FP4[usize::from(q >> 4)]) * d;
161        }
162    }
163}
164
165impl BlockNvfp4 {
166    fn decode(&self, y: &mut [f32]) {
167        for (s, scale) in self.d.iter().enumerate() {
168            let d = ue4m3_to_f32(*scale);
169            let yb = &mut y[s * QK_NVFP4_SUB..(s + 1) * QK_NVFP4_SUB];
170            let qs = &self.qs[s * QK_NVFP4_SUB / 2..(s + 1) * QK_NVFP4_SUB / 2];
171            for (j, q) in qs.iter().enumerate() {
172                yb[j] = f32::from(KVALUES_FP4[usize::from(q & 0x0F)]) * d;
173                yb[j + QK_NVFP4_SUB / 2] = f32::from(KVALUES_FP4[usize::from(q >> 4)]) * d;
174            }
175        }
176    }
177}
178
179macro_rules! decode_only_type {
180    ($blk:ty, $dtype:expr, $qk:expr, $vecdot:ty) => {
181        impl GgmlType for $blk {
182            const DTYPE: GgmlDType = $dtype;
183            const BLCK_SIZE: usize = $qk;
184
185            fn to_float(xs: &[Self], ys: &mut [f32]) {
186                for (b, y) in xs.iter().zip(ys.chunks_exact_mut($qk)) {
187                    b.decode(y);
188                }
189            }
190
191            fn from_float(_xs: &[f32], _ys: &mut [Self]) {
192                // Unreachable through the fallible entry points, which check
193                // `GgmlDType::supports_quantize` first.
194                panic!(
195                    "{:?} is decode-only: Fuel cannot quantize to it (GAP-211 C)",
196                    $dtype
197                );
198            }
199        }
200    };
201}
202
203decode_only_type!(BlockTq1_0, GgmlDType::Tq1_0, QK_K, BlockQ8K);
204decode_only_type!(BlockTq2_0, GgmlDType::Tq2_0, QK_K, BlockQ8K);
205decode_only_type!(BlockMxfp4, GgmlDType::Mxfp4, QK_MXFP4, BlockQ8_0);
206decode_only_type!(BlockNvfp4, GgmlDType::Nvfp4, QK_NVFP4, BlockQ8_0);
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    fn one<T: GgmlType>(b: T) -> Vec<f32> {
213        let mut y = vec![0f32; T::BLCK_SIZE];
214        T::to_float(&[b], &mut y);
215        y
216    }
217
218    #[test]
219    fn geometry_agrees_with_ggml_dtype() {
220        fn check<T: GgmlType>() {
221            assert_eq!(
222                std::mem::size_of::<T>(),
223                T::DTYPE.type_size(),
224                "{:?}",
225                T::DTYPE
226            );
227            assert_eq!(T::BLCK_SIZE, T::DTYPE.block_size(), "{:?}", T::DTYPE);
228            assert!(!T::DTYPE.supports_quantize());
229            assert!(T::DTYPE.check_quantize().is_err());
230        }
231        check::<BlockTq1_0>();
232        check::<BlockTq2_0>();
233        check::<BlockMxfp4>();
234        check::<BlockNvfp4>();
235        assert!(GgmlDType::Q4K.supports_quantize());
236    }
237
238    /// `ggml_e8m0_to_fp32_half`: code 127 is 2^0 halved; 0 and 1 are the denormals.
239    #[test]
240    fn e8m0_half_scale_cases() {
241        assert_eq!(e8m0_to_f32_half(127), 0.5);
242        assert_eq!(e8m0_to_f32_half(128), 1.0);
243        assert_eq!(e8m0_to_f32_half(1), f32::from_bits(0x0040_0000));
244        assert_eq!(e8m0_to_f32_half(0), f32::from_bits(0x0020_0000));
245    }
246
247    /// `ggml_ue4m3_to_fp32`, hand-evaluated: 0x38 = 1.0 * 0.5; 0x01 is the
248    /// smallest denormal 2^-9 * 0.5; 0x7E = 1.75 * 2^8 * 0.5; 0 and 0x7F are 0.
249    #[test]
250    fn ue4m3_scale_cases() {
251        assert_eq!(ue4m3_to_f32(0x38), 0.5);
252        assert_eq!(ue4m3_to_f32(0x01), 2f32.powi(-10));
253        assert_eq!(ue4m3_to_f32(0x7E), 224.0);
254        assert_eq!(ue4m3_to_f32(0x00), 0.0);
255        assert_eq!(ue4m3_to_f32(0x7F), 0.0);
256    }
257
258    #[test]
259    fn mxfp4_dequantizes_as_the_ggml_formula() {
260        let mut b = BlockMxfp4::zeros();
261        b.e = 128; // d = 1.0
262        b.qs[0] = 0x21; // low code 1 -> 1, high code 2 -> 2
263        b.qs[1] = 0xF8; // low code 8 -> 0, high code 15 -> -12
264        let y = one(b);
265        assert_eq!(y[0], 1.0);
266        assert_eq!(y[16], 2.0);
267        assert_eq!(y[1], 0.0);
268        assert_eq!(y[17], -12.0);
269        assert_eq!(y[2], 0.0); // untouched
270    }
271
272    #[test]
273    fn nvfp4_each_subblock_uses_its_own_scale() {
274        let mut b = BlockNvfp4::zeros();
275        b.d = [0x38, 0x38 + 8, 0x7F, 0x38]; // 0.5, 1.0, 0 (reserved), 0.5
276        b.qs[0] = 0x42; // sub-block 0: low code 2 -> 2, high code 4 -> 4
277        b.qs[8] = 0x42; // sub-block 1, same codes
278        b.qs[16] = 0x42; // sub-block 2, scale 0
279        let y = one(b);
280        assert_eq!(y[0], 1.0); // 2 * 0.5
281        assert_eq!(y[8], 2.0); // 4 * 0.5
282        assert_eq!(y[16], 2.0); // 2 * 1.0
283        assert_eq!(y[24], 4.0);
284        assert_eq!(y[32], 0.0);
285        assert_eq!(y[40], 0.0);
286    }
287
288    #[test]
289    fn tq2_0_dequantizes_as_the_ggml_formula() {
290        let mut b = BlockTq2_0::zeros();
291        b.d = f16::from_f32(2.0);
292        b.qs[0] = 0b11_10_01_00;
293        let y = one(b);
294        assert_eq!(y[0], -2.0); // q 0
295        assert_eq!(y[32], 0.0); // q 1
296        assert_eq!(y[64], 2.0); // q 2
297        assert_eq!(y[96], 4.0); // q 3
298        assert_eq!(y[1], -2.0); // untouched byte: q 0
299    }
300
301    /// Byte 171 is the trit run (2,0,0,0,0): ceil(162 * 256 / 243). Digit 0 is
302    /// +1, digit 1 is -1. Checked in the 32-run, the 16-run and the `qh` tail.
303    #[test]
304    fn tq1_0_dequantizes_as_the_ggml_formula() {
305        let mut b = BlockTq1_0::zeros();
306        b.d = f16::from_f32(2.0);
307        b.qs[0] = 171; // 32-run
308        b.qs[32] = 171; // 16-run
309        b.qh[0] = 171;
310        let y = one(b);
311        assert_eq!(y[0], 2.0);
312        assert_eq!(y[32], -2.0);
313        assert_eq!(y[160], 2.0); // 16-run starts after 5 * 32
314        assert_eq!(y[160 + 16], -2.0);
315        assert_eq!(y[240], 2.0); // qh tail starts after 160 + 80
316        assert_eq!(y[244], -2.0);
317        assert_eq!(y[1], -2.0); // zero byte: all digits -1
318    }
319}
fuel_quantized/tq_mx_types.rs

fuel_quantized/
tq_mx_types.rs
