re-enable rotary cache

modelling_RW.py

@@ -60,6 +60,7 @@ class RotaryEmbedding(torch.nn.Module):
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.head_dim = head_dim
         self.seq_len_cached = None
+        self.start_idx = None
         self.batch_size_cached = None
         self.cos_cached: torch.Tensor | None = None
         self.sin_cached: torch.Tensor | None = None
@@ -71,28 +72,23 @@ class RotaryEmbedding(torch.nn.Module):
         dtype=torch.bfloat16,
         start_idx: int = 0,
     ) -> torch.Tensor:
-        # if seq_len != self.seq_len_cached:
-        self.seq_len_cached = seq_len
-        t = torch.arange(start_idx, start_idx+seq_len, device=device).type_as(self.inv_freq)
-        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
-        emb = torch.cat((freqs, freqs), dim=-1).to(device)
+        if seq_len != self.seq_len_cached and self.start_idx != start_idx:
+            self.seq_len_cached = seq_len
+            t = torch.arange(start_idx, start_idx+seq_len, device=device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(device)
 
-        if dtype in [torch.float16, torch.bfloat16]:
-            emb = emb.float()
+            if dtype in [torch.float16, torch.bfloat16]:
+                emb = emb.float()
 
-        self.cos_cached = emb.cos()[None, :, :]
-        self.sin_cached = emb.sin()[None, :, :]
+            self.cos_cached = emb.cos()[None, :, :]
+            self.sin_cached = emb.sin()[None, :, :]
 
-        self.cos_cached = self.cos_cached.type(dtype)
-        self.sin_cached = self.sin_cached.type(dtype)
+            self.cos_cached = self.cos_cached.type(dtype)
+            self.sin_cached = self.sin_cached.type(dtype)
 
         return self.cos_cached, self.sin_cached
 
-    def forward(self, q, k, start_idx=0):
-        batch, seq_len, head_dim = q.shape
-        cos, sin = self.cos_sin(seq_len, q.device, q.dtype, start_idx=start_idx)
-        return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
-
 
 def _make_causal_mask(
     input_ids_shape: torch.Size, device: torch.device, past_key_values_length: int