SDNQ use inplace transpose and use view instead of reshape

modules/sdnq/__init__.py

@@ -159,12 +159,12 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
                 zero_point = zero_point.to(torch_dtype)
 
         if use_quantized_matmul:
-            scale = scale.transpose(0,1)
-            layer.weight.data = layer.weight.transpose(0,1)
+            scale.transpose_(0,1)
+            layer.weight.transpose_(0,1)
             if not dtype_dict[weights_dtype]["is_integer"]:
                 stride = layer.weight.stride()
                 if stride[0] > stride[1] and stride[1] == 1:
-                    layer.weight.data = layer.weight.t().contiguous().t()
+                    layer.weight.data = layer.weight.t_().contiguous().t_()
                 if not use_tensorwise_fp8_matmul:
                     scale = scale.to(torch.float32)
 

modules/sdnq/dequantizer.py

@@ -9,21 +9,21 @@ from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetr
 def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
     result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale).to(dtype=dtype)
     if result_shape is not None:
-        result = result.reshape(result_shape)
+        result = result.view(result_shape)
     return result
 
 
 def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
     result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
     if skip_quantized_matmul:
-        result = result.t()
+        result.t_()
     if result_shape is not None:
-        result = result.reshape(result_shape)
+        result = result.view(result_shape)
     return result
 
 
 def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatTensor, bias: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
-    return torch.addcmul(bias, weight.to(dtype=scale.dtype), scale).to(dtype=dtype).reshape(result_shape)
+    return torch.addcmul(bias, weight.to(dtype=scale.dtype), scale).to(dtype=dtype).view(result_shape)
 
 
 def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:

modules/sdnq/layers/conv/conv_fp8.py

@@ -25,24 +25,24 @@ def conv_fp8_matmul(
     input, input_scale = quantize_fp8_matmul_input(input)
 
     if groups == 1:
-        result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape)
+        result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
     else:
-        scale = scale.reshape(groups, 1, scale.shape[1] // groups)
-        input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1)
-        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
-        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        scale = scale.view(groups, 1, scale.shape[1] // groups)
+        input_scale = input_scale.view(groups, input_scale.shape[0] // groups, 1)
+        weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
+        input = input.view(input.shape[0], groups, input.shape[1] // groups)
         result = []
         if bias is not None:
-            bias = bias.reshape(groups, bias.shape[0] // groups)
+            bias = bias.view(groups, bias.shape[0] // groups)
             for i in range(groups):
-                result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype))
+                result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype))
         else:
             for i in range(groups):
-                result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype))
-        result = torch.cat(result, dim=-1).reshape(mm_output_shape)
+                result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype))
+        result = torch.cat(result, dim=-1).view(mm_output_shape)
 
     if conv_type == 1:
-        result = result.transpose(1,2)
+        result = result.transpose_(1,2)
     elif conv_type == 2:
         result = result.permute(0,3,1,2)
     elif conv_type == 3:

modules/sdnq/layers/conv/conv_fp8_tensorwise.py

@@ -29,11 +29,11 @@ def conv_fp8_matmul_tensorwise(
     if groups == 1:
         result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)
     else:
-        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
-        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
+        input = input.view(input.shape[0], groups, input.shape[1] // groups)
         result = []
         for i in range(groups):
-            result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype))
+            result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype))
         result = torch.cat(result, dim=-1)
     if bias is not None:
         dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape)
