import torch from typing import Optional, Tuple def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) @torch.jit.script def apply_rotary_pos_emb(x, cos, sin): # NOTE: This could probably be moved to Triton # Handle a possible sequence length mismatch in between q and k cos = cos[:, :, : x.shape[-2], :] sin = sin[:, :, : x.shape[-2], :] return (x * cos) + (rotate_half(x) * sin) class RotaryEmbedding(torch.nn.Module): """ Rotary position embeddings from RoFormer (Su et. al, 2021). """ def __init__(self, dim_model: int, *_, **__): super().__init__() # Generate and save the inverse frequency buffer (non trainable) inv_freq = 1.0 / (10000 ** (torch.arange(0, dim_model, 2).float() / dim_model)) self.register_buffer("inv_freq", inv_freq) self._seq_len_cached = None self._cos_cached = None self._sin_cached = None def update_cos_sin_tables(self, x, seq_dimension=1): seq_len = x.shape[seq_dimension] # Reset the tables if the sequence length has changed, # or if we're on a new device (possibly due to tracing for instance) if ( seq_len != self._seq_len_cached or self._cos_cached.device != x.device or self._cos_cached.dtype != x.dtype ): self._seq_len_cached = seq_len t = torch.arange( x.shape[seq_dimension], device=x.device, dtype=torch.float32 ) freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(x.dtype)) emb = torch.cat((freqs, freqs), dim=-1).to(x.device) self._cos_cached = emb.cos()[None, None, :, :].to(x.dtype) self._sin_cached = emb.sin()[None, None, :, :].to(x.dtype) return self._cos_cached, self._sin_cached def forward( self, q: torch.Tensor, k: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: self._cos_cached, self._sin_cached = self.update_cos_sin_tables( k, seq_dimension=-2 ) return ( apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached), ) def __test_rope__(): dtype=torch.float16 batch=4 seqlen=2048 dim=4096 num_heads=32 dim_key_head=dim // num_heads x=torch.randn(batch,seqlen,num_heads,dim_key_head).to(dtype=dtype).to('cuda') rpe=RotaryEmbedding(dim_key_head).to(dtype=dtype).to('cuda') q,k=rpe(q=x,k=x) #__test_rope__()