From 88fb6fa0014850615ca5b3e0ec1c018f67319237 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Thu, 15 Aug 2024 18:39:24 +0200
Subject: [PATCH] Improve Q8_0 and Q4_0, Add Q4_1 and Q5_1

Use FP32 to do the calculations while dequantizing to improve existing Q8_0 quants and to avoid double dtype conversion.
Add Q4_1 and Q5_1 dequant kernels
---
 dequant.py | 60 +++++++++++++++++++++++++++++++++++++++++++-----------
 ops.py     |  7 ++++---
 2 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/dequant.py b/dequant.py
index 2be51b8..920c237 100644
--- a/dequant.py
+++ b/dequant.py
@@ -17,13 +17,18 @@ def dequantize_tensor(tensor, dtype=torch.float16):
         return data.to(dtype)
     elif qtype in dequantize_functions:
         # this is the main pytorch op
-        return dequantize(data, qtype, oshape).to(dtype)
+        if False and dtype == torch.bfloat16: # disabled, lower quality
+            # dequant directly in native BF16
+            return dequantize(data, qtype, oshape, dtype) 
+        else:
+            # dequant in fp32, convert to target
+            return dequantize(data, qtype, oshape, dtype=torch.float32).to(dtype)
     else:
         # this is incredibly slow
         new = gguf.quants.dequantize(data.cpu().numpy(), qtype)
         return torch.from_numpy(new).to(data.device, dtype=dtype)
 
-def dequantize(data, qtype, oshape):
+def dequantize(data, qtype, oshape, dtype=None):
     """
     Dequantize tensor back to usable shape/dtype
     """
@@ -36,7 +41,7 @@ def dequantize(data, qtype, oshape):
 
     n_blocks = rows.numel() // type_size
     blocks = rows.reshape((n_blocks, type_size))
-    blocks = dequantize_blocks(blocks, block_size, type_size)
+    blocks = dequantize_blocks(blocks, block_size, type_size, dtype)
     return blocks.reshape(oshape)
 
 def to_uint32(x):
@@ -44,19 +49,36 @@ def to_uint32(x):
     x = x.view(torch.uint8).to(torch.int32) 
     return (x[:, 0] | x[:, 1] << 8 | x[:, 2] << 16 | x[:, 3] << 24).unsqueeze(1)
 
-def dequantize_blocks_Q8_0(blocks, block_size, type_size):
-    d = blocks[:, :2].view(torch.float16)
-    x = blocks[:, 2:].view(torch.int8).to(torch.float16)
-    return (x * d)
+def dequantize_blocks_Q8_0(blocks, block_size, type_size, dtype=None):
+    d = blocks[:, :2].view(torch.float16).to(dtype)
+    x = blocks[:, 2:].view(torch.int8)
+    return (d * x)
 
-def dequantize_blocks_Q5_0(blocks, block_size, type_size):
+def dequantize_blocks_Q5_1(blocks, block_size, type_size, dtype=None):
     n_blocks = blocks.shape[0]
 
-    d  = blocks[:,  :2]
+    d  = blocks[:,  :2].view(torch.float16).to(dtype)
+    m  = blocks[:, 2:4].view(torch.float16).to(dtype)
+    qh = blocks[:, 4:8]
+    qs = blocks[:, 8: ]
+
+    qh = to_uint32(qh)
+
+    qh = qh.reshape((n_blocks, 1)) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32)
+    ql = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2, 1)
+    qh = (qh & 1).to(torch.uint8)
+    ql = (ql & 0x0F).reshape((n_blocks, -1))
+
+    qs = (ql | (qh << 4))
+    return (d * qs) + m
+
+def dequantize_blocks_Q5_0(blocks, block_size, type_size, dtype=None):
+    n_blocks = blocks.shape[0]
+
+    d  = blocks[:,  :2].view(torch.float16).to(dtype)
     qh = blocks[:, 2:6]
     qs = blocks[:, 6: ]
 
-    d = d.view(torch.float16).to(torch.float32)
     qh = to_uint32(qh)
 
     qh = qh.reshape(n_blocks, 1) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32)
@@ -68,10 +90,22 @@ def dequantize_blocks_Q5_0(blocks, block_size, type_size):
     qs = (ql | (qh << 4)).to(torch.int8) - 16
     return (d * qs)
 
-def dequantize_blocks_Q4_0(blocks, block_size, type_size):
+def dequantize_blocks_Q4_1(blocks, block_size, type_size, dtype=None):
     n_blocks = blocks.shape[0]
 
-    d  = blocks[:,  :2].view(torch.float16)
+    d  = blocks[:,  :2].view(torch.float16).to(dtype)
+    m  = blocks[:, 2:4].view(torch.float16).to(dtype)
+    qs = blocks[:, 4: ]
+
+    qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2, 1)
+    qs = (qs & 0x0F).reshape(n_blocks, -1)
+
+    return (d * qs) + m
+
+def dequantize_blocks_Q4_0(blocks, block_size, type_size, dtype=None):
+    n_blocks = blocks.shape[0]
+
+    d  = blocks[:,  :2].view(torch.float16).to(dtype)
     qs = blocks[:,  2:]
 
     qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1))
@@ -80,6 +114,8 @@ def dequantize_blocks_Q4_0(blocks, block_size, type_size):
 
 dequantize_functions = {
     gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0,
+    gguf.GGMLQuantizationType.Q5_1: dequantize_blocks_Q5_1,
     gguf.GGMLQuantizationType.Q5_0: dequantize_blocks_Q5_0,
+    gguf.GGMLQuantizationType.Q4_1: dequantize_blocks_Q4_1,
     gguf.GGMLQuantizationType.Q4_0: dequantize_blocks_Q4_0,
 }
diff --git a/ops.py b/ops.py
index 7c79952..20ded45 100644
--- a/ops.py
+++ b/ops.py
@@ -72,7 +72,8 @@ class GGMLOps(comfy.ops.disable_weight_init):
             self.parameters_manual_cast = torch.float32
 
         def forward(self, x):
-            weight, bias = self.get_weights(x.dtype)
-            x = torch.nn.functional.linear(x, weight, bias)
-            del weight, bias
+            with torch.no_grad():
+                weight, bias = self.get_weights(x.dtype)
+                x = torch.nn.functional.linear(x, weight, bias)
+                del weight, bias
             return x