Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-04 12:00:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

SDNQ handle packed floats in fp mm

Browse Source
This commit is contained in:
Disty0
2025-12-27 16:29:18 +03:00
parent 43079794c8
commit db59d2b507
9 changed files with 104 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/layers/conv/conv_fp16.py
+11 -1
View File
@@ -5,6 +5,7 @@ from typing import List
import torch
from ...common import compile_func, fp_mm_func # noqa: TID252
from ...packed_float import unpack_float # noqa: TID252
from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
from .forward import get_conv_args, process_conv_input
@@ -24,6 +25,8 @@ def conv_fp16_matmul(
bias: torch.FloatTensor = None,
svd_up: torch.FloatTensor = None,
svd_down: torch.FloatTensor = None,
quantized_weight_shape: torch.Size = None,
weights_dtype: str = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
@@ -35,7 +38,10 @@ def conv_fp16_matmul(
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, scale = quantize_fp_mm_input_tensorwise(input, scale, matmul_dtype="float16")
weight = weight.to(dtype=torch.float16) # fp8 weights
if quantized_weight_shape is not None:
weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float16)
elif weight.dtype != torch.float16:
weight = weight.to(dtype=torch.float16) # fp8 weights
input, weight = check_mats(input, weight)
if groups == 1:
@@ -64,8 +70,10 @@ def conv_fp16_matmul(
def quantized_conv_forward_fp16_matmul(self, input) -> torch.FloatTensor:
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
return conv_fp16_matmul(
input, weight, scale,
@@ -76,6 +84,8 @@ def quantized_conv_forward_fp16_matmul(self, input) -> torch.FloatTensor:
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
modules/sdnq/layers/conv/conv_fp8.py
+9
View File
@@ -5,6 +5,7 @@ from typing import List
import torch
from ...common import compile_func # noqa: TID252
from ...packed_float import unpack_float # noqa: TID252
from .forward import get_conv_args, process_conv_input
from ..linear.linear_fp8 import quantize_fp_mm_input # noqa: TID252
@@ -23,6 +24,8 @@ def conv_fp8_matmul(
bias: torch.FloatTensor = None,
svd_up: torch.FloatTensor = None,
svd_down: torch.FloatTensor = None,
quantized_weight_shape: torch.Size = None,
weights_dtype: str = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
@@ -31,6 +34,8 @@ def conv_fp8_matmul(
svd_bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_fp_mm_input(input)
if quantized_weight_shape is not None:
weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn)
input, weight = check_mats(input, weight)
if groups == 1:
@@ -71,8 +76,10 @@ def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
return conv_fp8_matmul(
input, weight, scale,
@@ -83,6 +90,8 @@ def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
modules/sdnq/layers/conv/conv_fp8_tensorwise.py
+9
View File
@@ -5,6 +5,7 @@ from typing import List
import torch
from ...common import compile_func # noqa: TID252
from ...packed_float import unpack_float # noqa: TID252
from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
from .forward import get_conv_args, process_conv_input
@@ -24,6 +25,8 @@ def conv_fp8_matmul_tensorwise(
bias: torch.FloatTensor = None,
svd_up: torch.FloatTensor = None,
svd_down: torch.FloatTensor = None,
quantized_weight_shape: torch.Size = None,
weights_dtype: str = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
@@ -35,6 +38,8 @@ def conv_fp8_matmul_tensorwise(
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, scale = quantize_fp_mm_input_tensorwise(input, scale)
if quantized_weight_shape is not None:
weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn)
input, weight = check_mats(input, weight)
dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
@@ -66,8 +71,10 @@ def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTens
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
return conv_fp8_matmul_tensorwise(
input, weight, scale,
@@ -78,6 +85,8 @@ def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTens
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
modules/sdnq/layers/conv/conv_int8.py
+13 -12
View File
@@ -15,18 +15,18 @@ from ..linear.forward import check_mats # noqa: TID252
def conv_int8_matmul(
input: torch.FloatTensor,
weight: torch.CharTensor,
bias: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
svd_up: torch.FloatTensor,
svd_down: torch.FloatTensor,
quantized_weight_shape: torch.Size,
result_shape: torch.Size,
weights_dtype: str,
reversed_padding_repeated_twice: List[int],
padding_mode: str, conv_type: int,
groups: int, stride: List[int],
padding: List[int], dilation: List[int],
bias: torch.FloatTensor = None,
svd_up: torch.FloatTensor = None,
svd_down: torch.FloatTensor = None,
quantized_weight_shape: torch.Size = None,
weights_dtype: str = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
@@ -73,18 +73,19 @@ def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
else:
weight = self.weight
scale = self.scale
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
return conv_int8_matmul(
input, weight, self.bias,
scale, self.svd_up, self.svd_down,
quantized_weight_shape,
input, weight, scale,
self.sdnq_dequantizer.result_shape,
self.sdnq_dequantizer.weights_dtype,
self._reversed_padding_repeated_twice,
self.padding_mode, conv_type,
self.groups, stride, padding, dilation,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
modules/sdnq/layers/linear/linear_fp16.py
+17 -2
View File
@@ -3,6 +3,7 @@
import torch
from ...common import compile_func, fp_mm_func # noqa: TID252
from ...packed_float import unpack_float # noqa: TID252
from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
from .forward import check_mats
@@ -16,7 +17,13 @@ def fp16_matmul(
bias: torch.FloatTensor = None,
svd_up: torch.FloatTensor = None,
svd_down: torch.FloatTensor = None,
quantized_weight_shape: torch.Size = None,
weights_dtype: str = None,
) -> torch.FloatTensor:
if quantized_weight_shape is not None:
weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float16)
elif weight.dtype != torch.float16:
weight = weight.to(dtype=torch.float16) # fp8 weights
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if svd_up is not None:
@@ -26,7 +33,6 @@ def fp16_matmul(
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, scale = quantize_fp_mm_input_tensorwise(input, scale, matmul_dtype="float16")
weight = weight.to(dtype=torch.float16) # fp8 weights
input, weight = check_mats(input, weight)
if bias is not None:
return dequantize_symmetric_with_bias(fp_mm_func(input, weight), scale, bias, dtype=return_dtype, result_shape=output_shape)
@@ -37,9 +43,18 @@ def fp16_matmul(
def quantized_linear_forward_fp16_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
return fp16_matmul(input, weight, scale, bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down)
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
return fp16_matmul(
input, weight, scale,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
fp16_matmul = compile_func(fp16_matmul)
modules/sdnq/layers/linear/linear_fp8.py
+15 -1
View File
@@ -5,6 +5,7 @@ from typing import Tuple
import torch
from ...common import compile_func # noqa: TID252
from ...packed_float import unpack_float # noqa: TID252
from ...dequantizer import quantize_fp_mm # noqa: TID252
from .forward import check_mats
@@ -23,7 +24,11 @@ def fp8_matmul(
bias: torch.FloatTensor = None,
svd_up: torch.FloatTensor = None,
svd_down: torch.FloatTensor = None,
quantized_weight_shape: torch.Size = None,
weights_dtype: str = None,
) -> torch.FloatTensor:
if quantized_weight_shape is not None:
weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn)
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if svd_up is not None:
@@ -45,9 +50,18 @@ def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
return fp8_matmul(input, weight, scale, bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down)
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
return fp8_matmul(
input, weight, scale,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
fp8_matmul = compile_func(fp8_matmul)
modules/sdnq/layers/linear/linear_fp8_tensorwise.py
+15 -1
View File
@@ -5,6 +5,7 @@ from typing import Tuple
import torch
from ...common import compile_func # noqa: TID252
from ...packed_float import unpack_float # noqa: TID252
from ...dequantizer import quantize_fp_mm, dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
from .forward import check_mats
@@ -26,7 +27,11 @@ def fp8_matmul_tensorwise(
bias: torch.FloatTensor = None,
svd_up: torch.FloatTensor = None,
svd_down: torch.FloatTensor = None,
quantized_weight_shape: torch.Size = None,
weights_dtype: str = None,
) -> torch.FloatTensor:
if quantized_weight_shape is not None:
weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn)
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if svd_up is not None:
@@ -49,9 +54,18 @@ def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTenso
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
return fp8_matmul_tensorwise(input, weight, scale, bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down)
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
return fp8_matmul_tensorwise(
input, weight, scale,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
fp8_matmul_tensorwise = compile_func(fp8_matmul_tensorwise)
modules/sdnq/layers/linear/linear_int8.py
+2 -3
View File
@@ -55,8 +55,7 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
quantized_weight_shape = None
else:
weight = self.weight
scale = self.scale
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
return int8_matmul(
input, weight, scale,
@@ -64,7 +63,7 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc
svd_up=self.svd_up,
svd_down=self.svd_down,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
modules/sdnq/quantizer.py
+13 -4
View File
@@ -228,7 +228,15 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int
dtype_dict[weights_dtype]["is_unsigned"]
or dtype_dict[weights_dtype]["is_integer"] != dtype_dict[quantized_matmul_dtype]["is_integer"]
or dtype_dict[weights_dtype]["num_bits"] > dtype_dict[quantized_matmul_dtype]["num_bits"]
or (dtype_dict[weights_dtype]["is_packed"] and not dtype_dict[weights_dtype]["is_integer"])
or (
dtype_dict[weights_dtype]["is_packed"]
and not dtype_dict[weights_dtype]["is_integer"]
and not dtype_dict[quantized_matmul_dtype]["is_integer"]
and (
dtype_dict[weights_dtype]["num_bits"] >= dtype_dict[quantized_matmul_dtype]["num_bits"]
or dtype_dict[weights_dtype]["max"] > dtype_dict[quantized_matmul_dtype]["max"]
)
)
)
if layer_class_name in conv_types:
@@ -348,7 +356,7 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int
scale.t_()
weight.t_()
weight = prepare_weight_for_matmul(weight)
if not use_tensorwise_fp8_matmul and not dtype_dict[weights_dtype]["is_integer"]:
if not use_tensorwise_fp8_matmul and not dtype_dict[quantized_matmul_dtype]["is_integer"]:
scale = scale.to(dtype=torch.float32)
sdnq_dequantizer = SDNQDequantizer(
@@ -956,8 +964,9 @@ class SDNQConfig(QuantizationConfigMixin):
Args:
weights_dtype (`str`, *optional*, defaults to `"int8"`):
The target dtype for the weights after quantization. Supported values are:
("int16", "int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint16", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float16", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
The target dtype for the weights after quantization.
Check out `sdnq.common.accepted_weight_dtypes` for all the supported values.
These are some of the recommended values to use: ("int8", "int7", "int6", "uint5", "uint4", "uint3", "uint2", "float8_e4m3fn", "float7_e3m3fn", "float6_e3m2fn", "float5_e2m2fn", "float4_e2m1fn", "float3_e1m1fn", "float2_e1m0fn")
quantized_matmul_dtype (`str`, *optional*, defaults to `None`):
The target dtype for quantized matmul.
`None` will use "int8" with integer weight dtypes and "float8_e4m3fn" or "float16" with float weight dtypes.