From 5d14e0c90197c84009a235a0547821ef7fa93ea8 Mon Sep 17 00:00:00 2001
From: Haoming <hmstudy02@gmail.com>
Date: Mon, 1 Jun 2026 12:05:25 +0800
Subject: [PATCH] triton

---
 backend/args.py      | 1 +
 backend/quant_ops.py | 5 +++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/backend/args.py b/backend/args.py
index f86af015..917f90e3 100644
--- a/backend/args.py
+++ b/backend/args.py
@@ -103,6 +103,7 @@ parser.add_argument("--mmap-torch-files", action="store_true", help="Use mmap wh
 parser.add_argument("--disable-mmap", action="store_true", help="Don't use mmap when loading safetensors")
 
 parser.add_argument("--tiled-conv2d", type=int, default=0, metavar="TILE_SIZE", choices=[0, 64, 128, 256, 512], help="reduce VAE memory usage ; increase processing time")
+parser.add_argument("--enable-triton-backend", action="store_true", help="Enable the use of Triton backend in comfy-kitchen")
 
 
 class SageAttentionFuncs(enum.Enum):
diff --git a/backend/quant_ops.py b/backend/quant_ops.py
index d29f25c4..c3f8236e 100644
--- a/backend/quant_ops.py
+++ b/backend/quant_ops.py
@@ -22,9 +22,10 @@ else:
     if cuda_version < (13,):
         ck.registry.disable("cuda")
 
+from backend.args import args
 
-# https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/quant_ops.py#L24
-ck.registry.disable("triton")
+if not args.enable_triton_backend:
+    ck.registry.disable("triton")
 
 
 # region FP8 Layouts