vllm-prefix-lookup

vllm/vllm/v1/core/kv_cache_manager.py · ef9339e · L266–328 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

266    def get_computed_blocks(
267        self, request: Request, max_cache_hit_length: int | None = None
268    ) -> tuple[KVCacheBlocks, int, int]:
269        """Get the computed (cached) blocks for the request.
270        Note that the computed blocks must be full.
271
272        Args:
273            request: The request to get the computed blocks.
274            max_cache_hit_length: Upper bound on the hit length. Defaults to
275                ``request.get_max_cache_hit_length()``; explicit bounds can only
276                reduce the request's cache hit limit.
277
278        Returns:
279            A tuple containing:
280                - A list of blocks that are computed for the request.
281                - The number of computed tokens.
282                - ``shared_prefix_boundary``: the block-aligned token position of
283                  a shared prefix that a sparse-retention group (Mamba / sliding
284                  window) has not cached yet (Marconi-style APC), or 0 if none.
285                  Pinned so sparse prefix-cache retention does not drop
286                  the junction and defeat cross-request reuse.
287
288        """
289        # We skip finding the prefix cache hit when prefix caching is
290        # disabled or the request is marked as skipping kv cache read
291        # (which happens when the request requires prompt logprobs
292        # or calls a pooling model with all pooling).
293        if not self.prefix_cache_lookup_enabled(request):
294            return self.empty_kv_cache_blocks, 0, 0
295
296        # NOTE: When all tokens hit the cache, we must recompute the last token
297        # to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1.
298        # This can trigger recomputation of an entire block, rather than just
299        # the single last token, because allocate_slots() requires
300        # num_computed_tokens to be block-size aligned. Removing this limitation
301        # could slightly improve performance in the future.
302        max_cache_hit_length = request.get_max_cache_hit_length(max_cache_hit_length)
303        computed_blocks, num_new_computed_tokens, num_uncached = (
304            self.coordinator.find_longest_cache_hit(
305                request.block_hashes, max_cache_hit_length
306            )
307        )
308
309        # When kv_cache_report_mode is "full", emit BlockStored events
310        # for the reused prefix cache blocks so that external consumers
311        # (e.g. gateway) can learn about them.
312        if (
313            num_new_computed_tokens > 0
314            and self.enable_kv_cache_events
315            and getattr(request, "kv_cache_report_mode", "incremental") == "full"
316        ):
317            self.coordinator.emit_cached_block_events(request, computed_blocks)
318
319        # The junction to pin is where the lagging sparse-retention group stops
320        # (``num_new_computed_tokens``) plus the uncached shared prefix -- i.e.
321        # the longest single-group hit. Sub-block gaps are left to the mask,
322        # which floors to the alignment boundary (a no-op there).
323        shared_prefix_boundary = (
324            num_new_computed_tokens + num_uncached if num_uncached else 0
325        )
326
327        blocks = self.create_kv_cache_blocks(computed_blocks)
328        return blocks, num_new_computed_tokens, shared_prefix_boundary