| 45 | self.sin_cached: torch.Tensor | None = None |
| 46 | |
| 47 | def cos_sin( |
| 48 | self, |
| 49 | seq_len: int, |
| 50 | device="cuda", |
| 51 | dtype=torch.bfloat16, |
| 52 | ) -> torch.Tensor: |
| 53 | if seq_len != self.seq_len_cached: |
| 54 | self.seq_len_cached = seq_len |
| 55 | t = torch.arange(seq_len, device=device).type_as(self.inv_freq) |
| 56 | freqs = torch.einsum("i,j->ij", t, self.inv_freq) |
| 57 | emb = torch.cat((freqs, freqs), dim=-1).to(device) |
| 58 | |
| 59 | if dtype in [torch.float16, torch.bfloat16]: |
| 60 | emb = emb.float() |
| 61 | |
| 62 | self.cos_cached = emb.cos()[None, :, :] |
| 63 | self.sin_cached = emb.sin()[None, :, :] |
| 64 | |
| 65 | self.cos_cached = self.cos_cached.type(dtype) |
| 66 | self.sin_cached = self.sin_cached.type(dtype) |
| 67 | |
| 68 | return self.cos_cached, self.sin_cached |
| 69 | |
| 70 | def forward(self, _q, _k): |
| 71 | batch, seq_len, num_heads, head_dim = _q.shape |