SmartAttentionDispatcher

## Smart Attention Dispatcher Patches the model to use SageAttention kernels instead of PyTorch SDPA. Detects GPU architecture, installed libraries, and selects the correct kernel automatically. --- ### Inputs | Pin | Default | Description | |---|---|---| | `model` | — | Any ComfyUI model (Flux, SD3.5, SDXL, Qwen, ErnieImage, Z-Image, …). | | `sdpa_kernel` | False | Force PyTorch SDPA. Overrides all SA2/SA3 settings. Use to compare output or disable SA entirely. | | `sa2_kernel` | disable | SA2 kernel selection — see table below. | | `combine` | False | Dynamic mode: SA2 on boundary steps (first/last), SA3 on middle steps. Requires both sa2_kernel ≠ disable and sa3_kernel ≠ disable. | | `sa3_kernel` | disable | SA3 kernel selection — see table below. | **sa2_kernel options:** | Value | Description | |---|---| | `disable` | SA2 off. | | `auto` | Best kernel for the detected GPU: fp8 on Ada/Hopper/Blackwell, fp16 on Turing/Ampere. | | `fp16` | `sageattn_qk_int8_pv_fp16_cuda` — fp32 accumulator. Turing / Ampere. Blocked on Blackwell (SM120+). | | `fp8` | `sageattn_qk_int8_pv_fp8_cuda` — fp32+fp32 accumulator. Ada and newer. Bit-exact with SDPA on tested models. | | `fp8++` | `sageattn_qk_int8_pv_fp8_cuda` — fp32+fp16 accumulator. Slightly faster but results differ from SDPA. | | `triton` | `sageattn_qk_int8_pv_fp16_triton` — Triton fallback, all GPUs. | **sa3_kernel options:** | Value | Description | |---|---| | `disable` | SA3 off. | | `standard` | `sageattn3_blackwell` — FP4, Blackwell only (SM≥100, CUDA≥12.8, Python≥3.10). | | `per_block_mean` | Same kernel with `per_block_mean=True` — slightly different numerics, marginally faster in some cases. | --- ### Outputs | Pin | Type | Description | |---|---|---| | `model` | MODEL | Patched model with attention override applied. | --- ### Mode display (node status panel) The node shows the active mode after each run: | Display | Meaning | |---|---| | `SDPA` | PyTorch SDPA — baseline, no SA active. | | `SA2` | SageAttention2 active on all steps. | | `SA3` | SageAttention3 active on all steps. | | `SA2-SA3-SA2` | Dynamic: SA2 on first/last step, SA3 on middle steps. | | `SDPA-SA3-SDPA` | Dynamic: SDPA on first/last step, SA3 on middle steps. | | `SA3 (not installed) >>> SA2` | Fallback — reason shown in parentheses. | --- ### Tested models | Model | SA2 | SA3 | Notes | |---|---|---|---| | Flux.1 / Flux.2 / Flux.2 Klein | ✅ | ✅ | — | | SD3.5 | ✅ | ✅ | Cross-attention layers auto-fallback to SDPA. | | Z-Image (Lumina2) | ✅ | ✅ | — | | SDXL | ✅ | — | SA3 not tested on UNet. No speed gain observed. | | ErnieImage | ✅ | ✅ | Requires global sys.modules patch (applied automatically). | | Qwen-Image / Qwen-Edit | ✅ | ⚠️ | SA3 numerically unstable at long sequences (seq > 7000). | | LTX / Wan / HunyuanVideo | — | — | Not yet tested. | --- ### Performance notes - On **RTX 50xx (Blackwell) with PyTorch 2.8+ / CUDA 13.0** — no measurable speed gain over baseline SDPA. PyTorch SDPA is already optimized for SM120. SA kernels add overhead without benefit. - On **RTX 30xx / 40xx** — SA2 gives real speed improvement, especially at long sequences (Flux, SD3.5, Qwen). - SA3 on Qwen (seq 7000–14000) produces unpredictable results due to FP4 quantization error accumulation. --- ### Compatibility notes - Do **not** use `--use-sage-attention` launch flag together with this node. ComfyUI patches attention before the node loads — the node will restore the wrong baseline on SDPA mode. - `--fast` flag (fp16_accumulation, fp8_matrix_mult, etc.) is safe to use alongside this node. - Attention masks (inpainting, outpainting) automatically fall back to SDPA — SA2/SA3 do not support arbitrary masks.