eager: dequantize INT8 weights instead of torch._int_mm on AMD RDNA / CPU

comfy_kitchen/backends/eager/quantization.py

@@ -6,6 +6,9 @@
 #   Copyright (c) Meta Platforms, Inc. and affiliates.
 #   Licensed under the BSD 3-Clause License (see NOTICE file for details)
 
+import functools
+import os
+
 import torch
 
 from comfy_kitchen.float_utils import (
@@ -936,6 +939,33 @@ def dequantize_int8_simple_dtype(q: torch.Tensor, scale: torch.Tensor, output_dt
     return dequantize_int8_simple(q, scale).to(DTYPE_CODE_TO_DTYPE[output_dtype_code])
 
 
+_CDNA_ARCHS = frozenset({"gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx950"})
+
+
+def _use_int8_mm(device_type: str, device_index: int) -> bool:
+    mode = os.environ.get("COMFY_KITCHEN_INT8_EAGER", "auto").strip().lower()
+    if mode == "int_mm":
+        return True
+    if mode == "dequant":
+        return False
+    if mode == "auto":
+        return _use_int8_mm_auto(device_type, device_index)
+    raise ValueError("COMFY_KITCHEN_INT8_EAGER must be one of: auto, int_mm, dequant")
+
+
+@functools.lru_cache(maxsize=None)
+def _use_int8_mm_auto(device_type: str, device_index: int) -> bool:
+    if device_type != "cuda" or not torch.cuda.is_available():
+        return False
+    if getattr(torch.version, "hip", None) is None:
+        return True
+    try:
+        arch = torch.cuda.get_device_properties(device_index).gcnArchName.split(":")[0]
+    except Exception:
+        return False
+    return arch in _CDNA_ARCHS
+
+
 def int8_linear(
     x: torch.Tensor,
     weight: torch.Tensor,
@@ -983,6 +1013,15 @@ def int8_linear(
         h = _build_hadamard(convrot_groupsize, device=x.device, dtype=x.dtype)
         x = _rotate_activation(x, h, convrot_groupsize)
 
+    if not _use_int8_mm(x.device.type, x.device.index if x.device.index is not None else 0):
+        ws = weight_scale.float()
+        w = (weight.float() * (ws if ws.numel() == 1 else ws.reshape(-1, 1))).to(out_dtype)
+        return torch.nn.functional.linear(
+            x.to(out_dtype),
+            w,
+            None if bias is None else bias.to(device=x.device, dtype=out_dtype),
+        )
+
     orig_shape = x.shape
     x_2d = x.reshape(-1, x.shape[-1])