vllm.v1.core.kv_cache_utils ¶

class FreeKVCacheBlockQueue:
    """This class organizes a list of KVCacheBlock objects to a doubly linked
    list of free blocks. We implement this class instead of using Python
    builtin deque to support removing a block in the middle of the queue
    in O(1) time. To close the performance gap to the builtin deque which is
    implemented in C++, this class does not allocate any Python objects when
    manipulating the linked list. Instead, this class manipulates the
    prev_free_block and next_free_block attributes of the given blocks.

    The queue is ordered by block ID in the beginning. When a block is allocated
    and then freed, it will be appended back with the eviction order:
    1. The least recent used block is at the front (LRU).
    2. If two blocks have the same last accessed time (allocated by the
       same sequence), the one with more hash tokens (the tail of a block
       chain) is at the front.
    Note that we maintain this order by reversing the block order when free
    blocks of a request. This operation is outside of this class.

    Args:
        blocks: A list of KVCacheBlock objects.
    """

    def __init__(self, blocks: list[KVCacheBlock]) -> None:
        self.num_free_blocks = len(blocks)

        # Initialize doubly links of consecutive blocks
        for i in range(self.num_free_blocks):
            if i > 0:
                blocks[i].prev_free_block = blocks[i - 1]
            if i < self.num_free_blocks - 1:
                blocks[i].next_free_block = blocks[i + 1]

        # Create a fake head and a tail block for the doubly linked list to
        # reduce branching in the code
        #
        # The implementation guaranteed that the fake head and tail
        # are NEVER got popped, so we could safely assume each real blocks
        # in the queue has prev and next blocks.
        self.fake_free_list_head = KVCacheBlock(block_id=-1)
        self.fake_free_list_tail = KVCacheBlock(block_id=-1)
        if self.num_free_blocks > 0:
            # Connect fake_head and fake_tail to the first and last block
            # respectively.
            self.fake_free_list_head.next_free_block = blocks[0]
            blocks[0].prev_free_block = self.fake_free_list_head
            self.fake_free_list_tail.prev_free_block = blocks[-1]
            blocks[-1].next_free_block = self.fake_free_list_tail
        else:
            # For empty list, simply connect the fake head and tail.
            self.fake_free_list_head.next_free_block = self.fake_free_list_tail
            self.fake_free_list_tail.prev_free_block = self.fake_free_list_head

    def popleft(self) -> KVCacheBlock:
        """Pop the first free block and reduce num_free_blocks by 1.

        Returns:
            The first free block.
        """
        if (
            self.fake_free_list_head.next_free_block is self.fake_free_list_tail
            or self.fake_free_list_head.next_free_block is None
        ):
            assert self.num_free_blocks == 0, (
                f"num_free_blocks ({self.num_free_blocks}) is out of sync "
                "with the free list."
            )
            raise ValueError("No free blocks available")

        first_block: KVCacheBlock = self.fake_free_list_head.next_free_block

        if first_block.next_free_block is None:
            # This should not happen if the block is from the free list.
            # It indicates a bug in the caller's logic.
            raise RuntimeError(
                "Invalid block found in popleft() "
                "which doesn't have a valid next_free_block"
            )

        # Connect fake_head and the next block of first_block (i.e. second block
        # or fake tail).
        self.fake_free_list_head.next_free_block = first_block.next_free_block
        first_block.next_free_block.prev_free_block = self.fake_free_list_head

        # Remove the block from the linked list.
        first_block.prev_free_block = first_block.next_free_block = None

        self.num_free_blocks -= 1
        return first_block

    def popleft_n(self, n: int, **kwargs) -> list[KVCacheBlock]:
        """Pop the first n free blocks and reduce num_free_blocks by n.

        Args:
            n: The number of blocks to pop.

        Returns:
            A list of n free blocks.
        """
        if n == 0:
            return []
        assert self.num_free_blocks >= n
        self.num_free_blocks -= n

        curr_block = self.fake_free_list_head.next_free_block
        # Pop n blocks from the head of the list
        ret = []
        for _ in range(n):
            assert curr_block is not None
            ret.append(curr_block)
            last_block = curr_block
            curr_block = curr_block.next_free_block
            # Reset prev_free_block and next_free_block of all popped blocks
            last_block.prev_free_block = None
            last_block.next_free_block = None

        if curr_block is not None:
            # The queue is not empty, connect the fake head to
            # the new first block.
            self.fake_free_list_head.next_free_block = curr_block
            curr_block.prev_free_block = self.fake_free_list_head
        return ret

    def remove(self, block: KVCacheBlock) -> None:
        """Remove a block in the free list and reduce num_free_blocks by 1.

        Args:
            block: The block to remove.
        """
        if block.prev_free_block is None or block.next_free_block is None:
            # This should not happen if the block is from the free list.
            # It indicates a bug in the caller's logic.
            raise RuntimeError(f"remove() called on an invalid block: {block}")

        # Link the previous block to the next block.
        block.prev_free_block.next_free_block = block.next_free_block
        # Link the next block to the previous block.
        block.next_free_block.prev_free_block = block.prev_free_block

        # Remove the block from the linked list.
        block.prev_free_block = block.next_free_block = None
        self.num_free_blocks -= 1

    def append(self, block: KVCacheBlock) -> None:
        """Put a block back into the free list and increase
        num_free_blocks by 1.

        Args:
            block: The block to append.
        """
        if self.fake_free_list_tail.prev_free_block is None:
            raise RuntimeError(
                "prev_free_block of fake_free_list_tail should always exist"
            )
        last_block: KVCacheBlock = self.fake_free_list_tail.prev_free_block

        # Connect the new block after the last block.
        last_block.next_free_block = block
        block.prev_free_block = last_block

        # Connect the fake tail after the new block.
        block.next_free_block = self.fake_free_list_tail
        self.fake_free_list_tail.prev_free_block = block

        self.num_free_blocks += 1

    def append_n(self, blocks: list[KVCacheBlock]) -> None:
        """Put a list of blocks back into the free list

        Args:
            blocks: The blocks to append.
        """
        if len(blocks) == 0:
            return

        last_block = self.fake_free_list_tail.prev_free_block
        assert last_block is not None, (
            "prev_free_block of fake_free_list_tail should always exist"
        )
        # Add inter-connections between consecutive blocks
        for block in blocks:
            block.prev_free_block = last_block
            last_block.next_free_block = block
            last_block = block

        # Connect the last block of <blocks> to the fake tail
        last_block.next_free_block = self.fake_free_list_tail
        self.fake_free_list_tail.prev_free_block = last_block

        self.num_free_blocks += len(blocks)

    def get_all_free_blocks(self) -> list[KVCacheBlock]:
        """Get all free blocks in the free list. Mainly used for testing.

        Returns:
            A list of free blocks.
        """
        ret = []
        if self.fake_free_list_head.next_free_block is None:
            raise RuntimeError(
                "next_free_block of fake_free_list_head should always exist"
            )
        # Start from the first block
        curr_block: KVCacheBlock = self.fake_free_list_head.next_free_block
        # As long as next_free_block is available, we haven't reached to
        # the fake tail yet.
        while curr_block.next_free_block is not None:
            ret.append(curr_block)
            curr_block = curr_block.next_free_block
        return ret