@@ -41,7 +41,7 @@ def conv_fp8_matmul_tensorwise(
         dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
 
     if conv_type == 1:
-        result = result.transpose(1,2)
+        result = result.transpose_(1,2)
     elif conv_type == 2:
         result = result.permute(0,3,1,2)
     elif conv_type == 3:

modules/sdnq/layers/conv/conv_int8.py

@@ -33,11 +33,11 @@ def conv_int8_matmul(
     if groups == 1:
         result = torch._int_mm(input, weight)
     else:
-        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
-        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
+        input = input.view(input.shape[0], groups, input.shape[1] // groups)
         result = []
         for i in range(groups):
-            result.append(torch._int_mm(input[i], weight[i]))
+            result.append(torch._int_mm(input[:, i], weight[:, i]))
         result = torch.cat(result, dim=-1)
     if bias is not None:
         result = dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape)
@@ -45,7 +45,7 @@ def conv_int8_matmul(
         result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
 
     if conv_type == 1:
-        result = result.transpose(1,2)
+        result = result.transpose_(1,2)
     elif conv_type == 2:
         result = result.permute(0,3,1,2)
     elif conv_type == 3:

modules/sdnq/layers/linear/linear_fp8.py

@@ -21,10 +21,9 @@ def fp8_matmul(
     scale: torch.FloatTensor,
 ) -> torch.FloatTensor:
     return_dtype = input.dtype
-    output_shape = list(input.shape)
-    output_shape[-1] = weight.shape[-1]
+    output_shape = (*input.shape[:-1], weight.shape[-1])
     input, input_scale = quantize_fp8_matmul_input(input)
-    return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape)
+    return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).view(output_shape)
 
 
 def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:

modules/sdnq/layers/linear/linear_fp8_tensorwise.py

@@ -25,8 +25,7 @@ def fp8_matmul_tensorwise(
     scale: torch.FloatTensor,
 ) -> torch.FloatTensor:
     return_dtype = input.dtype
-    output_shape = list(input.shape)
-    output_shape[-1] = weight.shape[-1]
+    output_shape = (*input.shape[:-1], weight.shape[-1])
     dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
     input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
     if bias is not None:

modules/sdnq/layers/linear/linear_int8.py

@@ -30,8 +30,7 @@ def int8_matmul(
     if quantized_weight_shape is not None:
         weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8)
     return_dtype = input.dtype
-    output_shape = list(input.shape)
-    output_shape[-1] = weight.shape[-1]
+    output_shape = (*input.shape[:-1], weight.shape[-1])
     input, scale = quantize_int8_matmul_input(input, scale)
     if bias is not None:
         return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape)

modules/sdnq/packed_int.py

@@ -26,7 +26,7 @@ def unpack_int_asymetric(packed_tensor: torch.ByteTensor, shape: torch.Size, wei
 
 
 def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor:
-    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = tensor.contiguous().view(-1, 8)
     packed_tensor = torch.stack(
         (
             torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 1), 128)),
@@ -43,7 +43,7 @@ def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor:
 
 
 def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor:
-    packed_tensor = tensor.contiguous().reshape(-1, 4)
+    packed_tensor = tensor.contiguous().view(-1, 4)
     packed_tensor = torch.stack(
         (
             torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)),
@@ -56,7 +56,7 @@ def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor:
 
 
 def pack_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor:
-    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = tensor.contiguous().view(-1, 8)
     packed_tensor = torch.stack(
         (
             torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 5], 5)),
@@ -83,13 +83,13 @@ def pack_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor:
 
 
 def pack_uint4(tensor: torch.ByteTensor) -> torch.ByteTensor:
-    packed_tensor = tensor.contiguous().reshape(-1, 2)
+    packed_tensor = tensor.contiguous().view(-1, 2)
     packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4))
     return packed_tensor
 
 
 def pack_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor:
-    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = tensor.contiguous().view(-1, 8)
     packed_tensor = torch.stack(
         (
             torch.bitwise_or(
@@ -114,7 +114,7 @@ def pack_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor:
 
 
 def pack_uint2(tensor: torch.ByteTensor) -> torch.ByteTensor:
-    packed_tensor = tensor.contiguous().reshape(-1, 4)
+    packed_tensor = tensor.contiguous().view(-1, 4)
     packed_tensor = torch.bitwise_or(
         torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)),
         torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 4), torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
@@ -153,7 +153,7 @@ def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
             )
         ),
         dim=-1
-    ).reshape(shape)
+    ).view(shape)
     return result
 
 
@@ -172,7 +172,7 @@ def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
             )
         ),
         dim=-1
-    ).reshape(shape)
+    ).view(shape)
     return result
 
 
@@ -201,12 +201,12 @@ def unpack_uint5(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
             ),
         ),
         dim=-1
-    ).reshape(shape)
+    ).view(shape)
     return result
 
 
 def unpack_uint4(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
-    result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape)
+    result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).view(shape)
     return result
 
 
@@ -229,7 +229,7 @@ def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
             ),
         ),
         dim=-1
-    ).reshape(shape)
+    ).view(shape)
     return result
 
 
@@ -242,7 +242,7 @@ def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
             torch.bitwise_right_shift(packed_tensor, 6),
         ),
         dim=-1
-    ).reshape(shape)
+    ).view(shape)
     return result