[core] propagate sage attention updates.

.ai/skills/diffusers-cli/run.md

@@ -108,7 +108,7 @@ Each entry calls `pipeline.load_lora_weights(<lora_id>, adapter_name=<name>)`. A
 - `--cpu-offload {model, group}` — `model` uses `enable_model_cpu_offload`, `group` uses
   `enable_group_offload(offload_type="leaf_level", use_stream=True)`. Use `group` to fit a 9B+ model on a single
   A100. Onload target device comes from `--device-map` (must be a plain device string in this case).
-- `--attention-backend {default, flash_hub, flash_varlen_hub, flash_4_hub, sage_hub}` — hub-hosted kernels,
+- `--attention-backend {default, flash_hub, flash_varlen_hub, flash_4_hub, sage_hub, sage_blackwell_hub}` — hub-hosted kernels,
   auto-downloaded on first use. Failures (kernel not available, CUDA arch mismatch, network) raise a clear
   `SystemExit` listing the alternatives instead of silently reverting to the default. Only supported on
   transformer-based pipelines; UNet pipelines get a `logger.warning` and the flag is ignored.

docs/source/en/optimization/attention_backends.md

@@ -164,6 +164,7 @@ Refer to the table below for a complete list of available attention backends and
 | `_flash_3_varlen_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention-3 from kernels |
 | `sage` | [SageAttention](https://github.com/thu-ml/SageAttention) | Quantized attention (INT8 QK) |
 | `sage_hub` | [SageAttention](https://github.com/thu-ml/SageAttention) | Quantized attention (INT8 QK) from kernels |
+| `sage_blackwell_hub` | [SageAttention](https://github.com/thu-ml/SageAttention) | SageAttention3 FP4 attention for SM120 Blackwell GPUs from kernels |
 | `sage_varlen` | [SageAttention](https://github.com/thu-ml/SageAttention) | Variable length SageAttention |
 | `_sage_qk_int8_pv_fp8_cuda` | [SageAttention](https://github.com/thu-ml/SageAttention) | INT8 QK + FP8 PV (CUDA) |
 | `_sage_qk_int8_pv_fp8_cuda_sm90` | [SageAttention](https://github.com/thu-ml/SageAttention) | INT8 QK + FP8 PV (SM90) |

docs/source/en/using-diffusers/cli.md

@@ -134,7 +134,7 @@ Configure how the CLI loads model weights and custom pipeline code.
   `enable_auto_cpu_offload` as `memory_reserve_margin` (default `3GB`). Raise it when a large canvas runs out
   of memory mid-forward: the offloader keeps components resident while they fit, so on a high-VRAM card the
   default margin can leave too little room for the activations of a long video.
-- `--attention-backend {default, flash_hub, flash_varlen_hub, flash_4_hub, sage_hub}` — Hub-hosted attention
+- `--attention-backend {default, flash_hub, flash_varlen_hub, flash_4_hub, sage_hub, sage_blackwell_hub}` — Hub-hosted attention
   kernels, auto-downloaded on first use. Transformer-based pipelines only; ignored with a warning on legacy UNet
   pipelines. See [Attention backends](../optimization/attention_backends).
 - `--vae-tiling` / `--vae-slicing` — lower VAE decode VRAM. See

src/diffusers/models/attention_dispatch.py

@@ -242,6 +242,7 @@ class AttentionBackendName(str, Enum):
     # `sageattention`
     SAGE = "sage"
     SAGE_HUB = "sage_hub"
+    SAGE_BLACKWELL_HUB = "sage_blackwell_hub"
     SAGE_VARLEN = "sage_varlen"
     _SAGE_QK_INT8_PV_FP8_CUDA = "_sage_qk_int8_pv_fp8_cuda"
     _SAGE_QK_INT8_PV_FP8_CUDA_SM90 = "_sage_qk_int8_pv_fp8_cuda_sm90"
@@ -352,6 +353,11 @@ class _HubKernelConfig:
     AttentionBackendName.SAGE_HUB: _HubKernelConfig(
         repo_id="kernels-community/sage-attention",
         function_attr="sageattn",
+        version=3,
+    ),
+    AttentionBackendName.SAGE_BLACKWELL_HUB: _HubKernelConfig(
+        repo_id="kernels-community/sage-blackwell",
+        function_attr="sageattn3_blackwell",
         version=1,
     ),
     AttentionBackendName.FLASH_4_HUB: _HubKernelConfig(
@@ -473,6 +479,13 @@ def check_device_cuda(query: torch.Tensor, key: torch.Tensor, value: torch.Tenso
     return check_device_cuda
 
 
+def _check_head_dim_64_or_128(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> None:
+    # The SM120 SageAttention3 kernel rejects head dims below 64 outright, fails to compile its
+    # Triton pre-pass on non-power-of-two dims, and silently falls back to SDPA at 256 and above.
+    if query.shape[-1] not in (64, 128):
+        raise ValueError(f"Query, key, and value must have a head dimension of 64 or 128, got {query.shape[-1]}.")
+
+
 def _check_qkv_dtype_match(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> None:
     if query.dtype != key.dtype:
         raise ValueError("Query and key must have the same dtype.")
@@ -535,6 +548,7 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
         AttentionBackendName._FLASH_3_HUB,
         AttentionBackendName._FLASH_3_VARLEN_HUB,
         AttentionBackendName.SAGE_HUB,
+        AttentionBackendName.SAGE_BLACKWELL_HUB,
         AttentionBackendName.FLASH_4_HUB,
         AttentionBackendName.AITER_FA2_HUB,
     ]:
@@ -4103,6 +4117,40 @@ def _sage_attention_hub(
     return (out, lse) if return_lse else out
 
 
+@_AttentionBackendRegistry.register(
+    AttentionBackendName.SAGE_BLACKWELL_HUB,
+    constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_head_dim_64_or_128, _check_shape],
+)
+def _sage_attention_blackwell_hub(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None = None,
+    is_causal: bool = False,
+    scale: float | None = None,
+    return_lse: bool = False,
+    _parallel_config: "ParallelConfig" | None = None,
+) -> torch.Tensor:
+    if attn_mask is not None:
+        raise ValueError("`attn_mask` is not supported for sage attention")
+    if return_lse:
+        # `sageattn3_blackwell` returns the output only, so there is no LSE to hand back. This
+        # also rules out context parallelism, hence `supports_context_parallel` is not set above.
+        raise ValueError("`return_lse` is not supported by the `sage_blackwell_hub` backend.")
+    if scale is not None and scale != query.shape[-1] ** -0.5:
+        # The kernel derives the softmax scale from the head dimension internally and silently
+        # swallows unknown kwargs, so a custom scale would be ignored rather than applied.
+        raise ValueError("A custom `scale` is not supported by the `sage_blackwell_hub` backend.")
+
+    func = _HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_BLACKWELL_HUB].kernel_fn
+    # The kernel works on the HND layout, unlike the other Sage backends which take NHD. It also
+    # subtracts the per-token mean from `key` in place, so the transposed copies we build here
+    # double as protection for the caller's tensors.
+    query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value))
+    out = func(query, key, value, is_causal=is_causal)
+    return out.transpose(1, 2).contiguous()
+
+
 @_AttentionBackendRegistry.register(
     AttentionBackendName.SAGE_VARLEN,
     constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_shape],

tests/models/testing_utils/attention.py

@@ -36,6 +36,10 @@
 # ---------------------------------------------------------------------------
 
 _CUDA_AVAILABLE = torch.cuda.is_available()
+# Every build variant of `kernels-community/sage-blackwell` declares `archs: ["12.0a"]`, so the
+# kernel only loads on SM120 (consumer/workstation Blackwell). `a` targets are architecture
+# specific, so neither SM100 nor SM121 is covered.
+_IS_SM120 = _CUDA_AVAILABLE and torch.cuda.get_device_capability() == (12, 0)
 
 _PARAM_NATIVE_CUDNN = pytest.param(
     AttentionBackendName._NATIVE_CUDNN,
@@ -94,19 +98,53 @@
     ],
 )
 
+_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`.",
+        ),
+    ],
+)
+
+_PARAM_SAGE_BLACKWELL_HUB = pytest.param(
+    AttentionBackendName.SAGE_BLACKWELL_HUB,
+    id="sage_blackwell_hub",
+    marks=[
+        pytest.mark.skipif(
+            not _IS_SM120, reason="An SM120 Blackwell GPU is required for the sage_blackwell_hub backend."
+        ),
+        pytest.mark.skipif(
+            not is_kernels_available(),
+            reason="`kernels` package is required for sage_blackwell_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,
+    _PARAM_SAGE_BLACKWELL_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, AttentionBackendName.SAGE_BLACKWELL_HUB}
+
 
 def _skip_if_backend_requires_nondeterminism(backend):
     """Skip at runtime when torch.use_deterministic_algorithms(True) blocks the backend.
@@ -419,6 +457,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,10 @@
     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,
+    AttentionBackendName.SAGE_BLACKWELL_HUB,
 }
 
 

tests/models/transformers/test_models_transformer_flux.py

@@ -253,6 +253,14 @@ class TestFluxTransformerAttention(FluxTransformerTesterConfig, AttentionTesterM
 class TestFluxTransformerAttentionBackend(FluxTransformerTesterConfig, AttentionBackendTesterMixin):
     """Attention backend tests for Flux Transformer."""
 
+    def get_init_dict(self) -> dict[str, int | list[int]]:
+        # `sage_blackwell_hub` runs a kernel that only accepts head dims of 64 or 128, so widen the
+        # shared dummy config's `attention_head_dim` of 16. `axes_dims_rope` has to keep summing to it.
+        init_dict = super().get_init_dict()
+        init_dict["attention_head_dim"] = 64
+        init_dict["axes_dims_rope"] = [16, 16, 32]
+        return init_dict
+
 
 class TestFluxTransformerContextParallel(FluxTransformerTesterConfig, ContextParallelTesterMixin):
     """Context Parallel inference tests for Flux Transformer"""

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", "sage_blackwell_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 the sage hub backends do not support.
+    unsupported_attn_backends = ["flash_hub", "_flash_3_hub", "sage_hub", "sage_blackwell_hub"]
 
     def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
         inputs = super().get_dummy_inputs(batch_size=batch_size)