Update RotaryEmbedding caching

modelling_RW.py

@@ -56,13 +56,12 @@ class RotaryEmbedding(torch.nn.Module):
         base=10000,
     ):
         super().__init__()
-        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        inv_freq = 1.0 / (base
+                          **(torch.arange(0, head_dim, 2).float() / head_dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.head_dim = head_dim
-        self.seq_len_cached = None
-        self.batch_size_cached = None
-        self.cos_cached: torch.Tensor | None = None
-        self.sin_cached: torch.Tensor | None = None
+        self.cos_cache_dict: dict = {}
+        self.sin_cache_dict: dict = {}
 
     def cos_sin(
         self,
@@ -70,27 +69,24 @@ class RotaryEmbedding(torch.nn.Module):
         device="cuda",
         dtype=torch.bfloat16,
     ) -> torch.Tensor:
-        if seq_len != self.seq_len_cached:
+        if seq_len not in self.cos_cache_dict or seq_len not in self.sin_cache_dict:
             self.seq_len_cached = seq_len
             t = torch.arange(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()
+            self.cos_cache_dict[seq_len] = emb.cos()[None, :, :].type(dtype)
+            self.sin_cache_dict[seq_len] = emb.sin()[None, :, :].type(dtype)
 
-            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)
-
-        return self.cos_cached, self.sin_cached
+        return self.cos_cache_dict[seq_len], self.sin_cache_dict[seq_len]
 
     def forward(self, q, k):
         batch, seq_len, head_dim = q.shape
+        assert seq_len is not None, "seq_len must be known and not None"
         cos, sin = self.cos_sin(seq_len, q.device, q.dtype)
-        return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+        return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) *
+                                                                sin)
 
 
 def _make_causal_mask(