resolve rebase conflict

This commit is contained in:
JWLHS 2026-07-13 21:38:42 +08:00
parent 015e8a2676
commit 4c36e2d6f2
2 changed files with 19 additions and 34 deletions

View File

@ -3,22 +3,6 @@ import logging
from comfy.cli_args import args
def _rocm_kitchen_arch_supported():
"""comfy-kitchen's INT8 Triton kernels compile tl.dot to matrix-core instructions.
RDNA3/3.5/4 (gfx11xx/gfx12xx) have WMMA and CDNA (gfx9xx) has MFMA; RDNA1/RDNA2
(gfx10xx) have neither, so the INT8 path hangs the GPU there. Gates the automatic
ROCm default so those cards stay on the eager fallback (an explicit
--enable-triton-backend still forces it on any arch)."""
try:
arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0]
except Exception:
return False
if arch.startswith(("gfx11", "gfx12")):
return True
return arch in ("gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx950")
try:
import comfy_kitchen as ck
from comfy_kitchen.tensor import (
@ -42,13 +26,9 @@ try:
logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
# 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 AND a
# matrix-core GPU (RDNA3+ WMMA gfx11xx/gfx12xx, CDNA MFMA gfx9xx). RDNA1/RDNA2
# (gfx10xx) have no WMMA -> the INT8 tl.dot path hangs the GPU, so they stay eager.
# 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.disable_triton_backend:
ck.registry.disable("triton")
elif args.enable_triton_backend: # or (torch.version.hip is not None and _rocm_kitchen_arch_supported()):
if args.enable_triton_backend or torch.version.hip is not None:
try:
import triton
triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2])

View File

@ -1,19 +1,24 @@
"""
ComfyUI TorchAO INT4 weight-only quantization backend.
W4A16 scheme: weights quantized to INT4 via torchao,
activations remain FP16. Supports CUDA, Intel XPU, and CPU.
Includes:
- TINT4Linear: INT4 linear layer with LoRA forward injection + QuaRot
- quantize_model: torchao INT4 quantization entry point
- reconstruct_int4_state_dict: rebuild TINT4Linear from safetensors
- build_hadamard / rotate_weight: QuaRot Hadamard rotation
...
"""
import torch
from .linear import TINT4Linear
from .quantize import quantize_model, reconstruct_int4_state_dict
from .quarot import build_hadamard, rotate_weight
_AVAILABLE = False
try:
import torchao
from .linear import TINT4Linear
from .quantize import quantize_model, reconstruct_int4_state_dict
from .quarot import build_hadamard, rotate_weight
_AVAILABLE = True
except ImportError:
pass
if not _AVAILABLE:
raise ImportError(
"torchao is required for INT4 quantization. "
"Install: pip install torchao>=0.17.0"
)
__all__ = [
"TINT4Linear",