@dataclass
class KVCacheBlock:
    """KV-cache block metadata."""

    # Block ID, ranging from 0 to num_gpu_blocks - 1.
    block_id: int
    # Reference count.
    ref_cnt: int = 0
    # The hash key (block hash + group id) of the block, only available
    # when the block is full and cached.
    _block_hash: Optional[BlockHashWithGroupId] = None

    # Used to construct a doubly linked list for free blocks.
    # These two attributes should only be manipulated by FreeKVCacheBlockQueue.
    prev_free_block: Optional["KVCacheBlock"] = None
    next_free_block: Optional["KVCacheBlock"] = None

    # Whether the block is a null block that should never be cached.
    is_null: bool = False
    # Request type for the corresponding block
    _type_info: str = "default"
    # Last accesseed timestamp
    last_accessed: float = 0

    def incr_ref(self):
        self.ref_cnt += 1

    def decr_ref(self):
        self.ref_cnt -= 1

    @property
    def type_info(self) -> str:
        return self._type_info

    @type_info.setter
    def type_info(self, type_info: str):
        self._type_info = type_info

    @property
    def block_hash(self) -> Optional[BlockHashWithGroupId]:
        return self._block_hash

    @block_hash.setter
    def block_hash(self, block_hash: BlockHashWithGroupId):
        assert self.block_hash is None, (
            "The block already has a hash. This should not happen."
        )
        self._block_hash = block_hash

    def reset_hash(self):
        """Reset the block hash when the block is evicted."""
        self._block_hash = None

    def __repr__(self) -> str:
        # Use block_id instead of KVCacheBlock object to avoid calling __repr__
        # on KVCacheBlock object recursively.
        prev_block_id = self.prev_free_block.block_id if self.prev_free_block else None
        next_block_id = self.next_free_block.block_id if self.next_free_block else None
        return (
            f"KVCacheBlock(block_id={self.block_id}, "
            f"ref_cnt={self.ref_cnt}, "
            f"_block_hash={self._block_hash!r}, "
            f"prev_free_block={prev_block_id}, "
            f"next_free_block={next_block_id})"
        )

