SDQN add OpenVINO FP16 MM and add update_torch_dtype to SDNQQuantizer

modules/sdnq/common.py

@@ -380,8 +380,9 @@ int_mm_func = None
 
 if use_openvino_mm:
     try:
-        from .kernels.openvino_mm import openvino_int_mm
+        from .kernels.openvino_mm import openvino_int_mm, openvino_fp_mm
         int_mm_func = openvino_int_mm
+        fp_mm_func = openvino_fp_mm
     except Exception:
         use_openvino_mm = False
 elif use_triton_mm:

modules/sdnq/kernels/openvino_mm.py

@@ -6,35 +6,36 @@ from openvino.properties import hint as ov_hints
 
 core = ov.Core()
 
-NPU_MUL = 32 # NPU uses FP16 x INT8 -> FP16 instead of INT8 x INT8 -> INT32 and FP16 output overflows
 OV_DEVICE: str = os.environ.get("SDNQ_OPENVINO_DEVICE", "CPU")
-OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, tuple[int,int] | None], tuple[ov.InferRequest, str]] = {}
-core.set_property(OV_DEVICE, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY})
+OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, str, tuple[int,int] | None], tuple[ov.InferRequest, str]] = {}
+
+if OV_DEVICE == "NPU":
+    OV_DEVICE = "HETERO:NPU,CPU"
+for ov_device in core.get_available_devices():
+    core.set_property(ov_device, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY})
 
 
-def ov_int_mm(A: torch.CharTensor, B: torch.CharTensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor:
+def ov_mm(A: torch.CharTensor, B: torch.CharTensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor:
     C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32)
     infer_request.set_tensor("A", ov.Tensor(A.detach().contiguous().to("cpu").numpy(), shared_memory=True))
     infer_request.set_tensor("B", ov.Tensor(B.detach().contiguous().to("cpu").numpy(), shared_memory=True))
     infer_request.set_tensor(out_name, ov.Tensor(C.numpy(), shared_memory=True))
     infer_request.infer()
     C = C.to(A.device)
-    if OV_DEVICE == "NPU":
-        C.mul_(NPU_MUL**2)
     return C
 
 
 @torch.library.custom_op("sdnq::openvino_int_mm", mutates_args=())
 def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor:
-    if OV_DEVICE in {"NPU", "CPU"}:
-        cache_key = (OV_DEVICE, Tensor_A.shape, Tensor_B.shape)
+    if "GPU" not in OV_DEVICE:
+        cache_key = (OV_DEVICE, "int8", Tensor_A.shape, Tensor_B.shape)
     else:
-        cache_key = (OV_DEVICE, None, None)
+        cache_key = (OV_DEVICE, "int8", None, None)
     infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
     if infer_request is not None:
-        return ov_int_mm(Tensor_A, Tensor_B, infer_request, out_name)
+        return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
 
-    if OV_DEVICE in {"NPU", "CPU"}:
+    if "GPU" not in OV_DEVICE:
         shape_a = ov.Shape(Tensor_A.shape)
         shape_b = ov.Shape(Tensor_B.shape)
     else:
@@ -42,23 +43,100 @@ def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Ten
         shape_b = ov.PartialShape([-1,-1])
     input_a = ov_ops.parameter(shape_a, ov.Type.i8, name="A")
     input_b = ov_ops.parameter(shape_b, ov.Type.i8, name="B")
+    a = ov_ops.convert(input_a, ov.Type.f32)
+    b = ov_ops.convert(input_b, ov.Type.f32)
+
     low  = ov_ops.constant(-128.0, dtype=ov.Type.f32)
     high = ov_ops.constant(127.0,  dtype=ov.Type.f32)
+    a = ov_ops.fake_quantize(a, low, high, low, high, 256)
+    b = ov_ops.fake_quantize(b, low, high, low, high, 256)
 
-    a = ov_ops.fake_quantize(ov_ops.convert(input_a, ov.Type.f32), low, high, low, high, 256)
-    b = ov_ops.fake_quantize(ov_ops.convert(input_b, ov.Type.f32), low, high, low, high, 256)
-    if OV_DEVICE == "NPU":
-        a = ov_ops.divide(a, ov_ops.constant(NPU_MUL, dtype=ov.Type.f32))
-        b = ov_ops.divide(b, ov_ops.constant(NPU_MUL, dtype=ov.Type.f32))
+    # NPU uses FP16 x INT8 -> FP16 instead of INT8 x INT8 -> INT32 and FP16 output overflows
+    if "NPU" in OV_DEVICE:
+        fp16_scale = 0.25012213 * Tensor_B.shape[-2]
+        in_scale = ov_ops.constant(fp16_scale ** 0.5, dtype=ov.Type.f32)
+        out_scale = ov_ops.constant(fp16_scale, dtype=ov.Type.f32, name="out_scale_const")
+        a = ov_ops.divide(a, in_scale)
+        b = ov_ops.divide(b, in_scale)
+        out = ov_ops.matmul(a, b, False, False)
+        out = ov_ops.multiply(out, out_scale, name="out_scale")
+    else:
+        out = ov_ops.matmul(a, b, False, False)
 
-    ov_model = ov.Model([ov_ops.matmul(a, b, False, False)], [input_a, input_b], "ov_int8_mm")
+    ov_model = ov.Model([out], [input_a, input_b], "ov_int8_mm")
+    if "NPU" in OV_DEVICE: # NPU can't use FP32 for regular multiplications
+        for node in ov_model.get_ops():
+            if node.get_friendly_name() in {"out_scale", "out_scale_const"}:
+                node.get_rt_info()["affinity"] = "CPU"
+            else:
+                node.get_rt_info()["affinity"] = "NPU"
     ov_model = core.compile_model(ov_model, OV_DEVICE)
     infer_request = ov_model.create_infer_request()
     out_name = ov_model.outputs[0]
 
     OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
-    return ov_int_mm(Tensor_A, Tensor_B, infer_request, out_name)
+    return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
 
 @openvino_int_mm.register_fake
 def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
     return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32))
