diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py index 2eecaf6ec1..27a478acdf 100644 --- a/tinygrad/llm/kernels/amd.py +++ b/tinygrad/llm/kernels/amd.py @@ -55,11 +55,17 @@ class Linear(nn.Linear): self.in_features, self.out_features = in_features, out_features def set_quantized(self, decoded:Tensor): packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in QUANT_SIZES.items()} - raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None) + graph = decoded.uop.toposort() + raw = next((u for u in graph if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None) if raw is None: return + ggml_type = packed_sizes[prod(raw.shape)] + # the packed byte rate alone can't distinguish same-rate formats (Q4_0 vs Q4_K, Q5_0 vs Q5_K, MXFP4 vs IQ4_XS). + # the supported formats are 256-wide superblocks: their decode views the packed bytes at the superblock width + # (ggml_data_to_tensor reshapes to (-1, QUANT_SIZES[type])), while same-rate 32-wide formats reshape to 17-22 + if not any(u.op is Ops.RESHAPE and u.shape[-1:] == (QUANT_SIZES[ggml_type],) for u in graph): 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.ggml_type = packed_sizes[prod(raw.shape)] + self.ggml_type = ggml_type # store a typed buffer view: a lazy BITCAST is decomposed into byte-combining ALU before custom-kernel # scheduling and would copy the entire packed weight on every JIT graph if self.ggml_type == Q6_K: