Starts prefetching the next layer cache
(self, layer_idx: int)
| 26 | self.offload_kv_cache = offload_kv_cache |
| 27 | |
| 28 | def prefetch_layer(self, layer_idx: int): |
| 29 | "Starts prefetching the next layer cache" |
| 30 | if layer_idx < len(self): |
| 31 | with torch.cuda.stream(self.prefetch_stream): |
| 32 | # Prefetch next layer tensors to GPU |
| 33 | device = self.original_device[layer_idx] |
| 34 | self.key_cache[layer_idx] = self.key_cache[layer_idx].to(device, non_blocking=True) |
| 35 | self.value_cache[layer_idx] = self.value_cache[layer_idx].to(device, non_blocking=True) |
| 36 | |
| 37 | |
| 38 | def evict_previous_layer(self, layer_idx: int): |