[Blackwell][1/N] Add SageAttention3 diffusion backend on blackwell(GB200/B200/RTX5090/PRO6000/DGX Spark available) (#3015)

Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: lishunyang <lishunyang12@163.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

benchmarks/diffusion/bench_sage_attn3_matrix.sh

@@ -0,0 +1,108 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# 2x2 timing matrix (backend x compile mode) for SAGE_ATTN_3 vs TORCH_SDPA,
+# on HunyuanVideo-1.5 480p T2V and Wan 2.2 TI2V-5B. Run from the repo root:
+#
+#     bash benchmarks/diffusion/bench_sage_attn3_matrix.sh
+#
+# Requires: sageattn3 installed (for SAGE_ATTN_3 rows) and a Blackwell GPU.
+
+set -u
+
+STEPS=${STEPS:-50}
+SEED=${SEED:-42}
+PROMPT=${PROMPT:-"A dog running across a field of golden wheat."}
+OUT_DIR=${OUT_DIR:-/tmp/bench_sage3}
+mkdir -p "$OUT_DIR"
+
+declare -A TOTAL PER_STEP TEXT_ENC TRANSFORMER VAE
+
+run() {
+    local name="$1"; shift
+    local model="$1"; shift
+    local h="$1"; shift
+    local w="$1"; shift
+    local f="$1"; shift
+    local g="$1"; shift
+    local log="$OUT_DIR/${name}.log"
+    echo "[$(date +%T)] Running $name ..."
+    env "$@" python examples/offline_inference/text_to_video/text_to_video.py \
+        --model "$model" \
+        --prompt "$PROMPT" \
+        --height "$h" --width "$w" --num-frames "$f" \
+        --num-inference-steps "$STEPS" --seed "$SEED" --guidance-scale "$g" \
+        --enable-diffusion-pipeline-profiler \
+        --output "$OUT_DIR/${name}.mp4" > "$log" 2>&1 || {
+            echo "  FAILED — see $log"
+            TOTAL[$name]="FAIL"; PER_STEP[$name]="FAIL"
+            TEXT_ENC[$name]="-"; TRANSFORMER[$name]="-"; VAE[$name]="-"
+            return
+        }
+
+    local wait_ms
+    wait_ms=$(grep "add_req_and_wait=" "$log" | tail -1 | sed -nE 's/.*add_req_and_wait=([0-9.]+) ms.*/\1/p')
+    local total_s
+    total_s=$(grep -oE "Total generation time: [0-9.]+ seconds" "$log" | tail -1 | sed -nE 's/.*: ([0-9.]+) seconds.*/\1/p')
+    local text_enc_s
+    text_enc_s=$(grep -oE "text_encoder.forward took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
+    local vae_s
+    vae_s=$(grep -oE "vae.decode took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
+    local pipeline_s
+    pipeline_s=$(grep -oE "Pipeline\.forward took [0-9.]+s" "$log" | tail -1 | sed -nE 's/.*took ([0-9.]+)s.*/\1/p')
+    local transformer_s
+    transformer_s=$(awk "BEGIN {printf \"%.3f\", ${pipeline_s:-0} - ${vae_s:-0} - ${text_enc_s:-0}}")
+
+    TOTAL[$name]="${total_s:-0}"
+    PER_STEP[$name]=$(awk "BEGIN {printf \"%.3f\", ${wait_ms:-0} / 1000 / $STEPS}")
+    TEXT_ENC[$name]="$text_enc_s"
+    TRANSFORMER[$name]="$transformer_s"
+    VAE[$name]="$vae_s"
+
+    echo "  total=${TOTAL[$name]}s  per-step=${PER_STEP[$name]}s/it  transformer=${TRANSFORMER[$name]}s  vae=${VAE[$name]}s"
+}
+
+# ========== HunyuanVideo-1.5 480p T2V (480x832x33, CFG=6.0) ==========
+HV_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
+run hv15_sdpa_eager     "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1
+run hv15_sdpa_compiled  "$HV_MODEL" 480 832 33 6.0
+run hv15_sage3_eager    "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
+run hv15_sage3_compiled "$HV_MODEL" 480 832 33 6.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
+
+# ========== Wan 2.2 TI2V-5B (704x1280x33, CFG=5.0) ==========
+WAN_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers
+run wan22_sdpa_eager     "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1
+run wan22_sdpa_compiled  "$WAN_MODEL" 704 1280 33 5.0
+run wan22_sage3_eager    "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
+run wan22_sage3_compiled "$WAN_MODEL" 704 1280 33 5.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
+
+print_table() {
+    local title="$1"; shift
+    echo
+    echo "$title (${STEPS} steps, seed=${SEED})"
+    echo
+    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
+        "Config" "Total (s)" "Per-step (s/it)" "Text Enc (s)" "Transformer (s)" "VAE Decode (s)"
+    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
+        "-----------------------" "---------" "---------------" "------------" "---------------" "--------------"
+    for row in "$@"; do
+        case "$row" in
+            *_sdpa_eager)     label="SDPA + Eager" ;;
+            *_sdpa_compiled)  label="SDPA + Compiled" ;;
+            *_sage3_eager)    label="SAGE_ATTN_3 + Eager" ;;
+            *_sage3_compiled) label="SAGE_ATTN_3 + Compiled" ;;
+        esac
+        printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
+            "$label" "${TOTAL[$row]}" "${PER_STEP[$row]}" "${TEXT_ENC[$row]}" "${TRANSFORMER[$row]}" "${VAE[$row]}"
+    done
+}
+
+print_table "HunyuanVideo-1.5 480p T2V, 832x480x33" \
+    hv15_sdpa_eager hv15_sdpa_compiled hv15_sage3_eager hv15_sage3_compiled
+
+print_table "Wan 2.2 TI2V-5B, 1280x704x33" \
+    wan22_sdpa_eager wan22_sdpa_compiled wan22_sage3_eager wan22_sage3_compiled
+
+echo
+echo "Logs: $OUT_DIR/*.log"