class PrefixCachingMetrics:
    """Metrics for prefix caching with a hit rate of the max recent N requests.

    Args:
        max_recent_requests: The number of the max recent requests to aggregate.
            Defaults to 1000.
    """

    def __init__(self, max_recent_requests: int = 1000):
        self.max_recent_requests = max_recent_requests
        # The current aggregated values.
        self.aggregated_requests = 0
        self.aggregated_query_total = 0
        self.aggregated_query_hit = 0
        # A deque of (requests, queries, hits) for the most recent requests.
        self.query_queue: deque[tuple[int, int, int]] = deque()

    def observe(self, stats: PrefixCacheStats):
        """Observe the prefix caching for a set of requests.

        This function is called with information gathered when new requests
        are being scheduled and are looking for computed blocks.

        When there are more than `max_recent_requests` requests, the oldest set
        of requests are removed from the metrics.

        Args:
            stats: The prefix cache stats.
        """
        # reset_prefix_cache was invoked before the current update.
        # Reset the metrics before aggregating the current stats.
        if stats.reset:
            self.reset()

        # DO NOT appending empty stats to avoid helpful info get kicked out
        # due to sliding window.
        if stats.requests == 0:
            return

        # Update the metrics.
        self.query_queue.append((stats.requests, stats.queries, stats.hits))
        self.aggregated_requests += stats.requests
        self.aggregated_query_total += stats.queries
        self.aggregated_query_hit += stats.hits

        # Remove the oldest stats until number of requests does not exceed
        # the limit.
        # NOTE: We preserve the latest added stats regardless.
        while (
            len(self.query_queue) > 1
            and self.aggregated_requests > self.max_recent_requests
        ):
            old_requests, old_queries, old_hits = self.query_queue.popleft()
            self.aggregated_requests -= old_requests
            self.aggregated_query_total -= old_queries
            self.aggregated_query_hit -= old_hits

    def reset(self):
        """Reset the metrics."""
        self.aggregated_requests = 0
        self.aggregated_query_total = 0
        self.aggregated_query_hit = 0
        self.query_queue.clear()

    @property
    def hit_rate(self) -> float:
        """Calculate the hit rate for the past N requests."""
        if self.aggregated_query_total == 0:
            return 0.0
        return self.aggregated_query_hit / self.aggregated_query_total

