From a164f3e0c282eb38f9cfd86e16434b3ebc6678f3 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Sun, 5 Oct 2025 03:12:05 +0300
Subject: [PATCH] SDNQ Improve UINT3 and below quant speed

---
 modules/sdnq/packed_int.py | 93 ++++++++++++++++----------------------
 1 file changed, 39 insertions(+), 54 deletions(-)

diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py
index 31e8d9214..09a38efbc 100644
--- a/modules/sdnq/packed_int.py
+++ b/modules/sdnq/packed_int.py
@@ -219,79 +219,64 @@ def unpack_uint4(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
 
 
 def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
-    result = torch.cat(
-        (
-            torch.bitwise_and(
-                torch.cat(
-                    (
-                        packed_tensor[:, :3],
-                        torch.bitwise_right_shift(packed_tensor[:, :3], 3)
-                    ),
-                    dim=-1
-                ),
-                7
-            ),
-            torch.bitwise_or(
-                torch.bitwise_right_shift(packed_tensor[:, :2], 6),
-                torch.bitwise_and(
-                    torch.stack(
-                        (
-                            torch.bitwise_right_shift(packed_tensor[:, 2], 4),
-                            torch.bitwise_right_shift(packed_tensor[:, 2], 5),
+    result = torch.bitwise_and(
+        torch.cat(
+            (
+                packed_tensor[:, :3],
+                torch.bitwise_right_shift(packed_tensor[:, :3], 3),
+                torch.bitwise_or(
+                    torch.bitwise_right_shift(packed_tensor[:, :2], 6),
+                    torch.bitwise_and(
+                        torch.stack(
+                            (
+                                torch.bitwise_right_shift(packed_tensor[:, 2], 4),
+                                torch.bitwise_right_shift(packed_tensor[:, 2], 5),
+                            ),
+                            dim=-1
                         ),
-                        dim=-1
+                        4
                     ),
-                    4
                 ),
             ),
+            dim=-1
         ),
-        dim=-1
+        7
     ).view(shape)
     return result
 
 
 def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
-    result = torch.cat(
-        (
-            torch.bitwise_and(
-                torch.stack(
-                    (
-                        packed_tensor,
-                        torch.bitwise_right_shift(packed_tensor, 2),
-                        torch.bitwise_right_shift(packed_tensor, 4)
-                    ),
-                    dim=-1
-                ),
-                3
+    result = torch.bitwise_and(
+        torch.stack(
+            (
+                packed_tensor,
+                torch.bitwise_right_shift(packed_tensor, 2),
+                torch.bitwise_right_shift(packed_tensor, 4),
+                torch.bitwise_right_shift(packed_tensor, 6),
             ),
-            torch.bitwise_right_shift(packed_tensor, 6).unsqueeze(-1),
+            dim=-1
         ),
-        dim=-1
+        3
     ).view(shape)
     return result
 
 
 def unpack_uint1(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
-    result = torch.cat(
-        (
-            torch.bitwise_and(
-                torch.stack(
-                    (
-                        packed_tensor,
-                        torch.bitwise_right_shift(packed_tensor, 1),
-                        torch.bitwise_right_shift(packed_tensor, 2),
-                        torch.bitwise_right_shift(packed_tensor, 3),
-                        torch.bitwise_right_shift(packed_tensor, 4),
-                        torch.bitwise_right_shift(packed_tensor, 5),
-                        torch.bitwise_right_shift(packed_tensor, 6),
-                    ),
-                    dim=-1
-                ),
-                1
+    result = torch.bitwise_and(
+        torch.stack(
+            (
+                packed_tensor,
+                torch.bitwise_right_shift(packed_tensor, 1),
+                torch.bitwise_right_shift(packed_tensor, 2),
+                torch.bitwise_right_shift(packed_tensor, 3),
+                torch.bitwise_right_shift(packed_tensor, 4),
+                torch.bitwise_right_shift(packed_tensor, 5),
+                torch.bitwise_right_shift(packed_tensor, 6),
+                torch.bitwise_right_shift(packed_tensor, 7),
             ),
-            torch.bitwise_right_shift(packed_tensor, 7).unsqueeze(-1),
+            dim=-1
         ),
-        dim=-1
+        1
     ).view(shape)
     return result