diff --git a/test/unit/test_llm_amd.py b/test/unit/test_llm_amd.py index a62c0412dd..38e705e4e1 100644 --- a/test/unit/test_llm_amd.py +++ b/test/unit/test_llm_amd.py @@ -6,6 +6,18 @@ from tinygrad.llm.kernels.amd import q8_quantize, quantized_attention from tinygrad.llm.gguf import ggml_data_to_tensor class TestQ8Quantize(unittest.TestCase): + def test_word_quant_weights_use_typed_buffer_view(self): + for ggml_type, type_size in ((13, 176), (23, 136)): + with self.subTest(ggml_type=ggml_type): + raw = Tensor(np.zeros(type_size + 4, dtype=np.uint8), device="CPU").contiguous().realize()[4:] + decoded = ggml_data_to_tensor(raw, 256, ggml_type).reshape(1, 256) + linear = Linear(256, 1, bias=False) + linear.set_quantized(decoded) + self.assertEqual(linear.ggml_type, ggml_type) + self.assertEqual(linear.weight.dtype, dtypes.uint32) + self.assertEqual(linear.weight.nbytes(), type_size) + self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4) + def test_values_and_scales(self): if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required") x = np.linspace(-3.1, 2.7, 64, dtype=np.float32).reshape(2, 32) diff --git a/tinygrad/llm/kernels/__init__.py b/tinygrad/llm/kernels/__init__.py index a6e18f41c8..d4e1a7dbef 100644 --- a/tinygrad/llm/kernels/__init__.py +++ b/tinygrad/llm/kernels/__init__.py @@ -25,8 +25,12 @@ class Linear(nn.Linear): if raw is None: return raw_offset = raw.contiguous_view_offset() assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8 - self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(raw.max_numel(), dtypes.uint8, raw_offset))) self.ggml_type = packed_sizes[prod(raw.shape)] + # Q5_K and IQ4_XS kernels consume words. Store a typed buffer view directly: a lazy BITCAST is decomposed into + # byte-combining ALU before custom-kernel scheduling and would copy the entire packed weight on every JIT graph. + packed_dtype = dtypes.uint8 if self.ggml_type == 14 else dtypes.uint32 + self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(raw.max_numel() * raw.dtype.itemsize // packed_dtype.itemsize, + packed_dtype, raw_offset))) def __call__(self, x:Tensor) -> Tensor: static = isinstance(x.numel(), int) supported = self.use_custom_quant and amd_custom_kernels_supported(cast(str, self.weight.device)) diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py index 0a107391b2..fd550051af 100644 --- a/tinygrad/llm/kernels/amd.py +++ b/tinygrad/llm/kernels/amd.py @@ -432,14 +432,14 @@ def q8_linear(layer:Linear, x:Tensor) -> Tensor: out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device).uop if layer.ggml_type == Q5_K and use_wmma: - return run(_q5_linear_f16_wmma_kernel, out, raw.bitcast(dtypes.uint32), x.cast(dtypes.float16).contiguous().uop) + return run(_q5_linear_f16_wmma_kernel, out, raw, x.cast(dtypes.float16).contiguous().uop) if layer.ggml_type == IQ4_XS and use_wmma: - return run(_iq4_linear_f16_wmma_kernel, out, raw.bitcast(dtypes.uint32), x.cast(dtypes.float16).contiguous().uop, + return run(_iq4_linear_f16_wmma_kernel, out, raw, x.cast(dtypes.float16).contiguous().uop, iq4_half_lut(str(x.device)).uop) xq, xd = q8_quantize(x, tokens, layer.in_features) decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type) out = Tensor.empty(tokens, layer.out_features, (layer.in_features+1023)//1024, dtype=dtypes.float32, device=x.device).uop - return run(decode, out, raw if layer.ggml_type == Q6_K else raw.bitcast(dtypes.uint32), xq.uop, xd.uop) + return run(decode, out, raw, xq.uop, xd.uop) @functools.cache def iq4_half_lut(device:str) -> Tensor: