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

@@ -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):

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",

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

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()