diff --git a/.ai/skills/diffusers-cli/run.md b/.ai/skills/diffusers-cli/run.md index 5eb56d826535..2f1d0f6cae54 100644 --- a/.ai/skills/diffusers-cli/run.md +++ b/.ai/skills/diffusers-cli/run.md @@ -108,7 +108,7 @@ Each entry calls `pipeline.load_lora_weights(, adapter_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. diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md index 7602ddda1134..a33fbe89815f 100644 --- a/docs/source/en/optimization/attention_backends.md +++ b/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) | diff --git a/docs/source/en/using-diffusers/cli.md b/docs/source/en/using-diffusers/cli.md index 02be2bd2fd0e..8869e0bed032 100644 --- a/docs/source/en/using-diffusers/cli.md +++ b/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 diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 89bf25e4de8e..ee24bc075815 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/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], diff --git a/tests/models/testing_utils/attention.py b/tests/models/testing_utils/attention.py index f31323d1bf52..1436a20b9bf9 100644 --- a/tests/models/testing_utils/attention.py +++ b/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,6 +98,32 @@ ], ) +_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, @@ -101,12 +131,20 @@ _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 " diff --git a/tests/models/testing_utils/utils.py b/tests/models/testing_utils/utils.py index 07e4a38ddb21..5070c443887d 100644 --- a/tests/models/testing_utils/utils.py +++ b/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, } diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 53af9eedc50c..be76f892fc4c 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/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""" diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 7a03a8fe2353..5fcf37f6ff3f 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", "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)