diffsynth 2.1.12__py3-none-any.whl → 2.2.0__py3-none-any.whl

diffsynth/core/attention/__init__.py

@@ -1 +1,2 @@
-from .attention import attention_forward, FLEX_ATTN_AVAILABLE
+from .attention import attention_forward, list_attention_implementations
+from .flex_attention import FLEX_ATTN_AVAILABLE

diffsynth/core/attention/attention.py

@@ -1,255 +1,51 @@
-import torch, os, inspect
-from einops import rearrange, repeat
-
-if os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID") is not None:
-    try:
-        from kernels import get_kernel
-        kernel_module = get_kernel(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID"), version=int(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_VERSION")))
-        customized_flash_attn_func = kernel_module.flash_attn_func
-        CUSTOMIZED_FA_KERNEL_AVAILABLE = True
-    except Exception:
-        CUSTOMIZED_FA_KERNEL_AVAILABLE = False
-else:
-    CUSTOMIZED_FA_KERNEL_AVAILABLE = False
-
-try:
-    from flash_attn.cute import flash_attn_func as flash_attn_func_cute
-    FLASH_ATTN_4_AVAILABLE = True
-except Exception:
-    FLASH_ATTN_4_AVAILABLE = False
-
-try:
-    import flash_attn_interface
-    FLASH_ATTN_3_AVAILABLE = True
-except Exception:
-    FLASH_ATTN_3_AVAILABLE = False
-
-try:
-    import flash_attn
-    FLASH_ATTN_2_AVAILABLE = True
-except Exception:
-    FLASH_ATTN_2_AVAILABLE = False
-
-try:
-    from sageattention import sageattn
-    SAGE_ATTN_AVAILABLE = True
-except Exception:
-    SAGE_ATTN_AVAILABLE = False
-
-try:
-    import xformers.ops as xops
-    XFORMERS_AVAILABLE = True
-except Exception:
-    XFORMERS_AVAILABLE = False
-
-try:
-    from torch.nn.attention.flex_attention import flex_attention as flex_attention_func
-    flex_attention_func = torch.compile(flex_attention_func)
-    FLEX_ATTN_AVAILABLE = True
-except Exception:
-    FLEX_ATTN_AVAILABLE = False
-
-try:
-    if "enable_gqa" in inspect.signature(torch.nn.functional.scaled_dot_product_attention).parameters:
-        TORCH_SUPPORT_GQA = True
-    else:
-        TORCH_SUPPORT_GQA = False
-except Exception:
-    TORCH_SUPPORT_GQA = False
+import os
+from .registry import (
+    get_attention_implementation,
+    list_attention_implementations,
+    register_attention_implementation,
+)
+from .flash_attention import CUSTOMIZED_FA_KERNEL_AVAILABLE, FLASH_ATTN_4_AVAILABLE, FLASH_ATTN_3_AVAILABLE, FLASH_ATTN_2_AVAILABLE
+from .sage_attention import SAGE_ATTN_AVAILABLE
+from .xformers_attention import XFORMERS_AVAILABLE
+from .torch_attention import torch_sdpa, torch_sdpa_sliding_window
+from .flex_attention import flex_attention
 
 
 def initialize_attention_priority():
     if os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION') is not None:
-        return os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION').lower()
-    elif CUSTOMIZED_FA_KERNEL_AVAILABLE:
-        return "customized_fa_kernel"
-    elif FLASH_ATTN_4_AVAILABLE:
-        return "flash_attention_4"
-    elif FLASH_ATTN_3_AVAILABLE:
-        return "flash_attention_3"
-    elif FLASH_ATTN_2_AVAILABLE:
-        return "flash_attention_2"
-    elif SAGE_ATTN_AVAILABLE:
-        return "sage_attention"
-    elif XFORMERS_AVAILABLE:
-        return "xformers"
+        return os.environ.get('DIFFSYNTH_ATTENTION_IMPLEMENTATION').strip().lower()
     else:
+        for name in ["customized_fa_kernel", "flash_attention_4", "flash_attention_3", "flash_attention_2", "sage_attention", "xformers"]:
+            if name in list_attention_implementations():
+                return name
         return "torch"
 
 
 ATTENTION_IMPLEMENTATION = initialize_attention_priority()
 
 
