Delete models/bs_roformer/attend_sw.py

models/bs_roformer/attend_sw.py

@@ -1,88 +0,0 @@
-import logging
-import os
-
-import torch
-import torch.nn.functional as F
-from packaging import version
-from torch import Tensor, einsum, nn
-from torch.nn.attention import SDPBackend, sdpa_kernel
-
-logger = logging.getLogger(__name__)
-
-
-class Attend(nn.Module):
-    def __init__(self, dropout: float = 0.0, flash: bool = False, scale=None):
-        super().__init__()
-        self.scale = scale
-        self.dropout = dropout
-        self.attn_dropout = nn.Dropout(dropout)
-
-        self.flash = flash
-        assert not (
-            flash and version.parse(torch.__version__) < version.parse("2.0.0")
-        ), "expected pytorch >= 2.0.0 to use flash attention"
-
-        self.cpu_backends = [
-            SDPBackend.FLASH_ATTENTION,
-            SDPBackend.EFFICIENT_ATTENTION,
-            SDPBackend.MATH,
-        ]
-        self.cuda_backends: list | None = None
-
-        if not torch.cuda.is_available() or not flash:
-            return
-
-        device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
-        device_version = version.parse(
-            f"{device_properties.major}.{device_properties.minor}"
-        )
-
-        if device_version >= version.parse("8.0"):
-            if os.name == "nt":
-                cuda_backends = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
-                logger.info(f"windows detected, {cuda_backends=}")
-            else:
-                cuda_backends = [SDPBackend.FLASH_ATTENTION]
-                logger.info(f"gpu compute capability >= 8.0, {cuda_backends=}")
-        else:
-            cuda_backends = [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
-            logger.info(f"gpu compute capability < 8.0, {cuda_backends=}")
-
-        self.cuda_backends = cuda_backends
-
-    def flash_attn(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
-        _, _heads, _q_len, _, _k_len, is_cuda, _device = (
-            *q.shape,
-            k.shape[-2],
-            q.is_cuda,
-            q.device,
-        )  # type: ignore
-
-        if self.scale is not None:
-            default_scale = q.shape[-1] ** -0.5
-            q = q * (self.scale / default_scale)
-
-        backends = self.cuda_backends if is_cuda else self.cpu_backends
-        with sdpa_kernel(backends=backends):  # type: ignore
-            out = F.scaled_dot_product_attention(
-                q, k, v, dropout_p=self.dropout if self.training else 0.0
-            )
-
-        return out
-
-    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
-        _q_len, _k_len, _device = q.shape[-2], k.shape[-2], q.device
-
-        scale = self.scale or q.shape[-1] ** -0.5
-
-        if self.flash:
-            return self.flash_attn(q, k, v)
-
-        sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale
-
-        attn = sim.softmax(dim=-1)
-        attn = self.attn_dropout(attn)
-
-        out = einsum("b h i j, b h j d -> b h i d", attn, v)
-
-        return out