Archive of fixed ComfyUI for Flash Attention-3 & Sage-Attention
*As of August 19, 2025, the following article cannot be applied to RTX5060Ti 16GB + Pytorch 2.8.0 + cu129 or later environments, but it is being archived as information based on an RTX4070 12GB + Pytorch 2.7.1 + cu128 environment.
Introduction
Following the A1111-based WebUI, I have implemented the xformers kernel analysis function and the SA load-time log display function for ComfyUI as well.
In the A1111-based WebUI, everything related to Adetailer was running on Cutlass, but in ComfyUI, it can be seen that it is running on FA-3.
On the other hand, Ultimate Upscaler seems to run on Cutlass.
Objective
To reliably reproduce the final state every time: from the official state (no logs), use xformers (FA-3 priority) when the SA node is not used, and SageAttention when the SA node is used.
Only 3 files are changed. No unnecessary modifications.
Target files (3)
comfy\ldm\modules\attention.py
python_embeded\Lib\site-packages\xformers\ops\fmha\_init_.py
custom_nodes\ComfyUI-KJNodes\nodes\model_optimization_nodes.py
Code Modification
attention.py (change priority only)
Objective: Even if launched with --use-sage-attention, prioritize xformers (FA-3) unless KJ's SA node is used.
Change location: The if-elif chain that determines optimized_attention (replace only the relevant block)
[Official]
if model_management.sage_attention_enabled():
logging.info("Using sage attention")
optimized_attention = attention_sage
elif model_management.xformers_enabled():
logging.info("Using xformers attention")
optimized_attention = attention_xformers
elif model_management.flash_attention_enabled():
logging.info("Using Flash Attention")
optimized_attention = attention_flash
elif model_management.pytorch_attention_enabled():
logging.info("Using pytorch attention")
optimized_attention = attention_pytorch[Final State]
if model_management.xformers_enabled():
logging.info("Using xformers attention")
optimized_attention = attention_xformers
elif model_management.sage_attention_enabled():
logging.info("Using sage attention")
optimized_attention = attention_sage
elif model_management.flash_attention_enabled():
logging.info("Using Flash Attention")
optimized_attention = attention_flash
elif model_management.pytorch_attention_enabled():
logging.info("Using pytorch attention")
optimized_attention = attention_pytorchSupplementary Note
SDPA (PyTorch attention) has not been disabled. It is still selected normally depending on the conditions.
Only change this block. Do not touch anything else.
xformers/ops/fmha/init.py (add a log once immediately after dispatch)
Purpose: Output logger.info only once, specifically when the kernel selection (FA-3 / cutlass) switches.
Change location: Insert the following immediately after op = _dispatch_fw(inp, False) within the function _memory_efficient_attention_forward(...).
(Do not insert into explicit op paths or the _requires_grad side. Do not add extra handlers, change propagation, or call _set_use_fa3(True).)
[Insertion block]
try:
import logging
logger = logging.getLogger("xformers_attention_log")
if not hasattr(_memory_efficient_attention_forward, "_last_kernel"):
_memory_efficient_attention_forward._last_kernel = None
last_kernel = _memory_efficient_attention_forward._last_kernel
if getattr(op, "NAME", None) != last_kernel:
logger.info(f"[xformers] memory_efficient_attention: selected kernel = {op.NAME}")
_memory_efficient_attention_forward._last_kernel = getattr(op, "NAME", None)
except Exception:
print(f"[xformers] memory_efficient_attention: selected kernel = {getattr(op, 'NAME', str(op))}")Overall image (relevant parts only)
def _memory_efficient_attention_forward(inp: Inputs, op: Optional[Type[AttentionFwOpBase]]) -> torch.Tensor:
inp.validate_inputs()
output_shape = inp.normalize_bmhk()
if op is None:
op = _dispatch_fw(inp, False)
# ←ここに上記ブロックを挿入
else:
_ensure_op_supports_or_raise(ValueError, "memory_efficient_attention", op, inp)
out, *_ = op.apply(inp, needs_gradient=False)
return out.reshape(output_shape)Supplementary notes
logger.info method only (assumes INFO is output to the console via environment logger settings).
To output only during switching, hold the most recent name in the function attribute _last_kernel to suppress duplicates.
KJNodes/nodes/model_optimization_nodes.py (log once per generation)
Purpose: When using SA nodes, output a log for SA usage only once at the beginning of a generation (do not output at all if SA nodes are not used).
Addition (add if it does not exist. No need to edit if already present):
A) Flag (class initialization)
self._called_this_generation = False
self._called_sageattn_this_generation = FalseB) Only the first time for auto branching (when using SA nodes)
if sage_attention == "auto":
def func(q, k, v, is_causal=False, attn_mask=None, tensor_layout="NHD"):
if not self._called_sageattn_this_generation:
print("[SageAttention][DEBUG] sageattn (auto) called")
self._called_sageattn_this_generation = True
return sageattn(q, k, v, is_causal=is_causal, attn_mask=attn_mask, tensor_layout=tensor_layout)
return funcC) Only once per generation (when using SA nodes)
def attention_sage(...):
if not self._called_this_generation:
print("SageAttention kernel is being used for this generation.")
self._called_this_generation = True
...D) Reset at the start of a generation
def reset_attention_sage_flag(*args, **kwargs):
self._called_this_generation = False
self._called_sageattn_this_generation = False
model_clone.add_callback(CallbacksMP.ON_PRE_RUN, reset_attention_sage_flag)Supplementary notes
Start/end of replacement is handled by existing ON_PRE_RUN (replacement) and ON_CLEANUP (reversion).
When SA nodes are not used, the replacement itself does not occur = attention is xformers, and logs are only on the xformers side.
Operation check (expected logs)
SA node unused / normal generation:
FA-3 condition: Once at the beginning
[xformers] memory_efficient_attention: selected kernel = fa3F@...Unsupported conditions (mask/tile, etc.): Once upon switching
[xformers] memory_efficient_attention: selected kernel = cutlassF-pt
Using SA nodes:
Once at the beginning of every generation
SageAttention kernel is being used for this generation.Only on the first time 'auto' is specified
[SageAttention][DEBUG] sageattn (auto) called
Rollback
attention.py: Revert the if-elif chain to the original order (Sage → xformers → Flash → PyTorch)
xformers/init.py: Delete (or comment out) the inserted block
KJNodes: Delete (or comment out) the above A-D
That is all. No other changes are necessary. This allows you to reliably reproduce the final form from the official version every time.
Fixed CCSR
When applying the above, a bug occurred where FA-3 load logs were repeated during CCSR loading, so I fixed it.
Modified files
ComfyUI-CCSR xformers log control fix details
Overview
Fix to suppress continuous xformers logs during CCSR node execution and allow the progress bar to function normally
Modified file
`ComfyUI/custom_nodes/comfyui-ccsr/nodes.py`
Added imports
import sys
import loggingNewly added class: XFormersKernelOnce
Class Structure
class XFormersKernelOnce:
"""ロガー＋stdout/errを収集モードにし、終了時に1行だけ出す"""
def __init__(self):
self._filter = self._XFormersKernelFilter()
self._loggers = []
self._saved_out = None
self._saved_err = None
self._proxy_out = None
self._proxy_err = None
self._agg = self._KernelAggregator()Internal Class 1: Kernel Aggregator
class _KernelAggregator:
"""カーネル選択を記録・集約"""
def __init__(self):
self._kernels = []
self._fa2_seen = False
def record(self, kernel):
if "fa2" in kernel.lower():
self._fa2_seen = True
self._kernels.append(kernel)
def selected(self):
if self._fa2_seen:
# FA2が含まれていればFA2を優先
for k in self._kernels:
if "fa2" in k.lower():
return k
# 最後に観測したカーネルを返す
return self._kernels[-1] if self._kernels else NoneInternal Class 2: Logger Filter
class _XFormersKernelFilter(logging.Filter):
"""xformersのカーネル選択ログを捕捉して抑止"""
def filter(self, record):
try:
msg = record.getMessage()
except Exception:
return True
if "memory_efficient_attention: selected kernel" in msg:
return False # ログを抑止
return TrueInternal Class 3: Standard Output Proxy
class _StdoutProxy:
"""stdout/stderrをラップして対象行のみ収集して抑止"""
def __init__(self, underlying, agg):
self._u = underlying
self._agg = agg
def write(self, s):
try:
text = str(s)
except Exception:
text = s
if "memory_efficient_attention: selected kernel" in text:
# カーネル名を抽出
if "=" in text:
kernel = text.split("=")[-1].strip()
self._agg.record(kernel)
return len(s) # 書き込み長を返す（進行バー等の整合性を維持）
return self._u.write(s)
def flush(self):
return self._u.flush()
# 以降は透過委譲
def fileno(self): return self._u.fileno() if hasattr(self._u, "fileno") else 1
def isatty(self): return self._u.isatty() if hasattr(self._u, "isatty") else False
def readable(self): return self._u.readable() if hasattr(self._u, "readable") else False
def writable(self): return self._u.writable() if hasattr(self._u, "writable") else True
def seekable(self): return self._u.seekable() if hasattr(self._u, "seekable") else False
@property
def encoding(self): return getattr(self._u, "encoding", "utf-8")
@property
def errors(self): return getattr(self._u, "errors", None)
@property
def buffer(self): return getattr(self._u, "buffer", None)
def __getattr__(self, name): return getattr(self._u, name)At Context Manager Start
def __enter__(self):
# よく使われるロガーに一括装着
self._loggers = [
logging.getLogger(),
logging.getLogger("xformers"),
logging.getLogger("xformers.ops"),
logging.getLogger("xformers.ops.fmha"),
logging.getLogger("xformers_attention_log"),
]
for lg in self._loggers:
try:
lg.addFilter(self._filter)
except Exception:
pass
# stdout/errをプロキシに差し替え（より強力に）
self._saved_out, self._saved_err = sys.stdout, sys.stderr
self._proxy_out = self._StdoutProxy(self._saved_out, self._agg)
self._proxy_err = self._StdoutProxy(self._saved_err, self._agg)
# グローバルに設定
sys.stdout = self._proxy_out
sys.stderr = self._proxy_err
# さらに、xformersの内部ロガーも制御
try:
import xformers.ops.fmha
if hasattr(xformers.ops.fmha, '_memory_efficient_attention_forward'):
# 元の関数を保存
if not hasattr(xformers.ops.fmha, '_original_memory_efficient_attention_forward'):
xformers.ops.fmha._original_memory_efficient_attention_forward = xformers.ops.fmha._memory_efficient_attention_forward
# ログ出力を無効化した関数で置き換え
def _silent_memory_efficient_attention_forward(inp, op=None):
inp.validate_inputs()
output_shape = inp.normalize_bmhk()
if op is None:
op = xformers.ops.fmha._dispatch_fw(inp, False)
# ログ出力を無効化
if not hasattr(_silent_memory_efficient_attention_forward, "_last_kernel"):
_silent_memory_efficient_attention_forward._last_kernel = None
last_kernel = _silent_memory_efficient_attention_forward._last_kernel
current = getattr(op, "NAME", str(op))
if current != last_kernel:
# ログ出力を無効化し、集約器に記録のみ
if "fa2" in str(current).lower():
self._agg._fa2_seen = True
self._agg._kernels.append(str(current))
_silent_memory_efficient_attention_forward._last_kernel = current
else:
xformers.ops.fmha._ensure_op_supports_or_raise(ValueError, "memory_efficient_attention", op, inp)
out, *_ = op.apply(inp, needs_gradient=False)
return out.reshape(output_shape)
xformers.ops.fmha._memory_efficient_attention_forward = _silent_memory_efficient_attention_forward
# requires_grad版もパッチ
if hasattr(xformers.ops.fmha, '_memory_efficient_attention_forward_requires_grad'):
if not hasattr(xformers.ops.fmha, '_original_memory_efficient_attention_forward_requires_grad'):
xformers.ops.fmha._original_memory_efficient_attention_forward_requires_grad = xformers.ops.fmha._memory_efficient_attention_forward_requires_grad
def _silent_memory_efficient_attention_forward_requires_grad(inp, op=None):
inp.validate_inputs()
output_shape = inp.normalize_bmhk()
if op is None:
op = xformers.ops.fmha._dispatch_fw(inp, True)
# ログ出力を無効化
if not hasattr(_silent_memory_efficient_attention_forward_requires_grad, "_last_kernel"):
_silent_memory_efficient_attention_forward_requires_grad._last_kernel = None
last_kernel = _silent_memory_efficient_attention_forward_requires_grad._last_kernel
current = getattr(op, "NAME", str(op))
if current != last_kernel:
# ログ出力を無効化し、集約器に記録のみ
if "fa2" in str(current).lower():
self._agg._fa2_seen = True
self._agg._kernels.append(str(current))
_silent_memory_efficient_attention_forward_requires_grad._last_kernel = current
else:
xformers.ops.fmha._ensure_op_supports_or_raise(ValueError, "memory_efficient_attention", op, inp)
out = op.apply(inp, needs_gradient=True)
assert out[1] is not None
return (out[0].reshape(output_shape), out[1])
xformers.ops.fmha._memory_efficient_attention_forward_requires_grad = _silent_memory_efficient_attention_forward_requires_grad
except Exception as e:
print(f"Warning: Could not patch xformers: {e}")
return selfAt Context Manager End
def __exit__(self, exc_type, exc_val, exc_tb):
# まず復元（副作用を残さない）
if self._saved_out is not None:
sys.stdout = self._saved_out
if self._saved_err is not None:
sys.stderr = self._saved_err
for lg in self._loggers:
try:
lg.removeFilter(self._filter)
except Exception:
pass
# xformersのパッチを元に戻す
try:
import xformers.ops.fmha
if hasattr(xformers.ops.fmha, '_original_memory_efficient_attention_forward'):
xformers.ops.fmha._memory_efficient_attention_forward = xformers.ops.fmha._original_memory_efficient_attention_forward
if hasattr(xformers.ops.fmha, '_original_memory_efficient_attention_forward_requires_grad'):
xformers.ops.fmha._memory_efficient_attention_forward_requires_grad = xformers.ops.fmha._original_memory_efficient_attention_forward_requires_grad
except Exception:
pass
# 観測結果から1行だけ表示（FA2優先→なければ最後に観測）
selected = self._agg.selected()
if selected:
print(f"[xformers] memory_efficient_attention: selected kernel = {selected} (CCSR)")Application Points in CCSR Nodes
Before Fix
autocast_condition = dtype == torch.float16 and not mm.is_device_mps(device)
with torch.autocast(mm.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext():
for i in range(B):
# ... サンプリング処理After Fix
autocast_condition = dtype == torch.float16 and not mm.is_device_mps(device)
# xformersのログを制御（CCSRブロック内限定）
with XFormersKernelOnce():
with torch.autocast(mm.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext():
for i in range(B):
# ... サンプリング処理Technical Features
1. Multiple Control
Logger Filter: Controls logs via `logging`
Standard Output Proxy: Directly controls `print` statements
Function Patching: Replaces internal xformers implementation
2. Safety
No Side Effects: Always restores original state after processing completes
Exception Handling: Safely handles errors occurring at each stage
Progress Bar Protection: Maintains progress bar integrity by returning appropriate output length
3. Efficiency
FA2 Priority: Prioritizes recording when Flash Attention 2 is used
Deduplication: Do not record if the same kernel is selected consecutively
Aggregated output: Output as a single line after processing is complete
Effect
Continuous xformers logs during CCSR node execution are completely suppressed
Progress bar functions normally
Finally, output only one line: "selected kernel = [kernel name] (CCSR)"