[Attn] Add SageAttention Backend (#366)

fastvideo/v1/attention/backends/sage_attn.py

@@ -0,0 +1,66 @@
+from typing import List, Optional, Type
+
+import torch
+from sageattention import sageattn
+
+from fastvideo.v1.attention.backends.abstract import (
+    AttentionBackend)  # FlashAttentionMetadata,
+from fastvideo.v1.attention.backends.abstract import (AttentionImpl,
+                                                      AttentionMetadata)
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> List[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SAGE_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> Type["SageAttentionImpl"]:
+        return SageAttentionImpl
+
+    # @staticmethod
+    # def get_metadata_cls() -> Type["AttentionMetadata"]:
+    #     return FlashAttentionMetadata
+
+
+class SageAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: Optional[int] = 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:
+        output = sageattn(
+            query,
+            key,
+            value,
+            # since input is (batch_size, seq_len, head_num, head_dim)
+            tensor_layout="NHD",
+            is_causal=self.causal)
+        return output

fastvideo/v1/envs.py

@@ -170,7 +170,8 @@ environment_variables: Dict[str, Callable[[], Any]] = {
     # Available options:
     # - "TORCH_SDPA": use torch.nn.MultiheadAttention
     # - "FLASH_ATTN": use FlashAttention
-    # - "STA" : use sliding tile attention
+    # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention
+    # - "SAGE_ATTN": use Sage Attention
     "FASTVIDEO_ATTENTION_BACKEND":
     lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
 

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -437,7 +437,8 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
         lambda n, m: "refiner" in n and str.isdigit(n.split(".")[-1]),
     ]
     _supported_attention_backends = (_Backend.SLIDING_TILE_ATTN,
-                                     _Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+                                     _Backend.SAGE_ATTN, _Backend.FLASH_ATTN,
+                                     _Backend.TORCH_SDPA)
     _param_names_mapping = {
         # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
         r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":

fastvideo/v1/models/dits/wanvideo.py

@@ -354,7 +354,8 @@ class WanTransformer3DModel(BaseDiT):
         lambda n, m: "blocks" in n and str.isdigit(n.split(".")[-1]),
     ]
     _supported_attention_backends = (_Backend.SLIDING_TILE_ATTN,
-                                     _Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+                                     _Backend.SAGE_ATTN, _Backend.FLASH_ATTN,
+                                     _Backend.TORCH_SDPA)
     _param_names_mapping = {
         r"^patch_embedding\.(.*)$":
         r"patch_embedding.proj.\1",

fastvideo/v1/platforms/cuda.py

@@ -126,12 +126,25 @@ class CudaPlatformBase(Platform):
                 logger.info("Using Sliding Tile Attention backend.")
                 return "fastvideo.v1.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend"
             except ImportError as e:
-                # TODO(will): improve error message
                 logger.info(e)
                 logger.info(
                     "Sliding Tile Attention backend is not installed. Fall back to Flash Attention."
                 )
+        elif selected_backend == _Backend.SAGE_ATTN:
+            try:
+                from sageattention import sageattn  # noqa: F401
+
+                from fastvideo.v1.attention.backends.sage_attn import (  # noqa: F401
+                    SageAttentionBackend)
+                logger.info("Using Sage Attention backend.")
+                return "fastvideo.v1.attention.backends.sage_attn.SageAttentionBackend"
+            except ImportError as e:
+                logger.info(e)
+                logger.info(
+                    "Sage Attention backend is not installed. Fall back to Flash Attention."
+                )
         elif selected_backend == _Backend.TORCH_SDPA:
+            logger.info("Using Torch SDPA backend.")
             return "fastvideo.v1.attention.backends.sdpa.SDPABackend"
         elif selected_backend == _Backend.FLASH_ATTN or selected_backend is None:
             pass
@@ -174,9 +187,7 @@ class CudaPlatformBase(Platform):
                 target_backend = _Backend.TORCH_SDPA
 
         if target_backend == _Backend.TORCH_SDPA:
-            logger.info(
-                "Using torch.nn.functional.scaled_dot_product_attention backend."
-            )
+            logger.info("Using Torch SDPA backend.")
             return "fastvideo.v1.attention.backends.sdpa.SDPABackend"
 
         logger.info("Using Flash Attention backend.")

fastvideo/v1/platforms/interface.py

@@ -17,7 +17,7 @@ class _Backend(enum.Enum):
     FLASH_ATTN = enum.auto()
     SLIDING_TILE_ATTN = enum.auto()
     TORCH_SDPA = enum.auto()
-    # SAGE_ATTN = enum.auto()
+    SAGE_ATTN = enum.auto()
     NO_ATTENTION = enum.auto()