Patch Flash Attention KJ

Here's the subtle problem this solves. ComfyUI's default attention will try to use an efficient kernel, but if it can't, it quietly falls back to standard SDPA attention and never tells you. So you think you're running FlashAttention, you're not, and you're leaving speed and memory on the table without a clue. PatchFlashAttentionKJ patches the model to use FlashAttention explicitly, without that silent fallback. If it can't use flash attention, you find out, instead of unknowingly running the slow path.

It's an experimental node in kijai's KJNodes. Small, sharp, and honest, that "no silent fallback" behavior is the entire pitch.

How it works

It patches the attention of whatever model passes through it to use FlashAttention. The patch travels with the model, so anything downstream of this node runs with it; to turn it off you bypass or disconnect the node. Model in, patched model out.

The one dependency it can't fake: it requires the flash_attn library to be installed. No library, no patch, and unlike ComfyUI's default this node won't paper over that for you.

The inputs and outputs that matter

• model - the model to patch. Patched model comes out.
• allow_compile (optional, default off) - whether to permit compilation in the patched path. Leave it off unless you specifically want it and have a working compile toolchain; on the wrong setup, turning it on is a way to trade one problem for another.

Output is MODEL, wired onward into your sampler. That's the whole surface, this node does one thing.

Installing it

KJNodes is the easy half; flash_attn is the real work.

• KJNodes: ComfyUI Manager → search KJNodes for ComfyUI → install → restart. Or git clone https://github.com/kijai/ComfyUI-KJNodes into custom_nodes and pip install -r requirements.txt.
• flash_attn: this is not in KJNodes' requirements and is notoriously annoying to install. It typically compiles from source against your exact CUDA/PyTorch versions, which can take a long time and fail on mismatches; many people chase down a prebuilt wheel matching their setup instead. The same reputation dogs this whole family of optimized-attention libraries (SageAttention included) where the running joke is that nobody can get them installed on the first try. Budget time for it, or run somewhere the wheel is already present.

Common issues

ModuleNotFoundError: No module named 'flash_attn' (or the node errors on use). The library isn't installed in ComfyUI's Python. That's the expected failure when flash_attn is missing, and it's by design that this node surfaces it rather than falling back silently. Install flash_attn into the same environment ComfyUI runs (portable build: python_embeded\python.exe -m pip ...), then restart.

Install won't build. flash_attn is version-sensitive to CUDA and PyTorch. A source build fails when those don't line up; the pragmatic path is a precompiled wheel that matches your Python/CUDA/torch trio exactly.

No speedup. Flash attention's benefit depends on your GPU and the model. On hardware or a model where it doesn't apply cleanly, you won't see much; the value of this node over the default is mostly the honesty (you know what you're actually running), not a guaranteed win on every card.

allow_compile broke something. Turn it back off. It opts into a compilation path that needs a working toolchain; the default-off is the safe state.