diff --git a/benchmarks/diffusion/bench_sage_attn3_matrix.sh b/benchmarks/diffusion/bench_sage_attn3_matrix.sh new file mode 100644 index 000000000..f7a3f4816 --- /dev/null +++ b/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" diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 75915377f..abc63ebed 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/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: diff --git a/tests/diffusion/attention/test_sage_attn3.py b/tests/diffusion/attention/test_sage_attn3.py new file mode 100644 index 000000000..c29e4fab0 --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index 356285654..8c22cc2cd 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/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" diff --git a/vllm_omni/diffusion/attention/backends/sage_attn3.py b/vllm_omni/diffusion/attention/backends/sage_attn3.py new file mode 100644 index 000000000..51ac3a48a --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py index ab71e753b..f31e40e2d 100644 --- a/vllm_omni/diffusion/attention/backends/sdpa.py +++ b/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 ( diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index a5057739f..63002a98d 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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()