diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 0d7df5e13..e2e0d97ec 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -92,6 +92,7 @@ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE" parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.") parser.add_argument("--supports-fp8-compute", action="store_true", help="ComfyUI will act like if the device supports fp8 compute.") parser.add_argument("--enable-triton-backend", action="store_true", help="ComfyUI will enable the use of Triton backend in comfy-kitchen. Is disabled at launch by default.") +parser.add_argument("--disable-triton-backend", action="store_true", help="Force-disable the comfy-kitchen Triton backend, overriding the automatic ROCm/AMD default and --enable-triton-backend.") class LatentPreviewMethod(enum.Enum): NoPreviews = "none" diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py index 91b3e4fe9..b1aabdc93 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -3,6 +3,22 @@ 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 ( @@ -26,9 +42,13 @@ 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: + # 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. # 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: + 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()): try: import triton triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2])