+
+
+@torch.library.custom_op("sdnq::openvino_fp_mm", mutates_args=())
+def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor:
+    mm_dtype = "fp16" if Tensor_B.dtype == torch.float16 else "fp8"
+    if mm_dtype == "fp8":
+        Tensor_A = Tensor_A.to(dtype=torch.float16)
+        Tensor_B = Tensor_B.to(dtype=torch.float16)
+    if "GPU" not in OV_DEVICE:
+        cache_key = (OV_DEVICE, mm_dtype, Tensor_A.shape, Tensor_B.shape)
+    else:
+        cache_key = (OV_DEVICE, mm_dtype, None, None)
+    infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
+    if infer_request is not None:
+        return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
+
+    if "GPU" not in OV_DEVICE:
+        shape_a = ov.Shape(Tensor_A.shape)
+        shape_b = ov.Shape(Tensor_B.shape)
+    else:
+        shape_a = ov.PartialShape([-1,-1])
+        shape_b = ov.PartialShape([-1,-1])
+    input_a = ov_ops.parameter(shape_a, ov.Type.f16, name="A")
+    input_b = ov_ops.parameter(shape_b, ov.Type.f16, name="B")
+    a = ov_ops.convert(input_a, ov.Type.f32)
+    b = ov_ops.convert(input_b, ov.Type.f32)
+
+    if mm_dtype == "fp8":
+        low  = ov_ops.constant(-448.0, dtype=ov.Type.f32)
+        high = ov_ops.constant(448.0,  dtype=ov.Type.f32)
+        a = ov_ops.fake_quantize(a, low, high, low, high, 256)
+        b = ov_ops.fake_quantize(b, low, high, low, high, 256)
+        fp16_scale = 4 * Tensor_B.shape[-2]
+    else:
+        fp16_scale = 65536 * Tensor_B.shape[-2]
+
+    in_scale = ov_ops.constant(fp16_scale**0.5, dtype=ov.Type.f32)
+    out_scale = ov_ops.constant(fp16_scale, dtype=ov.Type.f32, name="out_scale_const")
+    a = ov_ops.convert(ov_ops.divide(a, in_scale), ov.Type.f16)
+    b = ov_ops.convert(ov_ops.divide(b, in_scale), ov.Type.f16)
+    out = ov_ops.matmul(a, b, False, False, name="fp_mm")
+    out = ov_ops.multiply(ov_ops.convert(out, ov.Type.f32), out_scale, name="out_scale")
+
+    ov_model = ov.Model([out], [input_a, input_b], "ov_fp_mm")
+    if "NPU" in OV_DEVICE: # NPU can't use FP32 for regular multiplications
+        for node in ov_model.get_ops():
+            if node.get_friendly_name() in {"out_scale", "out_scale_const"}:
+                node.get_rt_info()["affinity"] = "CPU"
+            else:
+                node.get_rt_info()["affinity"] = "NPU"
+    ov_model = core.compile_model(ov_model, OV_DEVICE)
+    infer_request = ov_model.create_infer_request()
+    out_name = ov_model.outputs[0]
+
+    OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
+    return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
+
+@openvino_fp_mm.register_fake
+def openvino_fp_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32))

