Starts prefetching the next layer cache
(self, layer_idx: int, device: torch.device)
| 25 | config: Phi3Config |
| 26 | """ |
| 27 | def prefetch_layer(self, layer_idx: int, device: torch.device): |
| 28 | "Starts prefetching the next layer cache" |
| 29 | with torch.cuda.stream(self.prefetch_stream): |
| 30 | # Prefetch next layer tensors to GPU |
| 31 | for name, param in self.layers[layer_idx].named_parameters(): |
| 32 | param.data = param.data.to(device, non_blocking=True) |
| 33 | |
| 34 | def evict_previous_layer(self, layer_idx: int): |
| 35 | "Moves the previous layer cache to the CPU" |