mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 10:16:07 +00:00
add GGUF IQ4_XS support (#15766)
* add GGUF IQ4_XS support * gguf 21 * gguf 21 * use plus * ggml_common autogen for constant arrays * fix * ggml_common in autogen * inline
This commit is contained in:
George Hotz
@@ -58,6 +58,7 @@ jobs:
|
||||
python3 -c "from tinygrad.runtime.autogen import avcodec"
|
||||
python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
|
||||
python3 -c "from tinygrad.runtime.autogen import mlx5"
|
||||
python3 -c "from tinygrad.runtime.autogen import ggml_common"
|
||||
REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
|
||||
- name: Check for differences
|
||||
run: |
|
||||
|
||||
+31
-2
@@ -1,12 +1,30 @@
|
||||
import os, struct, unittest
|
||||
from tinygrad import dtypes, Tensor, fetch, Device
|
||||
from tinygrad.nn.state import ggml_data_to_tensor, gguf_load
|
||||
from tinygrad.nn.state import _ggml_iq_grid, ggml_data_to_tensor, gguf_load
|
||||
from tinygrad.runtime.autogen import ggml_common as _ggml
|
||||
from tinygrad.device import is_dtype_supported
|
||||
import numpy as np
|
||||
from gguf import GGUFReader, GGUFValueType, GGMLQuantizationType, GGML_QUANT_SIZES, dequantize, quantize
|
||||
from gguf.quants import IQ2_S, IQ3_S, IQ3_XXS
|
||||
|
||||
ggml_test_block_count = 4
|
||||
|
||||
class TestGGUFTables(unittest.TestCase):
|
||||
def test_iq2_s_grid_matches_gguf_py(self):
|
||||
IQ2_S.init_grid()
|
||||
grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2s_grid, (1024, 8)).numpy()
|
||||
np.testing.assert_equal(grid, IQ2_S.grid.reshape(1024, 8))
|
||||
|
||||
def test_iq3_xxs_grid_matches_gguf_py(self):
|
||||
IQ3_XXS.init_grid()
|
||||
grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3xxs_grid, (256, 4)).numpy()
|
||||
np.testing.assert_equal(grid, IQ3_XXS.grid.reshape(256, 4))
|
||||
|
||||
def test_iq3_s_grid_matches_gguf_py(self):
|
||||
IQ3_S.init_grid()
|
||||
grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3s_grid, (512, 4)).numpy()
|
||||
np.testing.assert_equal(grid, IQ3_S.grid.reshape(512, 4))
|
||||
|
||||
@unittest.skipIf(any(not is_dtype_supported(t) for t in [ dtypes.uint8, dtypes.half ]), "Backend must support uint8 and half")
|
||||
class TestGGUF(unittest.TestCase):
|
||||
def test_load_tinyllama_q8_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q8_0.gguf?download=true")
|
||||
@@ -36,6 +54,10 @@ class TestGGUF(unittest.TestCase):
|
||||
def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
|
||||
def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
|
||||
def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
|
||||
def test_dequantization_iq3_xxs(self): self._test_dequantization(GGMLQuantizationType.IQ3_XXS)
|
||||
def test_dequantization_iq3_s(self): self._test_dequantization(GGMLQuantizationType.IQ3_S)
|
||||
def test_dequantization_iq2_s(self): self._test_dequantization(GGMLQuantizationType.IQ2_S)
|
||||
def test_dequantization_iq4_xs(self): self._test_dequantization(GGMLQuantizationType.IQ4_XS)
|
||||
def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16")
|
||||
def test_dequantization_bf16(self): self._test_dequantization(GGMLQuantizationType.BF16)
|
||||
@@ -140,7 +162,10 @@ class TestGGUFGEMV(unittest.TestCase):
|
||||
# generate random quantized blocks with valid fp16 scale fields (random bytes can produce NaN scales)
|
||||
q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
|
||||
scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
|
||||
if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0): q_data[:, :2] = scales[:, :2] # d at offset 0
|
||||
if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0,
|
||||
GGMLQuantizationType.IQ3_XXS,
|
||||
GGMLQuantizationType.IQ2_S,
|
||||
GGMLQuantizationType.IQ3_S, GGMLQuantizationType.IQ4_XS): q_data[:, :2] = scales[:, :2] # d at offset 0
|
||||
elif qtype in (GGMLQuantizationType.Q5_1, GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K):
|
||||
q_data[:, :4] = scales[:, :4] # d, m/dmin at offset 0
|
||||
elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2] # d at end
|
||||
@@ -173,6 +198,10 @@ class TestGGUFGEMV(unittest.TestCase):
|
||||
def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
|
||||
def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
|
||||
def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
|
||||
def test_gguf_gemv_iq3_xxs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_XXS)
|
||||
def test_gguf_gemv_iq3_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_S)
|
||||
def test_gguf_gemv_iq2_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_S)
|
||||
def test_gguf_gemv_iq4_xs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ4_XS)
|
||||
def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16")
|
||||
def test_gguf_gemv_bf16(self): self._test_gguf_gemv(GGMLQuantizationType.BF16)
|
||||
|
||||
+46
-6
@@ -34,6 +34,12 @@ safe_dtypes = {"BOOL":dtypes.bool, "I8":dtypes.int8, "U8":dtypes.uint8, "I16":dt
|
||||
"I64":dtypes.int64, "U64":dtypes.uint64, "F16":dtypes.float16, "BF16":dtypes.bfloat16, "F32":dtypes.float32, "F64":dtypes.float64}
|
||||
inverse_safe_dtypes = {v:k for k,v in safe_dtypes.items()}
|
||||
|
||||
# ggml packs each iq grid entry as N bytes (N=4 for uint32 grids, N=8 for uint64 grids) in a single word. See ggml-common.h.
|
||||
@functools.lru_cache(None)
|
||||
def _ggml_iq_grid(device: str, grid: tuple[int, ...], grid_shape: tuple[int, int]) -> Tensor:
|
||||
values = [float((w >> (8*i)) & 0xFF) for w in grid for i in range(grid_shape[1])]
|
||||
return Tensor(values, dtype=dtypes.float32, device=device).reshape(grid_shape)
|
||||
|
||||
def accept_filename(func: Callable[[Tensor], T]) -> Callable[[Tensor|str|pathlib.Path], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(fn: Tensor|str|pathlib.Path) -> T: return func(Tensor(pathlib.Path(fn)) if not isinstance(fn, Tensor) else fn)
|
||||
@@ -297,18 +303,18 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
|
||||
"""
|
||||
Converts ggml tensor data to a tinygrad tensor.
|
||||
|
||||
Supported native types: float32 (id: 0), float16 (id: 1), bfloat16 (id: 30),
|
||||
int8 (id: 16), int16 (id: 17), int32 (id: 18)
|
||||
Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 24),
|
||||
int16 (id: 25), int32 (id: 26), int64 (id: 27), float64 (id: 28), bfloat16 (id: 30)
|
||||
Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q5_0 (id: 6),
|
||||
Q5_1 (id: 7), Q8_0 (id: 8), Q4_K (id: 12), Q5_K (id: 13),
|
||||
Q6_K (id: 14), MXFP4 (id: 39), Q1_0 (id: 41)
|
||||
Q6_K (id: 14), IQ3_XXS (id: 18), IQ3_S (id: 21), IQ2_S (id: 22), IQ4_XS (id: 23), MXFP4 (id: 39), Q1_0 (id: 41)
|
||||
"""
|
||||
# https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356
|
||||
|
||||
# native types
|
||||
if (dtype := {
|
||||
0: dtypes.float32, 1: dtypes.float16, 16: dtypes.int8,
|
||||
17: dtypes.int16, 18: dtypes.int32, 30: dtypes.bfloat16,
|
||||
0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8,
|
||||
25: dtypes.int16, 26: dtypes.int32, 27: dtypes.int64, 28: dtypes.float64, 30: dtypes.bfloat16,
|
||||
}.get(ggml_type)) is not None:
|
||||
return t[:dtype.itemsize * n].contiguous().bitcast(dtype)
|
||||
|
||||
@@ -320,9 +326,10 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
|
||||
# map to (number of elements, number of bytes)
|
||||
if (nelements_nbytes := {
|
||||
2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
|
||||
12:(256,144), 13:(256,176), 14:(256,210), 39:(32,17),
|
||||
12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17),
|
||||
41:(128,18)
|
||||
}.get(ggml_type)) is not None:
|
||||
from tinygrad.runtime.autogen import ggml_common as _ggml
|
||||
blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
|
||||
if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
|
||||
if ggml_type == 3:
|
||||
@@ -351,6 +358,39 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
|
||||
scales = blocks[:,192:208].bitcast(dtypes.int8).unsqueeze(-1).expand((-1, 16, 16)).reshape((-1, 256))
|
||||
d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).expand((-1, 256))
|
||||
return d * (xl.bitwise_or(xh).bitcast(dtypes.int8) - 32).flatten(-2) * scales
|
||||
if ggml_type == 18:
|
||||
d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
|
||||
scale_words = blocks[:, 66:98].bitcast(dtypes.uint32)
|
||||
db = d * (scale_words.rshift(28).cast(dtypes.float32) + 0.5).reshape((-1, 8, 1, 1)) * 0.5
|
||||
sign_idx = scale_words.unsqueeze(-1).rshift(
|
||||
Tensor([0, 7, 14, 21], device=t.device, dtype=dtypes.uint32)).bitwise_and(0x7F).reshape((-1, 32)).cast(dtypes.int32)
|
||||
even_signs = Tensor([i | (0x80 if i.bit_count() % 2 else 0) for i in range(128)], dtype=dtypes.uint8, device=t.device)
|
||||
signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8))
|
||||
grid = _ggml_iq_grid(t.device, _ggml.iq3xxs_grid, (256, 4))[blocks[:, 2:66]].reshape((-1, 8, 4, 8))
|
||||
return (db * grid * signs).flatten(-3)
|
||||
if ggml_type == 21:
|
||||
d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
|
||||
scales = (1 + 2 * q_to_uint8(blocks[:, 106:110].reshape((-1, 4, 1)), 4).reshape((-1, 8))).cast(dtypes.float32).reshape((-1, 8, 1, 1))
|
||||
qh = q_to_uint8(blocks[:, 66:74].reshape((-1, 8, 1)), 1).reshape((-1, 64)).cast(dtypes.uint16)
|
||||
signs = (q_to_uint8(blocks[:, 74:106].reshape((-1, 32, 1)), 1).reshape((-1, 256)) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8))
|
||||
q = blocks[:, 2:66].cast(dtypes.uint16) + qh.lshift(8)
|
||||
return (d * scales * _ggml_iq_grid(t.device, _ggml.iq3s_grid, (512, 4))[q].reshape((-1, 8, 4, 8)) * signs).flatten(-3)
|
||||
if ggml_type == 22:
|
||||
d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
|
||||
db = d * (q_to_uint8(blocks[:, 74:82].reshape((-1, 8, 1)), 4).reshape((-1, 16)).cast(dtypes.float32) + 0.5).reshape((-1, 16, 1, 1)) * 0.25
|
||||
signs = (q_to_uint8(blocks[:, 34:66].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 16, 2, 8))
|
||||
qh = q_to_uint8(blocks[:, 66:74].reshape((-1, 8, 1)), 2).reshape((-1, 32)).cast(dtypes.uint16)
|
||||
q = blocks[:, 2:34].cast(dtypes.uint16) + qh.lshift(8)
|
||||
return (db * _ggml_iq_grid(t.device, _ggml.iq2s_grid, (1024, 8))[q].reshape((-1, 16, 2, 8)) * signs).flatten(-3)
|
||||
if ggml_type == 23:
|
||||
d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1))
|
||||
scale_shifts = Tensor([0, 2, 4, 6, 8, 10, 12, 14], device=t.device, dtype=dtypes.uint16)
|
||||
iq4_xs_lut = Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)
|
||||
scales_l = Tensor.stack((sl:=blocks[:, 4:8]).bitwise_and(0xF), sl.rshift(4), dim=2).reshape((-1, 8))
|
||||
scales_h = blocks[:, 2:4].bitcast(dtypes.uint16).unsqueeze(-1).rshift(scale_shifts).bitwise_and(0x03).reshape((-1, 8)).cast(dtypes.uint8)
|
||||
scales = (scales_l.bitwise_or(scales_h.lshift(4)).bitcast(dtypes.int8) - 32).cast(dtypes.float32).reshape((-1, 8, 1))
|
||||
q = (qs:=blocks[:, 8:].reshape((-1, 8, 16))).bitwise_and(0xF).cat(qs.rshift(4), dim=2)
|
||||
return (d * scales * iq4_xs_lut[q]).flatten(-2)
|
||||
if ggml_type == 39:
|
||||
e = blocks[:, 0].cast(dtypes.uint32)
|
||||
small_bits = Tensor([0x00200000, 0x00400000], dtype=dtypes.uint32, device=t.device)[e.clip(0, 1).cast(dtypes.int32)] # e = 0 or e = 1 case
|
||||
|
||||
@@ -9,6 +9,7 @@ rocr_src = "https://github.com/ROCm/rocm-systems/archive/refs/tags/rocm-7.1.1.ta
|
||||
linux_headers_deb = "https://snapshot.debian.org/archive/debian/20260207T145350Z/pool/main/l/linux/linux-libc-dev_6.18.9-1_all.deb"
|
||||
linux_headers_kern_deb = "https://snapshot.debian.org/archive/debian/20260207T145350Z/pool/main/l/linux/linux-headers-6.18.9+deb14-common_6.18.9-1_all.deb"
|
||||
liburing_src = "https://raw.githubusercontent.com/axboe/liburing/refs/tags/liburing-2.14/src/include/liburing.h"
|
||||
ggml_common_src = "https://raw.githubusercontent.com/ggml-org/ggml/d4fcfe88a8bcf5c9840be14be6c2fbf1f5b3b2db/src/ggml-common.h"
|
||||
macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"
|
||||
|
||||
llvm_lib = (r"'C:\\Program Files\\LLVM\\bin\\LLVM-C.dll' if WIN else '/opt/homebrew/opt/llvm@20/lib/libLLVM.dylib' if OSX else " +
|
||||
@@ -162,6 +163,9 @@ def __getattr__(nm):
|
||||
[f"{macossdk}/System/Library/Frameworks/CoreFoundation.framework/Headers/CF{s}.h" for s in ["String", "Data"]],
|
||||
args=["-isysroot", macossdk])
|
||||
case "llvm_qcom": return load("llvm_qcom", "'llvm-qcom'", [root/"extra/tinydreno.h"])
|
||||
case "ggml_common":
|
||||
return load("ggml_common", None, ["{}/ggml-common.h"], srcs=ggml_common_src,
|
||||
args=["-DGGML_COMMON_DECL_C", "-DGGML_COMMON_IMPL_C"], parse_macros=False)
|
||||
case "mlx5":
|
||||
kh = "{}/usr/src/linux-headers-6.18.9+deb14-common/include/linux/mlx5"
|
||||
return load("mlx5", None, [root/"extra/mlx_driver/mlx5.h", f"{kh}/mlx5_ifc.h"], srcs=linux_headers_kern_deb,
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -85,7 +85,7 @@ base_rules = [(r'\s*\\\n\s*', ' '), (r'\s*\n\s*', ' '), (r'//.*', ''), (r'/\*.*?
|
||||
(r'\((unsigned )?(char|uint64_t)\)', ''), (r'^.*\d+:\d+.*$', ''), (r'^.*\w##\w.*$', '')]
|
||||
|
||||
uints = (clang.CXType_Char_U, clang.CXType_UChar, clang.CXType_UShort, clang.CXType_UInt, clang.CXType_ULong, clang.CXType_ULongLong)
|
||||
ints = uints + (clang.CXType_Char_S, clang.CXType_Short, clang.CXType_Int, clang.CXType_Long, clang.CXType_LongLong)
|
||||
ints = uints + (clang.CXType_Char_S, clang.CXType_SChar, clang.CXType_Short, clang.CXType_Int, clang.CXType_Long, clang.CXType_LongLong)
|
||||
fps, specs = (clang.CXType_FunctionProto, clang.CXType_FunctionNoProto), (clang.CXCursor_ObjCSuperClassRef,) # this could include protocols
|
||||
|
||||
tmap = {clang.CXType_Void:"None", clang.CXType_Char_U:"ctypes.c_ubyte", clang.CXType_UChar:"ctypes.c_ubyte", clang.CXType_WChar:"ctypes.c_wchar",
|
||||
@@ -250,10 +250,12 @@ def gen(name, dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False,
|
||||
case clang.CXCursor_VarDecl if clang.clang_getCursorLinkage(c) == clang.CXLinkage_Internal:
|
||||
ty = clang.clang_getCursorType(c)
|
||||
if (ty.kind == clang.CXType_ConstantArray and clang.clang_getCanonicalType(clang.clang_getArrayElementType(ty)).kind in ints and
|
||||
(init:=children(c)[-1]).kind == clang.CXCursor_InitListExpr
|
||||
and all(re.match(r"\[.*\].*=", readext(f, extent(c))) for c in children(init))):
|
||||
(init:=children(c)[-1]).kind == clang.CXCursor_InitListExpr):
|
||||
cs = children(init)
|
||||
macros += [f"{nm(c)} = {{{','.join(f'{readext(f, extent(next(it:=iter(children(c)))))}:{readext(f, extent(next(it)))}' for c in cs)}}}"]
|
||||
if all(re.match(r"\[.*\].*=", readext(f, extent(ch))) for ch in cs):
|
||||
items = ','.join(f'{readext(f, extent(next(it:=iter(children(ch)))))}:{readext(f, extent(next(it)))}' for ch in cs)
|
||||
macros += [f"{nm(c)} = {{{items}}}"]
|
||||
else: macros += [f"{nm(c)} = ({','.join(readext(f, extent(ch)) for ch in cs)},)"]
|
||||
elif clang.clang_getCanonicalType(ty).kind in ints: macros += [f"{nm(c)} = {readext(f, extent(children(c)[-1]))}"]
|
||||
else: macros += [f"{nm(c)} = {tname(ty)}({readext(f, extent(children(c)[-1]))})"]
|
||||
case clang.CXCursor_VarDecl if clang.clang_getCursorLinkage(c) == clang.CXLinkage_External and dll:
|
||||
|
||||
Reference in New Issue
Block a user