Attention Optimizer

Attention is the single most expensive thing your diffusion model does - 40-70% of total generation time on transformer-based models like SDXL, Flux, and WAN. And the fastest way to run it isn't universal: PyTorch's built-in SDPA, Flash Attention, SageAttention, and xFormers all ship different kernels, and which one wins depends on your GPU, your model's head dimension, and your resolution. That's the roulette. Attention Optimizer is the fix: it benchmarks every backend you've actually got installed, on your hardware, and quietly applies the winner.

It's a model_patches/optimization node, so it sits right between your checkpoint loader and your first KSampler. Give it your model, run once, and it clones the model, measures each available attention backend against synthetic tensors, and hooks the fastest one in. First run costs you 5-10 seconds; every run after that reads a cache file and applies instantly.

How it works

The node starts by detecting your model's head dimension (SD 1.5 is 64, SDXL/Flux are 128, LTX-V is 160), its hash, and its dtype. Then it benchmarks each backend with random q/k/v tensors at the sequence length you specify - warming up, timing, and synchronizing CUDA properly, not just eyeballing it. It even validates each backend by running a real attention call, so a library that's installed but silently broken doesn't get picked.

Results are cached in benchmark_db.json inside the pack's folder, keyed by model hash + head_dim + seq_len + num_heads. That cache is per-machine, and that's the whole point: a 4090 and a 3090 often disagree about which kernel is fastest. The chosen backend is applied through ComfyUI's optimized_attention_override on a cloned model - reversible, per-model, and it doesn't mutate your original checkpoint.

Inputs and outputs that matter

The only required input is model. The rest are optional and you can mostly leave them alone:

• attention_backend - auto (default) benchmarks and picks the best. Or force one: sage_auto, sage_fp8_cuda_fast, flash, xformers, the plain ComfyUI basic/sub_quad/split, and the rest. Forcing skips the benchmark entirely.
• force_refresh - set to True if you changed hardware or installed a new backend and want to re-benchmark despite the cache.
• seq_len / num_heads - benchmark parameters (defaults 8192 and 24). Longer seq_len matches video models, where the speedups get big.
• auto_apply - if you just want the benchmark numbers and no patch, flip this off.

Outputs: a patched model plus metadata - best_attention (the winning backend), speedup and time_ms, head_dim, impl_type, and a full text report for pasting into a Discord argument. The kjnodes_mode string is a nice touch: it maps the winner to the exact mode string KJNodes' PatchSageAttention node expects, in case you'd rather wire it through that.

Installing it

No hard dependencies - pyproject.toml ships empty, and the node works out of the box with PyTorch SDPA. The speedups come from optional extras:

# in ComfyUI/custom_nodes
git clone https://github.com/D-Ogi/ComfyUI-Attention-Optimizer.git
# then, for the real gains:
pip install sageattention   # recommended for RTX 30xx/40xx
pip install flash-attn      # Ampere+ alternative, painful to build on Windows
pip install xformers

Or ComfyUI Manager → "Install via Git URL" with the same URL, then restart. The author also supports comfy node install comfyui-attention-optimizer.

Gotchas worth knowing

• The patch is global. The node applies its backend to the whole model the moment it runs, so you only need one instance per workflow - and place it right before your first KSampler to enforce execution order, since ComfyUI only guarantees order through graph dependencies. The author spells this out in his r/comfyui thread, and it's the #1 way people think they need multiple copies.
• "Backend not available" in the report just means the Python package isn't installed. pip install the one you want and hit force_refresh.
• No speedup? Check auto_apply is on, watch the console for [Benchmark] Applied: X, and try force_refresh=True - the cache can be stale after a ComfyUI upgrade.
• SageAttention on Windows is the known pain point. The Triton path notoriously needs hand-rolled builds (woctordho's community wheels are the usual fix), and it's a recurring complaint in the WAN ecosystem docs. On Windows, prefer the CUDA variants or stick with SDPA.
• Some models aren't affected. SeedVR2 and friends ship their own attention and ignore the override entirely.

Honest take: on consumer NVIDIA cards with SageAttention installed, 1.5-2x on SDXL-class models is real and the benchmark numbers match what the community sees on WAN. The 2-4x figures are long-sequence video territory, so don't chase them on a single 1024x1024 still. AMD/ROCm users get the baseline and not much more. It's a young project (zero search impressions, one author), so treat the benchmark as ground truth for your machine and re-run after updates - but the caching makes that a non-event.