Add Sage Attention 3 Backend (#815)

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?

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

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

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),
 

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(

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

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()