Skip to content

Attention Dispatcher #11368

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 14 commits into from
Closed

Attention Dispatcher #11368

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
ab25ee0
initial attention dispatcher support
a-r-r-o-w Apr 19, 2025
bb0628a
fix wan vae dtype
a-r-r-o-w Apr 19, 2025
03a7630
context manager for switching provider
a-r-r-o-w Apr 20, 2025
4b201df
fix flash/sage seqlen preparation when kv len does not match q len (c…
a-r-r-o-w Apr 21, 2025
9c4d4aa
fix flash-attn input shape bug; remove custom block mask code; update…
a-r-r-o-w Apr 21, 2025
60ac540
Merge branch 'main' into attention-dispatcher
a-r-r-o-w Apr 22, 2025
98320f6
Merge branch 'main' into attention-dispatcher
a-r-r-o-w May 14, 2025
a289d30
provider->backend
a-r-r-o-w May 14, 2025
0430966
refactor; flex attention fixes for compile compatibility
a-r-r-o-w May 15, 2025
6b614a3
Merge branch 'main' into attention-dispatcher
a-r-r-o-w May 16, 2025
66deeac
support more models
a-r-r-o-w May 16, 2025
45a809a
add support for FA3, NPU, XLA
a-r-r-o-w May 30, 2025
15ed1d1
Merge branch 'main' into attention-dispatcher
a-r-r-o-w Jun 4, 2025
29d8fcc
update
a-r-r-o-w Jun 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/diffusers/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,7 @@
[
"AllegroTransformer3DModel",
"AsymmetricAutoencoderKL",
"AttentionBackendName",
"AuraFlowTransformer2DModel",
"AutoencoderDC",
"AutoencoderKL",
Expand Down Expand Up @@ -215,6 +216,7 @@
"UVit2DModel",
"VQModel",
"WanTransformer3DModel",
"attention_backend",
]
)
_import_structure["optimization"] = [
Expand Down Expand Up @@ -748,6 +750,7 @@
from .models import (
AllegroTransformer3DModel,
AsymmetricAutoencoderKL,
AttentionBackendName,
AuraFlowTransformer2DModel,
AutoencoderDC,
AutoencoderKL,
Expand Down Expand Up @@ -819,6 +822,7 @@
UVit2DModel,
VQModel,
WanTransformer3DModel,
attention_backend,
)
from .optimization import (
get_constant_schedule,
Expand Down
2 changes: 2 additions & 0 deletions src/diffusers/models/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@

if is_torch_available():
_import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"]
_import_structure["attention_dispatch"] = ["AttentionBackendName", "attention_backend"]
_import_structure["auto_model"] = ["AutoModel"]
_import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"]
_import_structure["autoencoders.autoencoder_dc"] = ["AutoencoderDC"]
Expand Down Expand Up @@ -109,6 +110,7 @@
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
if is_torch_available():
from .adapter import MultiAdapter, T2IAdapter
from .attention_dispatch import AttentionBackendName, attention_backend
from .auto_model import AutoModel
from .autoencoders import (
AsymmetricAutoencoderKL,
Expand Down
Loading