feat: Sage Attention Algorithm#455

Open

Marius-Graml wants to merge 8 commits intomainfrom

Commits on Dec 10, 2025

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

Commits on Dec 17, 2025

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

Commits on Dec 24, 2025

Refactor sage_attn to use target_modules utilities
Marius-Graml
committed