-def rearrange_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", required_in_pattern="b n s d", dims=None):
-    dims = {} if dims is None else dims
-    if q_pattern != required_in_pattern:
-        q = rearrange(q, f"{q_pattern} -> {required_in_pattern}", **dims)
-    if k_pattern != required_in_pattern:
-        k = rearrange(k, f"{k_pattern} -> {required_in_pattern}", **dims)
-    if v_pattern != required_in_pattern:
-        v = rearrange(v, f"{v_pattern} -> {required_in_pattern}", **dims)
-    return q, k, v
-
-
-def rearrange_out(out: torch.Tensor, out_pattern="b n s d", required_out_pattern="b n s d", dims=None):
-    dims = {} if dims is None else dims
-    if out_pattern != required_out_pattern:
-        out = rearrange(out, f"{required_out_pattern} -> {out_pattern}", **dims)
-    return out
-
-
-def torch_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False):
-    required_in_pattern, required_out_pattern= "b n s d", "b n s d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-    if q.shape[1] != k.shape[1] or q.shape[1] != v.shape[1]:
-        # Grouped Query Attention
-        if TORCH_SUPPORT_GQA:
-            out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal, enable_gqa=True)
-        else:
-            # In low-version torch, `enable_gqa` is not supported.
-            k = repeat(k, "b n s d -> b (n m) s d", m=q.shape[1]//k.shape[1])
-            v = repeat(v, "b n s d -> b (n m) s d", m=q.shape[1]//v.shape[1])
-            out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal)
-    else:
-        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal)
-    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
-    return out
-
-
-def torch_sdpa_sliding_window(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sliding_window: int, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
-    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-
-    B, N, S, D = q.shape
-    W = sliding_window
-    chunk_size = W
-    num_chunks = (S + chunk_size - 1) // chunk_size
-
-    output = torch.empty_like(q)
-    dtype = q.dtype
-    device = q.device
-    min_val = torch.finfo(dtype).min
-
-    for i in range(num_chunks):
-        q_start = i * chunk_size
-        q_end = min(q_start + chunk_size, S)
-        actual_chunk_size = q_end - q_start
+def attention_forward(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False, compatibility_mode=False, window_size=None, use_flex=False, score_mod=None):
+    if use_flex or score_mod is not None:
+        return flex_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, score_mod=score_mod)
 
-        kv_start = max(0, q_start - W)
-        kv_end = min(S, q_end + W)
+    if compatibility_mode:
+        return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, is_causal=is_causal, window_size=window_size)
 
-        q_chunk = q[:, :, q_start:q_end, :]
-        k_chunk = k[:, :, kv_start:kv_end, :]
-        v_chunk = v[:, :, kv_start:kv_end, :]
+    backend = get_attention_implementation(ATTENTION_IMPLEMENTATION)
 
-        q_indices = torch.arange(q_start, q_end, device=device)
-        k_indices = torch.arange(kv_start, kv_end, device=device)
-        diff = q_indices.unsqueeze(1) - k_indices.unsqueeze(0)
-        valid = diff.abs() <= W
+    if attn_mask is not None and not backend.supports_mask:
+        return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, is_causal=is_causal)
+    if window_size is not None and not backend.supports_window:
+        assert is_causal == False and attn_mask is None
+        return torch_sdpa_sliding_window(q, k, v, window_size, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+    if is_causal and not backend.supports_causal:
+        return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, is_causal=is_causal)
 
