diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index e7cc20f580d4..d47ca9a44678 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -351,7 +351,7 @@ class _HubKernelConfig: AttentionBackendName.SAGE_HUB: _HubKernelConfig( repo_id="kernels-community/sage-attention", function_attr="sageattn", - version=1, + version=3, ), AttentionBackendName.FLASH_4_HUB: _HubKernelConfig( repo_id="kernels-community/flash-attn4", diff --git a/tests/models/testing_utils/attention.py b/tests/models/testing_utils/attention.py index f31323d1bf52..6ff2390cb41e 100644 --- a/tests/models/testing_utils/attention.py +++ b/tests/models/testing_utils/attention.py @@ -94,6 +94,18 @@ ], ) +_PARAM_SAGE_HUB = pytest.param( + AttentionBackendName.SAGE_HUB, + id="sage_hub", + marks=[ + pytest.mark.skipif(not _CUDA_AVAILABLE, reason="CUDA is required for sage_hub backend."), + pytest.mark.skipif( + not is_kernels_available(), + reason="`kernels` package is required for sage_hub backend. Install with `pip install kernels`.", + ), + ], +) + # All backends under test. _ALL_BACKEND_PARAMS = [ _PARAM_NATIVE_CUDNN, @@ -101,12 +113,19 @@ _PARAM_FLASH_3_HUB, _PARAM_FLASH_VARLEN_HUB, _PARAM_FLASH_3_VARLEN_HUB, + _PARAM_SAGE_HUB, ] # Backends that perform non-deterministic operations and therefore cannot run when # torch.use_deterministic_algorithms(True) is active (e.g. after enable_full_determinism()). _NON_DETERMINISTIC_BACKENDS = {AttentionBackendName._NATIVE_CUDNN} +# Backends whose kernel cannot be traced into a single graph. Sage dispatches on the compute +# capability on every call (`torch.cuda.device_count()` returns a non-Tensor, which Dynamo +# rejects) and its arch-specific paths reach a Triton quantizer and torch ops that have no +# registered fake implementations. +_NO_FULLGRAPH_COMPILE_BACKENDS = {AttentionBackendName.SAGE_HUB} + def _skip_if_backend_requires_nondeterminism(backend): """Skip at runtime when torch.use_deterministic_algorithms(True) blocks the backend. @@ -419,6 +438,9 @@ def test_compile(self, backend, atol=1e-2, rtol=1e-2): if getattr(self.model_class, "_repeated_blocks", None) is None: pytest.skip("Skipping tests as regional compilation is not supported.") + if backend in _NO_FULLGRAPH_COMPILE_BACKENDS: + pytest.skip(f"Backend '{backend.value}' does not support fullgraph compilation.") + if backend == AttentionBackendName.NATIVE and not is_torch_version(">=", "2.9.0"): pytest.xfail( "test_compile with the native backend requires torch >= 2.9.0 for stable " diff --git a/tests/models/testing_utils/utils.py b/tests/models/testing_utils/utils.py index 07e4a38ddb21..9f2499ddca73 100644 --- a/tests/models/testing_utils/utils.py +++ b/tests/models/testing_utils/utils.py @@ -9,6 +9,9 @@ AttentionBackendName.FLASH_VARLEN_HUB, AttentionBackendName._FLASH_3_HUB, AttentionBackendName._FLASH_3_VARLEN_HUB, + # Sage attention quantizes QK to INT8 and PV to FP8/FP16, so it only accepts + # fp16/bf16 inputs and rejects the fp32 the test models default to. + AttentionBackendName.SAGE_HUB, } diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 7a03a8fe2353..a301209bcf85 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -259,7 +259,7 @@ class TestQwenImageTransformerAttention(QwenImageTransformerTesterConfig, Attent class TestQwenImageTransformerAttentionBackend(QwenImageTransformerTesterConfig, AttentionBackendTesterMixin): """Attention backend tests for QwenImage Transformer.""" - unsupported_attn_backends = ["flash_hub", "_flash_3_hub"] + unsupported_attn_backends = ["flash_hub", "_flash_3_hub", "sage_hub"] def get_dummy_inputs(self, batch_size: int = 2): inputs = super().get_dummy_inputs(batch_size=batch_size) @@ -289,9 +289,9 @@ class TestQwenImageTransformerContextParallelAttnBackends( ): """Context Parallel inference x attention backends tests for QwenImage Transformer""" - # QwenImage always passes a joint attention mask (text + image), which flash_hub and - # _flash_3_hub do not support. - unsupported_attn_backends = ["flash_hub", "_flash_3_hub"] + # QwenImage always passes a joint attention mask (text + image), which flash_hub, + # _flash_3_hub and sage_hub do not support. + unsupported_attn_backends = ["flash_hub", "_flash_3_hub", "sage_hub"] def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: inputs = super().get_dummy_inputs(batch_size=batch_size)