docs/user_guide/diffusion/attention_backends.md

@@ -19,6 +19,8 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
 | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
+| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
+
 
 ## Configuration
 
@@ -185,6 +187,29 @@ Quick check:
 python -c "import sageattention; print(sageattention.__file__)"
 ```
 
+## SageAttention3 Installation
+
+vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
+
+Build from source:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd SageAttention/sageattention3_blackwell
+python setup.py install
+```
+
+Quick check:
+
+```bash
+python -c "import sageattn3; print(sageattn3.__file__)"
+```
+
+Notes:
+
+- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
+- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
+
 ## Usage Examples
 
 ### Default (auto-route)
@@ -235,6 +260,20 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
     --output outputs/wan22_sage.mp4
 ```
 
+### Enable SageAttention3
+
+Example:
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
+    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
+    --tensor-parallel-size 2 \
+    --output outputs/hv15_sage3.mp4
+```
+
 ### Mixed backends across roles
 
 Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:

tests/diffusion/attention/test_sage_attn3.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib
+import sys
+import types
+
+import pytest
+import torch
+from vllm.platforms.interface import DeviceCapability
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+from vllm_omni.platforms.cuda import platform as cuda_platform_module
+from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"
+
+
+def load_sage_attn3_module(monkeypatch: pytest.MonkeyPatch, kernel_impl):
+    fake_module = types.ModuleType("sageattn3")
+    fake_module.sageattn3_blackwell = kernel_impl
+    monkeypatch.setitem(sys.modules, "sageattn3", fake_module)
+    sys.modules.pop(SAGE_ATTN3_MODULE, None)
+    return importlib.import_module(SAGE_ATTN3_MODULE)
+
+
+def test_sage_attn3_forward_uses_blackwell_layout(monkeypatch: pytest.MonkeyPatch):
+    calls = {}
+
+    def fake_kernel(query, key, value, is_causal=False):
+        calls["query_shape"] = query.shape
+        calls["is_causal"] = is_causal
+        return query + key + value
+
+    backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
+    impl = backend_module.SageAttention3Impl(
+        num_heads=4,
+        head_size=64,
+        softmax_scale=1.0 / 8.0,
+        causal=False,
+    )
+
+    query = torch.randn(2, 8, 4, 64)
+    key = torch.randn(2, 8, 4, 64)
+    value = torch.randn(2, 8, 4, 64)
+
+    output = impl.forward_cuda(query, key, value)
+
+    assert calls["query_shape"] == (2, 4, 8, 64)
+    assert calls["is_causal"] is False
+    expected = (query.transpose(1, 2) + key.transpose(1, 2) + value.transpose(1, 2)).transpose(1, 2)
+    assert torch.allclose(output, expected)
+
+
+def test_sage_attn3_falls_back_to_sdpa_for_gqa(monkeypatch: pytest.MonkeyPatch):
+    def fake_kernel(*args, **kwargs):
+        raise AssertionError("sageattn3_blackwell should not be used for GQA")
+
+    backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
+    sdpa_calls = {}
+
+    def fake_sdpa(query, key, value, **kwargs):
+        sdpa_calls["query_shape"] = query.shape
+        sdpa_calls["key_shape"] = key.shape
+        sdpa_calls["enable_gqa"] = kwargs["enable_gqa"]
+        return query + 1
+
+    monkeypatch.setattr(backend_module.F, "scaled_dot_product_attention", fake_sdpa)
+
+    impl = backend_module.SageAttention3Impl(
+        num_heads=4,
+        head_size=64,
+        softmax_scale=1.0 / 8.0,
+        causal=False,
+    )
+
+    query = torch.randn(2, 8, 4, 64)
+    key = torch.randn(2, 8, 2, 64)
+    value = torch.randn(2, 8, 2, 64)
+
+    output = impl.forward_cuda(query, key, value)
+
+    assert sdpa_calls["query_shape"] == (2, 4, 8, 64)
+    assert sdpa_calls["key_shape"] == (2, 2, 8, 64)
+    assert sdpa_calls["enable_gqa"] is True
+    expected = (query.permute(0, 2, 1, 3) + 1).permute(0, 2, 1, 3)
+    assert torch.allclose(output, expected)
+
+
+def test_cuda_platform_selects_sage_attn3_alias(monkeypatch: pytest.MonkeyPatch):
+    original_import_module = importlib.import_module
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+    monkeypatch.setattr(
+        cuda_platform_module.importlib,
+        "import_module",
+        lambda module_name: object() if module_name == "sageattn3" else original_import_module(module_name),
+    )
+
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
+
+    assert backend_path == DiffusionAttentionBackendEnum.SAGE_ATTN_3.get_path()
+
+
+def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
+
+    assert backend_path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()

