Add sage attn

2025-04-30 01:55:30 +00:00
parent d9689cef95
commit 24c3ac6ea2
2 changed files with 3 additions and 2 deletions
@@ -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)

@@ -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 = [