mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 20:38:15 +08:00
resolve rebase conflict
This commit is contained in:
parent
015e8a2676
commit
4c36e2d6f2
@ -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])
|
||||
|
||||
@ -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",
|
||||
|
||||
Loading…
Reference in New Issue
Block a user