char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Sigbjørn Skjæretconvert : minor fixes for numpy 2.x (#23571)5d246a792

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
63.4 KiB1378 linesraw
1from __future__ import annotations
2from abc import ABC, abstractmethod
3from typing import Any, Callable, Sequence
4from math import log2, ceil
5
6from numpy.typing import DTypeLike
7
8from .constants import GGML_QUANT_SIZES, GGMLQuantizationType, QK_K
9from .lazy import LazyNumpyTensor
10
11import numpy as np
12
13
14def quant_shape_to_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]:
15    block_size, type_size = GGML_QUANT_SIZES[quant_type]
16    if shape[-1] % block_size != 0:
17        raise ValueError(f"Quantized tensor row size ({shape[-1]}) is not a multiple of {quant_type.name} block size ({block_size})")
18    return (*shape[:-1], shape[-1] // block_size * type_size)
19
20
21def quant_shape_from_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]:
22    block_size, type_size = GGML_QUANT_SIZES[quant_type]
23    if shape[-1] % type_size != 0:
24        raise ValueError(f"Quantized tensor bytes per row ({shape[-1]}) is not a multiple of {quant_type.name} type size ({type_size})")
25    return (*shape[:-1], shape[-1] // type_size * block_size)
26
27
28# This is faster than np.vectorize and np.apply_along_axis because it works on more than one row at a time
29def _apply_over_grouped_rows(func: Callable[[np.ndarray], np.ndarray], arr: np.ndarray, otype: DTypeLike, oshape: tuple[int, ...]) -> np.ndarray:
30    rows = arr.reshape((-1, arr.shape[-1]))
31    assert len(rows.shape)
32    osize = 1
33    for dim in oshape:
34        osize *= dim
35    out = np.empty(shape=osize, dtype=otype)
36    # compute over groups of 16 rows (arbitrary, but seems good for performance)
37    n_groups = (rows.shape[0] // 16) or 1
38    np.concatenate([func(group).ravel() for group in np.array_split(rows, n_groups)], axis=0, out=out)
39    return out.reshape(oshape)
40
41
42# round away from zero
43# ref: https://stackoverflow.com/a/59143326/22827863
44def np_roundf(n: np.ndarray) -> np.ndarray:
45    a = abs(n)
46    floored = np.floor(a)
47    b = floored + np.floor(2 * (a - floored))
48    return np.sign(n) * b
49
50
51class QuantError(Exception): ...
52
53
54_type_traits: dict[GGMLQuantizationType, type[__Quant]] = {}
55
56
57def quantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
58    if qtype == GGMLQuantizationType.F32:
59        return data.astype(np.float32, copy=False)
60    elif qtype == GGMLQuantizationType.F16:
61        return data.astype(np.float16, copy=False)
62    elif (q := _type_traits.get(qtype)) is not None:
63        return q.quantize(data)
64    else:
65        raise NotImplementedError(f"Quantization for {qtype.name} is not yet implemented")
66
67
68def dequantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
69    if qtype == GGMLQuantizationType.F32:
70        return data.view(np.float32)
71    elif qtype == GGMLQuantizationType.F16:
72        return data.view(np.float16).astype(np.float32)
73    elif (q := _type_traits.get(qtype)) is not None:
74        return q.dequantize(data)
75    else:
76        raise NotImplementedError(f"Dequantization for {qtype.name} is not yet implemented")
77
78
79class __Quant(ABC):
80    qtype: GGMLQuantizationType
81    block_size: int
82    type_size: int
83
84    grid: np.ndarray[Any, np.dtype[np.float32]] | None = None
85    grid_shape: tuple[int, int] = (0, 0)
86    grid_map: tuple[int | float, ...] = ()
87    grid_hex: bytes | None = None
88
89    def __init__(self):
90        return TypeError("Quant conversion classes can't have instances")
91
92    def __init_subclass__(cls, qtype: GGMLQuantizationType) -> None:
93        cls.qtype = qtype
94        cls.block_size, cls.type_size = GGML_QUANT_SIZES[qtype]
95        cls.__quantize_lazy: Any = LazyNumpyTensor._wrap_fn(
96            cls.__quantize_array,
97            meta_noop=(np.uint8, cls.__shape_to_bytes)
98        )
99        cls.__dequantize_lazy: Any = LazyNumpyTensor._wrap_fn(
100            cls.__dequantize_array,
101            meta_noop=(np.float32, cls.__shape_from_bytes)
102        )
103        assert qtype not in _type_traits
104        _type_traits[qtype] = cls
105
106    @classmethod
107    def init_grid(cls):
108        if cls.grid is not None or cls.grid_hex is None:
109            return
110
111        bits_per_elem = ceil(log2(len(cls.grid_map)))
112        assert bits_per_elem != 0, cls.qtype.name
113        elems_per_byte = 8 // bits_per_elem
114
115        grid = np.frombuffer(cls.grid_hex, dtype=np.uint8)
116        # decode hexadecimal chars from grid
117        grid = grid.reshape((-1, 2))
118        grid = (np.where(grid > 0x40, grid + 9, grid) & 0x0F) << np.array([4, 0], dtype=np.uint8).reshape((1, 2))
119        grid = grid[..., 0] | grid[..., 1]
120        # unpack the grid values
121        grid = grid.reshape((-1, 1)) >> np.array([i for i in range(0, 8, 8 // elems_per_byte)], dtype=np.uint8).reshape((1, elems_per_byte))
122        grid = (grid & ((1 << bits_per_elem) - 1)).reshape((-1, 1))
123        grid_map = np.array(cls.grid_map, dtype=np.float32).reshape((1, -1))
124        grid = np.take_along_axis(grid_map, grid, axis=-1)
125        cls.grid = grid.reshape((1, 1, *cls.grid_shape))
126
127    @classmethod
128    @abstractmethod
129    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
130        raise NotImplementedError
131
132    @classmethod
133    @abstractmethod
134    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
135        raise NotImplementedError
136
137    @classmethod
138    def quantize_rows(cls, rows: np.ndarray) -> np.ndarray:
139        rows = rows.astype(np.float32, copy=False)
140        shape = rows.shape
141        n_blocks = rows.size // cls.block_size
142        blocks = rows.reshape((n_blocks, cls.block_size))
143        blocks = cls.quantize_blocks(blocks)
144        assert blocks.dtype == np.uint8
145        assert blocks.shape[-1] == cls.type_size
146        return blocks.reshape(cls.__shape_to_bytes(shape))
147
148    @classmethod
149    def dequantize_rows(cls, rows: np.ndarray) -> np.ndarray:
150        rows = rows.view(np.uint8)
151        shape = rows.shape
152        n_blocks = rows.size // cls.type_size
153        blocks = rows.reshape((n_blocks, cls.type_size))
154        blocks = cls.dequantize_blocks(blocks)
155        assert blocks.dtype == np.float32
156        assert blocks.shape[-1] == cls.block_size
157        return blocks.reshape(cls.__shape_from_bytes(shape))
158
159    @classmethod
160    def __shape_to_bytes(cls, shape: Sequence[int]):
161        return quant_shape_to_byte_shape(shape, cls.qtype)
162
163    @classmethod
164    def __shape_from_bytes(cls, shape: Sequence[int]):
165        return quant_shape_from_byte_shape(shape, cls.qtype)
166
167    @classmethod
168    def __quantize_array(cls, array: np.ndarray) -> np.ndarray:
169        return _apply_over_grouped_rows(cls.quantize_rows, arr=array, otype=np.uint8, oshape=cls.__shape_to_bytes(array.shape))
170
171    @classmethod
172    def __dequantize_array(cls, array: np.ndarray) -> np.ndarray:
173        cls.init_grid()
174        return _apply_over_grouped_rows(cls.dequantize_rows, arr=array, otype=np.float32, oshape=cls.__shape_from_bytes(array.shape))
175
176    @classmethod
177    def __quantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any:
178        pass
179
180    @classmethod
181    def __dequantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any:
182        pass
183
184    @classmethod
185    def can_quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> bool:
186        return tensor.shape[-1] % cls.block_size == 0
187
188    @classmethod
189    def quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray:
190        if not cls.can_quantize(tensor):
191            raise QuantError(f"Can't quantize tensor with shape {tensor.shape} to {cls.qtype.name}")
192        if isinstance(tensor, LazyNumpyTensor):
193            return cls.__quantize_lazy(tensor)
194        else:
195            return cls.__quantize_array(tensor)
196
197    @classmethod
198    def dequantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray:
199        if isinstance(tensor, LazyNumpyTensor):
200            return cls.__dequantize_lazy(tensor)
201        else:
202            return cls.__dequantize_array(tensor)
203
204
205class BF16(__Quant, qtype=GGMLQuantizationType.BF16):
206    @classmethod
207    # same as ggml_compute_fp32_to_bf16 in ggml-impl.h
208    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
209        n = blocks.view(np.uint32)
210        # force nan to quiet
211        n = np.where((n & 0x7fffffff) > 0x7f800000, (n & np.uint32(0xffff0000)) | np.uint32(64 << 16), n)
212        # round to nearest even
213        n = (np.uint64(n) + (0x7fff + ((n >> 16) & 1))) >> 16
214        return n.astype(np.uint16).view(np.uint8)
215
216    @classmethod
217    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
218        return (blocks.view(np.int16).astype(np.int32) << 16).view(np.float32)
219
220
221class Q4_0(__Quant, qtype=GGMLQuantizationType.Q4_0):
222    @classmethod
223    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
224        n_blocks = blocks.shape[0]
225
226        imax = abs(blocks).argmax(axis=-1, keepdims=True)
227        max = np.take_along_axis(blocks, imax, axis=-1)
228
229        d = max / -8
230        with np.errstate(divide="ignore"):
231            id = np.where(d == 0, 0, 1 / d)
232        qs = np.trunc((blocks * id) + np.float32(8.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
233
234        qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
235        qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
236
237        d = d.astype(np.float16).view(np.uint8)
238
239        return np.concatenate([d, qs], axis=-1)
240
241    @classmethod
242    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
243        n_blocks = blocks.shape[0]
244
245        d, qs = np.hsplit(blocks, [2])
246
247        d = d.view(np.float16).astype(np.float32)
248
249        qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
250        qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.int8) - np.int8(8)
251
252        return (d * qs.astype(np.float32))
253
254
255class Q4_1(__Quant, qtype=GGMLQuantizationType.Q4_1):
256    @classmethod
257    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
258        n_blocks = blocks.shape[0]
259
260        max = blocks.max(axis=-1, keepdims=True)
261        min = blocks.min(axis=-1, keepdims=True)
262
263        d = (max - min) / 15
264        with np.errstate(divide="ignore"):
265            id = np.where(d == 0, 0, 1 / d)
266        qs = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
267
268        qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
269        qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
270
271        d = d.astype(np.float16).view(np.uint8)
272        m = min.astype(np.float16).view(np.uint8)
273
274        return np.concatenate([d, m, qs], axis=-1)
275
276    @classmethod
277    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
278        n_blocks = blocks.shape[0]
279
280        d, rest = np.hsplit(blocks, [2])
281        m, qs = np.hsplit(rest, [2])
282
283        d = d.view(np.float16).astype(np.float32)
284        m = m.view(np.float16).astype(np.float32)
285
286        qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
287        qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.float32)
288
289        return (d * qs) + m
290
291
292class Q5_0(__Quant, qtype=GGMLQuantizationType.Q5_0):
293    @classmethod
294    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
295        n_blocks = blocks.shape[0]
296
297        imax = abs(blocks).argmax(axis=-1, keepdims=True)
298        max = np.take_along_axis(blocks, imax, axis=-1)
299
300        d = max / -16
301        with np.errstate(divide="ignore"):
302            id = np.where(d == 0, 0, 1 / d)
303        q = np.trunc((blocks * id) + np.float32(16.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
304
305        qs = q.reshape((n_blocks, 2, cls.block_size // 2))
306        qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
307
308        qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4)
309
310        d = d.astype(np.float16).view(np.uint8)
311
312        return np.concatenate([d, qh, qs], axis=-1)
313
314    @classmethod
315    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
316        n_blocks = blocks.shape[0]
317
318        d, rest = np.hsplit(blocks, [2])
319        qh, qs = np.hsplit(rest, [4])
320
321        d = d.view(np.float16).astype(np.float32)
322        qh = qh.view(np.uint32)
323
324        qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32))
325        ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
326        qh = (qh & np.uint32(0x01)).astype(np.uint8)
327        ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1))
328
329        qs = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(16)
330
331        return (d * qs.astype(np.float32))
332
333
334class Q5_1(__Quant, qtype=GGMLQuantizationType.Q5_1):
335    @classmethod
336    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
337        n_blocks = blocks.shape[0]
338
339        max = blocks.max(axis=-1, keepdims=True)
340        min = blocks.min(axis=-1, keepdims=True)
341
342        d = (max - min) / 31
343        with np.errstate(divide="ignore"):
344            id = np.where(d == 0, 0, 1 / d)
345        q = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
346
347        qs = q.reshape((n_blocks, 2, cls.block_size // 2))
348        qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
349
350        qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4)
351
352        d = d.astype(np.float16).view(np.uint8)
353        m = min.astype(np.float16).view(np.uint8)
354
355        return np.concatenate([d, m, qh, qs], axis=-1)
356
357    @classmethod
358    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
359        n_blocks = blocks.shape[0]
360
361        d, rest = np.hsplit(blocks, [2])
362        m, rest = np.hsplit(rest, [2])
363        qh, qs = np.hsplit(rest, [4])
364
365        d = d.view(np.float16).astype(np.float32)
366        m = m.view(np.float16).astype(np.float32)
367        qh = qh.view(np.uint32)
368
369        qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32))
370        ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
371        qh = (qh & np.uint32(0x01)).astype(np.uint8)
372        ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1))
373
374        qs = (ql | (qh << np.uint8(4))).astype(np.float32)
375
376        return (d * qs) + m
377
378
379class Q8_0(__Quant, qtype=GGMLQuantizationType.Q8_0):
380    @classmethod
381    # Implementation of Q8_0 with bit-exact same results as reference implementation in ggml-quants.c
382    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
383
384        d = abs(blocks).max(axis=1, keepdims=True) / 127
385        with np.errstate(divide="ignore"):
386            id = np.where(d == 0, 0, 1 / d)
387        qs = np_roundf(blocks * id)
388
389        # (n_blocks, 2)
390        d = d.astype(np.float16).view(np.uint8)
391        # (n_blocks, block_size)
392        qs = qs.astype(np.int8).view(np.uint8)
393
394        return np.concatenate([d, qs], axis=1)
395
396    @classmethod
397    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
398        d, x = np.split(blocks, [2], axis=1)
399        d = d.view(np.float16).astype(np.float32)
400        x = x.view(np.int8).astype(np.float32)
401
402        return (x * d)
403
404
405class Q2_K(__Quant, qtype=GGMLQuantizationType.Q2_K):
406    @classmethod
407    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
408        n_blocks = blocks.shape[0]
409
410        scales, rest = np.hsplit(blocks, [QK_K // 16])
411        qs, rest = np.hsplit(rest, [QK_K // 4])
412        d, dmin = np.hsplit(rest, [2])
413
414        d = d.view(np.float16).astype(np.float32)
415        dmin = dmin.view(np.float16).astype(np.float32)
416
417        # (n_blocks, 16, 1)
418        dl = (d * (scales & 0xF).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1))
419        ml = (dmin * (scales >> 4).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1))
420
421        shift = np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
422
423        qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & np.uint8(3)
424
425        qs = qs.reshape((n_blocks, QK_K // 16, 16)).astype(np.float32)
426
427        qs = dl * qs - ml
428
429        return qs.reshape((n_blocks, -1))
430
431
432class Q3_K(__Quant, qtype=GGMLQuantizationType.Q3_K):
433    @classmethod
434    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
435        n_blocks = blocks.shape[0]
436
437        hmask, rest = np.hsplit(blocks, [QK_K // 8])
438        qs, rest = np.hsplit(rest, [QK_K // 4])
439        scales, d = np.hsplit(rest, [12])
440
441        d = d.view(np.float16).astype(np.float32)
442
443        # The scales are packed at 6-bit each in this pattern:
444        #  0: IIIIAAAA
445        #  1: JJJJBBBB
446        #  2: KKKKCCCC
447        #  3: LLLLDDDD
448        #  4: MMMMEEEE
449        #  5: NNNNFFFF
450        #  6: OOOOGGGG
451        #  7: PPPPHHHH
452        #  8: MMIIEEAA
453        #  9: NNJJFFBB
454        # 10: OOKKGGCC
455        # 11: PPLLHHDD
456        lscales, hscales = np.hsplit(scales, [8])
457        lscales = lscales.reshape((n_blocks, 1, 8)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1))
458        lscales = lscales.reshape((n_blocks, 16))
459        hscales = hscales.reshape((n_blocks, 1, 4)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 4, 1))
460        hscales = hscales.reshape((n_blocks, 16))
461        scales = (lscales & np.uint8(0x0F)) | ((hscales & np.uint8(0x03)) << np.uint8(4))
462        scales = (scales.astype(np.int8) - np.int8(32)).astype(np.float32)
463
464        dl = (d * scales).reshape((n_blocks, 16, 1))
465
466        ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
467        qh = hmask.reshape(n_blocks, -1, 1, 32) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1))
468        ql = ql.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(3)
469        qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(1))
470        qh = qh ^ np.uint8(1)  # strangely, the offset is zero when the bitmask is 1
471        q = (ql.astype(np.int8) - (qh << np.uint8(2)).astype(np.int8)).astype(np.float32)
472
473        return (dl * q).reshape((n_blocks, QK_K))
474
475
476class Q4_K(__Quant, qtype=GGMLQuantizationType.Q4_K):
477    K_SCALE_SIZE = 12
478
479    @staticmethod
480    def get_scale_min(scales: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
481        n_blocks = scales.shape[0]
482        scales = scales.view(np.uint8)
483        ### Unpacking the following: ###
484        #  0 EEAAAAAA
485        #  1 FFBBBBBB
486        #  2 GGCCCCCC
487        #  3 HHDDDDDD
488        #  4 eeaaaaaa
489        #  5 ffbbbbbb
490        #  6 ggcccccc
491        #  7 hhdddddd
492        #  8 eeeeEEEE
493        #  9 ffffFFFF
494        # 10 ggggGGGG
495        # 11 hhhhHHHH
496        scales = scales.reshape((n_blocks, 3, 4))
497        d, m, m_d = np.split(scales, 3, axis=-2)
498
499        sc = np.concatenate([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], axis=-1)
500        min = np.concatenate([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], axis=-1)
501
502        return (sc.reshape((n_blocks, 8)), min.reshape((n_blocks, 8)))
503
504    @classmethod
505    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
506        n_blocks = blocks.shape[0]
507
508        d, rest = np.hsplit(blocks, [2])
509        dmin, rest = np.hsplit(rest, [2])
510        scales, qs = np.hsplit(rest, [cls.K_SCALE_SIZE])
511
512        d = d.view(np.float16).astype(np.float32)
513        dmin = dmin.view(np.float16).astype(np.float32)
514
515        sc, m = Q4_K.get_scale_min(scales)
516
517        d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1))
518        dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1))
519
520        qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
521        qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 32)).astype(np.float32)
522
523        return (d * qs - dm).reshape((n_blocks, QK_K))
524
525
526class Q5_K(__Quant, qtype=GGMLQuantizationType.Q5_K):
527    @classmethod
528    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
529        n_blocks = blocks.shape[0]
530
531        d, rest = np.hsplit(blocks, [2])
532        dmin, rest = np.hsplit(rest, [2])
533        scales, rest = np.hsplit(rest, [Q4_K.K_SCALE_SIZE])
534        qh, qs = np.hsplit(rest, [QK_K // 8])
535
536        d = d.view(np.float16).astype(np.float32)
537        dmin = dmin.view(np.float16).astype(np.float32)
538
539        sc, m = Q4_K.get_scale_min(scales)
540
541        d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1))
542        dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1))
543
544        ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
545        qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1))
546        ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32))
547        qh = (qh & np.uint8(0x01)).reshape((n_blocks, -1, 32))
548        q = (ql | (qh << np.uint8(4))).astype(np.float32)
549
550        return (d * q - dm).reshape((n_blocks, QK_K))
551
552
553class Q6_K(__Quant, qtype=GGMLQuantizationType.Q6_K):
554    @classmethod
555    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
556        n_blocks = blocks.shape[0]
557
558        ql, rest = np.hsplit(blocks, [QK_K // 2])
559        qh, rest = np.hsplit(rest, [QK_K // 4])
560        scales, d = np.hsplit(rest, [QK_K // 16])
561
562        scales = scales.view(np.int8).astype(np.float32)
563        d = d.view(np.float16).astype(np.float32)
564        d = (d * scales).reshape((n_blocks, QK_K // 16, 1))
565
566        ql = ql.reshape((n_blocks, -1, 1, 64)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
567        ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32))
568        qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
569        qh = (qh & np.uint8(0x03)).reshape((n_blocks, -1, 32))
570        q = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(32)
571        q = q.reshape((n_blocks, QK_K // 16, -1)).astype(np.float32)
572
573        return (d * q).reshape((n_blocks, QK_K))
574
575
576class TQ1_0(__Quant, qtype=GGMLQuantizationType.TQ1_0):
577    @classmethod
578    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
579        n_blocks = blocks.shape[0]
580
581        d = abs(blocks).max(axis=-1, keepdims=True)
582        with np.errstate(divide="ignore"):
583            id = np.where(d == 0, 0, 1 / d)
584        qs = np_roundf(blocks * id)
585        qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8)
586
587        qs0, qs1, qh = qs[..., :(32 * 5)], qs[..., (32 * 5):(48 * 5)], qs[..., (48 * 5):]
588        qs0 = qs0.reshape((n_blocks, -1, 5, 32)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1))
589        qs0 = np.sum(qs0, axis=-2).reshape((n_blocks, -1))
590        qs1 = qs1.reshape((n_blocks, -1, 5, 16)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1))
591        qs1 = np.sum(qs1, axis=-2).reshape((n_blocks, -1))
592        qh = qh.reshape((n_blocks, -1, 4, 4)) * np.array([81, 27, 9, 3], dtype=np.uint8).reshape((1, 1, 4, 1))
593        qh = np.sum(qh, axis=-2).reshape((n_blocks, -1))
594        qs = np.concatenate([qs0, qs1, qh], axis=-1)
595        qs = (qs.astype(np.uint16) * 256 + (243 - 1)) // 243
596
597        qs = qs.astype(np.uint8)
598        d = d.astype(np.float16).view(np.uint8)
599
600        return np.concatenate([qs, d], axis=-1)
601
602    @classmethod
603    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
604        n_blocks = blocks.shape[0]
605
606        qs, rest = np.hsplit(blocks, [(QK_K - 4 * QK_K // 64) // 5])
607        qh, d = np.hsplit(rest, [QK_K // 64])
608
609        d = d.view(np.float16).astype(np.float32)
610
611        qs0, qs1 = qs[..., :32], qs[..., 32:]
612        qs0 = qs0.reshape((n_blocks, -1, 1, 32)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1))
613        qs0 = qs0.reshape((n_blocks, -1))
614        qs1 = qs1.reshape((n_blocks, -1, 1, 16)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1))
615        qs1 = qs1.reshape((n_blocks, -1))
616        qh = qh.reshape((n_blocks, -1, 1, 4)) * np.array([1, 3, 9, 27], dtype=np.uint8).reshape((1, 1, 4, 1))
617        qh = qh.reshape((n_blocks, -1))
618        qs = np.concatenate([qs0, qs1, qh], axis=-1)
619        qs = ((qs.astype(np.uint16) * 3) >> 8).astype(np.int8) - np.int8(1)
620
621        return (d * qs.astype(np.float32))
622
623
624class TQ2_0(__Quant, qtype=GGMLQuantizationType.TQ2_0):
625    @classmethod
626    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
627        n_blocks = blocks.shape[0]
628
629        d = abs(blocks).max(axis=-1, keepdims=True)
630        with np.errstate(divide="ignore"):
631            id = np.where(d == 0, 0, 1 / d)
632        qs = np_roundf(blocks * id)
633        qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8)
634
635        qs = qs.reshape((n_blocks, -1, 4, 32)) << np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
636        qs = qs[..., 0, :] | qs[..., 1, :] | qs[..., 2, :] | qs[..., 3, :]
637        qs = qs.reshape((n_blocks, -1))
638
639        d = d.astype(np.float16).view(np.uint8)
640
641        return np.concatenate([qs, d], axis=-1)
642
643    @classmethod
644    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
645        n_blocks = blocks.shape[0]
646
647        qs, d = np.hsplit(blocks, [QK_K // 4])
648
649        d = d.view(np.float16).astype(np.float32)
650
651        qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
652        qs = (qs & 0x03).reshape((n_blocks, -1)).astype(np.int8) - np.int8(1)
653
654        return (d * qs.astype(np.float32))
655
656
657class MXFP4(__Quant, qtype=GGMLQuantizationType.MXFP4):
658    # e2m1 values (doubled)
659    # ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
660    kvalues = (0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12)
661
662    @staticmethod
663    # see ggml_e8m0_to_fp32_half in ggml-impl.h
664    def e8m0_to_fp32_half(x: np.ndarray) -> np.ndarray:
665        bits = np.where(x < 2, np.uint32(0x00200000) << np.uint32(x), np.uint32(x - 1) << np.uint32(23))
666        return bits.view(np.float32)
667
668    @classmethod
669    def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
670        n_blocks = blocks.shape[0]
671
672        d = abs(blocks).max(axis=-1, keepdims=True)
673
674        with np.errstate(divide="ignore"):
675            e = np.where(d > 0, np.floor(np.log2(d)) - 2 + 127, 0).astype(np.uint8)
676
677        d = cls.e8m0_to_fp32_half(e)
678
679        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape((1, 1, 16))
680
681        errs = np.abs(d.reshape((n_blocks, 1, 1)) * kvalues.astype(np.float32) - blocks.reshape((n_blocks, cls.block_size, 1)))
682        best = np.argmin(errs, axis=-1, keepdims=True)
683
684        qs = best.reshape(n_blocks, 2, cls.block_size // 2).astype(np.uint8)
685        qs = qs[:, 0] | (qs[:, 1] << np.uint8(4))
686
687        qs = qs.reshape((n_blocks, cls.block_size // 2))
688
689        return np.concatenate([e, qs], axis=-1)
690
691    @classmethod
692    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
693        n_blocks = blocks.shape[0]
694
695        e, qs = np.hsplit(blocks, [1])
696
697        d = cls.e8m0_to_fp32_half(e)
698
699        qs = qs.reshape((n_blocks, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1))
700        qs = (qs & np.uint8(0x0F)).view(np.int8)
701
702        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
703        qs = np.take_along_axis(kvalues, qs, axis=-1).reshape((n_blocks, cls.block_size))
704
705        return (d * qs.astype(np.float32))
706
707
708class NVFP4(__Quant, qtype=GGMLQuantizationType.NVFP4):
709    # E2M1 values doubled (kvalues_mxfp4 convention)
710    kvalues = (0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12)
711
712    @staticmethod
713    def ue4m3_to_fp32(x: np.ndarray) -> np.ndarray:
714        """Decode unsigned E4M3 (bias=7) to float, with 0.5 factor for kvalues convention."""
715        exp = (x >> 3).astype(np.int32) & 0xF
716        man = (x & 0x7).astype(np.float32)
717        raw = np.where(
718            exp == 0,
719            man * 2**-9,
720            (1.0 + man / 8.0) * (2.0 ** (exp.astype(np.float32) - 7)))
721        return np.where((x == 0) | (x == 0x7F), 0.0, raw * 0.5)
722
723    @staticmethod
724    def fp32_to_ue4m3(x: np.ndarray) -> np.ndarray:
725        """Vectorized float32 to unsigned E4M3, matching ggml_fp32_to_ue4m3 in C."""
726        x = np.clip(x, 0.0, 448.0).astype(np.float32)
727        bits = x.view(np.uint32)
728        fp32_exp = ((bits >> 23) & 0xFF).astype(np.int32) - 127
729        fp32_man = ((bits >> 20) & 0x7).astype(np.int32)
730        ue4m3_exp = fp32_exp + 7
731
732        # Subnormal
733        sub_man = np.clip((x * 512.0 + 0.5).astype(np.int32), 0, 7)
734        sub_result = np.where(sub_man >= 1, sub_man, 0).astype(np.uint8)
735
736        # Normal with rounding
737        round_bit = ((bits >> 19) & 1).astype(np.int32)
738        man = fp32_man + round_bit
739        exp = ue4m3_exp.copy()
740        overflow = man > 7
741        man = np.where(overflow, 0, man)
742        exp = np.where(overflow, exp + 1, exp)
743        normal_result = np.where(exp >= 15, np.uint8(0x7E), ((exp << 3) | man).astype(np.uint8))
744
745        return np.where(x <= 0.0, np.uint8(0),
746                        np.where(ue4m3_exp <= 0, sub_result,
747                        np.where(ue4m3_exp >= 15, np.uint8(0x7E), normal_result)))
748
749    @classmethod
750    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
751        n_super = blocks.shape[0]
752
753        d_bytes, qs = np.hsplit(blocks, [4])
754        d = cls.ue4m3_to_fp32(d_bytes).reshape(n_super, 4, 1)  # (n_super, 4, 1)
755
756        qs = qs.reshape(n_super, 4, 8)
757        lo = (qs & np.uint8(0x0F)).view(np.int8)
758        hi = (qs >> np.uint8(4)).view(np.int8)
759        vals = np.concatenate([lo, hi], axis=-1)  # (n_super, 4, 16)
760
761        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
762        vals = np.take_along_axis(kvalues, vals, axis=-1)
763
764        return (d * vals.astype(np.float32)).reshape(n_super, 64)
765
766
767class IQ2_XXS(__Quant, qtype=GGMLQuantizationType.IQ2_XXS):
768    ksigns: bytes = (
769        b"\x00\x81\x82\x03\x84\x05\x06\x87\x88\x09\x0a\x8b\x0c\x8d\x8e\x0f"
770        b"\x90\x11\x12\x93\x14\x95\x96\x17\x18\x99\x9a\x1b\x9c\x1d\x1e\x9f"
771        b"\xa0\x21\x22\xa3\x24\xa5\xa6\x27\x28\xa9\xaa\x2b\xac\x2d\x2e\xaf"
772        b"\x30\xb1\xb2\x33\xb4\x35\x36\xb7\xb8\x39\x3a\xbb\x3c\xbd\xbe\x3f"
773        b"\xc0\x41\x42\xc3\x44\xc5\xc6\x47\x48\xc9\xca\x4b\xcc\x4d\x4e\xcf"
774        b"\x50\xd1\xd2\x53\xd4\x55\x56\xd7\xd8\x59\x5a\xdb\x5c\xdd\xde\x5f"
775        b"\x60\xe1\xe2\x63\xe4\x65\x66\xe7\xe8\x69\x6a\xeb\x6c\xed\xee\x6f"
776        b"\xf0\x71\x72\xf3\x74\xf5\xf6\x77\x78\xf9\xfa\x7b\xfc\x7d\x7e\xff"
777    )
778
779    # iq2xxs_grid, but with each byte of the original packed in 2 bits,
780    # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
781    grid_shape = (256, 8)
782    grid_map = (0x08, 0x19, 0x2b)
783    grid_hex = (
784        b"00000200050008000a00110014002000220028002a0041004400500058006100"
785        b"6400800082008a00a20001010401100115014001840198010002020222028202"
786        b"010404041004210424044004420448046004810484049004a404000502050805"
787        b"200546056905800591050906100640068406a406000805080808140828084108"
788        b"440850085208880804094009020a140a01100410101021104010601084109010"
789        b"951000110811201150115a118011241245120014081420142514491480141815"
790        b"6215001616160118041810184018811800190519a019511a002002200a204420"
791        b"6120802082202921482100220222012404241024402456240025412564259026"
792        b"082820289428442a014004401040184021402440404048405640604081408440"
793        b"9040004120416141804185410142104248425642684200440844204480449944"
794        b"124524450046014804481048404845480049584961498249454a904a00500850"
795        b"1150195020508050885004514251a4519152905492540a550156545600581158"
796        b"195864584059085a046010604060686000615561186260620064056410651265"
797        b"84654268008002800a8041808280048118814081118201840484108415844084"
798        b"608400854685948509864086608602880489118a0490109024904090a1901691"
799        b"8091459200942294449451958198209902a050a085a009a100a218a450a804a9"
800    )
801
802    @classmethod
803    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
804        n_blocks = blocks.shape[0]
805
806        d, qs = np.hsplit(blocks, [2])
807
808        d = d.view(np.float16).astype(np.float32)
809
810        qs = qs.view(np.uint32).reshape(n_blocks, -1, 2)
811
812        db = d * (np.float32(0.5) + (qs[..., 1] >> 28).astype(np.float32)) * np.float32(0.25)
813        db = db.reshape((n_blocks, -1, 1, 1))
814
815        # get the sign indices and unpack the bits
816        signs = qs[..., 1].reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4))
817        ksigns = np.frombuffer(cls.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128))
818        signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1))
819        signs = np.take_along_axis(ksigns, signs, axis=-1)
820        signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 1, 8))
821        signs = signs & np.uint8(0x01)
822        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
823        signs = signs.reshape((n_blocks, -1, 4, 8))
824
825        assert cls.grid is not None
826        grid = np.take_along_axis(cls.grid, qs[..., 0].copy().view(np.uint8).reshape((n_blocks, -1, 1, 1)), axis=-2)
827        grid = grid.reshape((n_blocks, -1, 4, 8))
828
829        return (db * grid * signs).reshape((n_blocks, -1))
830
831
832class IQ2_XS(__Quant, qtype=GGMLQuantizationType.IQ2_XS):
833    # iq2xs_grid, but with each byte of the original packed in 2 bits,
834    # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
835    grid_shape = (512, 8)
836    grid_map = (0x08, 0x19, 0x2b)
837    grid_hex = (
838        b"00000200050008000a0011001400160019002000220025002800410044004600"
839        b"49005000520055005800610064008000820085008800910094009900a0000101"
840        b"04010601090110011201150118011a0121012401400142014501480151015401"
841        b"6001680181018401900100020202050208021102140220024102440250025502"
842        b"80028a0201040404060409041004120415041804210424044004420445044804"
843        b"5104540456046004810484049004000502050505080511051405200541054405"
844        b"500561058005010604061006260640064206840600080208050808080a081108"
845        b"14082008250841084408500858088008a008aa08010904091009400981098909"
846        b"000a200a280a960aa00a01100410061009101010121015101810211024104010"
847        b"4210451048105110541060106a10811084109010001102110511081111111411"
848        b"2011411144115011801194119611011204120612101240126012001402140514"
849        b"0814111414142014411444144914501464148014011504151015401500161416"
850        b"49160118041810181218401854188618001905196619511aa91a002002200520"
851        b"08200a201120142020204120442050208020a020012104211021402148216521"
852        b"002222228022a82201240424102429244024002541255225992501261a26a626"
853        b"002808280a28202855288828a22868299029082a202a822a882a8a2a01400440"
854        b"0640094010401240154018402140244040404240454048404a40514054406040"
855        b"6540814084409040004102410541084111411441204141414441504180418541"
856        b"a241014204421042124229424042004402440544084411441444194420444144"
857        b"4444504480449444014504451045244540459a4500460a464446504601480448"
858        b"1048404845485448624800491149444950496949044a00500250055008501150"
859        b"145020502850415044505050805001510451105115514051425100524452aa52"
860        b"0154045410542154405460548154a154005508558055885521566856a1560058"
861        b"14584158505899581a5940594259855a0160046010604060546062608660a960"
862        b"006124624a62926200641664106540654565a46501686a682569066a546a626a"
863        b"00800280058008801180148020802a8041804480508080808280a880aa800181"
864        b"0481068110814081518159810082208280828282a082a8820184048410841284"
865        b"158440846084898400854485a58518866a860088088825885a8880888288a888"
866        b"0689228a808a888a968aa88a0190049010904090569084900091229164915692"
867        b"89920094059444945094589429959095929541965198a6984999159a609a00a0"
868        b"02a008a00aa020a02aa0a0a051a159a1a6a100a202a208a22aa280a2a0a240a4"
869        b"95a465a698a60aa820a822a828a8a0a8a8a804a984a986a928aa2aaa91aaaaaa"
870    )
871
872    @classmethod
873    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
874        n_blocks = blocks.shape[0]
875
876        d, rest = np.hsplit(blocks, [2])
877        qs, scales = np.hsplit(rest, [2 * QK_K // 8])
878
879        d = d.view(np.float16).astype(np.float32)
880        qs = qs.view(np.uint16)
881
882        scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
883        scales = (scales & 0x0F).reshape((n_blocks, -1))
884        db = d * (np.float32(0.5) + scales) * np.float32(0.25)
885        db = db.reshape((n_blocks, -1, 1, 1))
886
887        # get the sign indices and unpack the bits
888        signs = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape(1, 1, 128)
889        signs = np.take_along_axis(signs, (qs >> 9).reshape((n_blocks, -1, 1)), axis=-1)
890        signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
891        signs = signs & np.uint8(0x01)
892        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
893        signs = signs.reshape((n_blocks, -1, 2, 8))
894
895        assert cls.grid is not None
896        grid = np.take_along_axis(cls.grid, (qs & np.uint16(511)).reshape((n_blocks, -1, 1, 1)), axis=-2)
897        grid = grid.reshape((n_blocks, -1, 2, 8))
898
899        return (db * grid * signs).reshape((n_blocks, -1))
900
901
902class IQ2_S(__Quant, qtype=GGMLQuantizationType.IQ2_S):
903    # iq2s_grid, but with each byte of the original packed in 2 bits,
904    # by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
905    grid_shape = (1024, 8)
906    grid_map = (0x08, 0x19, 0x2b)
907    grid_hex = (
908        b"00000200050008000a0011001400160019002000220025002800410044004600"
909        b"490050005200550058006100640066006900800082008500880091009400a000"
910        b"a500aa0001010401060109011001120115011801210124014001420145014801"
911        b"510154015601590160016501680181018401900192019501a101a40100020202"
912        b"050208021102140220022a02410244024602490250025502800285028a029402"
913        b"a202010404040604090410041204150418042104240426042904400442044504"
914        b"48044a0451045404560459046004620465048104840486048904900495049804"
915        b"a104a40400050205050508050a05110514051605190520052505280541054405"
916        b"46054905500552055505580561056405800582058505880591059405a0050106"
917        b"0406060609061006150640064506480651065406600681068406900600080208"
918        b"050808081108140816081908200825082a084108440846084908500852085508"
919        b"580861086408800885089408aa08010904091009120915091809210940094509"
920        b"480951095409600981099009000a110a140a220a280a2a0a500a990a01100410"
921        b"0610091010101210151018102110241026104010421045104810511054105610"
922        b"59106010621065106810811084108610901095109810a110a410001102110511"
923        b"08110a1111111411161119112011221125112811411144114611491150115211"
924        b"5511581161116411801182118511881191119411011204120912101215122112"
925        b"2412401245125112541281128412901200140214051408141114141416141914"
926        b"2014251428144114441446144914501452145514581461146414801482148514"
927        b"881491149414a014011504150615091510151215151518152115241540154215"
928        b"4515481551155415601581158415901500160516081611161416201641164416"
929        b"50168016aa160118041806180918101815181818211840184218451848185118"
930        b"541860188118841800190219051908191119141920194119441950196919a219"
931        b"041a101a401a561a00200220052008201120142016201920202025202a204120"
932        b"4420502052205520642080208a209420aa200121042110211221152121214021"
933        b"4221452151215421602181218421902100220a22222228222a22442250228822"
934        b"8a22a82201240424062409241024152418242124242440244224452448245124"
935        b"5424602481248424902400250525082511251425202541254425502566258025"
936        b"0126042610264026592600280528112814284128442850288a28aa2801290429"
937        b"102995290a2a222a642a882a8a2a014004400640094010401240154018401a40"
938        b"21402440264040404240454048404a4051405440564059406040624065408140"
939        b"8440904095409840a140a4400041024105410841114114411641194120412241"
940        b"2541414144414641494150415241554158416141644180418241854188419141"
941        b"9441a04101420442104212421542184224424042454248425142544260428142"
942        b"844200440244054408440a441144144416441944204422442544284441444444"
943        b"46444944504452445544584461446444804482448544884491449444a0440145"
944        b"0445064509451045124515451845214524454045424545454845514554456045"
945        b"6a4581458445904500460246054608461146144620464146444650468046a546"
946        b"0148044809481048124815481848214824484048424845484848514854486048"
947        b"84489048004902490549084911491449204941494449504980499649014a044a"
948        b"104a404a00500250055008501150145016501950205022502550285041504450"
949        b"4650495050505250555058506150645080508250855088509150945001510451"
950        b"0651095110511251155118512151245140514251455148515151545160518151"
951        b"8451905100520552085211521452205241524452505269528052015404540654"
952        b"0954105412541554185421542454405442544554485451545454605481548454"
953        b"9054005502550555085511551455205541554455505580550156045610562656"
954        b"405600580258055808581158145820584158445850585a588058015904591059"
955        b"4059005a195a855aa85a01600460066010601260156018602160246040604560"
956        b"4860516054606060846090600061026105610861116114612061416144615061"
957        b"806199610462106240625662a162006405640864116414642064416444645064"
958        b"806401650465106540654a656865926500669466016804681068656898680069"
959        b"2a69426aa16a0080028005800880118014801980208025804180448050805280"
960        b"5580588061808080858091809480018104810981108112811581188121812481"
961        b"408142814581488151815481818184819081a981008205820a82118214824182"
962        b"4482508201840484068409841084128415841884218440844284458448845184"
963        b"5484608481848484908400850285058508851185148520854185448550858085"
964        b"8a85018604861086298640860088058811881488418844885088a28801890489"
965        b"40896589228a588a5a8a828aa28a019004900990109012901590189024904090"
966        b"4290459048905190549060908190849090900091059111911491419144915091"
967        b"5a910192049210924092a6920094029405940894119414942094419444945094"
968        b"8094969401950495109540959895a19500964696649601980498109826984098"
969        b"a998009949995299909a00a005a00aa014a022a02aa041a044a050a0a2a0aaa0"
970        b"40a165a102a20aa222a228a22aa282a288a28aa2a8a201a404a410a440a489a4"
971        b"a4a400a519a551a60aa828a8a2a854a986a908aa0aaa20aa22aa28aa88aaaaaa"
972    )
973
974    @classmethod
975    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
976        n_blocks = blocks.shape[0]
977
978        d, rest = np.hsplit(blocks, [2])
979        qs, rest = np.hsplit(rest, [QK_K // 8])
980        signs, rest = np.hsplit(rest, [QK_K // 8])
981        qh, scales = np.hsplit(rest, [QK_K // 32])
982
983        d = d.view(np.float16).astype(np.float32)
984
985        scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
986        scales = (scales & 0x0F).reshape((n_blocks, -1))
987        db = d * (np.float32(0.5) + scales) * np.float32(0.25)
988        db = db.reshape((n_blocks, -1, 1, 1))
989
990        # unpack the sign bits
991        signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
992        signs = signs & np.uint8(0x01)
993        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
994        signs = signs.reshape((n_blocks, -1, 2, 8))
995
996        qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4))
997        qs = qs.astype(np.uint16) | ((qh & 0x03).astype(np.uint16) << 8).reshape((n_blocks, -1))
998
999        assert cls.grid is not None
1000        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
1001        grid = grid.reshape((n_blocks, -1, 2, 8))
1002
1003        return (db * grid * signs).reshape((n_blocks, -1))
1004
1005
1006class IQ3_XXS(__Quant, qtype=GGMLQuantizationType.IQ3_XXS):
1007    grid_shape = (256, 4)
1008    grid_map = (0x04, 0x0c, 0x14, 0x1c, 0x24, 0x2c, 0x34, 0x3e)
1009    grid_hex = (
1010        b"0000020004001100130017002000220031004200730075000101030110011201"
1011        b"2101250130013201410154017001000202020402110220022202310233023702"
1012        b"5102570275020103070310031203250370031304370444045704730475040105"
1013        b"0705320552053506640610071407160743076107011003101010121021102310"
1014        b"3010321034104710501000110211111120112211011203121012121221123012"
1015        b"7212001302132013311346136613011405145014201524154615711505162217"
1016        b"4017002002201120132020202220262031204220012103210521102112212121"
1017        b"3021632167217021002202221122172220222222372240225522012310231423"
1018        b"7023742335245324032527254125742501270327162745270130103012302130"
1019        b"2330503065307230003102312031313144314631013203321032253252327232"
1020        b"1133333330344734723400350635223555351436363663363337603704401740"
1021        b"3540374053405740744120423742404260426642074345430444514464442545"
1022        b"4345704505471047124730471250415070500051065126515551145232527252"
1023        b"0253535310542354275472540255315550562457425724604460466064602161"
1024        b"6161176264623063366344640565526533660367216703700570077010703270"
1025        b"5270267140711272457252720073157333736073217441740075027524753076"
1026    )
1027
1028    @classmethod
1029    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
1030        n_blocks = blocks.shape[0]
1031
1032        d, rest = np.hsplit(blocks, [2])
1033        qs, scales = np.hsplit(rest, [QK_K // 4])
1034
1035        d = d.view(np.float16).astype(np.float32)
1036        scales = scales.view(np.uint32)
1037
1038        db = d * (np.float32(0.5) + (scales >> 28).astype(np.float32)) * np.float32(0.5)
1039        db = db.reshape((n_blocks, -1, 1, 1))
1040
1041        # get the sign indices and unpack the bits
1042        signs = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4))
1043        ksigns = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128))
1044        signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1))
1045        signs = np.take_along_axis(ksigns, signs, axis=-1)
1046        signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 1, 8))
1047        signs = signs & np.uint8(0x01)
1048        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
1049        signs = signs.reshape((n_blocks, -1, 4, 8))
1050
1051        assert cls.grid is not None
1052        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
1053        grid = grid.reshape((n_blocks, -1, 4, 8))
1054
1055        return (db * grid * signs).reshape((n_blocks, -1))
1056
1057
1058class IQ3_S(__Quant, qtype=GGMLQuantizationType.IQ3_S):
1059    grid_shape = (512, 4)
1060    grid_map = (0x01, 0x03, 0x05, 0x07, 0x09, 0x0b, 0x0d, 0x0f)
1061    grid_hex = (
1062        b"0000010002000500070010001100120014001600200021002500330040004200"
1063        b"4500470051005300600062007100740077000001010102010401100111011501"
1064        b"2001230127013101350144016101650172010002010205020702100213021602"
1065        b"2102250230023402420245024702510253027002730203031103150320032203"
1066        b"3103330336034403500352036703710375030004130417042104240432044004"
1067        b"4304510470040205040520052205260533054105450547056605730506061106"
1068        b"1306310652067106000702070407200722072607330750075407001001100210"
1069        b"0410101011101310151017102010221031103410361054105610611072100011"
1070        b"0111031106111011141121113011331141115011521170117611001212121512"
1071        b"1712201224123212401243125512601272120113041307131013131321132713"
1072        b"3013341341136213701303140514121414143114331442144614501454140115"
1073        b"1015131521153015321551152016241627164416461601170317101712172117"
1074        b"3517411762177017002001200320052007201020122014201620212023202720"
1075        b"3020322041204320452050205220672070207320752000210221102113211721"
1076        b"2221252131213421422151210122042207222122232230223722412253225722"
1077        b"7122742200230223052311232223242331233323422350236623012407242024"
1078        b"2324322435244124722475240425112522253725402553257025002602260726"
1079        b"2126552661260527112726273027432750270230113013301530173022303130"
1080        b"3330353042304430473051306330713001310331053114312131233140316031"
1081        b"7231763100321232203232323432503201331033143321332333273330334133"
1082        b"4333473355337333033411341634223431345234603464340135103512352535"
1083        b"3235443556357335163641360137033720372237353700400440124020402440"
1084        b"2740324041405040704002410741114113412241304135414341514155410142"
1085        b"0342104215422142334240425742624270420443114313432043224331433543"
1086        b"0044024424443744404471440545074521456245134634466046104715473047"
1087        b"4347514702501050145022504050445047505250665074500151035105511251"
1088        b"2151325172510052115223523052365253520253075310532753445351536553"
1089        b"7353015404542054325446541255265551555355425602570457225711601360"
1090        b"1560316033606060006120612761646112623462426255626262706200631463"
1091        b"2163406325644364626400650365346560650566406611671367007004700770"
1092        b"2070227036704070547062700271117124714371457101720472107216722172"
1093        b"3072517202733273357353730174057413742074507422754275027631760077"
1094    )
1095
1096    @classmethod
1097    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
1098        n_blocks = blocks.shape[0]
1099
1100        d, rest = np.hsplit(blocks, [2])
1101        qs, rest = np.hsplit(rest, [QK_K // 4])
1102        qh, rest = np.hsplit(rest, [QK_K // 32])
1103        signs, scales = np.hsplit(rest, [QK_K // 8])
1104
1105        d = d.view(np.float16).astype(np.float32)
1106
1107        scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
1108        scales = (scales & 0x0F).reshape((n_blocks, -1))
1109        db = d * (1 + 2 * scales)
1110        db = db.reshape((n_blocks, -1, 1, 1))
1111
1112        # unpack the sign bits
1113        signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
1114        signs = signs & np.uint8(0x01)
1115        signs = np.where(signs == 0, np.float32(1), np.float32(-1))
1116        signs = signs.reshape((n_blocks, -1, 4, 8))
1117
1118        qh = qh.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8)
1119        qh = (qh & 0x01).astype(np.uint16).reshape((n_blocks, -1))
1120        qs = qs.astype(np.uint16) | (qh << 8)
1121
1122        assert cls.grid is not None
1123        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
1124        grid = grid.reshape((n_blocks, -1, 4, 8))
1125
1126        return (db * grid * signs).reshape((n_blocks, -1))
1127
1128
1129class IQ1_S(__Quant, qtype=GGMLQuantizationType.IQ1_S):
1130    # iq1s_grid, with each byte packed into 2 bits
1131    # -1, 0, 1 <=> 0, 1, 2
1132    grid_shape = (2048, 8)
1133    grid_map = (-1, 0, 1)
1134    grid_hex = (
1135        b"00000200050008000a00110015002000220028002a0045005100540056006500"
1136        b"8000820088008a009500a000a200a800aa000401050111011401160119011a01"
1137        b"2501410146014901520155015a0161016401660168018501910194019601a501"
1138        b"0002020208020a0215022002220228022a024502510259026402690280028202"
1139        b"88028a02910295029902a002a202a802aa021104140416042504410449045504"
1140        b"5a046404650491049904a5040105040505050605150518051a05290540054505"
1141        b"4a0550055105540555055605590560056205650568056a058105910595059805"
1142        b"9a05a105a405a505a605a9051406190641064406500652065506580660066106"
1143        b"6606690685069106940699060008020808080a0815082008220828082a084508"
1144        b"5108560865088008820888088a089508a008a208a808aa080509110914091909"
1145        b"2409250941095009510955096109640969099109940996099909a509000a020a"
1146        b"080a0a0a150a200a220a280a2a0a450a510a590a610a650a800a820a850a880a"
1147        b"8a0a950aa00aa20aa80aaa0a1010111014101910241025104110441050105510"
1148        b"58106110641065106910911094109610a110a510011104110611091110111211"
1149        b"1511181121112411291145114a11501151115211541155115611591160116511"
1150        b"841192119511a111a41111121412161225124012461249125212551258125a12"
1151        b"641266128512911294129612a512011406140914141415141814191421142614"
1152        b"41144514461448144a1451145414551456145914621465146814841489149014"
1153        b"94149514981499149a14a114a414a514a914021505150a151115141515151615"
1154        b"191520152215251528152a154115441545154615511552155415551556155915"
1155        b"5a1561156415651566156915801582158415851588158a159015911594159515"
1156        b"961599159a15a015a215a51501160416051606161516161618161a1621162616"
1157        b"401642164416451648164a165116551656165816591661166416651668166916"
1158        b"6a1686168a1692169516a416a916111816182518411844184618491850185518"
1159        b"58185a1860186118641866186918851891189418a5181019121915191a192119"
1160        b"25194219441945194819511954195519561959195a19601965196a1989199119"
1161        b"921995199819a119a619a919091a161a241a261a441a461a491a501a521a551a"
1162        b"581a611a661a691a851a911a961a9a1a0020022008200a201520202022202520"
1163        b"28202a20452051205920612065208020822088208a209520a020a220a520a820"
1164        b"aa2005211121142119212521422144214921552158215a216121642165216621"
1165        b"8521902196219921a521012208220a22112215222022222228222a2245225122"
1166        b"562259226522812288228a2291229522a022a222a822aa220524142416241924"
1167        b"252444244524462449245224552458245a2466248524912494249924a124a524"
1168        b"0925152521252925402545254825512554255525592562256525682589259025"
1169        b"9425952598259a25a125a425a625a92505261026122619262526412649265526"
1170        b"6026612669268426862690269a260028022808280a2815282028222828282a28"
1171        b"45285128542865288028822888288a28a028a228a828aa280929112914291929"
1172        b"2529462949295229552961296429662969298529902996299929a429a529002a"
1173        b"022a082a0a2a202a222a282a2a2a452a512a562a592a652a802a822a882a8a2a"
1174        b"952aa02aa22aa82aaa2a054011401640254049405240554058405a4061406440"
1175        b"664094409940a140a6400041014104410641094112411541164118411a412141"
1176        b"26412941454148414a41514154415541564159415a41654168416a4181418441"
1177        b"8641904192419541a041a141a241054211421442164225424142524255425a42"
1178        b"6442694289429442a5420144154419442944454448444a445144544455445644"
1179        b"61446244654468446a44814486448944904492449544a044a144a94401450245"
1180        b"05450a4511451445154516451945204525452a45414544454545464549455045"
1181        b"5145544555455645584559456145644565456645694582458445854588459145"
1182        b"94459545964599459a45a545a845aa450146054609461446154618461a462146"
1183        b"2446294640464246454648465046514652465546564659466246654668468146"
1184        b"85468a4694469546a146a446a6460548114815481a4825484248494850485548"
1185        b"5848614864486648694885489148944896489948a5480149054906490a491049"
1186        b"144915491849214924492649404945494a495149524954495549564959496049"
1187        b"6249654966496a49864989499249954996499849a149a449a649a949164a444a"
1188        b"464a494a554a584a5a4a644a694a944aa54a0150045005500650095012501550"
1189        b"1a50215024502950405045504850515054505550565059506550685086508950"
1190        b"95509850a050a150a650a9500551085109510a51115114511551165118511951"
1191        b"20512551265128512a5141514451455146514951505151515251545155515651"
1192        b"585159515a51615164516551665169518251855191519451955196519951a051"
1193        b"a551aa5101520652125215521a5221522452425245524a525152545255525652"
1194        b"595262526552855290529252955299529a52a452045405541154145415541654"
1195        b"185419542154255428542a54415444544554465449544a545054515454545554"
1196        b"5654585459545a54615462546454655466546954805488548a54915494549554"
1197        b"96549954a154a454a554aa540155025504550555065509551055115512551455"
1198        b"1555165519551a55215524552555265529554055415542554455455546554855"
1199        b"4955505551555255545555555655585559555a55605561556455655566556855"
1200        b"69556a5581558455855589558a559055915594559555965598559955a155a455"
1201        b"a555a655a9550056015602560456065608560956115614561556185619562056"
1202        b"2156225624562556265628562956415645564656485649564a56505651565256"
1203        b"545655565656585659565a566156645665566956825685568656885689568a56"
1204        b"915695569a56a256a556a656a856a95604580558065809581058155818582158"
1205        b"2a58455848584a58515854585558565858585958605862586458655882588958"
1206        b"9058925895589858a158a9580159025905590a59115914591559165919592559"
1207        b"41594459455946594959505951595259545955595659585959595a5961596459"
1208        b"655966596959815985598959915994599559965998599959a559045a085a155a"
1209        b"1a5a205a255a265a295a455a485a495a515a555a565a585a595a625a655a685a"
1210        b"6a5a815a8a5a925a955a965a985a9a5aa15a0560146016601960256044605060"
1211        b"5560566058605a60616064606660696081609660a56001610461066109611261"
1212        b"15612161226126612961456149615161556156615961656166616a6184618a61"
1213        b"92619561a161a661a96111621662196240624162466255625662586260628562"
1214        b"91629662a56211641264156416641a6421642664296440644264456448644a64"
1215        b"516454645564566459645a646064626465648464856489649064926494649564"
1216        b"966498649a64a164a464a964056508650a651165156516651965446545654665"
1217        b"496550655165546555655665596561656465656566656965866589658a659165"
1218        b"9565966599659a65a265a565a665a86502660966156620662666286629664066"
1219        b"456648664a66516654665566566658665a666066656668668066826685668a66"
1220        b"9466966698669966a066a466a666aa661668196825684168526855685a686168"
1221        b"6968856891689868a66801690469106915692169246926692969406941694569"
1222        b"4669486951695469556956695969606965696a69826984698a699569a169a469"
1223        b"a569a969116a166a186a416a446a496a506a556a586a5a6a646a656a696a866a"
1224        b"946a986a9a6aa66a0080028008800a802080228028802a804580508051805480"
1225        b"5680598065808080828088808a809580a080a280a880aa800581118114811681"
1226        b"1981258141814481498150815281558156815881598164816681698185818981"
1227        b"948196819981a5810082028208820a8215822082228228822a82518254825982"
1228        b"65828082828288828a829582a082a282a882aa82148419844184448451845584"
1229        b"5a846184648469849484998401850985128515851a8526852985408541854585"
1230        b"4885518554855585568559855a856585668568856a8581858485868589859085"
1231        b"928595859885a68511861686198625864186448649864a865086558659865a86"
1232        b"618666866a86858691869a86a4860088028808880a8815882088228828882a88"
1233        b"41884588518854885988658869888088828888888a889588a088a288a888aa88"
1234        b"05890689118914891689258941894489468949895089528955895a8961896489"
1235        b"858996899989a589008a028a088a0a8a158a208a228a288a2a8a458a518a548a"
1236        b"568a808a828a888a8a8a958aa08aa28aa88aaa8a059011901690189019902590"
1237        b"419046904990559058905a9069906a9085909190949096909990a59001910491"
1238        b"069109911091159118911a912191249126912991409145915091519154915591"
1239        b"569159916291659184918691929195919891a191a491a691a991059211921492"
1240        b"19922592449246924992509252925592589266926992859294929692a9920194"
1241        b"04940694109415941894269440944a9451945494559456945894599460946194"
1242        b"62946594849486949294949495949894a194a9940095059508950a9510951195"
1243        b"14951595169519952195259529952a9541954495459546954995509551955295"
1244        b"549555955695589559955a956195649565956695699581958595889591959295"
1245        b"94959595969599959a95a095a295a595a895aa95019604961096159619962096"
1246        b"2696299645964896499651965296559656965996659668968296849689968a96"
1247        b"929694969596a496a696a9960598169819982598419846985098529855985698"
1248        b"5a98649865988598919896989998a59804990699099910991299159918991a99"
1249        b"209921992499269940994299459948994a995199549955995699599962996599"
1250        b"66996a99819984999099929995999a99a199a699059a159a259a449a469a499a"
1251        b"509a559a589a619a859a919a949a959a969a00a002a008a00aa015a020a022a0"
1252        b"28a02aa045a051a054a056a059a080a082a088a08aa095a0a0a0a2a0a8a0aaa0"
1253        b"05a109a111a114a116a119a11aa146a149a151a155a158a15aa161a164a185a1"
1254        b"90a192a196a199a102a208a20aa210a219a222a228a22aa245a251a256a259a2"
1255        b"65a280a282a288a28aa295a2a0a2a2a2a8a2aaa219a425a441a444a450a454a4"
1256        b"55a458a45aa461a465a466a468a469a485a406a509a510a512a515a518a526a5"
1257        b"29a542a545a551a554a555a556a559a565a56aa581a584a585a586a589a592a5"
1258        b"95a598a505a611a616a61aa621a625a644a646a64aa652a655a656a658a660a6"
1259        b"62a686a690a695a696a699a6a1a6a4a6a6a600a802a808a80aa820a822a828a8"
1260        b"2aa851a854a856a859a880a882a888a88aa895a8a0a8a2a8a8a8aaa805a914a9"
1261        b"19a921a925a941a950a955a95aa961a966a969a990a996a900aa02aa08aa0aaa"
1262        b"20aa22aa28aa2aaa51aa54aa56aa80aa82aa88aa8aaa95aaa0aaa2aaa8aaaaaa"
1263    )
1264
1265    delta = np.float32(0.125)
1266
1267    @classmethod
1268    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
1269        n_blocks = blocks.shape[0]
1270
1271        d, rest = np.hsplit(blocks, [2])
1272        qs, qh = np.hsplit(rest, [QK_K // 8])
1273
1274        d = d.view(np.float16).astype(np.float32)
1275        qh = qh.view(np.uint16)
1276
1277        dl = d * (2 * ((qh >> 12) & 7) + 1)
1278        dl = dl.reshape((n_blocks, -1, 1, 1))
1279        delta = np.where((qh & np.uint16(0x8000)) == 0, cls.delta, -cls.delta)
1280        delta = delta.reshape((n_blocks, -1, 1, 1))
1281
1282        qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4))
1283        qs = qs.astype(np.uint16) | ((qh & 7) << 8).reshape((n_blocks, -1))
1284
1285        assert cls.grid is not None
1286        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
1287        grid = grid.reshape((n_blocks, -1, 4, 8))
1288
1289        return (dl * (grid + delta)).reshape((n_blocks, -1))
1290
1291
1292class IQ1_M(__Quant, qtype=GGMLQuantizationType.IQ1_M):
1293    grid_shape = IQ1_S.grid_shape
1294    grid_map = IQ1_S.grid_map
1295    grid_hex = IQ1_S.grid_hex
1296
1297    delta = IQ1_S.delta
1298
1299    # Okay *this* type is weird. It's the only one which stores the f16 scales in multiple parts.
1300    @classmethod
1301    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
1302        n_blocks = blocks.shape[0]
1303
1304        qs, rest = np.hsplit(blocks, [QK_K // 8])
1305        qh, scales = np.hsplit(rest, [QK_K // 16])
1306
1307        # The f16 scale is packed across multiple bytes
1308        scales = scales.view(np.uint16)
1309        d = (scales.reshape((n_blocks, 4)) & np.uint16(0xF000)) >> np.array([12, 8, 4, 0], dtype=np.uint16).reshape((1, 4))
1310        d = d[..., 0] | d[..., 1] | d[..., 2] | d[..., 3]
1311        d = d.view(np.float16).astype(np.float32).reshape((n_blocks, 1))
1312
1313        scales = scales.reshape(n_blocks, -1, 1) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4))
1314        scales = (scales & 0x07).reshape((n_blocks, -1))
1315        dl = d * (2 * scales + 1)
1316        dl = dl.reshape((n_blocks, -1, 2, 1, 1))
1317
1318        qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
1319        qs = qs.astype(np.uint16) | ((qh & 0x07).astype(np.uint16) << 8).reshape((n_blocks, -1))
1320
1321        delta = np.where(qh & 0x08 == 0, cls.delta, -cls.delta)
1322        delta = delta.reshape((n_blocks, -1, 2, 2, 1))
1323
1324        assert cls.grid is not None
1325        grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
1326        grid = grid.reshape((n_blocks, -1, 2, 2, 8))
1327
1328        return (dl * (grid + delta)).reshape((n_blocks, -1))
1329
1330
1331class IQ4_NL(__Quant, qtype=GGMLQuantizationType.IQ4_NL):
1332    kvalues = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113)
1333
1334    @classmethod
1335    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
1336        n_blocks = blocks.shape[0]
1337
1338        d, qs = np.hsplit(blocks, [2])
1339
1340        d = d.view(np.float16).astype(np.float32)
1341
1342        qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
1343
1344        qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 1))
1345
1346        kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
1347        qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1))
1348
1349        return (d * qs)
1350
1351
1352class IQ4_XS(__Quant, qtype=GGMLQuantizationType.IQ4_XS):
1353    @classmethod
1354    def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
1355        n_blocks = blocks.shape[0]
1356
1357        d, rest = np.hsplit(blocks, [2])
1358        scales_h, rest = np.hsplit(rest, [2])
1359        scales_l, qs = np.hsplit(rest, [QK_K // 64])
1360
1361        d = d.view(np.float16).astype(np.float32)
1362        scales_h = scales_h.view(np.uint16)
1363
1364        scales_l = scales_l.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
1365        scales_h = scales_h.reshape((n_blocks, 1, -1)) >> np.array([2 * i for i in range(QK_K // 32)], dtype=np.uint16).reshape((1, -1, 1))
1366        scales_l = scales_l.reshape((n_blocks, -1)) & np.uint8(0x0F)
1367        scales_h = scales_h.reshape((n_blocks, -1)).astype(np.uint8) & np.uint8(0x03)
1368
1369        scales = (scales_l | (scales_h << np.uint8(4))).astype(np.int8) - np.int8(32)
1370        dl = (d * scales.astype(np.float32)).reshape((n_blocks, -1, 1))
1371
1372        qs = qs.reshape((n_blocks, -1, 1, 16)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
1373        qs = qs.reshape((n_blocks, -1, 32, 1)) & np.uint8(0x0F)
1374
1375        kvalues = np.array(IQ4_NL.kvalues, dtype=np.int8).reshape((1, 1, 1, -1))
1376        qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1, 32))
1377
1378        return (dl * qs).reshape((n_blocks, -1))