Use the type BlockTable (#1791)

This commit is contained in:
explainerauthors 2023-11-28 16:34:05 -08:00 committed by GitHub
parent 708e6c18b0
commit 6ed068a71a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -7,6 +7,10 @@ from vllm.sequence import Sequence, SequenceGroup, SequenceStatus
from vllm.utils import Device
# Mapping: logical block number -> physical block.
BlockTable = List[PhysicalTokenBlock]
class BlockAllocator:
"""Manages free physical token blocks for a device.
@ -26,7 +30,7 @@ class BlockAllocator:
self.num_blocks = num_blocks
# Initialize the free blocks.
self.free_blocks: List[PhysicalTokenBlock] = []
self.free_blocks: BlockTable = []
for i in range(num_blocks):
block = PhysicalTokenBlock(device=device,
block_number=i,
@ -51,10 +55,6 @@ class BlockAllocator:
return len(self.free_blocks)
# Mapping: logical block number -> physical block.
BlockTable = List[PhysicalTokenBlock]
class AllocStatus(enum.Enum):
"""Result for BlockSpaceManager.can_allocate