oops, actually add cosine sim flash attention

Author: lucidrains

memory_efficient_attention_pytorch/cosine_sim_flash_attention.py

@@ -0,0 +1,218 @@
+import math
+import torch
+from functools import partial
+from torch import nn, einsum
+import torch.nn.functional as F
+from torch.autograd.function import Function
+
+from einops import rearrange
+
+# constants
+
+EPSILON = 1e-6
+
+# helper functions
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def l2norm(t):
+    return F.normalize(t, dim = -1)
+
+# flash attention forwards and backwards
+
+class FlashAttentionFunction(Function):
+    @staticmethod
+    @torch.no_grad()
+    def forward(ctx, q, k, v, mask, scale, causal, q_bucket_size, k_bucket_size):
+        device = q.device
+        max_neg_value = -torch.finfo(q.dtype).max
+        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
+
+        k_len = k.shape[-2] # in cosine sim attention, row sums are bounded by key / values sequence length
+
+        o = torch.zeros_like(q)
+        all_row_sums = torch.zeros((*q.shape[:-1], 1), device = device)
+
+        q = q * scale
+
+        if not exists(mask):
+            mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
+        else:
+            mask = mask.split(q_bucket_size, dim = -2)
+
+        row_splits = zip(
+            q.split(q_bucket_size, dim = -2),
+            o.split(q_bucket_size, dim = -2),
+            mask,
+            all_row_sums.split(q_bucket_size, dim = -2),
+        )
+
+        for ind, (qc, oc, row_mask, row_sums) in enumerate(row_splits):
+            q_start_index = ind * q_bucket_size - qk_len_diff
+
+            col_splits = zip(
+                k.split(k_bucket_size, dim = -2),
+                v.split(k_bucket_size, dim = -2),
+            )
+
+            for k_ind, (kc, vc) in enumerate(col_splits):
+                k_start_index = k_ind * k_bucket_size
+
+                attn_weights = einsum('... i d, ... j d -> ... i j', qc, kc)
+
+                if exists(row_mask):
+                    attn_weights.masked_fill_(~row_mask, max_neg_value)
+
+                if causal and q_start_index < (k_start_index + k_bucket_size - 1):
+                    causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype = torch.bool, device = device).triu(q_start_index - k_start_index + 1)
+                    attn_weights.masked_fill_(causal_mask, max_neg_value)
+
+                attn_weights -= scale
+                exp_weights = torch.exp(attn_weights)
+
+                if exists(row_mask):
+                    exp_weights.masked_fill_(~row_mask, 0.)
+
+                block_row_sums = exp_weights.sum(dim = -1, keepdims = True).clamp(min = EPSILON)
+
+                exp_values = einsum('... i j, ... j d -> ... i d', exp_weights, vc)
+
+                oc.add_(exp_values / k_len)
+                row_sums.add_(block_row_sums)
+
+        ctx.args = (scale, causal, mask, q_bucket_size, k_bucket_size)
+        ctx.save_for_backward(q, k, v, o, all_row_sums)
+
+        o.mul_(k_len / all_row_sums)
+
+        return o
+
+    @staticmethod
+    @torch.no_grad()
+    def backward(ctx, do):
+        scale, causal, mask, q_bucket_size, k_bucket_size = ctx.args
+        q, k, v, o, l = ctx.saved_tensors
+
+        device = q.device
+
+        max_neg_value = -torch.finfo(q.dtype).max
+        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)
+
+        dq = torch.zeros_like(q)
+        dk = torch.zeros_like(k)
+        dv = torch.zeros_like(v)
+
+        row_splits = zip(
+            q.split(q_bucket_size, dim = -2),
+            o.split(q_bucket_size, dim = -2),
+            do.split(q_bucket_size, dim = -2),
+            mask,
+            l.split(q_bucket_size, dim = -2),
+            dq.split(q_bucket_size, dim = -2)
+        )
+
+        for ind, (qc, oc, doc, row_mask, lc, dqc) in enumerate(row_splits):
+            q_start_index = ind * q_bucket_size - qk_len_diff
+
+            col_splits = zip(
+                k.split(k_bucket_size, dim = -2),
+                v.split(k_bucket_size, dim = -2),
+                dk.split(k_bucket_size, dim = -2),
+                dv.split(k_bucket_size, dim = -2),
+            )
+
+            for k_ind, (kc, vc, dkc, dvc) in enumerate(col_splits):
+                k_start_index = k_ind * k_bucket_size
+
+                qc_scaled = qc * scale
+                attn_weights = einsum('... i d, ... j d -> ... i j', qc_scaled, kc)
+
+                if causal and q_start_index < (k_start_index + k_bucket_size - 1):
+                    causal_mask = torch.ones((qc.shape[-2], kc.shape[-2]), dtype = torch.bool, device = device).triu(q_start_index - k_start_index + 1)
+                    attn_weights.masked_fill_(causal_mask, max_neg_value)
+
+                exp_attn_weights = torch.exp(attn_weights)
+
+                if exists(row_mask):
+                    exp_attn_weights.masked_fill_(~row_mask, 0.)
+
+                p = exp_attn_weights / lc
+
+                dv_chunk = einsum('... i j, ... i d -> ... j d', p, doc)
+                dp = einsum('... i d, ... j d -> ... i j', doc, vc)
+
+                D = (doc * oc).sum(dim = -1, keepdims = True)
+                ds = p * scale * (dp - D)
+
+                dq_chunk = einsum('... i j, ... j d -> ... i d', ds, kc)
+                dk_chunk = einsum('... i j, ... i d -> ... j d', ds, qc)
+
+                dqc.add_(dq_chunk)
+                dkc.add_(dk_chunk)
+                dvc.add_(dv_chunk)
+
+        return dq, dk, dv, None, None, None, None, None
+
+# main class
+
+# flash attention for cosine sim attention
+# a bit less complicated, as no more need to worry about softmax numerical stability, and row sums are bounded
+
+class FlashAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        scale = 16,
+        heads = 8,
+        dim_head = 64,
+        causal = False,
+        q_bucket_size = 512,
+        k_bucket_size = 1024
+    ):
+        super().__init__()
+        self.heads = heads
+
+        self.scale = scale
+        self.causal = causal
+
+        inner_dim = heads * dim_head
+
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim)
+
+        # memory efficient attention related parameters
+        # can be overriden on forward
+        self.q_bucket_size = q_bucket_size
+        self.k_bucket_size = k_bucket_size
+
+    def forward(
+        self,
+        x,
+        context = None,
+        mask = None,
+        q_bucket_size = None,
+        k_bucket_size = None,
+    ):
+        q_bucket_size = default(q_bucket_size, self.q_bucket_size)
+        k_bucket_size = default(k_bucket_size, self.k_bucket_size)
+
+        h = self.heads
+        context = default(context, x)
+
+        q = self.to_q(x)
+        k, v = self.to_kv(context).chunk(2, dim = -1)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
+
+        q, k = map(l2norm, (q, k))
+
+        out = FlashAttentionFunction.apply(q, k, v, mask, self.scale, self.causal, q_bucket_size, k_bucket_size)
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'memory-efficient-attention-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.20',
+  version = '0.0.22',
   license='MIT',
   description = 'Memory Efficient Attention - Pytorch',
   long_description_content_type = 'text/markdown',