From f1098c77dc27f61b453feb50a7b2bf9a7cc01dd4 Mon Sep 17 00:00:00 2001 From: William Lin Date: Mon, 28 Apr 2025 12:40:07 -0600 Subject: [PATCH] [Attn] Add SageAttention Backend (#366) --- fastvideo/v1/attention/backends/sage_attn.py | 66 ++++++++++++++++++++ fastvideo/v1/envs.py | 3 +- fastvideo/v1/models/dits/hunyuanvideo.py | 3 +- fastvideo/v1/models/dits/wanvideo.py | 3 +- fastvideo/v1/platforms/cuda.py | 19 ++++-- fastvideo/v1/platforms/interface.py | 2 +- 6 files changed, 88 insertions(+), 8 deletions(-) create mode 100644 fastvideo/v1/attention/backends/sage_attn.py diff --git a/fastvideo/v1/attention/backends/sage_attn.py b/fastvideo/v1/attention/backends/sage_attn.py new file mode 100644 index 000000000..0a57d78c2 --- /dev/null +++ b/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 diff --git a/fastvideo/v1/envs.py b/fastvideo/v1/envs.py index 63b3d9252..6eff7fef5 100644 --- a/fastvideo/v1/envs.py +++ b/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), diff --git a/fastvideo/v1/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py index 9f819e24e..a97a104b6 100644 --- a/fastvideo/v1/models/dits/hunyuanvideo.py +++ b/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\.(.*)$": diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py index c0dea8eca..dc5b262e2 100644 --- a/fastvideo/v1/models/dits/wanvideo.py +++ b/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", diff --git a/fastvideo/v1/platforms/cuda.py b/fastvideo/v1/platforms/cuda.py index 56b9508d8..cc44d771e 100644 --- a/fastvideo/v1/platforms/cuda.py +++ b/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.") diff --git a/fastvideo/v1/platforms/interface.py b/fastvideo/v1/platforms/interface.py index b9900fffd..3813582bd 100644 --- a/fastvideo/v1/platforms/interface.py +++ b/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()