From 73e84d5ec8b943dcb42535229eb94ee7ab3abea1 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:57:09 -0700 Subject: [PATCH] Support convrot int4 models. (#14859) linear_dtype in comfy_quant metadata can be used to set if the int4 op does the matrix multiplication in int8 or int4, the default is int4 on GPUs that support it with fallback to int8 for GPUs that don't. --- comfy/ops.py | 26 +++++++++ comfy/quant_ops.py | 14 +++++ requirements.txt | 2 +- .../comfy_quant/test_mixed_precision.py | 54 ++++++++++++++++++- 4 files changed, 94 insertions(+), 2 deletions(-) diff --git a/comfy/ops.py b/comfy/ops.py index 35a1ee31e..0c6fe4cb4 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1104,6 +1104,21 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat scales["convrot_groupsize"] = int( layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256)) ) + elif module.quant_format == "convrot_w4a4": + scale = pop_scale("weight_scale") + if scale is None: + raise ValueError(f"Missing ConvRot W4A4 weight scale for layer {layer_name}") + params_conf = layer_conf.get("params", {}) + if not isinstance(params_conf, dict): + params_conf = {} + scales = { + "scale": scale, + "convrot_groupsize": int( + layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256)) + ), + "quant_group_size": 64, + "linear_dtype": layer_conf.get("linear_dtype", params_conf.get("linear_dtype", "int4")), + } else: raise ValueError(f"Unsupported quantization format: {module.quant_format}") @@ -1150,6 +1165,11 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False): quant_conf["convrot"] = True quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256) + elif module.quant_format == "convrot_w4a4": + quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256) + linear_dtype = getattr(params, "linear_dtype", "int4") + if linear_dtype != "int4": + quant_conf["linear_dtype"] = linear_dtype if extra_quant_conf: quant_conf.update(extra_quant_conf) sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8) @@ -1430,6 +1450,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec } if hasattr(params, "block_scale"): # NVFP4 kwargs["block_scale"] = params.block_scale[i] + if hasattr(params, "quant_group_size"): + kwargs["quant_group_size"] = params.quant_group_size + if hasattr(params, "convrot_groupsize"): + kwargs["convrot_groupsize"] = params.convrot_groupsize + if hasattr(params, "linear_dtype"): + kwargs["linear_dtype"] = params.linear_dtype return QuantizedTensor(weight._qdata[i], weight._layout_cls, type(params)(**kwargs)) def state_dict(self, *args, destination=None, prefix="", **kwargs): diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py index 44f25a97e..53a0cb603 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -10,6 +10,7 @@ try: QuantizedLayout, TensorCoreFP8Layout as _CKFp8Layout, TensorCoreNVFP4Layout as _CKNvfp4Layout, + TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout, TensorWiseINT8Layout as _CKTensorWiseINT8Layout, register_layout_op, register_layout_class, @@ -51,6 +52,9 @@ except ImportError as e: class _CKTensorWiseINT8Layout: pass + class _CKTensorCoreConvRotW4A4Layout: + pass + def register_layout_class(name, cls): pass @@ -179,6 +183,7 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase): # Backward compatibility alias - default to E4M3 TensorCoreFP8Layout = TensorCoreFP8E4M3Layout TensorWiseINT8Layout = _CKTensorWiseINT8Layout +TensorCoreConvRotW4A4Layout = _CKTensorCoreConvRotW4A4Layout # ============================================================================== @@ -190,6 +195,7 @@ register_layout_class("TensorCoreFP8E4M3Layout", TensorCoreFP8E4M3Layout) register_layout_class("TensorCoreFP8E5M2Layout", TensorCoreFP8E5M2Layout) register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout) register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout) +register_layout_class("TensorCoreConvRotW4A4Layout", _CKTensorCoreConvRotW4A4Layout) if _CK_MXFP8_AVAILABLE: register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout) @@ -227,6 +233,13 @@ QUANT_ALGOS["int8_tensorwise"] = { "quantize_input": False, } +QUANT_ALGOS["convrot_w4a4"] = { + "storage_t": torch.int8, + "parameters": {"weight_scale"}, + "comfy_tensor_layout": "TensorCoreConvRotW4A4Layout", + "quantize_input": False, +} + # ============================================================================== # Re-exports for backward compatibility @@ -239,6 +252,7 @@ __all__ = [ "TensorCoreFP8E4M3Layout", "TensorCoreFP8E5M2Layout", "TensorCoreNVFP4Layout", + "TensorCoreConvRotW4A4Layout", "TensorWiseINT8Layout", "QUANT_ALGOS", "register_layout_op", diff --git a/requirements.txt b/requirements.txt index e72f3045b..a8ea0eace 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,7 +22,7 @@ alembic SQLAlchemy>=2.0.0 filelock av>=16.0.0 -comfy-kitchen==0.2.16 +comfy-kitchen==0.2.17 comfy-aimdo==0.4.10 requests simpleeval>=1.0.0 diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py index 43b4b7ce9..7bbc96616 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/tests-unit/comfy_quant/test_mixed_precision.py @@ -15,7 +15,7 @@ if not has_gpu(): args.cpu = True from comfy import ops -from comfy.quant_ops import QuantizedTensor +from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor import comfy.utils @@ -283,7 +283,59 @@ class TestMixedPrecisionOps(unittest.TestCase): saved = model.state_dict() saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes()) self.assertTrue(saved_conf["convrot"]) + + def test_convrot_w4a4_loads_into_params(self): + """ConvRot W4A4 checkpoints must load as the dedicated kitchen layout.""" + if "convrot_w4a4" not in QUANT_ALGOS: + self.skipTest("comfy_kitchen does not provide ConvRot W4A4") + + torch.manual_seed(456) + layer_quant_config = { + "layer": { + "format": "convrot_w4a4", + "convrot_groupsize": 256, + "linear_dtype": "int8", + } + } + weight = torch.randn(16, 256, dtype=torch.bfloat16) + bias = torch.randn(16, dtype=torch.bfloat16) + q_weight = QuantizedTensor.from_float( + weight, + "TensorCoreConvRotW4A4Layout", + convrot_groupsize=256, + quant_group_size=64, + ) + state_dict = { + "layer.weight": q_weight._qdata, + "layer.bias": bias, + "layer.weight_scale": q_weight._params.scale, + } + + state_dict, _ = comfy.utils.convert_old_quants( + state_dict, + metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}, + ) + model = torch.nn.Module() + model.layer = ops.mixed_precision_ops({}).Linear(256, 16, device="cpu", dtype=torch.bfloat16) + model.load_state_dict(state_dict, strict=False) + + self.assertIsInstance(model.layer.weight, QuantizedTensor) + self.assertEqual(model.layer.weight._layout_cls, "TensorCoreConvRotW4A4Layout") + self.assertEqual(model.layer.weight._params.convrot_groupsize, 256) + self.assertEqual(model.layer.weight._params.quant_group_size, 64) + self.assertEqual(model.layer.weight._params.linear_dtype, "int8") + + input_tensor = torch.randn(4, 256, dtype=torch.bfloat16) + loaded_out = model.layer(input_tensor) + ref_out = torch.nn.functional.linear(input_tensor, q_weight, bias) + self.assertTrue(torch.equal(loaded_out, ref_out)) + + saved = model.state_dict() + saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes()) + self.assertEqual(saved_conf["format"], "convrot_w4a4") self.assertEqual(saved_conf["convrot_groupsize"], 256) + self.assertEqual(saved_conf["linear_dtype"], "int8") + self.assertNotIn("quant_group_size", saved_conf) if __name__ == "__main__": unittest.main()