[tests] refactor the pipeline-level attention processor tests

tests/models/autoencoders/test_models_autoencoder_kl.py

@@ -36,6 +36,7 @@
     torch_device,
 )
 from ..testing_utils import (
+    AttentionTesterMixin,
     BaseModelTesterConfig,
     MemoryTesterMixin,
     ModelTesterMixin,
@@ -203,6 +204,10 @@ class TestAutoencoderKLMemory(AutoencoderKLTesterConfig, MemoryTesterMixin):
     """Memory optimization tests for AutoencoderKL."""
 
 
+class TestAutoencoderKLAttention(AutoencoderKLTesterConfig, AttentionTesterMixin):
+    """Attention processor tests for AutoencoderKL."""
+
+
 class TestAutoencoderKLSlicingTiling(AutoencoderKLTesterConfig, AutoencoderTesterMixin):
     """Slicing and tiling tests for AutoencoderKL."""
 

tests/models/testing_utils/attention.py

@@ -15,13 +15,14 @@
 
 import gc
 import logging
+from typing import NamedTuple
 
 import pytest
 import torch
 
-from diffusers.models.attention import AttentionModuleMixin
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin
 from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry, attention_backend
-from diffusers.models.attention_processor import AttnProcessor
+from diffusers.models.attention_processor import Attention, AttnProcessor
 from diffusers.utils import is_kernels_available, is_torch_version
 
 from ...testing_utils import assert_tensors_close, backend_empty_cache, is_attention, is_torch_compile, torch_device
@@ -123,6 +124,40 @@ def _skip_if_backend_requires_nondeterminism(backend):
         )
 
 
