1use crate::iq_grids::*;
29use crate::k_quants::{BlockQ8K, GgmlType, QK_K};
30use fuel_ir::quantized::GgmlDType;
31use half::f16;
32
33const IQ1S_DELTA: f32 = 0.125;
34
35#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
36#[repr(C)]
37pub struct BlockIq2Xxs {
38 pub(crate) d: f16,
39 pub(crate) qs: [u16; QK_K / 8],
40}
41const _: () = assert!(std::mem::size_of::<BlockIq2Xxs>() == 66);
42
43#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
44#[repr(C)]
45pub struct BlockIq2Xs {
46 pub(crate) d: f16,
47 pub(crate) qs: [u16; QK_K / 8],
48 pub(crate) scales: [u8; QK_K / 32],
49}
50const _: () = assert!(std::mem::size_of::<BlockIq2Xs>() == 74);
51
52#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
53#[repr(C)]
54pub struct BlockIq2S {
55 pub(crate) d: f16,
56 pub(crate) qs: [u8; QK_K / 4],
57 pub(crate) qh: [u8; QK_K / 32],
58 pub(crate) scales: [u8; QK_K / 32],
59}
60const _: () = assert!(std::mem::size_of::<BlockIq2S>() == 82);
61
62#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
63#[repr(C)]
64pub struct BlockIq3Xxs {
65 pub(crate) d: f16,
66 pub(crate) qs: [u8; 3 * QK_K / 8],
67}
68const _: () = assert!(std::mem::size_of::<BlockIq3Xxs>() == 98);
69
70#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
71#[repr(C)]
72pub struct BlockIq3S {
73 pub(crate) d: f16,
74 pub(crate) qs: [u8; QK_K / 4],
75 pub(crate) qh: [u8; QK_K / 32],
76 pub(crate) signs: [u8; QK_K / 8],
77 pub(crate) scales: [u8; QK_K / 64],
78}
79const _: () = assert!(std::mem::size_of::<BlockIq3S>() == 110);
80
81#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
82#[repr(C)]
83pub struct BlockIq1S {
84 pub(crate) d: f16,
85 pub(crate) qs: [u8; QK_K / 8],
86 pub(crate) qh: [u16; QK_K / 32],
87}
88const _: () = assert!(std::mem::size_of::<BlockIq1S>() == 50);
89
90#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
91#[repr(C)]
92pub struct BlockIq1M {
93 pub(crate) qs: [u8; QK_K / 8],
94 pub(crate) qh: [u8; QK_K / 16],
95 pub(crate) scales: [u8; QK_K / 32],
96}
97const _: () = assert!(std::mem::size_of::<BlockIq1M>() == 56);
98
99fn sign(signs: u8, j: usize) -> f32 {
101 if signs & KMASK_IQ2XS[j] != 0 {
102 -1.0
103 } else {
104 1.0
105 }
106}
107
108impl BlockIq2Xxs {
109 pub(crate) fn decode(&self, y: &mut [f32]) {
111 let d = self.d.to_f32();
112 for ib32 in 0..QK_K / 32 {
113 let q = &self.qs[4 * ib32..4 * ib32 + 4];
114 let aux0 = u32::from(q[0]) | (u32::from(q[1]) << 16);
115 let aux1 = u32::from(q[2]) | (u32::from(q[3]) << 16);
116 let aux8 = aux0.to_le_bytes();
117 let db = d * (0.5 + (aux1 >> 28) as f32) * 0.25;
118 for l in 0..4 {
119 let grid = IQ2XXS_GRID[aux8[l] as usize].to_le_bytes();
120 let signs = KSIGNS_IQ2XS[((aux1 >> (7 * l)) & 127) as usize];
121 for j in 0..8 {
122 y[32 * ib32 + 8 * l + j] = db * f32::from(grid[j]) * sign(signs, j);
123 }
124 }
125 }
126 }
127}
128
129impl BlockIq2Xs {
130 pub(crate) fn decode(&self, y: &mut [f32]) {
132 let d = self.d.to_f32();
133 for ib32 in 0..QK_K / 32 {
134 let db = [
135 d * (0.5 + (self.scales[ib32] & 0xf) as f32) * 0.25,
136 d * (0.5 + (self.scales[ib32] >> 4) as f32) * 0.25,
137 ];
138 for l in 0..4 {
139 let q = self.qs[4 * ib32 + l];
140 let grid = IQ2XS_GRID[(q & 511) as usize].to_le_bytes();
141 let signs = KSIGNS_IQ2XS[(q >> 9) as usize];
142 for j in 0..8 {
143 y[32 * ib32 + 8 * l + j] = db[l / 2] * f32::from(grid[j]) * sign(signs, j);
144 }
145 }
146 }
147 }
148}
149
150impl BlockIq2S {
151 pub(crate) fn decode(&self, y: &mut [f32]) {
153 let d = self.d.to_f32();
154 let signs_base = QK_K / 8;
155 for ib32 in 0..QK_K / 32 {
156 let db = [
157 d * (0.5 + (self.scales[ib32] & 0xf) as f32) * 0.25,
158 d * (0.5 + (self.scales[ib32] >> 4) as f32) * 0.25,
159 ];
160 for l in 0..4 {
161 let dl = db[l / 2];
162 let idx = usize::from(self.qs[4 * ib32 + l])
163 | ((usize::from(self.qh[ib32]) << (8 - 2 * l)) & 0x300);
164 let grid = IQ2S_GRID[idx].to_le_bytes();
165 let signs = self.qs[signs_base + 4 * ib32 + l];
166 for j in 0..8 {
167 y[32 * ib32 + 8 * l + j] = dl * f32::from(grid[j]) * sign(signs, j);
168 }
169 }
170 }
171 }
172}
173
174impl BlockIq3Xxs {
175 pub(crate) fn decode(&self, y: &mut [f32]) {
177 let d = self.d.to_f32();
178 let scales_and_signs = QK_K / 4;
179 for ib32 in 0..QK_K / 32 {
180 let o = scales_and_signs + 4 * ib32;
181 let aux32 =
182 u32::from_le_bytes([self.qs[o], self.qs[o + 1], self.qs[o + 2], self.qs[o + 3]]);
183 let db = d * (0.5 + (aux32 >> 28) as f32) * 0.5;
184 for l in 0..4 {
185 let signs = KSIGNS_IQ2XS[((aux32 >> (7 * l)) & 127) as usize];
186 let g1 = IQ3XXS_GRID[self.qs[8 * ib32 + 2 * l] as usize].to_le_bytes();
187 let g2 = IQ3XXS_GRID[self.qs[8 * ib32 + 2 * l + 1] as usize].to_le_bytes();
188 for j in 0..4 {
189 y[32 * ib32 + 8 * l + j] = db * f32::from(g1[j]) * sign(signs, j);
190 y[32 * ib32 + 8 * l + j + 4] = db * f32::from(g2[j]) * sign(signs, j + 4);
191 }
192 }
193 }
194 }
195}
196
197impl BlockIq3S {
198 pub(crate) fn decode(&self, y: &mut [f32]) {
200 let d = self.d.to_f32();
201 for pair in 0..QK_K / 64 {
202 let db = [
203 d * (1 + 2 * i32::from(self.scales[pair] & 0xf)) as f32,
204 d * (1 + 2 * i32::from(self.scales[pair] >> 4)) as f32,
205 ];
206 for (half, dbh) in db.iter().enumerate() {
207 let ib32 = 2 * pair + half;
208 let qh = usize::from(self.qh[ib32]);
209 let qs = &self.qs[8 * ib32..8 * ib32 + 8];
210 for l in 0..4 {
211 let i1 = usize::from(qs[2 * l]) | ((qh << (8 - 2 * l)) & 256);
212 let i2 = usize::from(qs[2 * l + 1]) | ((qh << (7 - 2 * l)) & 256);
213 let g1 = IQ3S_GRID[i1].to_le_bytes();
214 let g2 = IQ3S_GRID[i2].to_le_bytes();
215 let signs = self.signs[4 * ib32 + l];
216 for j in 0..4 {
217 y[32 * ib32 + 8 * l + j] = dbh * f32::from(g1[j]) * sign(signs, j);
218 y[32 * ib32 + 8 * l + j + 4] = dbh * f32::from(g2[j]) * sign(signs, j + 4);
219 }
220 }
221 }
222 }
223 }
224}
225
226impl BlockIq1S {
227 pub(crate) fn decode(&self, y: &mut [f32]) {
229 let d = self.d.to_f32();
230 for ib in 0..QK_K / 32 {
231 let qh = self.qh[ib];
232 let dl = d * (2 * i32::from((qh >> 12) & 7) + 1) as f32;
233 let delta = if qh & 0x8000 != 0 {
234 -IQ1S_DELTA
235 } else {
236 IQ1S_DELTA
237 };
238 for l in 0..4 {
239 let idx =
240 usize::from(self.qs[4 * ib + l]) | ((usize::from(qh >> (3 * l)) & 7) << 8);
241 let grid = IQ1S_GRID[idx].to_le_bytes();
242 for j in 0..8 {
243 y[32 * ib + 8 * l + j] = dl * (f32::from(grid[j] as i8) + delta);
244 }
245 }
246 }
247 }
248}
249
250impl BlockIq1M {
251 fn sc(&self, i: usize) -> u16 {
252 u16::from_le_bytes([self.scales[2 * i], self.scales[2 * i + 1]])
253 }
254
255 pub(crate) fn decode(&self, y: &mut [f32]) {
257 let sc = [self.sc(0), self.sc(1), self.sc(2), self.sc(3)];
258 let scale =
259 (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
260 let d = f16::from_bits(scale).to_f32();
261 for ib in 0..QK_K / 32 {
262 let dl1 = d * (2 * i32::from((sc[ib / 2] >> (6 * (ib % 2))) & 0x7) + 1) as f32;
263 let dl2 = d * (2 * i32::from((sc[ib / 2] >> (6 * (ib % 2) + 3)) & 0x7) + 1) as f32;
264 let qs = &self.qs[4 * ib..4 * ib + 4];
265 let qh = &self.qh[2 * ib..2 * ib + 2];
266 let idx = [
267 usize::from(qs[0]) | ((usize::from(qh[0]) << 8) & 0x700),
268 usize::from(qs[1]) | ((usize::from(qh[0]) << 4) & 0x700),
269 usize::from(qs[2]) | ((usize::from(qh[1]) << 8) & 0x700),
270 usize::from(qs[3]) | ((usize::from(qh[1]) << 4) & 0x700),
271 ];
272 let bit = |b: u8, m: u8| if b & m != 0 { -IQ1S_DELTA } else { IQ1S_DELTA };
273 let delta = [
274 bit(qh[0], 0x08),
275 bit(qh[0], 0x80),
276 bit(qh[1], 0x08),
277 bit(qh[1], 0x80),
278 ];
279 for l in 0..4 {
280 let dl = if l < 2 { dl1 } else { dl2 };
281 let grid = IQ1S_GRID[idx[l]].to_le_bytes();
282 for j in 0..8 {
283 y[32 * ib + 8 * l + j] = dl * (f32::from(grid[j] as i8) + delta[l]);
284 }
285 }
286 }
287 }
288}
289
290fn dot_via_to_float<X: GgmlType, Y: GgmlType>(n: usize, xs: &[X], ys: &[Y]) -> f32 {
293 let mut a = vec![0f32; n];
294 let mut b = vec![0f32; n];
295 X::to_float(&xs[..n / X::BLCK_SIZE], &mut a);
296 Y::to_float(&ys[..n / Y::BLCK_SIZE], &mut b);
297 a.iter().zip(&b).map(|(x, y)| x * y).sum()
298}
299
300macro_rules! decode_only_type {
301 ($blk:ty, $dtype:expr) => {
302 impl GgmlType for $blk {
303 const DTYPE: GgmlDType = $dtype;
304 const BLCK_SIZE: usize = QK_K;
305 type VecDotType = BlockQ8K;
306
307 fn to_float(xs: &[Self], ys: &mut [f32]) {
308 for (b, y) in xs.iter().zip(ys.as_chunks_mut::<QK_K>().0.iter_mut()) {
309 b.decode(y);
310 }
311 }
312
313 fn from_float(_xs: &[f32], _ys: &mut [Self]) {
314 panic!(
317 "{:?} is decode-only: Fuel cannot quantize to it (GAP-211 B2)",
318 $dtype
319 );
320 }
321
322 fn vec_dot(n: usize, xs: &[Self], ys: &[Self::VecDotType]) -> f32 {
323 dot_via_to_float(n, xs, ys)
324 }
325
326 fn vec_dot_unopt(n: usize, xs: &[Self], ys: &[Self::VecDotType]) -> f32 {
327 dot_via_to_float(n, xs, ys)
328 }
329 }
330 };
331}
332
333decode_only_type!(BlockIq2Xxs, GgmlDType::Iq2Xxs);
334decode_only_type!(BlockIq2Xs, GgmlDType::Iq2Xs);
335decode_only_type!(BlockIq2S, GgmlDType::Iq2S);
336decode_only_type!(BlockIq3Xxs, GgmlDType::Iq3Xxs);
337decode_only_type!(BlockIq3S, GgmlDType::Iq3S);
338decode_only_type!(BlockIq1S, GgmlDType::Iq1S);
339decode_only_type!(BlockIq1M, GgmlDType::Iq1M);
340
341#[cfg(test)]
342mod tests {
343 use super::*;
344
345 fn one<T: GgmlType>(b: T) -> Vec<f32> {
346 let mut y = vec![0f32; QK_K];
347 T::to_float(&[b], &mut y);
348 y
349 }
350
351 #[test]
356 fn tables_have_the_structure_ggml_builds_them_with() {
357 for (i, s) in KSIGNS_IQ2XS.iter().enumerate() {
358 let parity = (i as u32).count_ones() & 1;
359 assert_eq!(u32::from(*s), i as u32 | (parity << 7), "ksigns[{i}]");
360 }
361 for (j, m) in KMASK_IQ2XS.iter().enumerate() {
362 assert_eq!(*m, 1u8 << j);
363 }
364 let all_in = |grid: &[u64], allowed: &[u8]| {
365 grid.iter()
366 .all(|e| e.to_le_bytes().iter().all(|b| allowed.contains(b)))
367 };
368 assert!(all_in(&IQ2XXS_GRID, &[0x08, 0x19, 0x2b]));
369 assert!(all_in(&IQ2XS_GRID, &[0x08, 0x19, 0x2b]));
370 assert!(all_in(&IQ2S_GRID, &[0x08, 0x19, 0x2b]));
371 assert!(all_in(&IQ1S_GRID, &[0xff, 0x00, 0x01]));
372 assert!(
373 IQ3S_GRID
374 .iter()
375 .all(|e| e.to_le_bytes().iter().all(|b| *b % 2 == 1 && *b <= 15))
376 );
377 }
378
379 #[test]
382 fn geometry_agrees_with_ggml_dtype() {
383 fn check<T: GgmlType>() {
384 assert_eq!(
385 std::mem::size_of::<T>(),
386 T::DTYPE.type_size(),
387 "{:?}",
388 T::DTYPE
389 );
390 assert_eq!(T::BLCK_SIZE, T::DTYPE.block_size(), "{:?}", T::DTYPE);
391 assert!(!T::DTYPE.supports_quantize());
392 assert!(T::DTYPE.check_quantize().is_err());
393 }
394 check::<BlockIq2Xxs>();
395 check::<BlockIq2Xs>();
396 check::<BlockIq2S>();
397 check::<BlockIq3Xxs>();
398 check::<BlockIq3S>();
399 check::<BlockIq1S>();
400 check::<BlockIq1M>();
401 assert!(GgmlDType::Q4K.supports_quantize());
402 }
403
404 #[test]
408 fn iq2_xxs_dequantizes_as_the_ggml_formula() {
409 let mut b = BlockIq2Xxs::zeros();
410 b.d = f16::from_f32(1.0);
411 b.qs[0] = 0x0100; b.qs[2] = 1; b.qs[3] = 0xF000; let y = one(b);
415 assert_eq!(y[0], -31.0);
416 assert_eq!(y[1], 31.0);
417 assert_eq!(y[7], -31.0);
418 assert_eq!(y[8], 3.875 * 43.0); assert_eq!(y[9], 31.0);
420 assert_eq!(y[32], 1.0); }
422
423 #[test]
426 fn iq2_xs_dequantizes_as_the_ggml_formula() {
427 let mut b = BlockIq2Xs::zeros();
428 b.d = f16::from_f32(1.0);
429 b.scales[0] = 0x21;
430 b.qs[0] = (1 << 9) | 1; let y = one(b);
432 assert_eq!(y[0], -(0.375 * 43.0));
433 assert_eq!(y[1], 3.0);
434 assert_eq!(y[7], -3.0);
435 assert_eq!(y[8], 3.0); assert_eq!(y[16], 5.0); }
438
439 #[test]
441 fn iq2_s_dequantizes_as_the_ggml_formula() {
442 let mut b = BlockIq2S::zeros();
443 b.d = f16::from_f32(1.0);
444 b.scales[0] = 0x21;
445 b.qs[0] = 1; b.qs[QK_K / 8] = 0x01; let y = one(b);
448 assert_eq!(y[0], -(0.375 * 43.0));
449 assert_eq!(y[1], 3.0);
450 assert_eq!(y[8], 3.0);
451 assert_eq!(y[16], 5.0);
452 }
453
454 #[test]
457 fn iq3_xxs_dequantizes_as_the_ggml_formula() {
458 let mut b = BlockIq3Xxs::zeros();
459 b.d = f16::from_f32(1.0);
460 b.qs[0] = 1;
461 let aux32: u32 = (3 << 28) | 1; b.qs[QK_K / 4..QK_K / 4 + 4].copy_from_slice(&aux32.to_le_bytes());
463 let y = one(b);
464 assert_eq!(y[0], -(1.75 * 20.0));
465 assert_eq!(y[1], 7.0);
466 assert_eq!(y[4], 7.0);
467 assert_eq!(y[7], -7.0);
468 }
469
470 #[test]
473 fn iq3_s_dequantizes_as_the_ggml_formula() {
474 let mut b = BlockIq3S::zeros();
475 b.d = f16::from_f32(1.0);
476 b.scales[0] = 0x21;
477 b.qs[0] = 1;
478 b.signs[0] = 0x01;
479 let y = one(b);
480 assert_eq!(y[0], -9.0);
481 assert_eq!(y[1], 3.0);
482 assert_eq!(y[4], 3.0);
483 assert_eq!(y[32], 5.0); }
485
486 #[test]
489 fn iq1_s_dequantizes_as_the_ggml_formula() {
490 let mut b = BlockIq1S::zeros();
491 b.d = f16::from_f32(1.0);
492 b.qs[0] = 1;
493 let y = one(b);
494 assert_eq!(y[0], 1.125);
495 assert_eq!(y[1], -0.875);
496 assert_eq!(y[8], -0.875);
497 b.qh[0] = 0x9000; let y = one(b);
499 assert_eq!(y[0], 3.0 * (1.0 - 0.125));
500 assert_eq!(y[1], 3.0 * (-1.0 - 0.125));
501 }
502
503 #[test]
507 fn iq1_m_dequantizes_as_the_ggml_formula() {
508 let mut b = BlockIq1M::zeros();
509 b.scales = [0x01, 0x00, 0x00, 0x00, 0x00, 0xc0, 0x00, 0x30];
510 b.qs[0] = 1;
511 let y = one(b);
512 assert_eq!(y[0], 3.0 * 1.125); assert_eq!(y[1], 3.0 * -0.875);
514 assert_eq!(y[8], 3.0 * -0.875);
515 assert_eq!(y[16], -0.875); assert_eq!(y[32], -0.875);
517 b.qh[0] = 0x08; let y = one(b);
519 assert_eq!(y[0], 3.0 * 0.875);
520 }
521
522 #[test]
524 fn quantizing_to_a_decode_only_type_is_a_typed_error_naming_it() {
525 let e = GgmlDType::Iq2Xxs.check_quantize().unwrap_err().to_string();
526 assert!(e.contains("Iq2Xxs") && e.contains("no quantizer"), "{e}");
527 }
528}