diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py index c0ecfa192..cfce8994a 100644 --- a/modules/sdnq/layers/conv/conv_fp8.py +++ b/modules/sdnq/layers/conv/conv_fp8.py @@ -25,6 +25,8 @@ def conv_fp8_matmul( input, input_scale = quantize_fp8_matmul_input(input) if groups == 1: + if bias is not None and bias.dtype != torch.bfloat16: + bias = bias.to(dtype=torch.bfloat16) result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(mm_output_shape).to(return_dtype) else: scale = scale.view(groups, 1, scale.shape[1] // groups) @@ -34,6 +36,8 @@ def conv_fp8_matmul( result = [] if bias is not None: bias = bias.view(groups, bias.shape[0] // groups) + if bias.dtype != torch.bfloat16: + bias = bias.to(dtype=torch.bfloat16) 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=torch.bfloat16)) else: diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py index 06e2a08aa..9b22f4696 100644 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ b/modules/sdnq/layers/linear/linear_fp8.py @@ -23,6 +23,8 @@ def fp8_matmul( return_dtype = input.dtype output_shape = (*input.shape[:-1], weight.shape[-1]) input, input_scale = quantize_fp8_matmul_input(input) + if bias is not None and bias.dtype != torch.bfloat16: + bias = bias.to(dtype=torch.bfloat16) return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(output_shape).to(return_dtype)