| 10 | |
| 11 | class Test(unittest.TestCase): |
| 12 | def setUp(self): |
| 13 | self.layer_num = 30 |
| 14 | self.block_num = 3000 |
| 15 | self.head_num = 4 |
| 16 | self.block_size = 64 |
| 17 | self.head_dim = 128 |
| 18 | |
| 19 | self.swap_block_num = 100 |
| 20 | self.cache_shape = [self.block_num, self.head_num, self.block_size, self.head_dim] |
| 21 | assert self.swap_block_num <= self.block_num |
| 22 | |
| 23 | # cache layout: layer_num * [block_num, head_num, block_size, head_dim] |
| 24 | # buffer layout: [block_num, layer_num, head_num, block_size, head_dim] |
| 25 | # self.gpu_cache_tensors = self._init_gpu_cache() |
| 26 | |
| 27 | self.block_bytes = self.head_num * self.block_size * self.head_dim * 2 |
| 28 | buffer_total_bytes = self.swap_block_num * self.layer_num * self.block_bytes |
| 29 | self.cpu_buffer = cuda_host_alloc(buffer_total_bytes) |
| 30 | |
| 31 | self.gpu_block_ids = random.sample(list(range(self.block_num)), self.swap_block_num) |
| 32 | self.cpu_block_ids = list(range(self.swap_block_num)) |
| 33 | |
| 34 | def tearDown(self) -> None: |
| 35 | cuda_host_free(self.cpu_buffer) |