From fc7b8d1eefcbe837a56b7c080509417fe5167e6c Mon Sep 17 00:00:00 2001 From: Alexander Matveev <59768536+alexm-neuralmagic@users.noreply.github.com> Date: Fri, 9 Aug 2024 11:49:36 -0400 Subject: [PATCH] [Performance] e2e overheads reduction: Small followup diff (#7364) --- vllm/core/block_manager_v1.py | 4 ++-- vllm/sequence.py | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm/core/block_manager_v1.py b/vllm/core/block_manager_v1.py index 622aca66a96d..ad26d3c516ff 100644 --- a/vllm/core/block_manager_v1.py +++ b/vllm/core/block_manager_v1.py @@ -336,9 +336,9 @@ def allocate(self, seq_group: SequenceGroup) -> None: # 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 diff --git a/vllm/sequence.py b/vllm/sequence.py index ba477efc54dd..fd2dc9656678 100644 --- a/vllm/sequence.py +++ b/vllm/sequence.py @@ -655,6 +655,9 @@ def get_unfinished_seqs(self) -> List[Sequence]: 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):