Sage Attention 3 public release and need to add a node to select attention type on steps
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 155
Description
### Feature Idea
After multiple months, Sage Attention 3 has been released to the public at https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell.
It employs the hardware-level FP4 cores in Blackwell series GPUs, enabling the attention to proceed x5 times faster than the fastest FlashAttention2 on these GPUs. Comparably to SageAttention2 it's around a x2 times boost.
The caveat is that this precision causes artifacts, especially in the MoE Wan2.2 model (where caching methods such as MagCache and EasyCache also give not great results). Likely it happens because of the MoE timesteps sensitivity, so a Attention Type steps node will probably be an essential addition
### Existing Solutions
There is this Pull Request, although it is dated and also needs renaming of the imports https://github.com/comfyanonymous/ComfyUI/pull/9047. Additionally, Kijai has already implemented Sage Attention 3 in his WanWrapper nodes https://github.com/kijai/ComfyUI-WanVideoWrapper/, with the vanilla Sage Attention 2 running on the first and the last steps
### Other
_No response_
Contributor guide
Research direction
Start by reviewing pull request 9047 and Kijai's ComfyUI-WanVideoWrapper implementation to understand the existing Sage Attention integrations. Then determine how an attention-type-by-step node should expose Sage Attention 3 alongside Sage Attention 2, especially for MoE Wan2.2 artifacts. Done means Sage Attention 3 is publicly usable with step-specific attention selection.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100