Automatically benchmark and optimize the attention mechanism in diffusion models for maximum generation speed.

Modern diffusion models (SDXL, Flux, WAN, LTX-V, Hunyuan Video) are based on transformer architecture. The core operation - attention - computes relationships between all elements in the image/video latent space. This is:

• The most expensive operation - attention takes 40-70% of total generation time
• O(n²) complexity - cost grows quadratically with resolution/frames
• GPU-dependent - different GPUs perform best with different implementations

Multiple optimized attention backends exist:

• PyTorch SDPA - built-in, always available
• Flash Attention - CUDA kernels, memory efficient
• SageAttention - INT8 quantization, up to 2-4x faster
• xFormers - memory efficient attention

But which one is fastest for YOUR specific GPU and model?

This plugin benchmarks all available backends and automatically applies the fastest one.

Tested on RTX 4090 with head_dim=128 (SDXL, Flux):

Result: 1.9x faster generation just by switching attention backend.

For video models (WAN, Hunyuan) with longer sequences, speedups can reach 2-4x.

1. Open ComfyUI Manager
2. Click "Install via Git URL"
3. Paste: https://github.com/D-Ogi/ComfyUI-Attention-Optimizer.git
4. Restart ComfyUI

cd ComfyUI/custom_nodes
git clone https://github.com/D-Ogi/ComfyUI-Attention-Optimizer.git

Restart ComfyUI.

The plugin works out-of-the-box with PyTorch SDPA. For better performance, install additional backends:

# SageAttention - recommended for RTX 30xx/40xx (1.5-2x speedup)
pip install sageattention

# Flash Attention - alternative for Ampere+ GPUs
pip install flash-attn

# xFormers - memory efficient option
pip install xformers

Note: On Windows, Flash Attention requires building from source or using prebuilt wheels. SageAttention is easier to install and often faster on consumer GPUs.

1. Add "Attention Optimizer" node to your workflow (category: model_patches)
2. Connect your model to the model input
3. Run - it benchmarks once, caches results, and auto-applies the fastest backend

┌─────────────────┐     ┌──────────────────────────┐     ┌─────────────┐
│ Load Checkpoint │────▶│ Attention Optimizer      │────▶│ KSampler    │
└─────────────────┘     │                          │     └─────────────┘
                        │ 1. Detect model params   │
                        │ 2. Check cache           │
                        │ 3. Benchmark (if needed) │
                        │ 4. Clone model & apply   │
                        │    attention override    │
                        └──────────────────────────┘

First run: Benchmarks all backends (~5-10 seconds), saves to cache. Subsequent runs: Loads from cache (instant), applies optimal backend.

=================================================================
BENCHMARK REPORT
=================================================================
dtype: float16 | head_dim: 128 | seq_len: 8192 | CUDA: 12.4 | Triton: 3.0.0
SageAttention: v2.1.1

>>> BEST: sage_fp8_cuda_fast (1.89x speedup) <<<
    impl: cuda | kjnodes mode: sageattn_qk_int8_pv_fp8_cuda++

Results (fastest first):
-----------------------------------------------------------------
 [v] sage_fp8_cuda_fast       2.671ms   1.89x  (cuda) <<<
 [v] sage_auto                2.679ms   1.88x  (auto)
 [v] sage_fp8_cuda            3.100ms   1.63x  (cuda)
 [v] sage_triton              3.446ms   1.47x  (triton)
 [v] sage_cuda                3.947ms   1.28x  (cuda)
 [v] pytorch                  5.049ms   1.00x  (pytorch)
 [v] xformers                 5.194ms   0.97x  (cuda/triton)
 [v] flash                    5.430ms   0.93x  (cuda)
 [ ] sage3                    ---       (N/A) Not installed
-----------------------------------------------------------------
[v] = validated (tested underlying library directly)
=================================================================

PyTorch SDPA uses cuDNN/cuBLAS - general purpose, always works.

Flash Attention fuses operations into single CUDA kernel, reducing memory bandwidth. Great for long sequences.

SageAttention quantizes Q/K to INT8, reducing memory and compute. Works best for head_dim ≤ 128.

xFormers similar to Flash Attention, good memory efficiency.

Models have different attention head dimensions:

• SD 1.5: head_dim=64
• SDXL, Flux: head_dim=128
• LTX-V: head_dim=160

SageAttention works best with head_dim ≤ 128. For larger dimensions, SDPA or Flash Attention may be faster.

Benchmark results are cached in benchmark_db.json based on:

• Model hash (architecture + weights)
• head_dim
• seq_len / num_heads parameters

Cache is per-machine - different GPUs will have different optimal backends.

Install the missing package:

pip install sageattention  # for sage_*
pip install flash-attn     # for flash
pip install xformers       # for xformers

1. Check if auto_apply is enabled
2. Try force_refresh=True to re-benchmark
3. Check console for [Benchmark] Applied: X message

Some models (like SeedVR2) use their own attention implementation and won't be affected by this plugin. Check the compatibility table above.

MIT License - see LICENSE

• SageAttention - THU-ML
• Flash Attention - Dao-AILab
• xFormers - Meta