Make SDNQ MatMul listen to the dequantize fp32 setting

modules/sdnq/layers/linear/linear_fp8.py

@@ -8,10 +8,9 @@ from ...common import use_torch_compile # noqa: TID252
 
 
 def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
-    input = input.flatten(0,-2).contiguous()
+    input = input.flatten(0,-2).contiguous().to(dtype=torch.float32)
     input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448)
     input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn)
-    input_scale = input_scale.to(dtype=torch.float32)
     return input, input_scale
 
 

modules/sdnq/layers/linear/linear_fp8_tensorwise.py

@@ -9,7 +9,7 @@ from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias
 
 
 def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
-    input = input.flatten(0,-2).contiguous()
+    input = input.flatten(0,-2).contiguous().to(dtype=scale.dtype)
     input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448)
     input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn)
     scale = torch.mul(input_scale, scale)

modules/sdnq/layers/linear/linear_int8.py

@@ -10,7 +10,7 @@ from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias
 
 
 def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]:
-    input = input.flatten(0,-2).contiguous()
+    input = input.flatten(0,-2).contiguous().to(dtype=scale.dtype)
     input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(127)
     input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(dtype=torch.int8)
     scale = torch.mul(input_scale, scale)