From 13d0aae706ef4d71c2f2e1b7f9bb35c354f50488 Mon Sep 17 00:00:00 2001 From: RandNMR73 <99706358+RandNMR73@users.noreply.github.com> Date: Wed, 24 Sep 2025 15:11:38 -0700 Subject: [PATCH] Add Sage Attention 3 Backend (#815) --- docs/source/inference/optimizations.md | 44 ++++++++++++- fastvideo/attention/backends/sage_attn3.py | 72 ++++++++++++++++++++++ fastvideo/configs/models/dits/base.py | 11 ++-- fastvideo/envs.py | 1 + fastvideo/pipelines/stages/denoising.py | 5 +- fastvideo/platforms/cuda.py | 15 +++++ fastvideo/platforms/interface.py | 1 + 7 files changed, 138 insertions(+), 11 deletions(-) create mode 100644 fastvideo/attention/backends/sage_attn3.py diff --git a/docs/source/inference/optimizations.md b/docs/source/inference/optimizations.md index 609f3f46f..0774bf500 100644 --- a/docs/source/inference/optimizations.md +++ b/docs/source/inference/optimizations.md @@ -1,32 +1,40 @@ (inference-optimizations)= + # Optimizations This page describes the various options for speeding up generation times in FastVideo. ## Table of Contents + - Optimized Attention Backends + - [Flash Attention](#optimizations-flash) - [Sliding Tile Attention](#optimizations-sta) - [Sage Attention](#optimizations-sage) + - [Sage Attention 3](#optimizations-sage3) - Caching Techniques - [TeaCache](#optimizations-teacache) (optimizations-backends)= + ## Attention Backends ### Available Backends + - Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` - Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN` - Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN` - Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN` - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN` +- Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE` ### Configuring Backends There are two ways to configure the attention backend in FastVideo. #### 1. In Python + In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this: ```python @@ -34,6 +42,7 @@ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN" ``` #### 2. In CLI + You can also set the environment variable on the command line: ```bash @@ -41,6 +50,7 @@ FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py ``` (optimizations-flash)= + ### Flash Attention **`FLASH_ATTN`** @@ -57,7 +67,7 @@ And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://git git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention cd hopper -pip install ninja +pip install ninja python setup.py install ``` @@ -66,7 +76,9 @@ FastVideo will automatically detect and use `FA3` if it is installed when using ::: (optimizations-sta)= + ### Sliding Tile Attention + **`SLIDING_TILE_ATTN`** ```bash @@ -76,7 +88,9 @@ pip install st_attn==0.0.4 Please see [this page](#sta-installation) for more installation instructions. (optimizations-vsa)= + ### Video Sparse Attention + **`VIDEO_SPARSE_ATTN`** ```bash @@ -87,19 +101,45 @@ python setup_vsa.py install Please see [this page](#vsa-installation) for more installation instructions. (optimizations-sage)= + ### Sage Attention + **`SAGE_ATTN`** To use [SageAttention](https://github.com/thu-ml/SageAttention) 2.1.1, please compile from source: ```bash git clone https://github.com/thu-ml/SageAttention.git -cd sageattention +cd sageattention python setup.py install # or pip install -e . ``` +(optimizations-sage3)= + +### Sage Attention 3 + +**`SAGE_ATTN_THREE`** + +[SageAttention 3](https://huggingface.co/jt-zhang/SageAttention3) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements. + +#### Hardware Requirements + +- RTX5090 + +#### Installation + +Note that Sage Attention 3 requires `python>=3.13`, `torch>=2.8.0`, `CUDA >=12.8`. If you are using `uv` and using `torch==2.8.0` make sure that `sentencepiece==0.2.1` in the pyproject.toml file. + +To use Sage Attention 3 in FastVideo, first get access to the SageAttention3 code, then move `sageattn/` and `setup.py` to the directory `fastvideo/attention/backends`, then install from using: + +```bash +python setup.py install +``` + (optimizations-teacache)= + ## Teacache + TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo. ### What is TeaCache? diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py new file mode 100644 index 000000000..d10e3634a --- /dev/null +++ b/fastvideo/attention/backends/sage_attn3.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +from fastvideo.attention.backends.sageattn.api import sageattn_blackwell + +from fastvideo.attention.backends.abstract import (AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +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_THREE" + + @staticmethod + def get_impl_cls() -> type["SageAttention3Impl"]: + return SageAttention3Impl + + @staticmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + raise NotImplementedError + + @staticmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + raise NotImplementedError + + # @staticmethod + # def get_metadata_cls() -> Type["AttentionMetadata"]: + # return FlashAttentionMetadata + + +class SageAttention3Impl(AttentionImpl): + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + 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( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + output = sageattn_blackwell(query, key, value, is_causal=self.causal) + output = output.transpose(1, 2) + return output diff --git a/fastvideo/configs/models/dits/base.py b/fastvideo/configs/models/dits/base.py index c9ba2dce4..148eb711b 100644 --- a/fastvideo/configs/models/dits/base.py +++ b/fastvideo/configs/models/dits/base.py @@ -15,13 +15,10 @@ class DiTArchConfig(ArchConfig): reverse_param_names_mapping: dict = field(default_factory=dict) lora_param_names_mapping: dict = field(default_factory=dict) _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( - AttentionBackendEnum.SLIDING_TILE_ATTN, - AttentionBackendEnum.SAGE_ATTN, - AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA, - AttentionBackendEnum.VIDEO_SPARSE_ATTN, - AttentionBackendEnum.VMOBA_ATTN, - ) + AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN, + AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, + AttentionBackendEnum.SAGE_ATTN_THREE) hidden_size: int = 0 num_attention_heads: int = 0 diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 25444854f..a42a5c8cd 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -175,6 +175,7 @@ environment_variables: dict[str, Callable[[], Any]] = { # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention # - "SAGE_ATTN": use Sage Attention + # - "SAGE_ATTN_THREE": use Sage Attention 3 "FASTVIDEO_ATTENTION_BACKEND": lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None), diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 9e486dd84..ed9a7d06e 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -85,8 +85,9 @@ class DenoisingStage(PipelineStage): AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, - AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA - ) # hack + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.SAGE_ATTN_THREE) # hack ) def forward( diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index edb717b1f..8bb58bb7a 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -115,6 +115,7 @@ class CudaPlatformBase(Platform): logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND) + logger.info("Selected backend: %s", selected_backend) if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN: try: from st_attn import sliding_tile_attention # noqa: F401 @@ -144,6 +145,20 @@ class CudaPlatformBase(Platform): logger.info( "Sage Attention backend is not installed. Fall back to Flash Attention." ) + elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE: + try: + from fastvideo.attention.backends.sageattn.api import sageattn_blackwell # noqa: F401 + + from fastvideo.attention.backends.sage_attn3 import ( # noqa: F401 + SageAttention3Backend) + logger.info("Using Sage Attention 3 backend.") + + return "fastvideo.attention.backends.sage_attn3.SageAttention3Backend" + except ImportError as e: + logger.info(e) + logger.info( + "Sage Attention 3 backend is not installed. Fall back to Flash Attention." + ) elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: try: from vsa import block_sparse_attn # noqa: F401 diff --git a/fastvideo/platforms/interface.py b/fastvideo/platforms/interface.py index c092d1c71..a8675325b 100644 --- a/fastvideo/platforms/interface.py +++ b/fastvideo/platforms/interface.py @@ -18,6 +18,7 @@ class AttentionBackendEnum(enum.Enum): SLIDING_TILE_ATTN = enum.auto() TORCH_SDPA = enum.auto() SAGE_ATTN = enum.auto() + SAGE_ATTN_THREE = enum.auto() VIDEO_SPARSE_ATTN = enum.auto() VMOBA_ATTN = enum.auto() NO_ATTENTION = enum.auto()