+class _FusionPath(NamedTuple):
+    """How one of the two `fuse_qkv_projections()` implementations behaves.
+
+    `attention_cls` is the attention base class that implementation walks: a model built on the *other* base class
+    gets fused right past, which is why `test_fuse_unfuse_qkv_projections` checks it before anything else.
+    `swaps_in_fused_processors` says what unfusing is expected to undo.
+    """
+
+    attention_cls: type
+    swaps_in_fused_processors: bool
+
+
+# `AttentionMixin.fuse_qkv_projections()`, inherited by Flux, Flux2, Chroma, Wan, ... It walks `AttentionModuleMixin`
+# modules and leaves the processors alone, so `unfuse_qkv_projections()` deletes the fused layers again.
+_SHARED_FUSION = _FusionPath(attention_cls=AttentionModuleMixin, swaps_in_fused_processors=False)
+
+# The per-model implementations on UNets, SD3, PixArt, AuraFlow, CogVideoX, AutoencoderKL, ... They walk `Attention`
+# modules, swap in a dedicated `Fused*` processor and stash the originals in `original_attn_processors`, so
+# `unfuse_qkv_projections()` only puts those processors back - the fused layers stay on the modules.
+_LEGACY_FUSION = _FusionPath(attention_cls=Attention, swaps_in_fused_processors=True)
+
+
+def _fusion_path(model):
+    """Which of the two `fuse_qkv_projections()` implementations this model inherits."""
+    if type(model).fuse_qkv_projections is AttentionMixin.fuse_qkv_projections:
+        return _SHARED_FUSION
+    return _LEGACY_FUSION
+
+
+def _fused_layer_name(module):
+    """Name of the layer `fuse_projections()` creates: cross-attention fuses K/V only, self-attention fuses Q/K/V."""
+    return "to_kv" if getattr(module, "is_cross_attention", False) else "to_qkv"
+
+
 @is_attention
 class AttentionTesterMixin:
     """
@@ -152,7 +187,7 @@ def teardown_method(self):
         backend_empty_cache(torch_device)
 
     @torch.no_grad()
-    def test_fuse_unfuse_qkv_projections(self, atol=1e-3, rtol=0):
+    def test_fuse_unfuse_qkv_projections(self, request, atol=1e-3, rtol=0):
         init_dict = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**init_dict)
@@ -162,46 +197,100 @@ def test_fuse_unfuse_qkv_projections(self, atol=1e-3, rtol=0):
         if not hasattr(model, "fuse_qkv_projections"):
             pytest.skip("Model does not support QKV projection fusion.")
 
+        fusion = _fusion_path(model)
+
+        attention_modules = [
+            module for module in model.modules() if isinstance(module, (AttentionModuleMixin, Attention))
+        ]
+        if not attention_modules:
+            pytest.skip("Model has no attention modules to fuse.")
+
+        walked_modules = [module for module in attention_modules if isinstance(module, fusion.attention_cls)]
+        stranded_modules = [module for module in attention_modules if not isinstance(module, fusion.attention_cls)]
+
+        if not walked_modules:
+            # Every attention module in this model is of the base class the model's `fuse_qkv_projections()` does not
+            # walk, so fusing silently does nothing. Marked xfail rather than skipped so that fixing the model shows
+            # up as an XPASS instead of staying quietly green.
+            request.node.add_marker(
+                pytest.mark.xfail(
+                    reason=(
+                        f"{type(model).__name__}.fuse_qkv_projections() only walks "
+                        f"`{fusion.attention_cls.__name__}` modules, but every attention module in this model is "
+                        f"a `{type(stranded_modules[0]).__name__}` instance, so nothing is ever fused."
+                    ),
+                    strict=True,
+                )
+            )
+
+        assert walked_modules, (
+            f"{type(model).__name__}.fuse_qkv_projections() does not reach any of this model's attention modules."
+        )
+
+        fusable_modules = [module for module in walked_modules if getattr(module, "_supports_qkv_fusion", True)]
+        if not fusable_modules:
+            pytest.skip("Model's attention modules do not support QKV projection fusion.")
+
         output_before_fusion = model(**inputs_dict, return_dict=False)[0]
+        processors_before_fusion = model.attn_processors
 
         model.fuse_qkv_projections()
 
-        has_fused_projections = False
-        for module in model.modules():
-            if isinstance(module, AttentionModuleMixin):
-                if hasattr(module, "to_qkv") or hasattr(module, "to_kv"):
-                    has_fused_projections = True
-                    assert module.fused_projections, "fused_projections flag should be True"
-                    break
-
-        if has_fused_projections:
-            output_after_fusion = model(**inputs_dict, return_dict=False)[0]
-
-            assert_tensors_close(
-                output_before_fusion,
-                output_after_fusion,
-                atol=atol,
-                rtol=rtol,
-                msg="Output should not change after fusing projections",
+        for module in fusable_modules:
+            assert module.fused_projections, "fused_projections flag should be True"
+            layer_name = _fused_layer_name(module)
+            assert getattr(module, layer_name, None) is not None, (
+                f"{type(module).__name__} should expose a fused `{layer_name}` layer after fusing."
             )
 
-            model.unfuse_qkv_projections()
+        processors_after_fusion = model.attn_processors
+        assert len(processors_after_fusion) == len(processors_before_fusion), (
+            "Fusing projections should not change the number of attention processors."
+        )
+        if fusion.swaps_in_fused_processors:
+            for name, processor in processors_after_fusion.items():
+                assert type(processor).__name__.startswith("Fused"), (
+                    f"Processor {name} should be a fused processor after fusing, got {type(processor).__name__}."
+                )
 
-            for module in model.modules():
-                if isinstance(module, AttentionModuleMixin):
-                    assert not hasattr(module, "to_qkv"), "to_qkv should be removed after unfusing"
-                    assert not hasattr(module, "to_kv"), "to_kv should be removed after unfusing"
-                    assert not module.fused_projections, "fused_projections flag should be False"
+        output_after_fusion = model(**inputs_dict, return_dict=False)[0]
 
-            output_after_unfusion = model(**inputs_dict, return_dict=False)[0]
+        assert_tensors_close(
+            output_before_fusion,
+            output_after_fusion,
+            atol=atol,
+            rtol=rtol,
+            msg="Output should not change after fusing projections",
+        )
+
+        model.unfuse_qkv_projections()
 
-            assert_tensors_close(
-                output_before_fusion,
-                output_after_unfusion,
-                atol=atol,
-                rtol=rtol,
-                msg="Output should match original after unfusing projections",
+        if fusion.swaps_in_fused_processors:
+            # This path only restores the original processors; the fused layers themselves are left in place.
+            processors_after_unfusion = model.attn_processors
+            assert len(processors_after_unfusion) == len(processors_before_fusion), (
+                "Unfusing projections should not change the number of attention processors."
             )
+            for name, processor in processors_after_unfusion.items():
+                assert type(processor) is type(processors_before_fusion[name]), (
+                    f"Processor {name} should be restored to {type(processors_before_fusion[name]).__name__} "
+                    f"after unfusing, got {type(processor).__name__}."
+                )
+        else:
+            for module in fusable_modules:
+                assert not hasattr(module, "to_qkv"), "to_qkv should be removed after unfusing"
+                assert not hasattr(module, "to_kv"), "to_kv should be removed after unfusing"
+                assert not module.fused_projections, "fused_projections flag should be False"
+
+        output_after_unfusion = model(**inputs_dict, return_dict=False)[0]
+
+        assert_tensors_close(
+            output_before_fusion,
+            output_after_unfusion,
+            atol=atol,
+            rtol=rtol,
+            msg="Output should match original after unfusing projections",
+        )
 
     def test_get_set_processor(self):
         init_dict = self.get_init_dict()

tests/models/transformers/test_models_transformer_chroma.py

@@ -20,6 +20,7 @@
 
 from ...testing_utils import enable_full_determinism, torch_device
 from ..testing_utils import (
+    AttentionTesterMixin,
     BaseModelTesterConfig,
     LoraHotSwappingForModelTesterMixin,
     LoraTesterMixin,
@@ -127,6 +128,10 @@ def test_deprecated_inputs_img_txt_ids_3d(self):
         )
 
 
+class TestChromaTransformerAttention(ChromaTransformerTesterConfig, AttentionTesterMixin):
+    """Attention processor tests for Chroma Transformer."""
+
+
 class TestChromaTransformerTraining(ChromaTransformerTesterConfig, TrainingTesterMixin):
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"ChromaTransformer2DModel"}

tests/models/transformers/test_models_transformer_hunyuan_dit.py

@@ -20,6 +20,7 @@
 
 from ...testing_utils import enable_full_determinism, torch_device
 from ..testing_utils import (
+    AttentionTesterMixin,
     BaseModelTesterConfig,
     ModelTesterMixin,
     TrainingTesterMixin,
@@ -125,6 +126,10 @@ def test_output(self, base_model_output):
         super().test_output(base_model_output, expected_output_shape=(batch_size,) + self.output_shape)
 
 
+class TestHunyuanDiTAttention(HunyuanDiTTesterConfig, AttentionTesterMixin):
+    """Attention processor tests for HunyuanDiT."""
+
+
 class TestHunyuanDiTTraining(HunyuanDiTTesterConfig, TrainingTesterMixin):
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {"HunyuanDiT2DModel"}

tests/models/transformers/test_models_transformer_ltx2.py

@@ -173,11 +173,16 @@ def test_gradient_checkpointing_is_applied(self):
 class TestLTX2TransformerAttention(LTX2TransformerTesterConfig, AttentionTesterMixin):
     """Attention processor tests for LTX2 Video Transformer."""
 
-    @pytest.mark.skip(
-        "LTX2Attention does not set is_cross_attention, so fuse_projections tries to fuse Q+K+V together even for cross-attention modules with different input dimensions."
+    @pytest.mark.xfail(
+        reason=(
+            "LTX2Attention does not set is_cross_attention, so fuse_projections tries to fuse Q+K+V together even "
+            "for cross-attention modules with different input dimensions."
+        ),
+        raises=RuntimeError,
+        strict=True,
     )
-    def test_fuse_unfuse_qkv_projections(self, atol=1e-3, rtol=0):
-        pass
+    def test_fuse_unfuse_qkv_projections(self, request, atol=1e-3, rtol=0):
+        super().test_fuse_unfuse_qkv_projections(request, atol=atol, rtol=rtol)
 
 
 class TestLTX2TransformerCompile(LTX2TransformerTesterConfig, TorchCompileTesterMixin):

tests/models/transformers/test_models_transformer_sd3.py

@@ -13,13 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
 import torch
 
 from diffusers import SD3Transformer2DModel
 from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
 from ..testing_utils import (
+    AttentionTesterMixin,
     BaseModelTesterConfig,
     BitsAndBytesTesterMixin,
     LoraTesterMixin,
@@ -117,6 +119,10 @@ def test_gradient_checkpointing_is_applied(self):
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
 
+class TestSD3TransformerAttention(SD3TransformerTesterConfig, AttentionTesterMixin):
+    """Attention processor tests for SD3 Transformer."""
+
+
 class TestSD3TransformerCompile(SD3TransformerTesterConfig, TorchCompileTesterMixin):
     pass
 
@@ -216,6 +222,22 @@ def test_gradient_checkpointing_is_applied(self):
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
 
+class TestSD35TransformerAttention(SD35TransformerTesterConfig, AttentionTesterMixin):
+    """Attention processor tests for SD3.5 Transformer."""
+
+    @pytest.mark.xfail(
+        reason=(
+            "fuse_qkv_projections() sets FusedJointAttnProcessor2_0 on every attention module, including the "
+            "self-attention-only `attn2` of the dual-attention layers, which is then called without "
+            "`encoder_hidden_states` and raises."
+        ),
+        raises=AttributeError,
+        strict=True,
+    )
+    def test_fuse_unfuse_qkv_projections(self, request, atol=1e-3, rtol=0):
+        super().test_fuse_unfuse_qkv_projections(request, atol=atol, rtol=rtol)
+
+
 class TestSD35TransformerCompile(SD35TransformerTesterConfig, TorchCompileTesterMixin):
     pass
 

tests/pipelines/aura_flow/test_pipeline_aura_flow.py

@@ -3,15 +3,12 @@
 
 from diffusers import AuraFlowPipeline, AuraFlowTransformer2DModel, AutoencoderKL, FlowMatchEulerDiscreteScheduler
 
-from ...testing_utils import assert_tensors_close
 from ..testing_utils import (
     BasePipelineTesterConfig,
     LoraMemoryTesterMixin,
     LoraTesterMixin,
     MemoryTesterMixin,
     PipelineTesterMixin,
-    check_qkv_fusion_matches_attn_procs_length,
-    check_qkv_fusion_processors_exist,
 )
 
 
@@ -83,52 +80,6 @@ def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=
         # AuraFlow pads the prompt embeddings to a common length, so batched and single runs diverge slightly more.
         super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
 
-    def test_fused_qkv_projections(self):
-        # Run on CPU to keep the device-dependent `torch.Generator` deterministic.
-        pipe = self.get_pipeline()
-
-        image = self.run_pipe(pipe)
-        original_image_slice = image[0, -1, -3:, -3:]
-
-        # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added
-        # to the pipeline level.
-        pipe.transformer.fuse_qkv_projections()
-        assert check_qkv_fusion_processors_exist(pipe.transformer), (
-            "Something wrong with the fused attention processors. Expected all the attention processors to be fused."
-        )
-        assert check_qkv_fusion_matches_attn_procs_length(
-            pipe.transformer, pipe.transformer.original_attn_processors
-        ), "Something wrong with the attention processors concerning the fused QKV projections."
-
-        image = self.run_pipe(pipe)
-        image_slice_fused = image[0, -1, -3:, -3:]
-
-        pipe.transformer.unfuse_qkv_projections()
-        image = self.run_pipe(pipe)
-        image_slice_disabled = image[0, -1, -3:, -3:]
-
-        assert_tensors_close(
-            original_image_slice,
-            image_slice_fused,
-            atol=1e-3,
-            rtol=1e-3,
-            msg="Fusion of QKV projections shouldn't affect the outputs.",
-        )
-        assert_tensors_close(
-            image_slice_fused,
-            image_slice_disabled,
-            atol=1e-3,
-            rtol=1e-3,
-            msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.",
-        )
-        assert_tensors_close(
-            original_image_slice,
-            image_slice_disabled,
-            atol=1e-2,
-            rtol=1e-2,
-            msg="Original outputs should match when fused QKV projections are disabled.",
-        )
-
 
 class TestAuraFlowPipelineMemory(AuraFlowPipelineTesterConfig, MemoryTesterMixin):
     """Memory optimization tests (CPU offload, group offload, layerwise casting) for the AuraFlow pipeline."""