fix useless copies

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)

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))

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: