From 5c67ac709b1a64abc6adaff7228b8863ed6b42b0 Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Sun, 18 Oct 2026 16:20:24 +0100 Subject: [PATCH 1/3] [core] download hub kernel in AttentionModuleMixin.set_attention_backend AttentionModuleMixin.set_attention_backend only stored the backend on the processor and never fetched the hub kernel, unlike ModelMixin.set_attention_backend. For hub backends (e.g. sage_hub) the kernel stayed unloaded and the call failed at the first forward pass. Call _maybe_download_kernel_for_backend before setting the backend and add a regression test using a stubbed kernels.get_kernel. Fixes #13284 Signed-off-by: Parth Vadera Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/diffusers/models/attention.py | 3 ++- tests/models/test_attention_dispatch.py | 27 ++++++++++++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 504f4afe3af3..7e19cb22b575 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -159,13 +159,14 @@ def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProce return self.processor def set_attention_backend(self, backend: str): - from .attention_dispatch import AttentionBackendName + from .attention_dispatch import AttentionBackendName, _maybe_download_kernel_for_backend available_backends = {x.value for x in AttentionBackendName.__members__.values()} if backend not in available_backends: raise ValueError(f"`{backend=}` must be one of the following: " + ", ".join(available_backends)) backend = AttentionBackendName(backend.lower()) + _maybe_download_kernel_for_backend(backend) self.processor._attention_backend = backend def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None: diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index 478fb2853c0f..eb877cfab99d 100644 --- a/tests/models/test_attention_dispatch.py +++ b/tests/models/test_attention_dispatch.py @@ -15,6 +15,7 @@ import contextlib import os +import types import pytest import torch @@ -22,9 +23,11 @@ import torch.multiprocessing as mp import torch.nn.functional as F +from diffusers.models import attention_dispatch from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig +from diffusers.models.attention_dispatch import AttentionBackendName, dispatch_attention_fn from diffusers.models.attention_dispatch import attention_backend as attention_backend_ctx -from diffusers.models.attention_dispatch import dispatch_attention_fn +from diffusers.models.transformers.transformer_flux import FluxAttention from ..testing_utils import ( is_attention, @@ -231,3 +234,25 @@ def make_qkv(seq_len): ) finally: torch.compiler.reset() + + +@is_attention +class TestAttentionModuleSetAttentionBackend: + @pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.") + def test_hub_kernel_is_loaded(self, monkeypatch): + import kernels + + def fake_sageattn(*args, **kwargs): + raise NotImplementedError + + fake_kernel = types.SimpleNamespace(__name__="fake_sage_attention", sageattn=fake_sageattn) + monkeypatch.setattr(kernels, "get_kernel", lambda *args, **kwargs: fake_kernel) + monkeypatch.setattr(attention_dispatch, "DIFFUSERS_TRUST_REMOTE_KERNELS", True) + config = attention_dispatch._HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB] + monkeypatch.setattr(config, "kernel_fn", None) + + attn = FluxAttention(query_dim=8, heads=2, dim_head=4) + attn.set_attention_backend("sage_hub") + + assert attn.processor._attention_backend == AttentionBackendName.SAGE_HUB + assert config.kernel_fn is fake_sageattn From 9424d9bf8b5b6c0eb4de50ed2500bfe778d2add9 Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Tue, 20 Oct 2026 10:38:51 +0100 Subject: [PATCH 2/3] Check backend requirements in AttentionModuleMixin.set_attention_backend Call _check_attention_backend_requirements(backend) before downloading the hub kernel in AttentionModuleMixin.set_attention_backend to match ModelMixin.set_attention_backend, and add a unit test covering missing backend requirements. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Parth Vadera --- src/diffusers/models/attention.py | 7 ++++++- tests/models/test_attention_dispatch.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 7e19cb22b575..3a7e810697b4 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -159,13 +159,18 @@ def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProce return self.processor def set_attention_backend(self, backend: str): - from .attention_dispatch import AttentionBackendName, _maybe_download_kernel_for_backend + from .attention_dispatch import ( + AttentionBackendName, + _check_attention_backend_requirements, + _maybe_download_kernel_for_backend, + ) available_backends = {x.value for x in AttentionBackendName.__members__.values()} if backend not in available_backends: raise ValueError(f"`{backend=}` must be one of the following: " + ", ".join(available_backends)) backend = AttentionBackendName(backend.lower()) + _check_attention_backend_requirements(backend) _maybe_download_kernel_for_backend(backend) self.processor._attention_backend = backend diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index eb877cfab99d..f4a320652f02 100644 --- a/tests/models/test_attention_dispatch.py +++ b/tests/models/test_attention_dispatch.py @@ -256,3 +256,16 @@ def fake_sageattn(*args, **kwargs): assert attn.processor._attention_backend == AttentionBackendName.SAGE_HUB assert config.kernel_fn is fake_sageattn + + def test_backend_requirements_are_checked(self, monkeypatch): + attn = FluxAttention(query_dim=8, heads=2, dim_head=4) + + monkeypatch.setattr(attention_dispatch, "_CAN_USE_FLASH_ATTN", False) + with pytest.raises(RuntimeError, match="Flash Attention backend 'flash' is not usable"): + attn.set_attention_backend("flash") + + monkeypatch.setattr(attention_dispatch, "is_kernels_available", lambda: False) + with pytest.raises(RuntimeError, match="Backend 'sage_hub' is not usable because the `kernels` package"): + attn.set_attention_backend("sage_hub") + + assert attn.processor._attention_backend is None From 82adbf2527a22ba923f5d6f10ba7c5692d918fa1 Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Fri, 30 Oct 2026 16:08:36 +0000 Subject: [PATCH 3/3] Resolve merge conflict in tests/models/test_attention_dispatch.py Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Parth Vadera --- tests/models/test_attention_dispatch.py | 84 +++++++++++++------------ 1 file changed, 45 insertions(+), 39 deletions(-) diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index f4a320652f02..7017bd3e05b4 100644 --- a/tests/models/test_attention_dispatch.py +++ b/tests/models/test_attention_dispatch.py @@ -15,7 +15,6 @@ import contextlib import os -import types import pytest import torch @@ -23,11 +22,9 @@ import torch.multiprocessing as mp import torch.nn.functional as F -from diffusers.models import attention_dispatch from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig -from diffusers.models.attention_dispatch import AttentionBackendName, dispatch_attention_fn from diffusers.models.attention_dispatch import attention_backend as attention_backend_ctx -from diffusers.models.transformers.transformer_flux import FluxAttention +from diffusers.models.attention_dispatch import dispatch_attention_fn from ..testing_utils import ( is_attention, @@ -189,6 +186,50 @@ def test_attn_backend_backward_parity(self, cp_type, attention_backend): ) +@is_attention +class TestAttentionModuleSetAttentionBackend: + @pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.") + def test_hub_kernel_is_loaded(self, monkeypatch): + import types + + import kernels + + from diffusers.models import attention_dispatch + from diffusers.models.attention_dispatch import AttentionBackendName + from diffusers.models.transformers.transformer_flux import FluxAttention + + def fake_sageattn(*args, **kwargs): + raise NotImplementedError + + fake_kernel = types.SimpleNamespace(__name__="fake_sage_attention", sageattn=fake_sageattn) + monkeypatch.setattr(kernels, "get_kernel", lambda *args, **kwargs: fake_kernel) + monkeypatch.setattr(attention_dispatch, "DIFFUSERS_TRUST_REMOTE_KERNELS", True) + config = attention_dispatch._HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB] + monkeypatch.setattr(config, "kernel_fn", None) + + attn = FluxAttention(query_dim=8, heads=2, dim_head=4) + attn.set_attention_backend("sage_hub") + + assert attn.processor._attention_backend == AttentionBackendName.SAGE_HUB + assert config.kernel_fn is fake_sageattn + + def test_backend_requirements_are_checked(self, monkeypatch): + from diffusers.models import attention_dispatch + from diffusers.models.transformers.transformer_flux import FluxAttention + + attn = FluxAttention(query_dim=8, heads=2, dim_head=4) + + monkeypatch.setattr(attention_dispatch, "_CAN_USE_FLASH_ATTN", False) + with pytest.raises(RuntimeError, match="Flash Attention backend 'flash' is not usable"): + attn.set_attention_backend("flash") + + monkeypatch.setattr(attention_dispatch, "is_kernels_available", lambda: False) + with pytest.raises(RuntimeError, match="Backend 'sage_hub' is not usable because the `kernels` package"): + attn.set_attention_backend("sage_hub") + + assert attn.processor._attention_backend is None + + @is_attention class TestVarlenAttentionCompile: """Op-level tests for the varlen backends under `torch.compile` with dynamic shapes, model independent.""" @@ -234,38 +275,3 @@ def make_qkv(seq_len): ) finally: torch.compiler.reset() - - -@is_attention -class TestAttentionModuleSetAttentionBackend: - @pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.") - def test_hub_kernel_is_loaded(self, monkeypatch): - import kernels - - def fake_sageattn(*args, **kwargs): - raise NotImplementedError - - fake_kernel = types.SimpleNamespace(__name__="fake_sage_attention", sageattn=fake_sageattn) - monkeypatch.setattr(kernels, "get_kernel", lambda *args, **kwargs: fake_kernel) - monkeypatch.setattr(attention_dispatch, "DIFFUSERS_TRUST_REMOTE_KERNELS", True) - config = attention_dispatch._HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB] - monkeypatch.setattr(config, "kernel_fn", None) - - attn = FluxAttention(query_dim=8, heads=2, dim_head=4) - attn.set_attention_backend("sage_hub") - - assert attn.processor._attention_backend == AttentionBackendName.SAGE_HUB - assert config.kernel_fn is fake_sageattn - - def test_backend_requirements_are_checked(self, monkeypatch): - attn = FluxAttention(query_dim=8, heads=2, dim_head=4) - - monkeypatch.setattr(attention_dispatch, "_CAN_USE_FLASH_ATTN", False) - with pytest.raises(RuntimeError, match="Flash Attention backend 'flash' is not usable"): - attn.set_attention_backend("flash") - - monkeypatch.setattr(attention_dispatch, "is_kernels_available", lambda: False) - with pytest.raises(RuntimeError, match="Backend 'sage_hub' is not usable because the `kernels` package"): - attn.set_attention_backend("sage_hub") - - assert attn.processor._attention_backend is None