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 1/5] 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 35a1ee31ec7..0c6fe4cb4b9 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1104,6 +1104,21 @@ def pop_scale(name, dtype=None): 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 _expert_qt_from(self, weight: QuantizedTensor, i: int) -> QuantizedTensor: } 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 44f25a97ef3..53a0cb603a8 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -10,6 +10,7 @@ QuantizedLayout, TensorCoreFP8Layout as _CKFp8Layout, TensorCoreNVFP4Layout as _CKNvfp4Layout, + TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout, TensorWiseINT8Layout as _CKTensorWiseINT8Layout, register_layout_op, register_layout_class, @@ -51,6 +52,9 @@ class _CKNvfp4Layout: 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 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase): 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 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase): "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 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase): "TensorCoreFP8E4M3Layout", "TensorCoreFP8E5M2Layout", "TensorCoreNVFP4Layout", + "TensorCoreConvRotW4A4Layout", "TensorWiseINT8Layout", "QUANT_ALGOS", "register_layout_op", diff --git a/requirements.txt b/requirements.txt index e72f3045bc1..a8ea0eacee2 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 43b4b7ce97b..7bbc9661661 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/tests-unit/comfy_quant/test_mixed_precision.py @@ -15,7 +15,7 @@ def 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 @@ def test_int8_convrot_metadata_loads_into_params(self): 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() From b7a648ca2011489ba40eaacf01a5d6f4e9fab539 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:39:01 -0700 Subject: [PATCH 2/5] Try to fix the model reloading issue some people have. (#14822) --- comfy/model_management.py | 11 +++++++++++ comfy_execution/caching.py | 25 +++++++++++++++++-------- execution.py | 17 +++++++++++------ 3 files changed, 39 insertions(+), 14 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index b15d08ba15b..222005b6f4d 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -616,6 +616,8 @@ def get_torch_device_name(device): #Freeing registerables on pressure does imply a GPU sync, so go big on #the hysteresis so each expensive sync gives us back a good chunk. REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024 +WINDOWS_PIN_EVICTION_SWAP_PERCENT = 5.0 +WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE = 512 * 1024 ** 2 def module_size(module): module_mem = 0 @@ -642,6 +644,15 @@ def free_pins(size, evict_active=False): size -= freed return freed_total +def should_free_pins_for_ram_pressure(shortfall): + if shortfall <= 0: + return False + if not WINDOWS: + return True + if psutil.virtual_memory().available < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE: + return True + return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT + def ensure_pin_budget(size, evict_active=False): if args.high_ram: return True diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index ad75a0e5019..6bd99b68f30 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -503,6 +503,8 @@ async def ensure_subcache_for(self, node_id, children_ids): RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3 +RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2 + def all_outputs_dynamic(outputs): if outputs is None: @@ -517,7 +519,6 @@ def all_outputs_dynamic(outputs): return True - class RAMPressureCache(LRUCache): def __init__(self, key_class, enable_providers=False): @@ -539,9 +540,9 @@ def set_local(self, node_id, value): self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time() super().set_local(node_id, value) - def ram_release(self, target, free_active=False): + def ram_release(self, target, free_active=False, min_entry_size=0): if psutil.virtual_memory().available >= target: - return + return 0 clean_list = [] @@ -555,8 +556,9 @@ def ram_release(self, target, free_active=False): oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key]) ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE + oom_ram_usage = ram_usage def scan_list_for_ram_usage(outputs): - nonlocal ram_usage + nonlocal ram_usage, oom_ram_usage if outputs is None: return for output in outputs: @@ -564,19 +566,26 @@ def scan_list_for_ram_usage(outputs): scan_list_for_ram_usage(output) elif isinstance(output, torch.Tensor) and output.device.type == 'cpu': ram_usage += output.numel() * output.element_size() + oom_ram_usage += output.numel() * output.element_size() elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation: #old ModelPatchers are the first to go - ram_usage = 1e30 + oom_ram_usage = 1e30 scan_list_for_ram_usage(cache_entry.outputs) - oom_score *= ram_usage + if ram_usage < min_entry_size: + continue + + oom_score *= oom_ram_usage #In the case where we have no information on the node ram usage at all, #break OOM score ties on the last touch timestamp (pure LRU) - bisect.insort(clean_list, (oom_score, self.timestamps[key], key)) + bisect.insort(clean_list, (oom_score, self.timestamps[key], key, ram_usage)) + freed = 0 while psutil.virtual_memory().available < target and clean_list: - _, _, key = clean_list.pop() + _, _, key, ram_usage = clean_list.pop() del self.cache[key] self.used_generation.pop(key, None) self.timestamps.pop(key, None) self.children.pop(key, None) + freed += ram_usage + return freed diff --git a/execution.py b/execution.py index c45317593c0..19b8cdd68be 100644 --- a/execution.py +++ b/execution.py @@ -29,6 +29,7 @@ HierarchicalCache, LRUCache, RAMPressureCache, + RAM_CACHE_LARGE_INTERMEDIATE, ) from comfy_execution.graph import ( DynamicPrompt, @@ -794,12 +795,16 @@ async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs= if self.cache_type == CacheType.RAM_PRESSURE: ram_release_callback(ram_inactive_headroom) ram_shortfall = ram_headroom - psutil.virtual_memory().available - freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2)) - if freed < ram_shortfall: - if freed > 64 * (1024 ** 2): - # AIMDO MEM_DECOMMIT can outrun psutil.available catching up. - time.sleep(0.05) - ram_release_callback(ram_headroom, free_active=True) + if ram_shortfall > 0: + freed = ram_release_callback(ram_headroom, free_active=True, min_entry_size=RAM_CACHE_LARGE_INTERMEDIATE) + ram_shortfall -= freed + if comfy.model_management.should_free_pins_for_ram_pressure(ram_shortfall): + freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2)) + if freed < ram_shortfall: + if freed > 64 * (1024 ** 2): + # AIMDO MEM_DECOMMIT can outrun psutil.available catching up. + time.sleep(0.05) + ram_release_callback(ram_headroom, free_active=True) else: # Only execute when the while-loop ends without break # Send cached UI for intermediate output nodes that weren't executed From 62e025a4f34d16eeedfc1c93e50a48a69098df7f Mon Sep 17 00:00:00 2001 From: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:30:26 +0800 Subject: [PATCH 3/5] Fix FP8 activation quantization for >2D activations in mixed_precision_ops (#14643) mixed_precision_ops.Linear.forward only quantized activations that were 2D, or 3D (reshaped to 2D). Inputs with rank >= 4 (e.g. Anima's MLP activations, which are not reshaped to 3D the way the attention path is) fell through the `input_reshaped.ndim == 2` guard and reached scaled_mm as bf16, silently dispatching a bf16 kernel instead of FP8. Since MLP is roughly half the compute, the FP8 speedup was far below expectation. Generalize the existing 3D->2D reshape to any rank >= 3 (flatten the leading dims, keep the contraction dim) and reshape the output back to the original leading dims. 2D and 3D inputs are handled exactly as before; only rank >= 4 inputs change (now quantized instead of skipped). This matches the rank-agnostic handling already used by the training path (flatten(0, -2) / unflatten). Fixes #14595. --- comfy/ops.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/comfy/ops.py b/comfy/ops.py index 0c6fe4cb4b9..13c2604fb74 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1257,7 +1257,7 @@ def forward(self, input, *args, **kwargs): run_every_op() input_shape = input.shape - reshaped_3d = False + reshaped_nd = False #If cast needs to apply lora, it should be done in the compute dtype compute_dtype = input.dtype @@ -1294,12 +1294,12 @@ def forward(self, input, *args, **kwargs): # Inference path (unchanged) if _use_quantized and quantize_input: - # Reshape 3D tensors to 2D for quantization (needed for NVFP4 and others) - input_reshaped = input.reshape(-1, input_shape[2]) if input.ndim == 3 else input + # Reshape >=3D tensors to 2D for quantization (needed for NVFP4 and others) + input_reshaped = input.reshape(-1, input_shape[-1]) if input.ndim >= 3 else input # Fall back to non-quantized for non-2D tensors if input_reshaped.ndim == 2: - reshaped_3d = input.ndim == 3 + reshaped_nd = input.ndim >= 3 # dtype is now implicit in the layout class scale = getattr(self, 'input_scale', None) if scale is not None: @@ -1314,9 +1314,9 @@ def forward(self, input, *args, **kwargs): weight_only_quant=weight_only_quant, ) - # Reshape output back to 3D if input was 3D - if reshaped_3d: - output = output.reshape((input_shape[0], input_shape[1], self.weight.shape[0])) + # Reshape output back to original rank if input was >2D + if reshaped_nd: + output = output.reshape((*input_shape[:-1], self.weight.shape[0])) return output From 099522f85bcd8586eac4133c02f31c70dafe85d2 Mon Sep 17 00:00:00 2001 From: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> Date: Fri, 10 Jul 2026 11:11:52 +0800 Subject: [PATCH 4/5] Enable comfy-kitchen Triton backend by default on ROCm/AMD (#14862) On AMD/ROCm the CUDA backend is unavailable, so Triton is the only accelerated comfy-kitchen backend. It was disabled by default (opt-in --enable-triton-backend), leaving AMD on the slow eager path. Enable it by default when torch.version.hip is set AND Triton is >= 3.7 -- older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path, so on Triton < 3.7 it stays disabled with a log line. NVIDIA behavior is unchanged; the explicit --enable-triton-backend flag still works as an override. Fixes #14861 --- comfy/quant_ops.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py index 53a0cb603a8..91b3e4fe928 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -25,10 +25,18 @@ ck.registry.disable("cuda") logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.") - if args.enable_triton_backend: + # On ROCm/AMD the CUDA backend is unavailable, so Triton is the only accelerated + # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7: + # older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path. + if args.enable_triton_backend or torch.version.hip is not None: try: import triton - logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__) + triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2]) + if args.enable_triton_backend or triton_version >= (3, 7): + logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__) + else: + logging.info("Triton %s is too old for the ROCm INT8 path (needs >= 3.7); comfy-kitchen triton backend disabled.", triton.__version__) + ck.registry.disable("triton") except ImportError as e: logging.error(f"Failed to import triton, Error: {e}, the comfy-kitchen triton backend will not be available.") ck.registry.disable("triton") From e2a6e30d892402ffcf01d6280c8e2744a4448b9d Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 9 Jul 2026 20:17:06 -0700 Subject: [PATCH 5/5] Fix black image on turing when using int4 models. (#14864) --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index a8ea0eacee2..790ef49408f 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.17 +comfy-kitchen==0.2.18 comfy-aimdo==0.4.10 requests simpleeval>=1.0.0