diff --git a/cacheflow/master/block_manager.py b/cacheflow/master/block_manager.py index 488049a2..594da808 100644 --- a/cacheflow/master/block_manager.py +++ b/cacheflow/master/block_manager.py @@ -165,10 +165,11 @@ class BlockSpaceManager: for cpu_block in block_table: if cpu_block in mapping: - new_block_table.append(mapping[cpu_block]) - continue - gpu_block = self.gpu_allocator.allocate() - mapping[cpu_block] = gpu_block + gpu_block = mapping[cpu_block] + else: + gpu_block = self.gpu_allocator.allocate() + mapping[cpu_block] = gpu_block + new_block_table.append(gpu_block) # Free the CPU block swapped in to GPU. self.cpu_allocator.free(cpu_block) self.block_tables[seq.seq_id] = new_block_table @@ -194,12 +195,11 @@ class BlockSpaceManager: for gpu_block in block_table: if gpu_block in mapping: - new_block_table.append(mapping[gpu_block]) - continue - cpu_block = self.cpu_allocator.allocate() + cpu_block = mapping[gpu_block] + else: + cpu_block = self.cpu_allocator.allocate() + mapping[gpu_block] = cpu_block new_block_table.append(cpu_block) - - mapping[gpu_block] = cpu_block # Free the GPU block swapped out to CPU. self.gpu_allocator.free(gpu_block) self.block_tables[seq.seq_id] = new_block_table