fix: optimize priority queue cleanup operation

vllm-project · Aug 7, 2024 · fd520b2 · fd520b2
1 parent bbd8f45
commit fd520b2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 17 deletions.
diff --git a/vllm/core/evictor_v1.py b/vllm/core/evictor_v1.py
@@ -104,15 +104,14 @@ def _cleanup_if_necessary(self):
             self._cleanup()
 
     def _cleanup(self):
-        new_priority_queue: List[Tuple[int, int, int]] = []
-        for last_accessed, neg_num_hashed_tokens, block_hash in (
-                self.priority_queue):
-            if (block_hash in self.free_table
-                    and self.free_table[block_hash].last_accessed
-                    == last_accessed):
-                heapq.heappush(
-                    new_priority_queue,
-                    (last_accessed, neg_num_hashed_tokens, block_hash))
+        new_priority_queue: List[Tuple[float, int, int]] = []
+
+        for block_hash, block in self.free_table.items():
+            new_priority_queue.append(
+                (block.last_accessed, -block.num_hashed_tokens,
+                 block.block_hash))
+        heapq.heapify(new_priority_queue)
+
         self.priority_queue = new_priority_queue
 
     @property

diff --git a/vllm/core/evictor_v2.py b/vllm/core/evictor_v2.py
@@ -119,14 +119,14 @@ def _cleanup_if_necessary(self):
             self._cleanup()
 
     def _cleanup(self):
-        new_priority_queue: List[Tuple[int, int, int, int]] = []
-        for last_accessed, neg_num_hashed_tokens, content_hash, block_id in (
-                self.priority_queue):
-            if (block_id in self.free_table and
-                    self.free_table[block_id].last_accessed == last_accessed):
-                heapq.heappush(new_priority_queue,
-                               (last_accessed, neg_num_hashed_tokens,
-                                content_hash, block_id))
+        new_priority_queue: List[Tuple[float, int, int, int]] = []
+
+        for block_id, block in self.free_table.items():
+            new_priority_queue.append(
+                (block.last_accessed, -block.num_hashed_tokens,
+                 block.content_hash, block_id))
+        heapq.heapify(new_priority_queue)
+
         self.priority_queue = new_priority_queue
 
     def remove(self, block_id: int):