[Performance] e2e overheads reduction: Small followup diff (#7364)

This commit is contained in:
Alexander Matveev 2024-08-09 11:49:36 -04:00 committed by GitHub
parent 67abdbb42f
commit fc7b8d1eef
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 5 additions and 2 deletions

View File

@ -336,9 +336,9 @@ class BlockSpaceManagerV1(BlockSpaceManager):
# Assign the self-attention block tables for each sequence.
if len(wait_seqs) == 1:
self.block_tables[wait_seqs[0].seq_id] = block_table
self.block_tables[seq.seq_id] = block_table
else:
for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
for seq in wait_seqs:
self.block_tables[seq.seq_id] = block_table.copy()
# Allocate encoder sequence

View File

@ -655,6 +655,9 @@ class SequenceGroup:
return [seq for seq in self.seqs if not seq.is_finished()]
def get_finished_seqs(self) -> List[Sequence]:
if self.is_single_seq:
return self.seqs if self.seqs[0].is_finished() else []
return [seq for seq in self.seqs if seq.is_finished()]
def update_num_computed_tokens(self, num_new_computed_tokens: int):