(self, idx: torch.Tensor)
| 210 | ) |
| 211 | |
| 212 | def build_mask_cache(self, idx: torch.Tensor) -> torch.Tensor: |
| 213 | ones = torch.ones( |
| 214 | (self.config.block_size, self.config.block_size), |
| 215 | device=idx.device, |
| 216 | dtype=torch.bool, |
| 217 | ) |
| 218 | return torch.tril(ones).unsqueeze(0).unsqueeze(0) |
| 219 | |
| 220 | def build_kv_caches( |
| 221 | self, idx: torch.Tensor, max_seq_length: int, rope_cache_length: int |