feat: Sage Attention Algorithm#455

Open

Marius-Graml wants to merge 12 commits intomainfrom

Commits on Jan 19, 2026

Add sage attention algorithm to pruna framework by using diffusers attention backend

Marius Graml
authored and
Marius-Graml
committed
Add compatibility for sage-attn with torch-compile

Marius Graml
authored and
Marius-Graml
committed
Add tests für sage attn algorithm
Marius-Graml
committed
Change formatting using ruff
Marius-Graml
committed
Quick commit, add sage_attn2++ as reference paper, add cachers and quantizers as compatible after and before, add sage_attn in corresponding cachers and quantizers algorithms as compatible, add dty…
Marius-Graml
committed
Add target modules including hyperparameter for excluding first and last attention block per attention component. Remove dtype gaurd as dtypes of q, k, and v per attn module is implicitly checked b…
Marius-Graml
committed
Add doc strings to functions and methods
Marius-Graml
committed
Refactor sage_attn to use target_modules utilities
Marius-Graml
committed
Remove incorrect comment, add all transformers (i.e., transformer*) as default target module, remove warning print
Marius-Graml
committed
Remove unused code
Marius-Graml
committed
Update comments
Marius-Graml
committed
ruff check fix
Marius-Graml
committed