diff --git a/README.md b/README.md index c4be05a..5949b2f 100644 --- a/README.md +++ b/README.md @@ -1,54 +1,109 @@ # ComfyUI-QuantOps -Extended quantization layouts for ComfyUI, enabling loading and inference with models quantized by [convert_to_quant](https://github.com/silveroxides/convert_to_quant). +Additional quantized tensor layouts and model loaders for ComfyUI. The metadata formats are produced by [convert_to_quant](https://github.com/silveroxides/convert_to_quant). -This is experimental and due to lack of proper support and merging of PR in ComfyUI, do not expect this to work without putting in the effort. -I don't have the time or the energy to keep this up and will close entire project if i keep getting bunch of low effort issues posted expecting me go serve a fix up on a silver platter. +## Automatic routing +QuantOps records ComfyUI's native quantization formats before registering its own layouts. Normal checkpoint and diffusion-model loaders then route each model from its embedded metadata: -### tl;dr Go complain at ComfyOrg. Not here. +1. Models using only ComfyUI-native formats stay on the ComfyUI path. +2. Models containing a QuantOps-only format receive QuantOps custom operations. +3. Unquantized and unknown formats are left to ComfyUI instead of being claimed by QuantOps. -### The following is the last update I make regarding this. +This is automatic in normal ComfyUI workflows and in SwarmUI-generated workflows. Swarm presets and dtype overrides continue through the original loader arguments. -In order to use int8_tensorwise(RTX 30xx-series or newer GPU) you will need the following: +The log reports every decision as `decision=native-bypass`, `decision=quantops`, `decision=no-quantization`, or `decision=core-pass-through`. -- torch 2.10+cu130 or higher -- installed the latest of my custom comfy-kitchen fork wheels with the int8-tensorwise support -- enable the use of triton backend by using --enable-triton-backend launch argument in ComfyUI +## QuantOps-only metadata formats -Step 1: Install Triton -Activate your virtual environment used by ComfyUI and install triton. -For Windows you need to use this but linux can install latest triton as usual. -``` -# for torch 2.10 and 2.11 -pip install -U "triton-windows<3.7" -# for torch 2.12 -pip install -U "triton-windows<3.8" -``` +These canonical metadata formats are not native to the current ComfyUI quantization registry and have been verified through the automatic loader path: -Step 3: Install my comfy-kitchen -Download the latest uploaded version matching you python of my pre-compiled .whl file from my [HuggingFace repository](https://huggingface.co/silveroxides/comfy-kitchen-int8-wheels/tree/main) (Latest as of 13 June 2026) +| Metadata format | Storage and scaling | Default execution | +| --- | --- | --- | +| `int8_blockwise` | INT8 weights with a 2D block-scale grid, normally 128 x 128 blocks | PyTorch fallback by default; optional Triton backend in the Quantized loaders | +| `float8_e4m3fn_rowwise` | FP8 E4M3 weights with one scale per output row | QuantOps Triton kernel on compatible CUDA systems; dequantized fallback otherwise | +| `float8_e4m3fn_blockwise` | FP8 E4M3 weights with a 2D block-scale grid, normally 64 x 64 blocks | QuantOps Triton kernel on compatible CUDA systems; dequantized fallback otherwise | -Install it directly pointing to the file path: -``` -pip install --no-deps --force-reinstall --no-cache-dir "path/to/comfy-kitchen.whl" +The legacy metadata name `int8` is accepted as a blockwise INT8 compatibility alias when the scale shape identifies a block layout. New files should use `int8_blockwise`. + +### Conditional format + +`hybrid_mxfp8` describes 32-element MXFP8 blocks with an additional weight scalar. It is usable only when the installed `comfy-kitchen` exports `HybridMXFP8Layout`. The installed `comfy-kitchen` 0.2.17 in the tested environment does not export that class, so this format is not available there. + +## Dedicated BNB 4-bit formats + +The `BNB4bitUNETLoader` node supports bitsandbytes-compatible `NF4` and `FP4` safetensors without requiring bitsandbytes at runtime. It uses pure PyTorch dequantization and currently targets Flux, Flux 2, Chroma, Chroma Radiance, and Chroma Radiance X0 model structures. + +NF4 and FP4 use their dedicated loader. They are not automatic `QUANT_ALGOS` metadata formats. + +## Formats already native to ComfyUI + +QuantOps does not replace ComfyUI operations for these formats: + +- `convrot_w4a4` +- `float8_e4m3fn` +- `float8_e5m2` +- `int8_tensorwise`, including per-channel files whose names may say rowwise +- `mxfp8` +- `nvfp4` + +For these models the expected log is `decision=native-bypass` followed by `path=core result=completed`. Embedded metadata is authoritative; filenames are not used to select a quantization layout. + +## GGUF integration + +GGUF is delegated to an installed provider that registers `UnetLoaderGGUF`, such as [ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF). QuantOps does not duplicate or replace its GGUF tensor implementation. + +When the provider is present, QuantOps adds automatic image-architecture recognition for: + +- `krea2` +- `ideogram` + +SwarmUI already selects `UnetLoaderGGUF` for `.gguf` diffusion models. QuantOps extends the provider in memory and logs the detected architecture, GGUF tensor types, state-key count, and acceptance or error result. If no provider is installed, GGUF integration is disabled without affecting other formats. + +## Loader coverage + +- Stock ComfyUI diffusion-model loaders: automatic native or QuantOps routing. +- Stock ComfyUI checkpoint loaders: automatic routing for the diffusion-model portion. +- `QuantizedModelLoader` and `QuantizedUNETLoader`: explicit format and backend controls. +- `QuantizedCLIPLoader` and `QuantizedDualCLIPLoader`: quantized text-encoder loading. +- `BNB4bitUNETLoader`: dedicated NF4 and FP4 loading. + +## Runtime diagnostics + +Useful successful log sequences include: + +```text +ComfyUI-QuantOps diagnostic: loader=diffusion decision=native-bypass ... +ComfyUI-QuantOps diagnostic: loader=diffusion path=core result=completed + +ComfyUI-QuantOps diagnostic: loader=diffusion decision=quantops ... +ComfyUI-QuantOps FP8 blockwise: path=triton-dynamic ... +ComfyUI-QuantOps diagnostic: loader=diffusion path=quantops result=completed + +ComfyUI-QuantOps GGUF diagnostic: decision=gguf-provider ... +ComfyUI-QuantOps GGUF diagnostic: result=accepted architecture=krea2 ... ``` -Step 4: Install/Update ComfyUI-QuantOps -You just need to ensure it's fully up to date to read the new model formats. -Run these commands: +Repeated kernel-path messages are rate-limited. If an FP8 Triton kernel fails, QuantOps disables that kernel path for the session, emits one concise warning, and uses the dequantized fallback. + +## Triton + +Use a Triton build compatible with the PyTorch and CUDA versions in the ComfyUI virtual environment. On Windows this normally means `triton-windows`. The ComfyUI launch option below enables its `comfy-kitchen` Triton backend: +```text +--enable-triton-backend ``` + +QuantOps FP8 kernels are selected when a compatible Triton installation and CUDA device are available. The explicit Quantized loaders expose the blockwise INT8 backend choice. + +## Updating + +```text cd custom_nodes/ComfyUI-QuantOps git pull ``` -When launching Comfyui add launch argument: -``` ---enable-triton-backend -``` - -You can get most of the models here: https://huggingface.co/silveroxides +Quantized model releases are available from [silveroxides on Hugging Face](https://huggingface.co/silveroxides). ## License @@ -56,5 +111,5 @@ MIT License ## Acknowledgements -- [lyogavin](https://github.com/lyogavin) for [PR #10864](https://github.com/comfyanonymous/ComfyUI/pull/10864) to ComfyUI. -- [Clybius](https://github.com/Clybius) for inspiring me to take on quantization and his [Learned-Rounding](https://github.com/Clybius/Learned-Rounding) repository. +- [lyogavin](https://github.com/lyogavin) for [ComfyUI PR #10864](https://github.com/comfyanonymous/ComfyUI/pull/10864). +- [Clybius](https://github.com/Clybius) for the [Learned-Rounding](https://github.com/Clybius/Learned-Rounding) project. diff --git a/__init__.py b/__init__.py index 48a9869..92a3beb 100644 --- a/__init__.py +++ b/__init__.py @@ -3,7 +3,6 @@ This custom node extends ComfyUI's quantization system with additional layouts: - INT8 blockwise (with optional Triton acceleration) -- INT8 tensorwise (uses torch._int_mm with dynamic activation quant) - Row-wise and Block-wise FP8 variants All layouts are lazy-loaded to avoid import errors when optional dependencies @@ -12,6 +11,12 @@ import logging +import torch +from comfy.quant_ops import QUANT_ALGOS, register_layout_class + + +_NATIVE_COMFY_FORMATS = frozenset(QUANT_ALGOS) + # ============================================================================= # Module-level state for comfy-kitchen backend integration # ============================================================================= @@ -33,8 +38,6 @@ def is_ck_triton_available() -> bool: def _setup_comfy_kitchen_backends(): """ Configure comfy-kitchen backends for QuantOps. - - 1. Re-enable triton backend (ComfyUI disables it by default) """ global _CK_AVAILABLE, _CK_TRITON_AVAILABLE @@ -47,23 +50,23 @@ def _setup_comfy_kitchen_backends(): _CK_TRITON_AVAILABLE = False return - # Step 1: Re-enable triton backend (ComfyUI disables it) try: - ck.enable_backend("triton") - backends = ck.list_backends() triton_info = backends.get("triton", {}) if triton_info.get("available") and not triton_info.get("disabled"): _CK_TRITON_AVAILABLE = True - logging.info("ComfyUI-QuantOps: Enabled comfy-kitchen triton backend") + logging.info("ComfyUI-QuantOps: comfy-kitchen triton backend available") + elif triton_info.get("available"): + logging.info("ComfyUI-QuantOps: comfy-kitchen triton backend disabled") + _CK_TRITON_AVAILABLE = False else: unavail_reason = triton_info.get("unavailable_reason", "unknown") logging.info(f"ComfyUI-QuantOps: comfy-kitchen triton unavailable: {unavail_reason}") _CK_TRITON_AVAILABLE = False except Exception as e: - logging.warning(f"ComfyUI-QuantOps: Failed to enable ck triton backend: {e}") + logging.warning(f"ComfyUI-QuantOps: Failed to inspect ck triton backend: {e}") _CK_TRITON_AVAILABLE = False @@ -75,37 +78,40 @@ def _setup_comfy_kitchen_backends(): def _register_layouts(): """Register our custom layouts into ComfyUI's layout registry and QUANT_ALGOS dict.""" try: - from comfy.quant_ops import QUANT_ALGOS, register_layout_class - import torch - - # Import our layouts (this also registers their operation handlers) - from .quant_layouts.int8_layout import BlockWiseINT8Layout - from .quant_layouts.fp8_variants import RowWiseFP8Layout, BlockWiseFP8Layout - - # Register layouts using the new comfy_kitchen API - register_layout_class("BlockWiseINT8Layout", BlockWiseINT8Layout) - register_layout_class("RowWiseFP8Layout", RowWiseFP8Layout) - register_layout_class("BlockWiseFP8Layout", BlockWiseFP8Layout) - - # Tensorwise INT8 from comfy_kitchen - try: - from comfy_kitchen.tensor.int8 import TensorWiseINT8Layout - register_layout_class("TensorWiseINT8Layout", TensorWiseINT8Layout) - # Load our patch for per-channel scale support - from .quant_layouts import tensorwise_int8_layout - logging.info("ComfyUI-QuantOps: Registered TensorWiseINT8Layout") - except ImportError: - logging.debug("ComfyUI-QuantOps: TensorWiseINT8Layout not available") + registered = [] + if "int8_blockwise" not in QUANT_ALGOS: + from .quant_layouts.int8_layout import BlockWiseINT8Layout + + register_layout_class("BlockWiseINT8Layout", BlockWiseINT8Layout) + registered.append("BlockWiseINT8Layout") + if "float8_e4m3fn_rowwise" not in QUANT_ALGOS: + from .quant_layouts.fp8_variants import RowWiseFP8Layout + + register_layout_class("RowWiseFP8Layout", RowWiseFP8Layout) + registered.append("RowWiseFP8Layout") + if "float8_e4m3fn_blockwise" not in QUANT_ALGOS: + from .quant_layouts.fp8_variants import BlockWiseFP8Layout + + register_layout_class("BlockWiseFP8Layout", BlockWiseFP8Layout) + registered.append("BlockWiseFP8Layout") + + # Tensorwise INT8 is native in current ComfyUI. Only register it on + # older installs that do not already expose the format. + if "int8_tensorwise" not in QUANT_ALGOS: + try: + from comfy_kitchen.tensor.int8 import TensorWiseINT8Layout + register_layout_class("TensorWiseINT8Layout", TensorWiseINT8Layout) + QUANT_ALGOS["int8_tensorwise"] = { + "storage_t": torch.int8, + "parameters": {"weight_scale", "input_scale"}, + "comfy_tensor_layout": "TensorWiseINT8Layout", + } + registered.append("TensorWiseINT8Layout") + logging.info("ComfyUI-QuantOps: Registered TensorWiseINT8Layout") + except ImportError: + logging.debug("ComfyUI-QuantOps: TensorWiseINT8Layout not available") # Register QUANT_ALGOS - QUANT_ALGOS.setdefault( - "int8_tensorwise", - { - "storage_t": torch.int8, - "parameters": {"weight_scale", "input_scale"}, # Keep input_scale if checkpoints have it - "comfy_tensor_layout": "TensorWiseINT8Layout", # Must match the class name above - } - ) QUANT_ALGOS.setdefault( "int8_blockwise", { @@ -137,8 +143,10 @@ def _register_layouts(): # MXFP8 from comfy_kitchen try: from comfy_kitchen.tensor import TensorCoreMXFP8Layout - register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout) - logging.info("ComfyUI-QuantOps: Registered TensorCoreMXFP8Layout") + if "mxfp8" not in QUANT_ALGOS: + register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout) + registered.append("TensorCoreMXFP8Layout") + logging.info("ComfyUI-QuantOps: Registered TensorCoreMXFP8Layout") except ImportError: logging.debug("ComfyUI-QuantOps: TensorCoreMXFP8Layout not available") @@ -153,12 +161,14 @@ def _register_layouts(): ) # Hybrid MXFP8 from comfy_kitchen - try: - from comfy_kitchen.tensor import HybridMXFP8Layout - register_layout_class("HybridMXFP8Layout", HybridMXFP8Layout) - logging.info("ComfyUI-QuantOps: Registered HybridMXFP8Layout") - except ImportError: - logging.debug("ComfyUI-QuantOps: HybridMXFP8Layout not available") + if "hybrid_mxfp8" not in QUANT_ALGOS: + try: + from comfy_kitchen.tensor import HybridMXFP8Layout + register_layout_class("HybridMXFP8Layout", HybridMXFP8Layout) + registered.append("HybridMXFP8Layout") + logging.info("ComfyUI-QuantOps: Registered HybridMXFP8Layout") + except ImportError: + logging.debug("ComfyUI-QuantOps: HybridMXFP8Layout not available") QUANT_ALGOS.setdefault( "hybrid_mxfp8", @@ -182,7 +192,6 @@ def _register_layouts(): ) # Verify registration - registered = ["BlockWiseINT8Layout", "TensorWiseINT8Layout", "RowWiseFP8Layout", "BlockWiseFP8Layout", "TensorCoreMXFP8Layout"] logging.info(f"ComfyUI-QuantOps: Registered layouts: {registered}") except Exception as e: @@ -195,9 +204,33 @@ def _register_layouts(): # Setup backends first (enables ck triton, registers our kernels) _setup_comfy_kitchen_backends() +logging.info( + "ComfyUI-QuantOps diagnostic: native formats before extension: %s", + ", ".join(sorted(_NATIVE_COMFY_FORMATS)), +) # Register layouts _register_layouts() +logging.info( + "ComfyUI-QuantOps diagnostic: formats added to registry: %s", + ", ".join(sorted(set(QUANT_ALGOS) - _NATIVE_COMFY_FORMATS)) or "none", +) + +# Patch stock ComfyUI loaders so QuantOps-only metadata works from normal loaders. +try: + from .auto_patch import install_auto_patch + + install_auto_patch(_NATIVE_COMFY_FORMATS) +except Exception as e: + logging.warning(f"ComfyUI-QuantOps: failed to install stock-loader auto patch: {e}") + +# Extend an installed GGUF provider without replacing its loader implementation. +try: + from .gguf_integration import install_gguf_integration + + install_gguf_integration() +except Exception as e: + logging.warning(f"ComfyUI-QuantOps: failed to install GGUF integration: {e}") # Import nodes for ComfyUI discovery from .nodes.loader_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS diff --git a/auto_patch.py b/auto_patch.py new file mode 100644 index 0000000..42e17c7 --- /dev/null +++ b/auto_patch.py @@ -0,0 +1,515 @@ +""" +Automatic stock-loader integration for ComfyUI-QuantOps. + +ComfyUI core can detect QuantOps metadata after this custom node registers +extra QUANT_ALGOS entries, but the stock mixed-precision loader only knows how +to instantiate ComfyUI-native formats. This patch makes stock loaders use +QuantOps custom_operations when a model contains QuantOps-only formats. +""" + +import dataclasses +import json +import logging +from typing import Iterable, Optional, Tuple + +import torch + + +NATIVE_COMFY_FORMATS = { + "convrot_w4a4", + "float8_e4m3fn", + "float8_e5m2", + "int8_tensorwise", + "mxfp8", + "nvfp4", +} + +QUANTOPS_FORMATS = { + "int8", + "int8_blockwise", + "int8_tensorwise", + "float8_e4m3fn_rowwise", + "float8_e4m3fn_blockwise", + "mxfp8", + "hybrid_mxfp8", + "nvfp4", +} + + +def _metadata_formats(quant_metadata: Optional[dict]) -> set[str]: + if not quant_metadata: + return set() + layers = quant_metadata.get("layers", {}) + return {conf.get("format") for conf in layers.values() if conf.get("format")} + + +def _format_summary(quant_metadata: Optional[dict]) -> tuple[int, str]: + counts = {} + if quant_metadata: + for conf in quant_metadata.get("layers", {}).values(): + quant_format = conf.get("format") + if quant_format: + counts[quant_format] = counts.get(quant_format, 0) + 1 + summary = ", ".join(f"{name}={count}" for name, count in sorted(counts.items())) + return sum(counts.values()), summary or "none" + + +def _quantops_formats( + quant_metadata: Optional[dict], native_formats: Iterable[str] +) -> set[str]: + formats = _metadata_formats(quant_metadata) + return (formats & QUANTOPS_FORMATS) - set(native_formats) + + +def _needs_quantops( + quant_metadata: Optional[dict], + native_formats: Iterable[str] = NATIVE_COMFY_FORMATS, +) -> bool: + return bool(_quantops_formats(quant_metadata, native_formats)) + + +def _format_names(formats: Iterable[str]) -> str: + return ", ".join(sorted(formats)) if formats else "none" + + +def _backend_details(formats: Iterable[str]) -> str: + details = [] + formats = set(formats) + if formats & {"int8", "int8_blockwise"}: + try: + from .quant_layouts.int8_layout import BlockWiseINT8Layout + + backend = "triton" if BlockWiseINT8Layout.use_triton else "pytorch fallback" + except Exception: + backend = "unknown" + details.append(f"blockwise INT8 backend={backend}") + return "; ".join(details) + + +def _prepare_quantops_options( + state_dict: dict, + metadata: Optional[dict], + model_options: Optional[dict], + model_prefix: str = "", + native_formats: Iterable[str] = NATIVE_COMFY_FORMATS, + loader_name: str = "stock", +) -> Tuple[dict, dict, dict, Optional[dict], bool]: + """Inject .comfy_quant tensors and QuantOps custom_operations if needed.""" + from .utils.safetensors_loader import convert_old_quants + + metadata = dict(metadata or {}) + input_key_count = len(state_dict) + if "_quantization_metadata" in metadata: + metadata_source = "file-metadata" + elif f"{model_prefix}scaled_fp8" in state_dict: + metadata_source = "legacy-scaled-fp8" + elif any(key.endswith(".comfy_quant") for key in state_dict): + metadata_source = "comfy-quant-tensors" + else: + metadata_source = "dtype-scale-inference" + state_dict, metadata, quant_metadata = convert_old_quants( + state_dict, + model_prefix=model_prefix, + metadata=metadata, + ) + if not quant_metadata: + metadata_source = "none" + + quantops_formats = _quantops_formats(quant_metadata, native_formats) + formats = _metadata_formats(quant_metadata) + native = formats & set(native_formats) + quant_layers, format_summary = _format_summary(quant_metadata) + existing_custom_ops = bool(model_options and model_options.get("custom_operations") is not None) + dtype_override = (model_options or {}).get("dtype", "default") + fp8_optimizations = bool(model_options and model_options.get("fp8_optimizations", False)) + if not quantops_formats: + if not formats: + decision = "no-quantization" + elif formats <= set(native_formats): + decision = "native-bypass" + else: + decision = "core-pass-through" + logging.info( + "ComfyUI-QuantOps diagnostic: loader=%s decision=%s prefix=%s state_keys=%d metadata_source=%s quant_layers=%d formats=%s dtype_override=%s fp8_optimizations=%s existing_custom_operations=%s", + loader_name, + decision, + model_prefix or "", + input_key_count, + metadata_source, + quant_layers, + format_summary, + dtype_override, + fp8_optimizations, + existing_custom_ops, + ) + return state_dict, metadata, dict(model_options or {}), quant_metadata, False + + from .unified_ops import make_quant_ops + + patched_options = dict(model_options or {}) + base_ops = patched_options.get("custom_operations", None) + patched_options["custom_operations"] = make_quant_ops(base_ops) + patched_options.setdefault("quantization_metadata", {"mixed_ops": True}) + + details = _backend_details(quantops_formats) + logging.info( + "ComfyUI-QuantOps diagnostic: loader=%s decision=quantops prefix=%s state_keys=%d metadata_source=%s quant_layers=%d formats=%s quantops_formats=%s native_formats=%s dtype_override=%s fp8_optimizations=%s existing_custom_operations=%s backend=%s", + loader_name, + model_prefix or "", + input_key_count, + metadata_source, + quant_layers, + format_summary, + _format_names(quantops_formats), + _format_names(native), + dtype_override, + fp8_optimizations, + existing_custom_ops, + details or "default", + ) + return state_dict, metadata, patched_options, quant_metadata, True + + +def _dtype_from_metadata(value, default=torch.bfloat16): + if isinstance(value, torch.dtype): + return value + if isinstance(value, str): + return { + "torch.bfloat16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "torch.float16": torch.float16, + "float16": torch.float16, + "torch.float32": torch.float32, + "float32": torch.float32, + }.get(value, default) + return default + + +def _parse_comfy_quant_tensor(tensor) -> Optional[dict]: + try: + text = tensor.detach().cpu().numpy().tobytes().decode("utf-8").strip() + if text.startswith("{{") and text.endswith("}}"): + text = text[1:-1] + return json.loads(text) + except Exception: + return None + + +def _quant_metadata_from_state_dict(state_dict: dict, quant_metadata: Optional[dict]) -> dict: + if quant_metadata is not None: + return quant_metadata + + layers = {} + for key, value in state_dict.items(): + if not key.endswith(".comfy_quant"): + continue + layer_name = key[: -len(".comfy_quant")] + layer_conf = _parse_comfy_quant_tensor(value) + if layer_conf: + layers[layer_name] = layer_conf + return {"layers": layers} if layers else {"layers": {}} + + +def _get_first_tensor(state_dict: dict, keys: Iterable[str]): + for key in keys: + value = state_dict.get(key) + if value is not None: + return value + return None + + +def _manual_dequantize_int8(weight: torch.Tensor, scale: Optional[torch.Tensor], dtype: torch.dtype) -> torch.Tensor: + if scale is None: + return weight.to(dtype) + + scale = scale.to(torch.float32) + if scale.ndim == 0 or (scale.ndim == 1 and scale.numel() == 1): + return (weight.to(torch.float32) * scale.reshape(())).to(dtype) + if scale.ndim == 1 and scale.numel() == weight.shape[0]: + return (weight.to(torch.float32) * scale.reshape(-1, 1)).to(dtype) + if scale.ndim == 2 and scale.shape[0] == weight.shape[0] and scale.shape[1] == 1: + return (weight.to(torch.float32) * scale.to(torch.float32)).to(dtype) + + raise RuntimeError(f"Unsupported INT8 fallback scale shape: {tuple(scale.shape)}") + + +def _dequantize_layer(state_dict: dict, layer_name: str, layer_conf: dict) -> bool: + from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class + + weight_key = f"{layer_name}.weight" + weight = state_dict.get(weight_key) + if weight is None: + return False + + quant_format = layer_conf.get("format") + if not quant_format: + return False + + scale = _get_first_tensor( + state_dict, + (f"{layer_name}.weight_scale", f"{layer_name}.scale_weight"), + ) + scale_2 = state_dict.get(f"{layer_name}.weight_scale_2") + target_dtype = _dtype_from_metadata(layer_conf.get("orig_dtype"), torch.bfloat16) + orig_shape = tuple(layer_conf.get("orig_shape", tuple(weight.shape))) + + try: + qconfig = QUANT_ALGOS[quant_format] + layout_name = qconfig["comfy_tensor_layout"] + layout_cls = get_layout_class(layout_name) + + params_kwargs = { + "scale": scale.to(torch.float32) if scale is not None else torch.tensor(1.0), + "orig_dtype": target_dtype, + "orig_shape": orig_shape, + "block_size": int(layer_conf.get("group_size") or qconfig.get("group_size") or 128), + "is_weight": True, + } + if scale_2 is not None: + params_kwargs["block_scale"] = scale + params_kwargs["scale"] = scale_2.to(torch.float32) + + field_names = {field.name for field in dataclasses.fields(layout_cls.Params)} + params_kwargs = { + key: value for key, value in params_kwargs.items() if key in field_names + } + params = layout_cls.Params(**params_kwargs) + quantized = QuantizedTensor( + weight.to(qconfig["storage_t"]), + layout_name, + params, + ) + state_dict[weight_key] = quantized.dequantize().to(target_dtype).cpu() + except Exception as exc: + if weight.dtype == torch.int8: + logging.warning( + "ComfyUI-QuantOps: generic fallback failed for %s (%s); using INT8 dequant fallback", + layer_name, + exc, + ) + state_dict[weight_key] = _manual_dequantize_int8(weight, scale, target_dtype).cpu() + else: + logging.warning( + "ComfyUI-QuantOps: could not dequantize %s fallback layer %s: %s", + quant_format, + layer_name, + exc, + ) + return False + + for suffix in ( + ".weight_scale", + ".scale_weight", + ".weight_scale_2", + ".weight_scalar", + ".input_scale", + ".scale_input", + ".comfy_quant", + ): + state_dict.pop(f"{layer_name}{suffix}", None) + return True + + +def _dequantize_for_native_fallback( + state_dict: dict, + metadata: Optional[dict], + quant_metadata: Optional[dict], + native_formats: Iterable[str] = NATIVE_COMFY_FORMATS, +) -> Tuple[dict, dict, int]: + fallback_sd = dict(state_dict) + fallback_metadata = dict(metadata or {}) + quant_metadata = _quant_metadata_from_state_dict(fallback_sd, quant_metadata) + + quantops_formats = _quantops_formats(quant_metadata, native_formats) + converted = 0 + for layer_name, layer_conf in quant_metadata.get("layers", {}).items(): + if layer_conf.get("format") not in quantops_formats: + continue + if _dequantize_layer(fallback_sd, layer_name, layer_conf): + converted += 1 + + fallback_metadata.pop("_quantization_metadata", None) + return fallback_sd, fallback_metadata, converted + + +def _is_quantops_load_failure(exc: BaseException) -> bool: + text = str(exc) + return ( + "Unsupported quantization format" in text + or "TensorWiseINT8Layout" in text + or "BlockWiseINT8Layout" in text + or "RowWiseFP8Layout" in text + or "BlockWiseFP8Layout" in text + or "HybridMXFP8Layout" in text + ) + + +def install_auto_patch( + native_formats: Iterable[str] = NATIVE_COMFY_FORMATS, +) -> None: + import comfy.model_detection + import comfy.sd + + native_formats = frozenset(native_formats) + + if getattr(comfy.sd, "_quantops_auto_patch_installed", False): + return + + original_load_diffusion_model_state_dict = comfy.sd.load_diffusion_model_state_dict + original_load_state_dict_guess_config = comfy.sd.load_state_dict_guess_config + + def load_diffusion_model_state_dict_quantops( + sd, + model_options={}, + metadata=None, + disable_dynamic=False, + ): + prepared_sd, prepared_metadata, patched_options, quant_metadata, did_patch = _prepare_quantops_options( + dict(sd), + metadata, + model_options, + native_formats=native_formats, + loader_name="diffusion", + ) + if not did_patch: + result = original_load_diffusion_model_state_dict( + sd, + model_options=model_options, + metadata=metadata, + disable_dynamic=disable_dynamic, + ) + logging.info("ComfyUI-QuantOps diagnostic: loader=diffusion path=core result=completed") + return result + + try: + result = original_load_diffusion_model_state_dict( + dict(prepared_sd), + model_options=patched_options, + metadata=prepared_metadata, + disable_dynamic=disable_dynamic, + ) + logging.info("ComfyUI-QuantOps diagnostic: loader=diffusion path=quantops result=completed") + return result + except Exception as exc: + if not _is_quantops_load_failure(exc): + raise + logging.warning( + "ComfyUI-QuantOps: auto custom load failed (%s). Falling back to dequantized BF16 load.", + exc, + ) + fallback_sd, fallback_metadata, converted = _dequantize_for_native_fallback( + prepared_sd, + prepared_metadata, + quant_metadata, + native_formats=native_formats, + ) + if converted == 0: + raise + fallback_options = dict(model_options or {}) + fallback_options.pop("custom_operations", None) + result = original_load_diffusion_model_state_dict( + fallback_sd, + model_options=fallback_options, + metadata=fallback_metadata, + disable_dynamic=disable_dynamic, + ) + logging.info( + "ComfyUI-QuantOps diagnostic: loader=diffusion path=bf16-fallback result=completed converted_layers=%d", + converted, + ) + return result + + def load_state_dict_guess_config_quantops( + sd, + output_vae=True, + output_clip=True, + output_clipvision=False, + embedding_directory=None, + output_model=True, + model_options={}, + te_model_options={}, + metadata=None, + disable_dynamic=False, + ): + diffusion_model_prefix = comfy.model_detection.unet_prefix_from_state_dict(sd) + prepared_sd, prepared_metadata, patched_options, quant_metadata, did_patch = _prepare_quantops_options( + dict(sd), + metadata, + model_options, + model_prefix=diffusion_model_prefix, + native_formats=native_formats, + loader_name="checkpoint", + ) + if not did_patch: + result = original_load_state_dict_guess_config( + sd, + output_vae=output_vae, + output_clip=output_clip, + output_clipvision=output_clipvision, + embedding_directory=embedding_directory, + output_model=output_model, + model_options=model_options, + te_model_options=te_model_options, + metadata=metadata, + disable_dynamic=disable_dynamic, + ) + logging.info("ComfyUI-QuantOps diagnostic: loader=checkpoint path=core result=completed") + return result + + try: + result = original_load_state_dict_guess_config( + dict(prepared_sd), + output_vae=output_vae, + output_clip=output_clip, + output_clipvision=output_clipvision, + embedding_directory=embedding_directory, + output_model=output_model, + model_options=patched_options, + te_model_options=te_model_options, + metadata=prepared_metadata, + disable_dynamic=disable_dynamic, + ) + logging.info("ComfyUI-QuantOps diagnostic: loader=checkpoint path=quantops result=completed") + return result + except Exception as exc: + if not _is_quantops_load_failure(exc): + raise + logging.warning( + "ComfyUI-QuantOps: auto checkpoint load failed (%s). Falling back to dequantized BF16 UNet load.", + exc, + ) + fallback_sd, fallback_metadata, converted = _dequantize_for_native_fallback( + prepared_sd, + prepared_metadata, + quant_metadata, + native_formats=native_formats, + ) + if converted == 0: + raise + fallback_options = dict(model_options or {}) + fallback_options.pop("custom_operations", None) + result = original_load_state_dict_guess_config( + fallback_sd, + output_vae=output_vae, + output_clip=output_clip, + output_clipvision=output_clipvision, + embedding_directory=embedding_directory, + output_model=output_model, + model_options=fallback_options, + te_model_options=te_model_options, + metadata=fallback_metadata, + disable_dynamic=disable_dynamic, + ) + logging.info( + "ComfyUI-QuantOps diagnostic: loader=checkpoint path=bf16-fallback result=completed converted_layers=%d", + converted, + ) + return result + + comfy.sd.load_diffusion_model_state_dict = load_diffusion_model_state_dict_quantops + comfy.sd.load_state_dict_guess_config = load_state_dict_guess_config_quantops + comfy.sd._quantops_auto_patch_installed = True + comfy.sd._quantops_original_load_diffusion_model_state_dict = original_load_diffusion_model_state_dict + comfy.sd._quantops_original_load_state_dict_guess_config = original_load_state_dict_guess_config + + logging.info("ComfyUI-QuantOps: installed stock-loader auto patch") diff --git a/gguf_integration.py b/gguf_integration.py new file mode 100644 index 0000000..a6d1469 --- /dev/null +++ b/gguf_integration.py @@ -0,0 +1,108 @@ +"""Optional integration with an installed ComfyUI GGUF loader.""" + +import functools +import logging +import os +import sys + + +GGUF_IMAGE_ARCHITECTURES = {"ideogram", "krea2"} + + +def _provider_name(loader_module) -> str: + module_file = getattr(loader_module, "__file__", "") + if module_file: + return os.path.basename(os.path.dirname(module_file)) + return loader_module.__name__ + + +def install_gguf_integration() -> bool: + """Enable supported image architectures and diagnostics on the GGUF provider.""" + import nodes + + loader_class = nodes.NODE_CLASS_MAPPINGS.get("UnetLoaderGGUF") + if loader_class is None: + logging.info( + "ComfyUI-QuantOps GGUF diagnostic: provider=unavailable status=disabled " + "reason=UnetLoaderGGUF-not-registered" + ) + return False + + method_globals = loader_class.load_unet.__globals__ + gguf_loader = method_globals.get("gguf_sd_loader") + loader_module = sys.modules.get(getattr(gguf_loader, "__module__", "")) + architectures = getattr(loader_module, "IMG_ARCH_LIST", None) + if gguf_loader is None or loader_module is None or not isinstance(architectures, set): + logging.warning( + "ComfyUI-QuantOps GGUF diagnostic: provider=%s status=disabled " + "reason=incompatible-provider-api", + loader_class.__module__, + ) + return False + + provider = _provider_name(loader_module) + already_supported = architectures & GGUF_IMAGE_ARCHITECTURES + added = GGUF_IMAGE_ARCHITECTURES - architectures + architectures.update(GGUF_IMAGE_ARCHITECTURES) + + if getattr(gguf_loader, "_quantops_diagnostic_wrapper", False): + logging.info( + "ComfyUI-QuantOps GGUF diagnostic: provider=%s status=active " + "architectures_added=none architectures_supported=%s", + provider, + ",".join(sorted(GGUF_IMAGE_ARCHITECTURES)), + ) + return True + + @functools.wraps(gguf_loader) + def logged_gguf_loader(path, handle_prefix="model.diffusion_model.", is_text_model=False): + path_text = os.fspath(path) if path is not None else "" + model_kind = "text" if is_text_model else "diffusion" + logging.info( + "ComfyUI-QuantOps GGUF diagnostic: decision=gguf-provider provider=%s " + "model_kind=%s file=%s path=%s", + provider, + model_kind, + os.path.basename(path_text), + path_text, + ) + try: + state_dict, extra = gguf_loader(path, handle_prefix, is_text_model) + except Exception as error: + logging.error( + "ComfyUI-QuantOps GGUF diagnostic: decision=gguf-provider provider=%s " + "model_kind=%s file=%s result=error error_type=%s error=%s", + provider, + model_kind, + os.path.basename(path_text), + type(error).__name__, + error, + ) + raise + + metadata = extra.get("metadata", {}) + logging.info( + "ComfyUI-QuantOps GGUF diagnostic: decision=gguf-provider provider=%s " + "model_kind=%s file=%s result=accepted architecture=%s state_keys=%d " + "metadata_fields=%d", + provider, + model_kind, + os.path.basename(path_text), + extra.get("arch_str", "unknown"), + len(state_dict), + len(metadata), + ) + return state_dict, extra + + logged_gguf_loader._quantops_diagnostic_wrapper = True + loader_module.gguf_sd_loader = logged_gguf_loader + method_globals["gguf_sd_loader"] = logged_gguf_loader + + logging.info( + "ComfyUI-QuantOps GGUF diagnostic: provider=%s status=active " + "architectures_added=%s architectures_already_supported=%s", + provider, + ",".join(sorted(added)) or "none", + ",".join(sorted(already_supported)) or "none", + ) + return True diff --git a/kernels/fp8_kernels.py b/kernels/fp8_kernels.py index ef28a38..0aa01aa 100644 --- a/kernels/fp8_kernels.py +++ b/kernels/fp8_kernels.py @@ -176,9 +176,8 @@ def fp8_gemm_blockwise_kernel( a_s_ptrs = a_s_ptr + offs_m * a_s_k_blocks # Weight scale pointers: shape [N//input_block_size, K//input_block_size] - # For N tile pid_n, we need scales[pid_n, :] across K iterations b_s_k_blocks = tl.cdiv(K, input_block_size) - b_s_base = b_s_ptr + pid_n * b_s_k_blocks + b_s_rows = offs_n // input_block_size # Accumulator accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) @@ -189,8 +188,8 @@ def fp8_gemm_blockwise_kernel( # Load FP8 tiles and cast to float32 mask_k = offs_k < K - k_start - a_fp8 = tl.load(a_ptrs, mask=mask_k[None, :], other=0) - b_fp8 = tl.load(b_ptrs, mask=mask_k[:, None], other=0) + a_fp8 = tl.load(a_ptrs, mask=mask_k[None, :], other=0.0) + b_fp8 = tl.load(b_ptrs, mask=mask_k[:, None], other=0.0) # Cast to float32 for computation a_f32 = a_fp8.to(tl.float32) @@ -205,10 +204,10 @@ def fp8_gemm_blockwise_kernel( # Load scales a_s = tl.load(a_s_ptrs + k_scale_idx) # [BLOCK_SIZE_M] - b_s = tl.load(b_s_base + k_scale_idx) # scalar (for this N tile, K block) + b_s = tl.load(b_s_ptr + b_s_rows * b_s_k_blocks + k_scale_idx) # Apply scales: result = dot * a_scale[:, None] * b_scale - accumulator += dot_result * a_s[:, None] * b_s + accumulator += dot_result * a_s[:, None] * b_s[None, :] # Advance pointers a_ptrs += BLOCK_SIZE_K @@ -328,7 +327,7 @@ def fp8_addmm_blockwise_kernel( a_s_k_blocks = tl.cdiv(K, input_block_size) a_s_ptrs = a_s_ptr + offs_m * a_s_k_blocks b_s_k_blocks = tl.cdiv(K, input_block_size) - b_s_base = b_s_ptr + pid_n * b_s_k_blocks + b_s_rows = offs_n // input_block_size # Accumulator accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) @@ -338,8 +337,8 @@ def fp8_addmm_blockwise_kernel( k_start = k_idx * BLOCK_SIZE_K mask_k = offs_k < K - k_start - a_fp8 = tl.load(a_ptrs, mask=mask_k[None, :], other=0) - b_fp8 = tl.load(b_ptrs, mask=mask_k[:, None], other=0) + a_fp8 = tl.load(a_ptrs, mask=mask_k[None, :], other=0.0) + b_fp8 = tl.load(b_ptrs, mask=mask_k[:, None], other=0.0) a_f32 = a_fp8.to(tl.float32) b_f32 = b_fp8.to(tl.float32) @@ -348,9 +347,9 @@ def fp8_addmm_blockwise_kernel( k_scale_idx = k_start // input_block_size a_s = tl.load(a_s_ptrs + k_scale_idx) - b_s = tl.load(b_s_base + k_scale_idx) + b_s = tl.load(b_s_ptr + b_s_rows * b_s_k_blocks + k_scale_idx) - accumulator += dot_result * a_s[:, None] * b_s + accumulator += dot_result * a_s[:, None] * b_s[None, :] a_ptrs += BLOCK_SIZE_K b_ptrs += BLOCK_SIZE_K @@ -487,8 +486,8 @@ def fp8_gemm_rowwise_kernel( k_start = k_idx * BLOCK_SIZE_K mask_k = offs_k < K - k_start - a_fp8 = tl.load(a_ptrs, mask=mask_k[None, :], other=0) - b_fp8 = tl.load(b_ptrs, mask=mask_k[:, None], other=0) + a_fp8 = tl.load(a_ptrs, mask=mask_k[None, :], other=0.0) + b_fp8 = tl.load(b_ptrs, mask=mask_k[:, None], other=0.0) a_f32 = a_fp8.to(tl.float32) b_f32 = b_fp8.to(tl.float32) diff --git a/nodes/loader_nodes.py b/nodes/loader_nodes.py index 8d7078e..c91006e 100644 --- a/nodes/loader_nodes.py +++ b/nodes/loader_nodes.py @@ -19,6 +19,8 @@ import comfy.latent_formats import comfy.conds +from ..auto_patch import NATIVE_COMFY_FORMATS + # Try to import UnifiedSafetensorsLoader for aimdo-free loading try: from unifiedefficientloader import UnifiedSafetensorsLoader @@ -206,6 +208,45 @@ def _configure_int8_backend(kernel_backend): logging.warning(f"Failed to configure comfy_kitchen backend: {e}") +def _quant_metadata_formats(quant_metadata): + if not quant_metadata: + return set() + return { + conf.get("format") + for conf in quant_metadata.get("layers", {}).values() + if conf.get("format") + } + + +def _selected_formats(quant_metadata, quant_format): + formats = _quant_metadata_formats(quant_metadata) + if not formats and quant_format != "auto": + formats = {quant_format} + return formats + + +def _format_names(formats): + return ", ".join(sorted(formats)) if formats else "none" + + +def _backend_details(formats): + details = [] + if formats & {"int8", "int8_blockwise"}: + try: + from ..quant_layouts.int8_layout import BlockWiseINT8Layout + + backend = "triton" if BlockWiseINT8Layout.use_triton else "pytorch fallback" + except Exception: + backend = "unknown" + details.append(f"blockwise INT8 backend={backend}") + return "; ".join(details) + + +def _quantops_only_formats(quant_metadata, quant_format): + formats = _selected_formats(quant_metadata, quant_format) + return {fmt for fmt in formats if fmt not in NATIVE_COMFY_FORMATS} + + def _build_model_options( quant_format, sd, @@ -250,7 +291,7 @@ def _build_model_options( for conf in quant_metadata.get("layers", {}).values() if conf.get("format") } - has_int8 = any(fmt in ("int8", "int8_tensorwise") for fmt in layer_formats) + has_int8 = any(fmt in ("int8", "int8_blockwise") for fmt in layer_formats) else: has_int8 = any( k.endswith(".weight") and sd[k].dtype == torch.int8 @@ -259,7 +300,7 @@ def _build_model_options( ) if has_int8: _configure_int8_backend(kernel_backend) - elif quant_format in ("int8", "int8_tensorwise"): + elif quant_format in ("int8", "int8_blockwise"): _configure_int8_backend(kernel_backend) # Forward text-encoder quantization metadata into model_options @@ -280,14 +321,28 @@ def _build_model_options( if quant_metadata is not None and "quantization_metadata" not in model_options: model_options["quantization_metadata"] = {"mixed_ops": True} - # Attach unified custom operations dynamically - try: - from ..unified_ops import make_quant_ops + quantops_formats = _quantops_only_formats(quant_metadata, quant_format) + if quantops_formats: + try: + from ..unified_ops import make_quant_ops - base_ops = model_options.get("custom_operations", None) - model_options["custom_operations"] = make_quant_ops(base_ops) - except ImportError as e: - logging.warning(f"unified_ops not available: {e}") + base_ops = model_options.get("custom_operations", None) + model_options["custom_operations"] = make_quant_ops(base_ops) + details = _backend_details(quantops_formats) + logging.info( + "ComfyUI-QuantOps: used by QuantOps loader; custom operations active for formats: %s%s", + _format_names(quantops_formats), + f" ({details})" if details else "", + ) + except ImportError as e: + logging.warning(f"unified_ops not available: {e}") + else: + formats = _selected_formats(quant_metadata, quant_format) + if formats: + logging.info( + "ComfyUI-QuantOps: not used by QuantOps loader; native ComfyUI handles formats: %s", + _format_names(formats), + ) return model_options diff --git a/quant_layouts/fp8_variants.py b/quant_layouts/fp8_variants.py index 40f04c8..717db0e 100644 --- a/quant_layouts/fp8_variants.py +++ b/quant_layouts/fp8_variants.py @@ -259,6 +259,43 @@ def get_plain_tensors(cls, qtensor) -> Tuple[torch.Tensor, torch.Tensor, int]: # ============================================================================== +_fp8_path_counts = {} +_FP8_LOG_LIMIT = 3 + + +def _disable_fp8_kernels(error): + global _HAS_FP8_KERNELS + + _HAS_FP8_KERNELS = False + lines = str(error).strip().splitlines() + summary = lines[-1] if lines else type(error).__name__ + logging.warning( + "ComfyUI-QuantOps FP8: disabled Triton kernels after %s: %s", + type(error).__name__, + summary, + ) + + +def _log_fp8_path(layout, path, input_shape, weight_shape): + key = (layout, path) + count = _fp8_path_counts.get(key, 0) + _fp8_path_counts[key] = count + 1 + if count < _FP8_LOG_LIMIT: + logging.info( + "ComfyUI-QuantOps FP8 %s: path=%s input=%s weight=%s", + layout, + path, + input_shape, + weight_shape, + ) + elif count == _FP8_LOG_LIMIT: + logging.info( + "ComfyUI-QuantOps FP8 %s: suppressing further path=%s logs", + layout, + path, + ) + + @register_layout_op(torch.ops.aten.linear.default, RowWiseFP8Layout) def rowwise_fp8_linear(func, args, kwargs): """Row-wise FP8 linear operation with native kernel support.""" @@ -286,11 +323,6 @@ def rowwise_fp8_linear(func, args, kwargs): dtype=w_qdata.dtype, ) - logging.debug( - f"FP8 rowwise: Native kernel (dynamic quant), " - f"input={a_qdata.shape}, weight={w_qdata.shape}" - ) - # For rowwise, bias needs manual addition (no fused kernel yet) result = fp8_gemm_rowwise( a_qdata, @@ -305,12 +337,13 @@ def rowwise_fp8_linear(func, args, kwargs): device=result.device, dtype=result.dtype ) + _log_fp8_path("rowwise", "triton-dynamic", a_qdata.shape, w_qdata.shape) return result.to(orig_dtype) except Exception as e: - logging.warning(f"FP8 rowwise native kernel failed: {e}") + _disable_fp8_kernels(e) # Fallback: dequantize - logging.debug("FP8 rowwise: Using dequant fallback") + _log_fp8_path("rowwise", "dequant-fallback", input_tensor.shape, weight.shape) if isinstance(weight, QuantizedTensor): weight = weight.dequantize() if isinstance(input_tensor, QuantizedTensor): @@ -330,7 +363,7 @@ def rowwise_fp8_mm(func, args, kwargs): if isinstance(input_tensor, QuantizedTensor): input_tensor = input_tensor.dequantize() - return func(input_tensor, weight) + return torch.mm(input_tensor, weight) @register_layout_op(torch.ops.aten.addmm.default, RowWiseFP8Layout) @@ -347,7 +380,7 @@ def rowwise_fp8_addmm(func, args, kwargs): if isinstance(weight, QuantizedTensor): weight = weight.dequantize() - return func(bias, input_tensor, weight, **kwargs) + return torch.addmm(bias, input_tensor, weight, **kwargs) @register_layout_op(torch.ops.aten.view.default, RowWiseFP8Layout) @@ -356,12 +389,17 @@ def rowwise_fp8_func(func, args, kwargs): """Handle view/transpose for row-wise FP8 tensors.""" input_tensor = args[0] if isinstance(input_tensor, QuantizedTensor): - plain_input, scale = RowWiseFP8Layout.get_plain_tensors(input_tensor) - ar = list(args) - ar[0] = plain_input - # Use _copy_with to preserve params - return input_tensor._copy_with(qdata=func(*ar, **kwargs)) - return func(*args, **kwargs) + path = "transpose-dequant" if len(args) == 1 else "view-dequant" + _log_fp8_path("rowwise", path, None, input_tensor.shape) + plain_input = input_tensor.dequantize() + if len(args) == 1: + return torch.t(plain_input) + shape = args[1] if len(args) == 2 and isinstance(args[1], (tuple, list)) else args[1:] + return plain_input.view(*shape) + if len(args) == 1: + return torch.t(input_tensor) + shape = args[1] if len(args) == 2 and isinstance(args[1], (tuple, list)) else args[1:] + return input_tensor.view(*shape) @register_layout_op(torch.ops.aten.linear.default, BlockWiseFP8Layout) @@ -385,11 +423,6 @@ def blockwise_fp8_linear(func, args, kwargs): input_tensor ) - logging.debug( - f"FP8 blockwise: Native kernel (both quantized), " - f"input={a_qdata.shape}, weight={w_qdata.shape}, block_size={w_block_size}" - ) - try: if bias is not None: result = fp8_addmm_blockwise( @@ -408,11 +441,10 @@ def blockwise_fp8_linear(func, args, kwargs): w_scale, input_block_size=w_block_size, ) + _log_fp8_path("blockwise", "triton-quantized", a_qdata.shape, w_qdata.shape) return result.to(orig_dtype) except Exception as e: - logging.warning( - f"FP8 native kernel failed, falling back to dequant: {e}" - ) + _disable_fp8_kernels(e) # Input is not quantized - quantize it dynamically elif input_tensor.dtype in [torch.float16, torch.bfloat16, torch.float32]: @@ -424,11 +456,6 @@ def blockwise_fp8_linear(func, args, kwargs): dtype=w_qdata.dtype, ) - logging.debug( - f"FP8 blockwise: Native kernel (dynamic quant), " - f"input={a_qdata.shape}, weight={w_qdata.shape}" - ) - if bias is not None: result = fp8_addmm_blockwise( a_qdata, @@ -446,14 +473,13 @@ def blockwise_fp8_linear(func, args, kwargs): w_scale, input_block_size=w_block_size, ) + _log_fp8_path("blockwise", "triton-dynamic", a_qdata.shape, w_qdata.shape) return result.to(orig_dtype) except Exception as e: - logging.warning( - f"FP8 dynamic quant failed, falling back to dequant: {e}" - ) + _disable_fp8_kernels(e) # Fallback: dequantize - logging.debug("FP8 blockwise: Using dequant fallback") + _log_fp8_path("blockwise", "dequant-fallback", input_tensor.shape, weight.shape) if isinstance(weight, QuantizedTensor): weight = weight.dequantize() if isinstance(input_tensor, QuantizedTensor): @@ -473,7 +499,7 @@ def blockwise_fp8_mm(func, args, kwargs): if isinstance(input_tensor, QuantizedTensor): input_tensor = input_tensor.dequantize() - return func(input_tensor, weight) + return torch.mm(input_tensor, weight) @register_layout_op(torch.ops.aten.addmm.default, BlockWiseFP8Layout) @@ -490,7 +516,7 @@ def blockwise_fp8_addmm(func, args, kwargs): if isinstance(weight, QuantizedTensor): weight = weight.dequantize() - return func(bias, input_tensor, weight, **kwargs) + return torch.addmm(bias, input_tensor, weight, **kwargs) @register_layout_op(torch.ops.aten.view.default, BlockWiseFP8Layout) @@ -499,11 +525,14 @@ def blockwise_fp8_func(func, args, kwargs): """Handle view/transpose for block-wise FP8 tensors.""" input_tensor = args[0] if isinstance(input_tensor, QuantizedTensor): - plain_input, scale, block_size = BlockWiseFP8Layout.get_plain_tensors( - input_tensor - ) - ar = list(args) - ar[0] = plain_input - # Use _copy_with to preserve params - return input_tensor._copy_with(qdata=func(*ar, **kwargs)) - return func(*args, **kwargs) + path = "transpose-dequant" if len(args) == 1 else "view-dequant" + _log_fp8_path("blockwise", path, None, input_tensor.shape) + plain_input = input_tensor.dequantize() + if len(args) == 1: + return torch.t(plain_input) + shape = args[1] if len(args) == 2 and isinstance(args[1], (tuple, list)) else args[1:] + return plain_input.view(*shape) + if len(args) == 1: + return torch.t(input_tensor) + shape = args[1] if len(args) == 2 and isinstance(args[1], (tuple, list)) else args[1:] + return input_tensor.view(*shape) diff --git a/quant_layouts/int8_layout.py b/quant_layouts/int8_layout.py index 06abe64..81d4994 100644 --- a/quant_layouts/int8_layout.py +++ b/quant_layouts/int8_layout.py @@ -258,6 +258,7 @@ def dequantize(qdata, params) -> torch.Tensor: if BlockWiseINT8Layout.use_triton and qdata.dim() == 2 and qdata.is_cuda: try: weight_dequant = _get_triton_function("weight_dequant") + _log_int8_path("TRITON_WEIGHT_DEQUANT", qdata.shape, None) return weight_dequant( qdata, scale, block_size=block_size, output_dtype=output_dt ) @@ -277,11 +278,12 @@ def dequantize(qdata, params) -> torch.Tensor: raise RuntimeError( f"Weight scale shape mismatch: scale.shape={scale.shape}, expected {expected_scale_shape}" ) + _log_int8_path("PYTORCH_WEIGHT_DEQUANT", qdata.shape, None) qdata_blocked = qdata.reshape( M // block_size, block_size, N // block_size, block_size ) qdata_blocked = qdata_blocked.permute(0, 2, 1, 3) - scale_broadcast = scale.unsqueeze(-1).unsqueeze(-1) + scale_broadcast = scale.to(dtype=output_dt, device=qdata_blocked.device).unsqueeze(-1).unsqueeze(-1) dequant = qdata_blocked.to(output_dt) * scale_broadcast dequant = dequant.permute(0, 2, 1, 3).reshape(M, N) else: @@ -289,6 +291,7 @@ def dequantize(qdata, params) -> torch.Tensor: if BlockWiseINT8Layout.use_triton and qdata.is_cuda: try: act_dequant = _get_triton_function("act_dequant") + _log_int8_path("TRITON_ACT_DEQUANT", qdata.shape, None) return act_dequant( qdata, scale, block_size=block_size, output_dtype=output_dt ) @@ -311,8 +314,9 @@ def dequantize(qdata, params) -> torch.Tensor: raise RuntimeError( f"Activation scale shape mismatch: scale.shape={scale.shape}, expected {expected_scale_shape}" ) + _log_int8_path("PYTORCH_ACT_DEQUANT", qdata.shape, None) qdata_blocked = qdata.reshape(*batch_shape, K // block_size, block_size) - scale_broadcast = scale.unsqueeze(-1) + scale_broadcast = scale.to(dtype=output_dt, device=qdata_blocked.device).unsqueeze(-1) dequant = qdata_blocked.to(output_dt) * scale_broadcast dequant = dequant.reshape(qdata.shape) @@ -334,7 +338,17 @@ def get_plain_tensors(cls, qtensor): # ============================================================================== # Call counters to avoid log spam -_int8_path_counts = {"NATIVE_TRITON": 0, "PYTORCH_BOTH_QUANT": 0, "DEQUANT_FALLBACK": 0, "DYNAMIC_ACT_QUANT": 0} +_int8_path_counts = { + "NATIVE_TRITON": 0, + "PYTORCH_BOTH_QUANT": 0, + "PYTORCH_DYNAMIC_FALLBACK": 0, + "PYTORCH_WEIGHT_DEQUANT": 0, + "PYTORCH_ACT_DEQUANT": 0, + "TRITON_WEIGHT_DEQUANT": 0, + "TRITON_ACT_DEQUANT": 0, + "DEQUANT_FALLBACK": 0, + "DYNAMIC_ACT_QUANT": 0, +} _LOG_LIMIT = 3 # Log first N occurrences per path, then summary @@ -348,24 +362,44 @@ def _log_int8_path(path_type, input_shape, weight_shape, reason=None): if count < _LOG_LIMIT: if path_type == "NATIVE_TRITON": logging.info( - f"INT8: Native Triton matmul - input={input_shape}, weight={weight_shape}" + f"ComfyUI-QuantOps INT8 blockwise: backend=triton used native Triton matmul - input={input_shape}, weight={weight_shape}" ) elif path_type == "PYTORCH_BOTH_QUANT": logging.info( - f"INT8: PyTorch fallback (both quantized) - input={input_shape}, weight={weight_shape}" + f"ComfyUI-QuantOps INT8 blockwise: backend=pytorch fallback used dequantized linear (both quantized) - input={input_shape}, weight={weight_shape}" + ) + elif path_type == "PYTORCH_DYNAMIC_FALLBACK": + logging.info( + f"ComfyUI-QuantOps INT8 blockwise: backend=pytorch fallback used dequantized linear after dynamic activation quant - input={input_shape}, weight={weight_shape}" + ) + elif path_type == "PYTORCH_WEIGHT_DEQUANT": + logging.info( + f"ComfyUI-QuantOps INT8 blockwise: backend=pytorch fallback dequantized weight for generic torch op - qdata={input_shape}. Native INT8 matmul was not used for this op." + ) + elif path_type == "PYTORCH_ACT_DEQUANT": + logging.info( + f"ComfyUI-QuantOps INT8 blockwise: backend=pytorch fallback dequantized activation for generic torch op - qdata={input_shape}. Native INT8 matmul was not used for this op." + ) + elif path_type == "TRITON_WEIGHT_DEQUANT": + logging.info( + f"ComfyUI-QuantOps INT8 blockwise: backend=triton dequantized weight for generic torch op - qdata={input_shape}. Native INT8 matmul was not used for this op." + ) + elif path_type == "TRITON_ACT_DEQUANT": + logging.info( + f"ComfyUI-QuantOps INT8 blockwise: backend=triton dequantized activation for generic torch op - qdata={input_shape}. Native INT8 matmul was not used for this op." ) elif path_type == "DYNAMIC_ACT_QUANT": logging.info( - f"INT8: Dynamic activation quant - input={input_shape}, weight={weight_shape}" + f"ComfyUI-QuantOps INT8 blockwise: dynamically quantized activation - input={input_shape}, weight={weight_shape}" ) elif path_type == "DEQUANT_FALLBACK": logging.warning( - f"INT8: Dequant fallback ({reason}) - input={input_shape}, weight={weight_shape}. " + f"ComfyUI-QuantOps INT8 blockwise: dequant fallback ({reason}) - input={input_shape}, weight={weight_shape}. " f"Native INT8 matmul NOT used - only memory savings, no compute speedup." ) elif count == _LOG_LIMIT: logging.info( - f"INT8: Suppressing further '{path_type}' logs (limit={_LOG_LIMIT})" + f"ComfyUI-QuantOps INT8 blockwise: suppressing further '{path_type}' logs (limit={_LOG_LIMIT})" ) @@ -544,6 +578,7 @@ def int8_linear(func, args, kwargs): if result.dtype != out_dtype: result = result.to(out_dtype) + _log_int8_path("NATIVE_TRITON", a_int8.shape, b_int8.shape) return result.reshape(*orig_shape[:-1], weight.shape[0]) except Exception as e: logging.warning( @@ -551,6 +586,7 @@ def int8_linear(func, args, kwargs): ) # PyTorch fallback for dynamic quant path + _log_int8_path("PYTORCH_DYNAMIC_FALLBACK", a_int8.shape, b_int8.shape) output = _int8_gemm_pytorch_fallback( a_int8, a_scale, b_int8, b_scale, b_block_size, bias ) @@ -589,7 +625,7 @@ def int8_mm(func, args, kwargs): if isinstance(input_tensor, QuantizedTensor): input_tensor = input_tensor.dequantize() - return func(input_tensor, weight) + return torch.mm(input_tensor, weight) @register_layout_op(torch.ops.aten.addmm.default, BlockWiseINT8Layout) @@ -606,7 +642,7 @@ def int8_addmm(func, args, kwargs): if isinstance(weight, QuantizedTensor): weight = weight.dequantize() - return func(bias, input_tensor, weight, **kwargs) + return torch.addmm(bias, input_tensor, weight, **kwargs) @register_layout_op(torch.ops.aten.view.default, BlockWiseINT8Layout) @@ -615,10 +651,18 @@ def int8_func(func, args, kwargs): """Handle view/transpose for INT8 tensors.""" input_tensor = args[0] if isinstance(input_tensor, QuantizedTensor): - qdata = input_tensor._qdata - ar = list(args) - ar[0] = qdata - new_qdata = func(*ar, **kwargs) - # Use _copy_with to preserve params - return input_tensor._copy_with(qdata=new_qdata) - return func(*args, **kwargs) + path = ( + "TRITON_WEIGHT_DEQUANT" + if BlockWiseINT8Layout.use_triton + else "PYTORCH_WEIGHT_DEQUANT" + ) + _log_int8_path(path, input_tensor.shape, None) + plain_input = input_tensor.dequantize() + if len(args) == 1: + return torch.t(plain_input) + shape = args[1] if len(args) == 2 and isinstance(args[1], (tuple, list)) else args[1:] + return plain_input.view(*shape) + if len(args) == 1: + return torch.t(input_tensor) + shape = args[1] if len(args) == 2 and isinstance(args[1], (tuple, list)) else args[1:] + return input_tensor.view(*shape) diff --git a/unified_ops.py b/unified_ops.py index c150244..eb7fb6e 100644 --- a/unified_ops.py +++ b/unified_ops.py @@ -247,6 +247,20 @@ def _is_per_channel_scale(s, weight_n): ) if self.block_size is None: self.block_size = qconfig.get("group_size", None) + if self.layout_type in [ + "TensorCoreFP8Layout", + "TensorCoreFP8E4M3Layout", + "TensorCoreFP8E5M2Layout", + ] and scale is not None: + if scale.ndim == 1 and scale.numel() == weight_tensor.shape[0]: + self.layout_type = "RowWiseFP8Layout" + elif scale.ndim == 2 and scale.numel() > 1: + self.layout_type = "BlockWiseFP8Layout" + if self.block_size is None: + M, N = weight_tensor.shape + scale_M, scale_N = scale.shape + if M % scale_M == 0 and N % scale_N == 0: + self.block_size = M // scale_M else: if scale is not None: if scale.ndim == 0 or ( diff --git a/utils/safetensors_loader.py b/utils/safetensors_loader.py index 63b47d5..e8efdaa 100644 --- a/utils/safetensors_loader.py +++ b/utils/safetensors_loader.py @@ -397,9 +397,9 @@ def convert_old_quants(state_dict, model_prefix="", metadata=None): seen.add(layer_name) weight = state_dict[k] - scale = state_dict.get(layer_name + ".weight_scale") or state_dict.get( - layer_name + ".scale_weight" - ) + scale = state_dict.get(layer_name + ".weight_scale") + if scale is None: + scale = state_dict.get(layer_name + ".scale_weight") scale_2 = state_dict.get(layer_name + ".weight_scale_2") layer_conf = _infer_layer_format(weight, scale, scale_2)