modules/sdnq/layers/linear/linear_fp8_tensorwise.py

@@ -48,9 +48,9 @@ def fp8_matmul_tensorwise(
     input, input_scale = quantize_fp_mm_input_tensorwise(input, dtype=scale.dtype)
     input, weight = check_mats(input, weight, allow_contiguous_mm=False)
     if bias is not None:
-        return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).to(dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape)
+        return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape)
     else:
-        return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).to(dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape)
+        return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape)
 
 
 def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor:

modules/sdnq/quantizer.py

@@ -589,18 +589,6 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
                 self.quantization_config.modules_to_not_convert.append(param_name)
         return False
 
-    def check_quantized_param(self, *args, **kwargs) -> bool:
-        """
-        needed for transformers compatibility, returns self.check_if_quantized_param
-        """
-        return self.check_if_quantized_param(*args, **kwargs)
-
-    def param_needs_quantization(self, model, param_name: str, *args, **kwargs) -> bool:
-        """
-        needed for transformers compatibility, returns self.check_if_quantized_param
-        """
-        return self.check_if_quantized_param(model, None, param_name, *args, **kwargs)
-
     @devices.inference_context()
     def create_quantized_param( # pylint: disable=arguments-differ
         self,
@@ -656,16 +644,6 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
         parent_module, tensor_name = get_module_from_name(model, param_name.removesuffix(tensor_name).removesuffix("."))
         setattr(parent_module, tensor_name, layer)
 
-    def get_quantize_ops(self):
-        return SDNQQuantize(self)
-
-    def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
-        max_memory = {key: val * 0.80 for key, val in max_memory.items()}
-        return max_memory
-
-    def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed
-        return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"]
-
     def _process_model_before_weight_loading( # pylint: disable=arguments-differ
         self,
         model: torch.nn.Module,
@@ -726,6 +704,20 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
         devices.torch_gc(force=True, reason="sdnq")
         return model
 
+    def get_quantize_ops(self):
+        return SDNQQuantize(self)
+
+    def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
+        max_memory = {key: val * 0.80 for key, val in max_memory.items()}
+        return max_memory
+
+    def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed
+        return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"]
+
+    def update_torch_dtype(self, torch_dtype: torch.dtype) -> torch.dtype:
+        self.torch_dtype = torch_dtype
+        return torch_dtype
+
     def get_state_dict_and_metadata(self, state_dict: dict | torch.nn.Module, **kwargs) -> tuple[dict | None, dict]: # pylint: disable=unused-argument, arguments-differ
         # transformers
         if isinstance(state_dict, torch.nn.Module):
@@ -736,12 +728,6 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
     def get_accelerator_warm_up_factor(self):
         return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"]
 
-    def get_cuda_warm_up_factor(self):
-        """
-        needed for transformers compatibility, returns self.get_accelerator_warm_up_factor
-        """
-        return self.get_accelerator_warm_up_factor()
-
     def _dequantize(self, model):
         return dequantize_sdnq_model(model)
 
@@ -764,6 +750,30 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
     def is_compileable(self):
         return True
 
+    def check_quantized_param(self, *args, **kwargs) -> bool:
+        """
+        needed for transformers compatibility, returns self.check_if_quantized_param
+        """
+        return self.check_if_quantized_param(*args, **kwargs)
+
+    def param_needs_quantization(self, model, param_name: str, *args, **kwargs) -> bool:
+        """
+        needed for transformers compatibility, returns self.check_if_quantized_param
+        """
+        return self.check_if_quantized_param(model, None, param_name, *args, **kwargs)
+
+    def get_cuda_warm_up_factor(self):
+        """
+        needed for transformers compatibility, returns self.get_accelerator_warm_up_factor
+        """
+        return self.get_accelerator_warm_up_factor()
+
+    def update_dtype(self, dtype: torch.dtype) -> torch.dtype:
+        """
+        needed for transformers compatibility, returns self.update_torch_dtype
+        """
+        return self.update_torch_dtype(dtype)
+
 
 @dataclass
 class SDNQConfig(QuantizationConfigMixin):