diff --git a/comfy/ops.py b/comfy/ops.py index 7cf077eea..51b3e985f 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1312,6 +1312,11 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec return _quantized_weight_state_dict(self, sd, prefix, extra_quant_params=("input_scale", "pre_quant_scale")) def _forward(self, input, weight, bias): + if getattr(self, "layout_type", None) == "TensorCoreNVFP4Layout" and isinstance(input, QuantizedTensor) and isinstance(weight, QuantizedTensor): + output = torch.mm(input, weight.t()) + if bias is not None: + output = output + bias + return output return torch.nn.functional.linear(input, weight, bias) def forward_comfy_cast_weights( diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py index 7bbc96616..8aa4b7e17 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/tests-unit/comfy_quant/test_mixed_precision.py @@ -1,4 +1,5 @@ import unittest +from unittest import mock import torch import sys import os @@ -15,7 +16,7 @@ if not has_gpu(): args.cpu = True from comfy import ops -from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor +from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, TensorCoreNVFP4Layout import comfy.utils @@ -37,6 +38,27 @@ class SimpleModel(torch.nn.Module): class TestMixedPrecisionOps(unittest.TestCase): + @staticmethod + def _nvfp4_tensor(rows, columns, fill, scale=1.0): + params = TensorCoreNVFP4Layout.Params( + scale=None if scale is None else torch.tensor(scale), + orig_dtype=torch.float32, + orig_shape=(rows, columns), + block_scale=torch.ones(1), + ) + return QuantizedTensor( + torch.full((rows, columns // 2), fill, dtype=torch.uint8), + "TensorCoreNVFP4Layout", + params, + ) + + @staticmethod + def _nvfp4_linear(in_features=16, out_features=8): + layer = ops.mixed_precision_ops({}).Linear(in_features, out_features, device="cpu", dtype=torch.float32) + layer.layout_type = "TensorCoreNVFP4Layout" + layer.quant_format = "nvfp4" + return layer + def test_all_layers_standard(self): """Test that model with no quantization works normally""" # Create model @@ -337,5 +359,137 @@ class TestMixedPrecisionOps(unittest.TestCase): self.assertEqual(saved_conf["linear_dtype"], "int8") self.assertNotIn("quant_group_size", saved_conf) + def test_nvfp4_quantized_mm_dispatch_with_bias_and_rank_restore(self): + input_tensor = torch.randn(2, 3, 16) + quantized_input = self._nvfp4_tensor(6, 16, 1) + quantized_weight = self._nvfp4_tensor(8, 16, 2) + bias = torch.arange(8, dtype=torch.float32) + kernel_output = torch.arange(48, dtype=torch.float32).reshape(6, 8) + layer = self._nvfp4_linear() + layer.weight = quantized_weight + + with ( + mock.patch.object(QuantizedTensor, "from_float", return_value=quantized_input) as quantize, + mock.patch("comfy_kitchen.scaled_mm_nvfp4", return_value=kernel_output) as scaled_mm, + mock.patch.object(QuantizedTensor, "dequantize", autospec=True) as dequantize, + mock.patch.object(ops, "cast_bias_weight", return_value=(quantized_weight, bias, None)), + mock.patch.object(ops, "uncast_bias_weight"), + ): + output = layer(input_tensor) + + self.assertTrue(torch.equal(output, (kernel_output + bias).reshape(2, 3, 8))) + self.assertEqual(quantize.call_args.args[0].shape, (6, 16)) + self.assertEqual(scaled_mm.call_count, 1) + self.assertEqual(dequantize.call_count, 0) + + def test_nvfp4_quantized_mm_dispatch_without_bias(self): + quantized_input = self._nvfp4_tensor(5, 16, 1) + quantized_weight = self._nvfp4_tensor(8, 16, 2) + kernel_output = torch.randn(5, 8) + layer = self._nvfp4_linear() + + with ( + mock.patch("comfy_kitchen.scaled_mm_nvfp4", return_value=kernel_output) as scaled_mm, + mock.patch.object(QuantizedTensor, "dequantize", autospec=True) as dequantize, + ): + output = layer._forward(quantized_input, quantized_weight, None) + + self.assertIs(output, kernel_output) + self.assertEqual(scaled_mm.call_count, 1) + self.assertEqual(dequantize.call_count, 0) + + def test_nvfp4_mm_backend_fallbacks_dequantize(self): + input_tensor = torch.randn(5, 16) + weight = torch.randn(8, 16) + bias = torch.randn(8) + expected = torch.nn.functional.linear(input_tensor, weight, bias) + quantized_input = self._nvfp4_tensor(5, 16, 1) + quantized_weight = self._nvfp4_tensor(8, 16, 2) + dequantized = { + (5, 16): input_tensor, + (8, 16): weight, + } + layer = self._nvfp4_linear() + + def dequantize(tensor): + value = dequantized[tensor._params.orig_shape if not tensor._params.transposed else tuple(reversed(tensor._params.orig_shape))] + return value.t() if tensor._params.transposed else value + + for error in (RuntimeError("no supported backend"), TypeError("unsupported kernel signature")): + with self.subTest(error=type(error).__name__): + with ( + mock.patch("comfy_kitchen.scaled_mm_nvfp4", side_effect=error) as scaled_mm, + mock.patch.object(QuantizedTensor, "dequantize", autospec=True, side_effect=dequantize) as dequantize_call, + ): + output = layer._forward(quantized_input, quantized_weight, bias) + + self.assertTrue(torch.equal(output, expected)) + self.assertEqual(scaled_mm.call_count, 1) + self.assertEqual(dequantize_call.call_count, 2) + + def test_nvfp4_mm_unexpected_failures_remain_loud(self): + class Cancellation(BaseException): + pass + + quantized_input = self._nvfp4_tensor(5, 16, 1) + quantized_weight = self._nvfp4_tensor(8, 16, 2) + layer = self._nvfp4_linear() + + for error in (ValueError("invalid metadata"), Cancellation("cancelled")): + with self.subTest(error=type(error).__name__): + with ( + mock.patch("comfy_kitchen.scaled_mm_nvfp4", side_effect=error), + mock.patch.object(QuantizedTensor, "dequantize", autospec=True) as dequantize, + ): + with self.assertRaises(type(error)): + layer._forward(quantized_input, quantized_weight, None) + self.assertEqual(dequantize.call_count, 0) + + def test_nvfp4_mm_invalid_quantization_metadata_remains_loud(self): + quantized_input = self._nvfp4_tensor(5, 16, 1) + layer = self._nvfp4_linear() + + missing_params = self._nvfp4_tensor(8, 16, 2) + del missing_params._params + with self.assertRaises(AttributeError): + layer._forward(quantized_input, missing_params, None) + + unknown_layout = self._nvfp4_tensor(8, 16, 2) + unknown_layout._layout_cls = "UnknownLayout" + with self.assertRaises(KeyError): + layer._forward(quantized_input, unknown_layout, None) + + missing_scale = self._nvfp4_tensor(8, 16, 2, scale=None) + with mock.patch("comfy_kitchen.scaled_mm_nvfp4", side_effect=ValueError("missing scale")): + with self.assertRaisesRegex(ValueError, "missing scale"): + layer._forward(quantized_input, missing_scale, None) + + def test_nvfp4_mm_guard_preserves_other_linear_paths(self): + quantized_input = self._nvfp4_tensor(5, 16, 1) + quantized_weight = self._nvfp4_tensor(8, 16, 2) + regular_input = torch.randn(5, 16) + regular_weight = torch.randn(8, 16) + bias = torch.randn(8) + sentinel = torch.randn(5, 8) + layer = self._nvfp4_linear() + + cases = ( + ("TensorCoreFP8E4M3Layout", quantized_input, quantized_weight), + ("TensorCoreNVFP4Layout", regular_input, quantized_weight), + ("TensorCoreNVFP4Layout", quantized_input, regular_weight), + ) + for layout_type, input_value, weight_value in cases: + with self.subTest(layout_type=layout_type, input_type=type(input_value).__name__, weight_type=type(weight_value).__name__): + with ( + mock.patch.object(torch.nn.functional, "linear", return_value=sentinel) as functional_linear, + mock.patch.object(torch, "mm") as mm, + ): + layer.layout_type = layout_type + output = layer._forward(input_value, weight_value, bias) + + self.assertIs(output, sentinel) + functional_linear.assert_called_once_with(input_value, weight_value, bias) + mm.assert_not_called() + if __name__ == "__main__": unittest.main()