From 24c3ac6ea2d382d6f2ff4940dd46e5209d21667d Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Wed, 30 Apr 2025 01:55:30 +0000 Subject: [PATCH] Add sage attn --- fastvideo/v1/configs/models/dits/base.py | 1 + fastvideo/v1/models/dits/base.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/fastvideo/v1/configs/models/dits/base.py b/fastvideo/v1/configs/models/dits/base.py index c247ca7ed..e5fd1d829 100644 --- a/fastvideo/v1/configs/models/dits/base.py +++ b/fastvideo/v1/configs/models/dits/base.py @@ -12,6 +12,7 @@ class DiTArchConfig(ArchConfig): _param_names_mapping: dict = field(default_factory=dict) _supported_attention_backends: Tuple[_Backend, ...] = (_Backend.SLIDING_TILE_ATTN, + _Backend.SAGE_ATTN, _Backend.FLASH_ATTN, _Backend.TORCH_SDPA) diff --git a/fastvideo/v1/models/dits/base.py b/fastvideo/v1/models/dits/base.py index 477c71e4d..3f405d426 100644 --- a/fastvideo/v1/models/dits/base.py +++ b/fastvideo/v1/models/dits/base.py @@ -17,8 +17,8 @@ class BaseDiT(nn.Module, ABC): num_attention_heads: int num_channels_latents: int # always supports torch_sdpa - _supported_attention_backends: Tuple[_Backend, - ...] = (_Backend.TORCH_SDPA, ) + _supported_attention_backends: Tuple[ + _Backend, ...] = DiTConfig()._supported_attention_backends def __init_subclass__(cls) -> None: required_class_attrs = [