vllm-block-pool
vllm/vllm/v1/core/block_pool.py · ef9339e · L681–830 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
681 def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
682 """Get new blocks from the free block pool.
683
684 Note that we do not check block cache in this function.
685
686 Args:
687 num_blocks: The number of blocks to allocate.
688
689 Returns:
690 A list of new block.
691
692 """
693 if num_blocks > self.get_num_free_blocks():
694 raise ValueError(f"Cannot get {num_blocks} free blocks from the pool")
695
696 ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)
697
698 if self._reuse_watchers:
699 self._notify_reuse(ret)
700
701 # In order to only iterate the list once, we duplicated code a bit
702 if self.enable_caching:
703 for block in ret:
704 self._maybe_evict_cached_block(block)
705 assert block.ref_cnt == 0
706 block.ref_cnt += 1
707 if self.metrics_collector:
708 self.metrics_collector.on_block_allocated(block)
709 else:
710 for block in ret:
711 assert block.ref_cnt == 0
712 block.ref_cnt += 1
713 if self.metrics_collector:
714 self.metrics_collector.on_block_allocated(block)
715 return ret
716
717 def _notify_reuse(self, blocks: Iterable[KVCacheBlock]) -> None:
718 for block in blocks:
719 watcher = self._reuse_watchers.pop(block.block_id, None)
720 if watcher is not None:
721 watcher(block)
722
723 def unpin_blocks(
724 self,
725 blocks: Iterable[KVCacheBlock],
726 on_reuse: Callable[[KVCacheBlock], None],
727 ) -> None:
728 """Release references to blocks that stay readable until reused.
729
730 The blocks become last-resort eviction candidates regardless of prefix
731 caching: they join the tail of the free queue and count as free.
732 ``on_reuse`` fires when ``get_new_blocks`` hands a block out (or the
733 cache is reset), so the caller can stop reading it.
734 """
735 released: list[KVCacheBlock] = []
736 for block in blocks:
737 assert block.ref_cnt > 0 and not block.is_null
738 block.ref_cnt -= 1
739 self._reuse_watchers[block.block_id] = on_reuse
740 if block.ref_cnt == 0:
741 released.append(block)
742 self.free_block_queue.append_n(released)
743
744 def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
745 """If a block is cached in `cached_block_hash_to_block`, we reset its hash
746 metadata and evict it from the cache.
747
748 Args:
749 block: The block to evict.
750
751 Returns:
752 True if the block is evicted, False otherwise.
753
754 """
755 # Clean up metrics tracking first to prevent leaks
756 if self.metrics_collector:
757 self.metrics_collector.on_block_evicted(block)
758
759 evicted_hashes = self._remove_cached_block_hashes(block)
760 if not evicted_hashes:
761 # The block doesn't have hash, eviction is not needed
762 return False
763
764 self._emit_block_removed_events(evicted_hashes)
765 return True
766
767 def touch(self, blocks: Sequence[KVCacheBlock]) -> None:
768 """Touch a block increases its reference count by 1, and may remove
769 the block from the free queue. This is used when a block is hit by
770 another request with the same prefix.
771
772 Args:
773 blocks: A list of blocks to touch.
774
775 """
776 for block in blocks:
777 if block.is_null:
778 continue
779 # ref_cnt=0 means this block is in the free list (i.e. eviction
780 # candidate), so remove it.
781 if block.ref_cnt == 0:
782 self.free_block_queue.remove(block)
783 block.ref_cnt += 1
784 if self.metrics_collector:
785 self.metrics_collector.on_block_accessed(block)
786
787 def is_block_writable(self, block: KVCacheBlock) -> bool:
788 """Return whether a block can be mutated by its sole owner."""
789 return not block.is_null and block.ref_cnt == 1 and block.block_hash is None
790
791 def free_blocks(
792 self, ordered_blocks: Iterable[KVCacheBlock], evict_first: bool = False
793 ) -> None:
794 """Free a list of blocks. The blocks should be ordered by their
795 eviction priority, where the first block will be evicted first.
796
797 Args:
798 ordered_blocks: A list of blocks to free ordered by their eviction
799 priority.
800 evict_first: Also put freed cached blocks at the front of the free
801 queue, so they are evicted before other cached blocks. For
802 blocks whose contents another tier already holds.
803
804 """
805 # Identify blocks with hash (LRU cache) and without it (never match APC)
806 blocks_to_evict_last = []
807 blocks_to_evict_first = []
808 other_pools: dict[BlockPool, list[KVCacheBlock]] = {}
809 for block in ordered_blocks:
810 if block.is_null: # never pooled; its ref_cnt is not maintained
811 continue
812 if block.pool is not None and block.pool is not self._weak_self:
813 pool = block.pool()
814 if pool is not None:
815 other_pools.setdefault(pool, []).append(block)
816 continue
817 block.ref_cnt -= 1
818 if block.ref_cnt == 0:
819 if block.block_hash is None or not self.enable_caching or evict_first:
820 # LIFO reuse of non-cached blocks for better GPU locality,
821 # and of cached ones when asked to evict them first.
822 blocks_to_evict_first.append(block)
823 else:
824 # FIFO reuse of cached blocks for LRU eviction behavior.
825 blocks_to_evict_last.append(block)
826
827 # Blocks to reuse first are prepended to the front of the free queue.
828 self.free_block_queue.prepend_n(blocks_to_evict_first)
829 # Blocks to reuse last are appended to the end of the free queue.
830 self.free_block_queue.append_n(blocks_to_evict_last)