This commit is contained in:
Woosuk Kwon 2023-02-14 09:05:39 +00:00
parent be1e2163c9
commit 263e91b386

View File

@ -210,8 +210,8 @@ class BlockSpaceManager:
} }
return block_number_mapping return block_number_mapping
def _free_blocks(self, blocks: Iterable[PhysicalTokenBlock]) -> None: def _free_block_table(self, block_table: BlockTable) -> None:
for block in blocks: for block in block_table:
if block.device == Device.GPU: if block.device == Device.GPU:
self.gpu_allocator.free(block) self.gpu_allocator.free(block)
else: else:
@ -219,10 +219,10 @@ class BlockSpaceManager:
def free(self, seq: Sequence) -> None: def free(self, seq: Sequence) -> None:
block_table = self.block_tables[seq.seq_id] block_table = self.block_tables[seq.seq_id]
self._free_blocks(block_table) self._free_block_table(block_table)
del self.block_tables[seq.seq_id] del self.block_tables[seq.seq_id]
def reset(self) -> None: def reset(self) -> None:
for block_table in self.block_tables.values(): for block_table in self.block_tables.values():
self._free_blocks(block_table) self._free_block_table(block_table)
self.block_tables.clear() self.block_tables.clear()