Patch Flash Attention DN

Ever set up Flash Attention, hit ComfyUI's --use-flash-attention startup flag, and quietly suspected it was doing nothing? You're probably right - the flag doesn't reliably force FA2 in every build, and that's the whole reason this node exists. Patch Flash Attention DN is a tiny inline node that patches your model's attention to run through Flash Attention 2 (flash_attn_func) and makes the patch explicit: on, off, and visible in the graph.

It sits between your loader and sampler and hands you back a patched MODEL - same idea as Patch Sage Attention KJ from KJNodes, which the author openly credits. If you're already running KJNodes' sage-attention patch, you don't need this. If you're on AMD (where sage attention is more of a fight) or you just want the FA2 path forced and confirmable, this is the one to reach for. Note the honesty in the README: it's untested on NVIDIA, "but should theoretically work." It's marked experimental, so treat it as a speed patch, not a correctness guarantee.

How it works

Under the hood it's almost embarrassingly simple - which is why it's reliable. The node clones the model and injects a wrapped flash_attn_func into the model's optimized_attention_override transformer option. That override is the same hook ComfyUI's built-in optimizers use, so instead of hoping the sampler picks FA2, you've replaced the attention call outright.

The wrapper handles the reshuffling between ComfyUI's (batch, heads, seq, dim_head) layout and Flash Attention's (batch, seq, heads, dim_head), casts any fp32 tensors to fp16 (and casts the output back so the rest of the graph never notices), and runs with causal=False. One real limitation: flash_attn_func doesn't take attention masks, so if your model feeds one in, it's dropped with a warning. Fine for most sampling runs, but don't reach for this node on masked workflows and expect identical results.

The inputs that matter

Two inputs, one output - this isn't a node you configure so much as place.

• model - the MODEL from your checkpoint loader or UNET loader. Wire the patched output into your KSampler; the whole point is that you connect model in and MODEL out, never the original.
• enabled (default True) - flip to False and the node just passes the model through unchanged. Handy for A/B testing whether FA2 is actually buying you anything on a given model.

Output is a single MODEL, exactly what your sampler expects.

Installing it (the real hurdle)

The node itself is a 30-second install - the dependency is the actual project:

cd ComfyUI/custom_nodes
git clone https://github.com/0xDELUXA/ComfyUI-DN_PatchFlashAttention

Restart ComfyUI and you'll find it under DN > FlashAttention. (It's in the ComfyUI Manager registry as "ComfyUI-DN_PatchFlashAttention" if you prefer that route.)

What this does not do is install flash_attn for you. That's on you, in the same Python environment ComfyUI runs in. On NVIDIA the shortest path is pip install flash_attn --no-build-isolation - or grab a prebuilt wheel, because compiling Flash Attention from source is a slow afternoon. On AMD, the README's path is building from source with FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE set, plus pip install triton-windows and a HIP_PATH env var pointing at your ROCm SDK on Windows. If the node errors with flash_attn is not installed, that's the checklist to go through.

Troubleshooting & gotchas

• The "flash_attn is not installed" RuntimeError - the single most common failure, and it means the dependency, not the node, is broken. See above.
• Windows users: set FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" before launching ComfyUI or FA2 silently won't engage with triton-windows. Community installs of Flash Attention on Windows are notoriously fiddly; if you're pulling your hair out, look for prebuilt-wheel helper installers before attempting a source build.
• Different results with masks - masks get ignored with a warning, so masked/attention-heavy workflows may differ slightly.
• fp32 models get silently cast to fp16 - expected, that's the design, and output is cast back so downstream nodes don't care.

Tested by the author on Flux, Qwen, and SDXL; it'll generally work wherever your sampler would otherwise use default attention. If you're on AMD and tired of watching VRAM burn on the stock attention path, this is one of the few flash-attention patches that's actually built around the AMD experience.