[wip][core] propagate sage attention updates.

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",

tests/models/testing_utils/attention.py

@@ -94,19 +94,38 @@
     ],
 )
 
+_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,
     _PARAM_FLASH_HUB,
     _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 "

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,
 }
 
 

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)