From dec460e6655ab5df10de3fe37d858058c8167f36 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 27 May 2025 03:02:36 +0300 Subject: [PATCH] SDNQ use torch.bitwise ops instead of python --- modules/model_quant_sdnq.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py index cc8b62433..0110b2db6 100644 --- a/modules/model_quant_sdnq.py +++ b/modules/model_quant_sdnq.py @@ -433,15 +433,14 @@ def pack_uint4(tensor: torch.Tensor) -> torch.Tensor: if tensor.dtype != torch.uint8: raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.") packed_tensor = tensor.contiguous().reshape(-1, 2) - packed_tensor = torch.bitwise_and(packed_tensor[..., ::2], 15) | packed_tensor[..., 1::2] << 4 + packed_tensor = torch.bitwise_or(torch.bitwise_and(packed_tensor[:, 0], 15), torch.bitwise_left_shift(packed_tensor[:, 1], 4)) return packed_tensor def pack_int4(tensor: torch.Tensor) -> torch.Tensor: if tensor.dtype != torch.int8: raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.int8 type is supported.") - tensor = tensor + 8 - return pack_uint4(tensor.to(dtype=torch.uint8)) + return pack_uint4((tensor + 8).to(dtype=torch.uint8)) def unpack_uint4(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: