content: import torch
  from torch import nn, einsum
  import torch.nn.functional as F
  
  from collections import namedtuple
  from functools import wraps
  from packaging import version
  
  from einops import rearrange, repeat
  
  # constants
  
  Config = namedtuple('Config', ['enable_flash', 'enable_math', 'enable_mem_efficient'])
  
  # helpers
  
  def exists(val):
      return val is not None
  
  def once(fn):
      called = False
      @wraps(fn)
      def inner(x):
          nonlocal called
          if called:
              return
          called = True
          return fn(x)
      return inner
  
  print_once = once(print)
  
  # main class
  
  class Attend(nn.Module):
      def __init__(
          self,
          dropout = 0.,
          causal = False,
          flash = False
      ):
          super().__init__()
          self.dropout = dropout
          self.attn_dropout = nn.Dropout(dropout)
  
          self.causal = causal
          self.register_buffer("mask", None, persistent=False)
  
          self.flash = flash
          assert not (flash and version.parse(torch.__version__) < version.parse('2.0.0')), 'in order to use flash attention, you must be using pytorch 2.0 or above'
  
          # determine efficient attention configs for cuda and cpu
  
          self.cpu_config = Config(True, True, True)
          self.cuda_config = None
  
          if not torch.cuda.is_available() or not flash:
              return
  
          device_properties = torch.cuda.get_device_properties(torch.device('cuda'))
  
          if device_properties.major == 8 and device_properties.minor == 0:
              print_once('A100 GPU detected, using flash attention if input tensor is on cuda')
              self.cuda_config = Config(True, False, False)
          else:
              print_once('Non-A100 GPU detected, using math or mem efficient attention if input tensor is on cuda')
              self.cuda_config = Config(False, True, True)
  
      def flash_attn(self, q, k, v, mask = None):
          _, heads, q_len, _, k_len, is_cuda = *q.shape, k.shape[-2], q.is_cuda
  
          k = repeat(k, 'b ... -> b h ...', h = heads)
          v = repeat(v, 'b ... -> b h ...', h = heads)
  
          causal = self.causal
  
          if exists(mask):
              mask = rearrange(mask, 'b j -> b 1 1 j')
              mask = mask.expand(-1, heads, q_len, -1)
  
              if causal:
                  causal_mask = torch.ones((q_len, k_len), device = q.device, dtype = torch.bool).triu(k_len - q_len + 1)
                  mask = mask & ~causal_mask                
                  causal = False
  
          config = self.cuda_config if is_cuda else self.cpu_config
  
          with torch.backends.cuda.sdp_kernel(**config._asdict()):
              out = F.scaled_dot_product_attention(
                  q, k, v,
                  attn_mask = mask,
                  dropout_p = self.dropout if self.training else 0., 
                  is_causal = causal
              )
  
          return out
  
      def forward(self, q, k, v, mask = None, attn_bias = None):
          """
          einstein notation
          b - batch
          h - heads
          n, i, j - sequence length (base sequence length, source, target)
          d - feature dimension
          """
  
          n, device = q.shape[-2], q.device
  
          scale = q.shape[-1] ** -0.5
  
          if self.flash:
              assert not exists(attn_bias), 'attention bias not supported for flash attention'
              return self.flash_attn(q, k, v, mask = mask)
  
          # similarity
  
          sim = einsum("b h i d, b j d -> b h i j", q, k) * scale
  
          # attention bias
  
          if exists(attn_bias):
              sim = sim + attn_bias
  
          # key padding mask
  
          if exists(mask):
              mask = rearrange(mask, 'b j -> b 1 1 j')
              sim = sim.masked_fill(~mask, -torch.finfo(sim.dtype).max)
  
          # causal mask
  
          if self.causal:
              i, j = sim.shape[-2:]
              causal_mask = torch.ones((i, j), device = sim.device, dtype = torch.bool).triu(j - i + 1)
              sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)
  
          # attention
  
          attn = sim.softmax(dim=-1)
          attn = self.attn_dropout(attn)
  
          # aggregate values
  
          out = einsum("b h i j, b j d -> b h i d", attn, v)
  
          return out
  