79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
import enum
|
|
from typing import List, Optional
|
|
|
|
from cacheflow.block import LogicalTokenBlock
|
|
from cacheflow.decoding import DecodingParams
|
|
|
|
|
|
class SequenceStatus(enum.Enum):
|
|
PENDING = enum.auto()
|
|
RUNNING = enum.auto()
|
|
SWAPPED = enum.auto()
|
|
FINISHED = enum.auto()
|
|
|
|
|
|
class Sequence:
|
|
|
|
def __init__(
|
|
self,
|
|
seq_id: int,
|
|
token_ids: List[int],
|
|
max_context_len: int,
|
|
block_size: int,
|
|
) -> None:
|
|
self.seq_id = seq_id
|
|
self.max_context_len = max_context_len
|
|
self.block_size = block_size
|
|
|
|
self.logical_token_blocks: List[LogicalTokenBlock] = []
|
|
# Initialize the logical token blocks with the given token ids.
|
|
self.append(token_ids)
|
|
|
|
self.status = SequenceStatus.PENDING
|
|
|
|
def add_block(self) -> None:
|
|
block = LogicalTokenBlock(
|
|
block_number=len(self.logical_token_blocks),
|
|
block_size=self.block_size,
|
|
)
|
|
self.logical_token_blocks.append(block)
|
|
|
|
def append(self, token_ids: List[int]) -> None:
|
|
while token_ids:
|
|
if not self.logical_token_blocks:
|
|
self.add_block()
|
|
|
|
last_block = self.logical_token_blocks[-1]
|
|
if last_block.is_full():
|
|
self.add_block()
|
|
last_block = self.logical_token_blocks[-1]
|
|
|
|
num_empty_slots = last_block.get_num_empty_slots()
|
|
last_block.append(token_ids[:num_empty_slots])
|
|
token_ids = token_ids[num_empty_slots:]
|
|
|
|
def get_token_ids(self) -> List[int]:
|
|
token_ids: List[int] = []
|
|
for block in self.logical_token_blocks:
|
|
token_ids.extend(block.get_token_ids())
|
|
return token_ids
|
|
|
|
|
|
class SequenceGroup:
|
|
|
|
def __init__(
|
|
self,
|
|
group_id: int,
|
|
seqs: List[Sequence],
|
|
decoding_params: DecodingParams,
|
|
) -> None:
|
|
self.group_id = group_id
|
|
self.seqs = seqs
|
|
self.decoding_params = decoding_params
|
|
|
|
def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
|
|
if status is None:
|
|
return len(self.seqs)
|
|
else:
|
|
return len([seq for seq in self.seqs if seq.status == status])
|