ComfyUI Attention Selector

Simple attention backend selector for ComfyUI.

Supported Backends

• basic: Basic einsum-based attention
• sub_quad: Sub-quadratic attention (memory efficient)
• split: Split attention (for low VRAM)
• pytorch: PyTorch's scaled_dot_product_attention (SDPA)
• xformers: xFormers memory efficient attention
• sage: SageAttention (latest v2.x / GitHub version, supports HND/NHD layouts, smooth_k, pv_accum_dtype)
• flash: Flash Attention (requires flash-attn package)

SageAttention Update (v2.x / Latest from GitHub)

The old SageAttention APIs (sage_auto, sage_cuda, sage_triton, sage_fp8_*, sage3) have been deprecated in favor of the unified sageattn() API which automatically selects the optimal implementation based on your GPU compute capability and tensor shapes.

The latest SageAttention (v2.x / GitHub master) provides:

Multiple Implementations (Auto-selected by GPU):

• sageattn: Unified API that auto-selects the best implementation
• sageattn_qk_int8_pv_fp16_cuda: SM80 GPUs (A100, A30, A40, RTX 3090/3080)
• sageattn_qk_int8_pv_fp16_triton: SM86 GPUs (some laptop/datacenter GPUs)
• sageattn_qk_int8_pv_fp8_cuda: SM89 GPUs (RTX 4090/4080/4070)
• sageattn_qk_int8_pv_fp8_cuda_sm90: SM90 GPUs (H100, H200)
• sageattn_qk_int8_pv_fp8_cuda: SM120/SM121 GPUs (B200, GB200)

Key Features:

• INT8 quantization for Q and K matrices (reduced memory bandwidth)
• FP8/FP16 PV accumulation for improved performance
• smooth_k option for better accuracy (subtracts key mean)
• Multiple tensor layouts: HND and NHD support
• Flexible head dimensions: Supports 64, 96, 128
• Multiple dtypes: float16, bfloat16
• pv_accum_dtype options:
  • "fp32": Full FP32 accumulation (slowest, most accurate)
  • "fp32+fp32": FP32 with periodic buffer flush (balanced)
  • "fp32+fp16": FP16 accumulation added to FP32 buffer (fastest)

Installation (Latest from GitHub):

# Install latest version from GitHub (recommended for v2.x features)
pip install git+https://github.com/thu-ml/SageAttention.git

# Or install from PyPI (may be older version)
pip install sageattention

Note: Building from source requires CUDA toolkit and a compatible GPU.

Installation

git clone https://github.com/huagusam/comfyui-attention-selector.git

Installing Flash Attention

Flash Attention requires CUDA to be installed and a compatible GPU:

pip install flash-attn --no-build-isolation

Note: Flash Attention must be built from source and requires CUDA development tools.

Usage

Add the "Attention Selector" node to your workflow and select the desired backend from the dropdown menu. The node will automatically detect available backends and fall back to PyTorch if the selected backend is not available.

License

MIT License - see LICENSE

Credits

• SageAttention - THU-ML
• Flash Attention - Dao-AILab
• xFormers - Meta
• ComfyUI Attention Optimizer - D-Ogi

https://github.com/huagusam/comfyui-attention-selector.git

## License

MIT License - see [LICENSE]

## Credits

- [SageAttention](https://github.com/thu-ml/SageAttention) - THU-ML

- [Flash Attention](https://github.com/Dao-AILab/flash-attention) - Dao-AILab

- [xFormers](https://github.com/facebookresearch/xformers) - Meta

- [ComfyUI Attention Optimizer](https://github.com/D-Ogi/ComfyUI-Attention-Optimizer) - D-Ogi