File size: 3,752 Bytes
2c20074 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """Both ops against `gguf.quants.dequantize`, the reference implementation of the block layouts.
The blocks are random bytes rather than a quantized tensor: `gguf` can unpack every type but only
pack a couple of them, and unpacking is defined for any byte pattern, so this needs no quantizer and
no checkpoint. The one constraint is that a block's scales are fp16 fields — masked below so a random
pattern cannot land on an exponent of all ones and make the whole block inf/nan.
"""
import numpy as np
import pytest
import torch
from gguf_kernels import MAX_GEMV_ROWS, dequantize, mul_mat_vec
gguf = pytest.importorskip("gguf", reason="the reference unpacker comes from the `gguf` package")
# name -> (ggml type id, values per block, bytes per block)
QUANT_TYPES = {
"Q4_K": (12, 256, 144),
"Q5_K": (13, 256, 176),
"Q6_K": (14, 256, 210),
"Q8_0": (8, 32, 34),
}
def random_blocks(rows: int, cols: int, ggml_name: str, device="cuda"):
"""Random blocks and the values `gguf` reads out of them."""
_, block_values, block_bytes = QUANT_TYPES[ggml_name]
generator = np.random.default_rng(0)
packed = generator.integers(0, 256, (rows, cols // block_values * block_bytes), dtype=np.uint8)
# every fp16 scale sits at an even offset in its block, so clearing bit 6 of each odd byte keeps
# every possible fp16 field finite whatever the rest of the pattern is
packed[:, 1::2] &= 0xBF
quant_type = getattr(gguf.GGMLQuantizationType, ggml_name)
reference = gguf.quants.dequantize(packed.reshape(-1), quant_type).reshape(rows, cols)
return torch.from_numpy(packed).to(device), torch.from_numpy(reference).to(device)
@pytest.mark.kernels_ci
@pytest.mark.parametrize("ggml_name", QUANT_TYPES)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_dequantize_matches_reference(ggml_name, dtype):
ggml_type = QUANT_TYPES[ggml_name][0]
rows, cols = 64, 512
blocks, reference = random_blocks(rows, cols, ggml_name)
out = dequantize(blocks, ggml_type, rows, cols, dtype)
assert out.shape == (rows, cols) and out.dtype == dtype
# the kernel writes `dtype` directly, so the tolerance is that dtype's own resolution
torch.testing.assert_close(out.float(), reference, rtol=torch.finfo(dtype).eps * 4, atol=0)
@pytest.mark.kernels_ci
@pytest.mark.parametrize("ggml_name", QUANT_TYPES)
@pytest.mark.parametrize("n_rows", [1, MAX_GEMV_ROWS])
def test_mul_mat_vec_matches_matmul(ggml_name, n_rows):
ggml_type = QUANT_TYPES[ggml_name][0]
out_features, in_features = 128, 512
blocks, reference = random_blocks(out_features, in_features, ggml_name)
x = torch.randn(n_rows, in_features, dtype=torch.bfloat16, device="cuda")
out = mul_mat_vec(blocks, x, ggml_type, out_features)
assert out.shape == (n_rows, out_features) and out.dtype == torch.float32
# the kernel quantizes the activations to q8_1, so this is close to a matmul, not equal to one
expected = x.float() @ reference.T
torch.testing.assert_close(out, expected, rtol=2e-2, atol=2e-2 * expected.abs().max())
@pytest.mark.kernels_ci
def test_gemv_is_compileable():
"""A graph break here would cost more than the kernel saves, so the fake has to be right."""
ggml_type = QUANT_TYPES["Q4_K"][0]
out_features, in_features = 128, 512
blocks, reference = random_blocks(out_features, in_features, "Q4_K")
x = torch.randn(1, in_features, dtype=torch.bfloat16, device="cuda")
compiled = torch.compile(
lambda t: mul_mat_vec(blocks, t, ggml_type, out_features), fullgraph=True
)
expected = x.float() @ reference.T
torch.testing.assert_close(compiled(x), expected, rtol=2e-2, atol=2e-2 * expected.abs().max())
|