chore: import upstream snapshot with attribution

Diff-Transformer/multihead_attention.py

@@ -0,0 +1,95 @@
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from kernel.rotary import apply_rotary_emb
+from flash_attn import flash_attn_func
+try:
+    from apex.normalization import FusedRMSNorm as RMSNorm 
+except ModuleNotFoundError:
+    print("No fused RMSNorm")
+    from rms_norm import RMSNorm
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """torch.repeat_interleave(x, dim=1, repeats=n_rep)"""
+    bs, n_kv_heads, slen, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, None, :, :]
+        .expand(bs, n_kv_heads, n_rep, slen, head_dim)
+        .reshape(bs, n_kv_heads * n_rep, slen, head_dim)
+    )
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        depth,
+        num_heads,
+        num_kv_heads=None,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        
+        self.num_heads = num_heads
+        
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.n_rep = self.num_heads // self.num_kv_heads
+        
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** -0.5
+        
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.k_proj = nn.Linear(embed_dim, embed_dim // self.n_rep, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim // self.n_rep, bias=False)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+    
+    def forward(
+        self,
+        x,
+        rel_pos,
+        attn_mask=None,
+    ):
+        bsz, tgt_len, embed_dim = x.size()
+        src_len = tgt_len
+
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+
+        q = q.view(bsz, tgt_len, self.num_heads, self.head_dim)
+        k = k.view(bsz, src_len, self.num_kv_heads, self.head_dim)
+        v = v.view(bsz, src_len, self.num_kv_heads, self.head_dim)
+
+        q = apply_rotary_emb(q, *rel_pos, interleaved=True)
+        k = apply_rotary_emb(k, *rel_pos, interleaved=True)
+
+        offset = src_len - tgt_len
+        q = q.transpose(1, 2)
+        k = repeat_kv(k.transpose(1, 2), self.n_rep)
+        v = repeat_kv(v.transpose(1, 2), self.n_rep)
+        q *= self.scaling
+        attn_weights = torch.matmul(q, k.transpose(-1, -2))
+        if attn_mask is None:
+            attn_mask = torch.triu(
+                torch.zeros([tgt_len, src_len])
+                .float()
+                .fill_(float("-inf"))
+                .type_as(attn_weights),
+                1 + offset,
+            )
+        attn_weights = torch.nan_to_num(attn_weights)
+        attn_weights += attn_mask   
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(
+            attn_weights
+        )
+
+        attn = torch.matmul(attn_weights, v)
+        attn = attn.transpose(1, 2).reshape(bsz, tgt_len, self.num_heads * self.head_dim)
+
+        attn = self.out_proj(attn)
+        return attn