class WAFreeQueue(FreeKVCacheBlockQueue):
    """

    This class (WAQ) is an enhanced version of FreeKVCacheBlockQueue (FeQ).

    The main difference between WAQ and FeQ is that
    WAQ organizes multiple double-linked queues,
    totaling `num_free_blocks` blocks.
    Each double-linked queue represents a workload category;
    all blocks in a given queue belong to the same category.

    Initially, there is only one default queue
    since blocks have no assigned category.
    During allocation, each `KVCacheBlock` is assigned a
    `type_info` tag indicating its workload category.

    When `append` is called,
    WAQ inserts the KVCacheBlock into its category's queue.
    If the queue doesn't exist, it will be created.

    When `remove` is called,
    the block is removed from its category's queue in O(1) time.

    When `popLeft` is called, WAQ considers all queues and
    selects the block with the lowest reuse probability.
    Since blocks within each queue remain sorted by last accessed time (LRU),
    only the first block in each queue needs evaluation.
    The block with the minimal reuse probability is chosen as the victim.

    When a block is allocated and then freed,
    it will also be appended back with the eviction order:
    1. Insert to the corresponding workload queue.
    2. The least recent used block is at the front (LRU).
    3. If two blocks have the same last accessed time (allocated by the
       same sequence), the one with more hash tokens (the tail of a block
       chain) is at the front.
    Note that we maintain this order by reversing the block order when free
    blocks of a request. This operation is outside of this class.

    Args:
        blocks: A list of KVCacheBlock objects.
        wa_offline_param_path: The path of the offline parameter.

    """

    def __init__(
        self, blocks: list[KVCacheBlock], wa_offline_param_path: Optional[str]
    ) -> None:
        self.num_free_blocks = len(blocks)
        try:
            if wa_offline_param_path is None:
                raise FileNotFoundError
            with open(wa_offline_param_path) as file:
                self.offline_params = json.load(file)
        except FileNotFoundError:
            if wa_offline_param_path == "TEST":  # only for test
                print(
                    "[WABlockQueue] Use empty offline_params, note that "
                    "this should be called when testing"
                )
                self.offline_params = {"a": {}, "p": {}}
            else:
                print(
                    f"[WABlockQueue]offline param file {wa_offline_param_path} "
                    "not found. Please generate it by profiler"
                )
        self.life_range = 5

        print(
            f"[WABlockQueue] initialize with {wa_offline_param_path=}"
            f"blocks_number={len(blocks)}"
        )
        self.workload_to_head_tail: dict[str, list[Optional[KVCacheBlock]]] = dict()
        self.block_to_workload: dict[int, str] = dict()

        # Initialize the doubly linked list of free blocks.
        # Construct the original default linked list
        self.default_free_list_head: Optional[KVCacheBlock] = blocks[0]
        self.default_free_list_tail: Optional[KVCacheBlock] = blocks[-1]
        for i in range(self.num_free_blocks):
            if i > 0:
                blocks[i].prev_free_block = blocks[i - 1]
            if i < self.num_free_blocks - 1:
                blocks[i].next_free_block = blocks[i + 1]
            self.block_to_workload[i] = "default"
        self.workload_to_head_tail["default"] = [
            self.default_free_list_head,
            self.default_free_list_tail,
        ]

    def _get_priority(self, workload, cur_time, last_accessed, time_range):
        if (
            workload not in self.offline_params["a"]
            or workload not in self.offline_params["p"]
        ):
            a = 0
            p = 1
        else:
            a = self.offline_params["a"][workload]
            p = self.offline_params["p"][workload]
        priority = math.exp(-a * (cur_time - last_accessed))
        pt = math.exp(-a * (cur_time - last_accessed + time_range))
        return (priority - pt) * p

    def popleft(self) -> KVCacheBlock:
        """Pop the least priority free block and reduce num_free_blocks by 1.

        Returns:
            The first free block.
        """
        selected_block = None
        min_prioirty = 1e9
        cur_time = time.time()
        for workload, (list_head, _) in self.workload_to_head_tail.items():
            # if empty in list head, just skip it
            if list_head is None:
                continue
            assert workload == list_head.type_info, (
                f"{workload=}, {list_head.type_info=}, mismatch"
            )
            if workload == "default":  # new blocks, first
                selected_block = list_head
                break
            cur_priority = self._get_priority(
                workload, cur_time, list_head.last_accessed, self.life_range
            )
            if cur_priority < min_prioirty:
                min_prioirty = cur_priority
                selected_block = list_head

        if selected_block is None:
            raise ValueError("No free blocks available")
        self.remove(selected_block)
        return selected_block

    def remove(self, block: KVCacheBlock) -> None:
        """Remove a block in the free list and reduce num_free_blocks by 1.

        Args:
            block: The block to remove.
        """
        workload = block.type_info
        assert workload in self.workload_to_head_tail

        if block.prev_free_block is not None:
            # Link the previous block to the next block.
            block.prev_free_block.next_free_block = block.next_free_block
        if block.next_free_block is not None:
            # Link the next block to the previous block.
            block.next_free_block.prev_free_block = block.prev_free_block

        if block == self.workload_to_head_tail[workload][0]:
            # Update the head if the block is the head.
            self.workload_to_head_tail[workload][0] = block.next_free_block
        if block == self.workload_to_head_tail[workload][1]:
            # Update the tail if the block is the tail.
            self.workload_to_head_tail[workload][1] = block.prev_free_block

        # Remove the block from the linked list.
        block.prev_free_block = block.next_free_block = None
        self.num_free_blocks -= 1

    def append(self, block: KVCacheBlock) -> None:
        """Put a block back into the free list and increase
        num_free_blocks by 1.

        Args:
            block: The block to append.
        """
        workload = block.type_info
        assert workload != "default", (
            f"{workload=}, should not be default when first append"
        )

        # try get the last block in it's workload
        last_block = self.workload_to_head_tail.get(workload, [None, None])[1]

        if last_block is not None:
            # case1: workload exists and not empty, append new blocks
            # Link the last block to the new block.
            last_block.next_free_block = block
            block.prev_free_block = last_block
            self.workload_to_head_tail[workload][1] = block
        else:
            # case2: workload not exists, create new
            # Create a new list for the workload.
            self.workload_to_head_tail[workload] = [block, block]
            block.prev_free_block = None

        block.next_free_block = None
        self.num_free_blocks += 1

    def append_n(self, blocks: list[KVCacheBlock]) -> None:
        """Put a list of blocks back into the free list

        Args:
            blocks: The blocks to append.
        """
        if len(blocks) == 0:
            return
        for block in blocks:
            self.append(block)

    def get_all_free_blocks(self) -> list[KVCacheBlock]:
        """Get all free blocks in the free list. Mainly used for testing.

        Returns:
            A list of free blocks.
        """
        ret = []
        for workload, (list_head, _) in self.workload_to_head_tail.items():
            curr_block = list_head
            while curr_block is not None:
                ret.append(curr_block)
                curr_block = curr_block.next_free_block
        return ret

    def popleft_n(self, n: int, **kwargs) -> list[KVCacheBlock]:
        """Pop the first n free blocks and reduce num_free_blocks by n.

        Args:
            n: The number of blocks to pop.

        Returns:
            A list of n free blocks.
        """
        if n == 0:
            return []
        assert self.num_free_blocks >= n
        type_info = kwargs.get("type_info", "")
        res = []
        for _ in range(n):
            block = self.popleft()
            block.type_info = type_info
            res.append(block)
        return res