vllm_omni/diffusion/attention/backends/registry.py

@@ -59,6 +59,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"
     TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"
     SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"
+    SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
     CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
     FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
 

vllm_omni/diffusion/attention/backends/sage_attn3.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from sageattn3 import sageattn3_blackwell  # noqa: F401
+except ImportError:
+    logger.warning(
+        "SageAttention3Backend is not available. Install `sageattn3` from "
+        "https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell"
+    )
+    raise ImportError
+
+
+# Wrapping sageattn3_blackwell as a torch.library custom op keeps it opaque to
+# torch.compile. Otherwise Dynamo graph-breaks on the raw pybind11 kernel and
+# Inductor fails scheduling with KeyError: 'op5'. The hasattr guard keeps this
+# idempotent across test re-imports that pop the module from sys.modules.
+if not hasattr(torch.ops.vllm_omni, "sageattn3_blackwell"):
+
+    @torch.library.custom_op("vllm_omni::sageattn3_blackwell", mutates_args=())
+    def _sageattn3_blackwell_op(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        is_causal: bool,
+    ) -> torch.Tensor:
+        from sageattn3 import sageattn3_blackwell as _kernel
+
+        return _kernel(query, key, value, is_causal=is_causal)
+
+    @_sageattn3_blackwell_op.register_fake
+    def _(query, key, value, is_causal):
+        return torch.empty_like(query)
+
+
+_sageattn3_blackwell_op = torch.ops.vllm_omni.sageattn3_blackwell
+
+
+class SageAttention3Backend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SAGE_ATTN_3"
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttention3Impl"]:
+        return SageAttention3Impl
+
+
+class SageAttention3Impl(AttentionImpl):
+    _warned_gqa_fallback_global: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+
+        if key.shape[1] != query.shape[1]:
+            if query.shape[1] % key.shape[1] != 0:
+                raise ValueError(
+                    "GQA/MQA requires query heads to be a multiple of KV heads, "
+                    f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
+                )
+            if not type(self)._warned_gqa_fallback_global:
+                logger.warning("SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA.")
+                type(self)._warned_gqa_fallback_global = True
+            output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                is_causal=self.causal,
+                dropout_p=self.dropout,
+                scale=self.softmax_scale,
+                enable_gqa=True,
+            )
+        else:
+            output = _sageattn3_blackwell_op(query, key, value, self.causal)
+
+        return output.transpose(1, 2).contiguous()

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -30,10 +30,6 @@ def _maybe_reshape_attn_mask(
       - broadcast_k: [batch_size, 1, 1, seq_len_k]
       - full_qk: [batch_size, 1, seq_len_q, seq_len_k]
     """
-    # Skip Attention Mask if all values are 1, `None` mask can speedup the computation
-    if attn_mask is not None and torch.all(attn_mask != 0):
-        attn_mask = None
-
     # Reshape Attention Mask
     # 2D [batch_size, seq_len_k] mask only.
     if (

vllm_omni/platforms/cuda/platform.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import importlib
+
 import torch
 import vllm.envs as envs
 from vllm.config import VllmConfig
@@ -129,6 +131,22 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
                     logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.")
                 logger.debug("Defaulting to diffusion attention backend SDPA")
                 return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+            if backend_upper == "SAGE_ATTN_3":
+                sage_attn3_supported = compute_capability is not None and compute_capability.major >= 10
+                if not sage_attn3_supported:
+                    logger.warning(
+                        "SageAttention3 requires a Blackwell-class GPU with compute capability >= 10.0. "
+                        "Falling back to TORCH_SDPA backend."
+                    )
+                    return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+                try:
+                    importlib.import_module("sageattn3")
+                except ImportError:
+                    logger.warning(
+                        "SageAttention3 package not available. Install it from "
+                        "SageAttention/sageattention3_blackwell. Falling back to TORCH_SDPA backend."
+                    )
+                    return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
             backend = DiffusionAttentionBackendEnum[backend_upper]
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()