-        local_mask = torch.zeros(actual_chunk_size, kv_end - kv_start, dtype=dtype, device=device)
-        local_mask.masked_fill_(~valid, min_val)
-        local_mask = local_mask.unsqueeze(0).unsqueeze(0)
-
-        out_chunk = torch_sdpa(
-            q_chunk, k_chunk, v_chunk, attn_mask=local_mask, scale=scale
-        )
-        output[:, :, q_start:q_end, :] = out_chunk
-
-    output = rearrange_out(output, out_pattern, required_out_pattern, dims)
-    return output
-
-
-def flash_attention_4(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
-    required_in_pattern, required_out_pattern= "b s n d", "b s n d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
-    out = flash_attn_func_cute(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
-    if isinstance(out, tuple):
-        out = out[0]
-    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
-    return out
-
-
-def flash_attention_3(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
-    required_in_pattern, required_out_pattern= "b s n d", "b s n d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
-    out = flash_attn_interface.flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
-    if isinstance(out, tuple):
-        out = out[0]
-    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
-    return out
-
-
-def flash_attention_2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
-    required_in_pattern, required_out_pattern= "b s n d", "b s n d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
-    out = flash_attn.flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
-    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
-    return out
-
-
-def flash_attention_customized_kernel(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
-    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
-    out = customized_flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
-    if isinstance(out, tuple):
-        out = out[0]
-    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
-    return out
-
-
-def sage_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
-    required_in_pattern, required_out_pattern= "b n s d", "b n s d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-    out = sageattn(q, k, v, sm_scale=scale)
-    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
-    return out
-
-
-def xformers_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
-    required_in_pattern, required_out_pattern= "b s n d", "b s n d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-    out = xops.memory_efficient_attention(q, k, v, scale=scale)
-    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
-    return out
-
-
-def flex_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, score_mod=None):
-    assert FLEX_ATTN_AVAILABLE, "Flex Attention is not available. Please upgrade torch to 2.5.0 or later."
-    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
-    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
-    out = flex_attention_func(query=q, key=k, value=v, block_mask=attn_mask, scale=scale, score_mod=score_mod)
-    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
-    return out
-
-
-def attention_forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False, compatibility_mode=False, window_size=None, use_flex=False, score_mod=None):
-    if compatibility_mode or (attn_mask is not None) or ATTENTION_IMPLEMENTATION == "torch":
-        if use_flex or score_mod is not None:
-            return flex_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, score_mod=score_mod)
-        elif window_size is None:
-            return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale, is_causal=is_causal)
-        else:
-            # Sliding Window Attention is not compatible with `is_causal` and `attn_mask`.
-            assert is_causal == False and attn_mask is None
-            return torch_sdpa_sliding_window(q, k, v, window_size, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
-    elif ATTENTION_IMPLEMENTATION == "customized_fa_kernel":
-        return flash_attention_customized_kernel(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
-    elif ATTENTION_IMPLEMENTATION == "flash_attention_4":
-        return flash_attention_4(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
-    elif ATTENTION_IMPLEMENTATION == "flash_attention_3":
-        return flash_attention_3(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
-    elif ATTENTION_IMPLEMENTATION == "flash_attention_2":
-        return flash_attention_2(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale, is_causal=is_causal, window_size=window_size)
-    elif ATTENTION_IMPLEMENTATION == "sage_attention":
-        if window_size is not None or is_causal: return attention_forward(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask, scale, is_causal, compatibility_mode=True, window_size=window_size)
-        return sage_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
-    elif ATTENTION_IMPLEMENTATION == "xformers":
-        if window_size is not None or is_causal: return attention_forward(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask, scale, is_causal, compatibility_mode=True, window_size=window_size)
-        return xformers_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
-    else:
-        raise NotImplementedError("No available attention implementation.")
+    kwargs = {"scale": scale}
+    if backend.supports_causal:
+        kwargs["is_causal"] = is_causal
+    if backend.supports_mask:
+        kwargs["attn_mask"] = attn_mask
+    if backend.supports_window:
+        kwargs["window_size"] = window_size
+    return backend.func(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, **kwargs)

diffsynth/core/attention/common.py

@@ -0,0 +1,27 @@
+import inspect
+
+import torch
+from einops import rearrange
+
+try:
+    TORCH_SUPPORT_GQA = "enable_gqa" in inspect.signature(torch.nn.functional.scaled_dot_product_attention).parameters
+except Exception:
+    TORCH_SUPPORT_GQA = False
+
+
+def rearrange_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", required_in_pattern="b n s d", dims=None):
+    dims = {} if dims is None else dims
+    if q_pattern != required_in_pattern:
+        q = rearrange(q, f"{q_pattern} -> {required_in_pattern}", **dims)
+    if k_pattern != required_in_pattern:
+        k = rearrange(k, f"{k_pattern} -> {required_in_pattern}", **dims)
+    if v_pattern != required_in_pattern:
+        v = rearrange(v, f"{v_pattern} -> {required_in_pattern}", **dims)
+    return q, k, v
+
+
+def rearrange_out(out: torch.Tensor, out_pattern="b n s d", required_out_pattern="b n s d", dims=None):
+    dims = {} if dims is None else dims
+    if out_pattern != required_out_pattern:
+        out = rearrange(out, f"{required_out_pattern} -> {out_pattern}", **dims)
+    return out

diffsynth/core/attention/flash_attention.py

@@ -0,0 +1,118 @@
+import os
+
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+if os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID") is not None:
+    try:
+        from kernels import get_kernel
+        kernel_module = get_kernel(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_REPO_ID"), version=int(os.environ.get("DIFFSYNTH_FLASH_ATTN_KERNEL_VERSION")))
+        customized_flash_attn_func = kernel_module.flash_attn_func
+        CUSTOMIZED_FA_KERNEL_AVAILABLE = True
+    except Exception:
+        CUSTOMIZED_FA_KERNEL_AVAILABLE = False
+else:
+    CUSTOMIZED_FA_KERNEL_AVAILABLE = False
+
+try:
+    from flash_attn.cute import flash_attn_func as flash_attn_func_cute
+    FLASH_ATTN_4_AVAILABLE = True
+except Exception:
+    FLASH_ATTN_4_AVAILABLE = False
+
+try:
+    from flash_attn_3 import flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except Exception:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = hasattr(flash_attn, "flash_attn_func")
+except Exception:
+    FLASH_ATTN_2_AVAILABLE = False
+
+
+def _require(flag, name):
+    if not flag:
+        raise RuntimeError(f"{name} is not available.")
+
+
+def flash_attention_4(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_4_AVAILABLE, "FlashAttention-4")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = flash_attn_func_cute(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    if isinstance(out, tuple):
+        out = out[0]
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+def flash_attention_3(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_3_AVAILABLE, "FlashAttention-3")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = flash_attn_interface.flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    if isinstance(out, tuple):
+        out = out[0]
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+def flash_attention_2(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(FLASH_ATTN_2_AVAILABLE, "FlashAttention-2")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = flash_attn.flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+def flash_attention_customized_kernel(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None, is_causal=False, window_size=None):
+    _require(CUSTOMIZED_FA_KERNEL_AVAILABLE, "customized FlashAttention kernel")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    window_size = (window_size, window_size) if window_size is not None else (-1, -1)
+    out = customized_flash_attn_func(q, k, v, softmax_scale=scale, causal=is_causal, window_size=window_size)
+    if isinstance(out, tuple):
+        out = out[0]
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+if CUSTOMIZED_FA_KERNEL_AVAILABLE:
+    register_attention_implementation(
+        "customized_fa_kernel",
+        flash_attention_customized_kernel,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_4_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_4",
+        flash_attention_4,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_3_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_3",
+        flash_attention_3,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )
+if FLASH_ATTN_2_AVAILABLE:
+    register_attention_implementation(
+        "flash_attention_2",
+        flash_attention_2,
+        supports_causal=True,
+        supports_window=True,
+        supports_gqa=True,
+    )

diffsynth/core/attention/flex_attention.py

@@ -0,0 +1,28 @@
+import torch
+
+from .common import rearrange_qkv, rearrange_out
+
+try:
+    from torch.nn.attention.flex_attention import flex_attention as _flex_attention_func
+    FLEX_ATTN_AVAILABLE = True
+except Exception:
+    FLEX_ATTN_AVAILABLE = False
+
+# Compile flex attention lazily
+_compiled_flex_attention_func = None
+
+
+def get_flex_attention_func():
+    global _compiled_flex_attention_func
+    if _compiled_flex_attention_func is None:
+        _compiled_flex_attention_func = torch.compile(_flex_attention_func)
+    return _compiled_flex_attention_func
+
+
+def flex_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, score_mod=None):
+    assert FLEX_ATTN_AVAILABLE, "Flex Attention is not available. Please upgrade torch to 2.5.0 or later."
+    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    out = get_flex_attention_func()(query=q, key=k, value=v, block_mask=attn_mask, scale=scale, score_mod=score_mod)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out

diffsynth/core/attention/registry.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+
+@dataclass(frozen=True)
+class AttentionBackend:
+    func: Callable
+    supports_causal: bool = False
+    supports_window: bool = False
+    supports_mask: bool = False
+    supports_gqa: bool = False
+
+
+_ATTENTION_BACKENDS = {}
+
+
+def register_attention_implementation(
+    name: str,
+    func: Optional[Callable] = None,
+    supports_causal: bool = False,
+    supports_window: bool = False,
+    supports_mask: bool = False,
+    supports_gqa: bool = False,
+):
+    key = name.strip().lower()
+    _ATTENTION_BACKENDS[key] = AttentionBackend(
+        func=func,
+        supports_causal=supports_causal,
+        supports_window=supports_window,
+        supports_mask=supports_mask,
+        supports_gqa=supports_gqa,
+    )
+
+
+def get_attention_implementation(name: str) -> AttentionBackend:
+    key = name.strip().lower()
+    if key not in _ATTENTION_BACKENDS:
+        available = ", ".join(list_attention_implementations()) or "none"
+        raise ValueError(
+            f"Unknown attention implementation: {name!r}. "
+            f"Available implementations: {available}. "
+            f"Please check the environment variable `DIFFSYNTH_ATTENTION_IMPLEMENTATION`."
+        )
+    return _ATTENTION_BACKENDS[key]
+
+
+def list_attention_implementations():
+    return sorted(_ATTENTION_BACKENDS)

diffsynth/core/attention/sage_attention.py

@@ -0,0 +1,20 @@
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+try:
+    from sageattention import sageattn
+    SAGE_ATTN_AVAILABLE = True
+except Exception:
+    SAGE_ATTN_AVAILABLE = False
+
+def sage_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+    if not SAGE_ATTN_AVAILABLE:
+        raise RuntimeError("SageAttention is not available.")
+    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    out = sageattn(q, k, v, sm_scale=scale)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+if SAGE_ATTN_AVAILABLE:
+    register_attention_implementation("sage_attention", sage_attention)

diffsynth/core/attention/torch_attention.py

@@ -0,0 +1,82 @@
+import torch
+from einops import repeat
+
+from .common import TORCH_SUPPORT_GQA, rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+
+def torch_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, attn_mask=None, scale=None, is_causal=False, window_size=None):
+    if window_size is not None:
+        if is_causal or attn_mask is not None:
+            raise ValueError("Sliding window attention does not support `is_causal=True` or `attn_mask`.")
+        return torch_sdpa_sliding_window(q, k, v, window_size, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale)
+
+    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    if q.shape[1] != k.shape[1] or q.shape[1] != v.shape[1]:
+        # Grouped Query Attention
+        if TORCH_SUPPORT_GQA:
+            out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal, enable_gqa=True)
+        else:
+            # In low-version torch, `enable_gqa` is not supported.
+            k = repeat(k, "b n s d -> b (n m) s d", m=q.shape[1]//k.shape[1])
+            v = repeat(v, "b n s d -> b (n m) s d", m=q.shape[1]//v.shape[1])
+            out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal)
+    else:
+        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask, scale=scale, is_causal=is_causal)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+def torch_sdpa_sliding_window(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sliding_window: int, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+    required_in_pattern, required_out_pattern = "b n s d", "b n s d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+
+    B, N, S, D = q.shape
+    W = sliding_window
+    chunk_size = W
+    num_chunks = (S + chunk_size - 1) // chunk_size
+
+    output = torch.empty_like(q)
+    dtype = q.dtype
+    device = q.device
+    min_val = torch.finfo(dtype).min
+
+    for i in range(num_chunks):
+        q_start = i * chunk_size
+        q_end = min(q_start + chunk_size, S)
+        actual_chunk_size = q_end - q_start
+
+        kv_start = max(0, q_start - W)
+        kv_end = min(S, q_end + W)
+
+        q_chunk = q[:, :, q_start:q_end, :]
+        k_chunk = k[:, :, kv_start:kv_end, :]
+        v_chunk = v[:, :, kv_start:kv_end, :]
+
+        q_indices = torch.arange(q_start, q_end, device=device)
+        k_indices = torch.arange(kv_start, kv_end, device=device)
+        diff = q_indices.unsqueeze(1) - k_indices.unsqueeze(0)
+        valid = diff.abs() <= W
+
+        local_mask = torch.zeros(actual_chunk_size, kv_end - kv_start, dtype=dtype, device=device)
+        local_mask.masked_fill_(~valid, min_val)
+        local_mask = local_mask.unsqueeze(0).unsqueeze(0)
+
+        out_chunk = torch_sdpa(
+            q_chunk, k_chunk, v_chunk, attn_mask=local_mask, scale=scale
+        )
+        output[:, :, q_start:q_end, :] = out_chunk
+
+    output = rearrange_out(output, out_pattern, required_out_pattern, dims)
+    return output
+
+
+register_attention_implementation(
+    "torch",
+    torch_sdpa,
+    supports_window=True,
+    supports_causal=True,
+    supports_mask=True,
+    supports_gqa=True,
+)

diffsynth/core/attention/xformers_attention.py

@@ -0,0 +1,22 @@
+from .common import rearrange_qkv, rearrange_out
+from .registry import register_attention_implementation
+
+try:
+    import xformers.ops as xops
+    XFORMERS_AVAILABLE = True
+except Exception:
+    XFORMERS_AVAILABLE = False
+
+
+def xformers_attention(q, k, v, q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d", dims=None, scale=None):
+    if not XFORMERS_AVAILABLE:
+        raise RuntimeError("xformers is not available.")
+    required_in_pattern, required_out_pattern = "b s n d", "b s n d"
+    q, k, v = rearrange_qkv(q, k, v, q_pattern, k_pattern, v_pattern, required_in_pattern, dims)
+    out = xops.memory_efficient_attention(q, k, v, scale=scale)
+    out = rearrange_out(out, out_pattern, required_out_pattern, dims)
+    return out
+
+
+if XFORMERS_AVAILABLE:
+    register_attention_implementation("xformers", xformers_attention)