Prefix Caching:KV block 跨請求複用

源码版本v0.25.1

職責

很多請求會共享一段相同的前綴:system prompt、chat template、函數定義、少樣本示例 (few-shot examples)。這些 token 的 K/V 已經算過了,如果下一條請求能直接複用,就省掉一整段 prefill。前綴快取 (prefix caching) 就是這件事:把每個 KV block 按「內容 hash」做指紋,下一條請求來了先查 hash,命中就把已有 block 引用過來,不重算。

啟用開關是 CacheConfig.enable_prefix_caching,預設 True(cache.py:93)。EngineCore.__init__ 在 enable_prefix_caching 或有 KV connector 時構造 request_block_hasher(core.py:210-219),它本質是 get_request_block_hasher(kv_cache_utils.py:673)返回的閉包,對 request 的 token sequence 按 hash_block_size 切片、鏈式哈希。每個 block 的 hash 把「父 block 的 hash + 當前 block 的 token ids + extra_keys(LoRA / 多模態)」一起餵給哈希函數(kv_cache_utils.py:577-604),所以一個 block 的指紋就是「從開頭到這個 block 的整條前綴」的指紋。

KVCacheManager 把命中查詢委託給 KVCacheCoordinator.find_longest_cache_hit(kv_cache_coordinator.py:364),coordinator 再下發給單類型 manager。命中後 allocate_slots 複用這些 block,只對未命中部分分配新 block。Scheduler.schedule 在排程每個請求時先調 get_computed_blocks(kv_cache_manager.py:206),命中就少分配 token 預算。

設計動機

• block 級而不是 token 級複用:KV cache 按 PagedAttention 的 block 切分,自然以 block 為最小單位,hash 也按 block 算,粒度大、查錶快。
• 鏈式哈希保證前綴語義:hash_block_tokens(parent_hash, tokens, extra_keys)(kv_cache_utils.py:577)讓每個 block 的指紋隱含整條前綴,不會出現「中間不同但末尾相同」的誤命中。
• 多種 hash 算法可選:prefix_caching_hash_algo 支持 sha256 / sha256_cbor / xxhash / xxhash_cbor(cache.py:95-110),xxhash 更快但不是密碼學安全的,多租戶場景要謹慎。
• 關閉路徑有專門 coordinator:KVCacheCoordinatorNoPrefixCache(kv_cache_coordinator.py:377)讓 find_longest_cache_hit 直接返回 ((), 0),避免 enable_caching=False 時還在做 hash 計算。
• 多模態 / LoRA extra_keys:generate_block_hash_extra_keys(kv_cache_utils.py:714)把多模態輸入和 LoRA ID 作為額外哈希鍵,避免不同圖片佔同一 block 槽位。
• 非因果注意力層強制關閉:EngineCore._initialize_kv_caches 檢測到 non_causal=True 的 spec 時把 enable_prefix_caching 改 False(core.py:255-269),因為 Prefix LM 的雙向注意力會讓 prefix 命中語義錯亂。
• 跳過讀取 prefix cache:Request.get_skip_reading_prefix_cache(request.py:266-277)讓 prompt logprobs / pooling all 請求顯式不讀 prefix,因為它們需要完整重算。
• watermark 防抖動:KVCacheManager.watermark_blocks(kv_cache_manager.py:164-167)留一小撮空閒 block,避免高頻 preemption。

關鍵檔案

• enable_prefix_caching:93 — 預設 True,啟動時是否啟用前綴快取。
• prefix_caching_hash_algo:95 — sha256 / sha256_cbor / xxhash / xxhash_cbor 四選一。
• EngineCore 构造 request_block_hasher:210-219 — 選 hash 函數 + init_none_hash + get_request_block_hasher。
• 非因果注意力强制关闭:255-269 — Prefix LM / 雙向注意力時 enable_prefix_caching = False。
• Request.update_block_hashes:237-240 — 呼叫 _block_hasher(self) 給新滿 block 算 hash。
• get_skip_reading_prefix_cache:266-277 — prompt logprobs / pooling 顯式跳過命中。
• hash_block_tokens:577 — 鏈式哈希:hash_function((parent_hash, token_ids_tuple, extra_keys))。
• get_request_block_hasher:673 — 返回閉包,只對滿 block 算 hash,父 hash 串起來。
• get_computed_blocks:206 — 入口:請求來 → 查最長命中 → 返回 (KVCacheBlocks, num_computed_tokens)。
• allocate_slots:248 — 命中部分複用,未命中部分從 BlockPool 取新 block。
• KVCacheCoordinatorNoPrefixCache:377 — 關閉路徑的 coordinator stub。
• UnitaryKVCacheCoordinator.find_longest_cache_hit:480 — 單 KV group 模型的命中查詢。
• Scheduler 调用 get_computed_blocks:724-726 — 排程時拿命中數,減少 num_new_tokens。

資料流

請求進來時 Scheduler 先調 get_computed_blocks,這一步在 coordinator 內部走 find_longest_cache_hit,按 block_hashes 列表逐個查 BlockPool 是否已有同樣 hash 的 block:

# vllm/v1/core/kv_cache_manager.py L206-L246
def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]:
    """Get the computed (cached) blocks for the request.
    Note that the computed blocks must be full.
    """
    # We skip finding the prefix cache hit when prefix caching is
    # disabled or the request is marked as skipping kv cache read
    # (which happens when the request requires prompt logprobs
    # or calls a pooling model with all pooling).
    if not self.enable_caching or request.skip_reading_prefix_cache:
        return self.empty_kv_cache_blocks, 0

    # NOTE: When all tokens hit the cache, we must recompute the last token
    # to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1.
    # This can trigger recomputation of an entire block, rather than just
    # the single last token, because allocate_slots() requires
    # num_computed_tokens to be block-size aligned. Removing this limitation
    # could slightly improve performance in the future.
    max_cache_hit_length = request.num_tokens - 1
    computed_blocks, num_new_computed_tokens = (
        self.coordinator.find_longest_cache_hit(
            request.block_hashes, max_cache_hit_length
        )
    )

    if self.log_stats:
        assert self.prefix_cache_stats is not None
        self.prefix_cache_stats.record(
            num_tokens=request.num_tokens,
            num_hits=num_new_computed_tokens,
            preempted=request.num_preemptions > 0,
        )

    return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens

每個 block 的 hash 是在 request 自身的 update_block_hashes 裡算的,request_block_hasher 是閉包,只對滿 block 算 hash,父 hash 串起來:

# vllm/v1/core/kv_cache_utils.py L687-L728
def request_block_hasher(request: Request) -> list[BlockHash]:
    start_token_idx = len(request.block_hashes) * hash_block_size
    num_tokens = request.num_tokens

    if start_token_idx + hash_block_size > num_tokens:
        # Early stop when there no new full blocks created.
        return []

    curr_mm_idx = 0
    if start_token_idx > 0:
        # Set curr_mm_idx = -1 to indicate the last mm input.
        # Note that since we reach to this branch only when the block is
        # completed with generated tokens, we only need to consider the
        # last mm input.
        curr_mm_idx = -1

    prev_block_hash_value = (
        request.block_hashes[-1] if request.block_hashes else None
    )
    new_block_hashes: list[BlockHash] = []
    while True:
        end_token_idx = start_token_idx + hash_block_size
        if end_token_idx > num_tokens:
            # We only hash full blocks
            break

        # MM and LoRA requests need extra keys for block-hash computation.
        extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
            request, start_token_idx, end_token_idx, curr_mm_idx
        )

        # Compute the hash of the current block
        block_tokens = request.all_token_ids[start_token_idx:end_token_idx]
        block_hash = hash_block_tokens(
            caching_hash_fn, prev_block_hash_value, block_tokens, extra_keys
        )

        new_block_hashes.append(block_hash)
        start_token_idx += hash_block_size
        prev_block_hash_value = block_hash

    return new_block_hashes

Scheduler 拿到命中數後,把 num_new_tokens 減掉命中部分,allocate_slots 在 block pool 裡只取未命中的 block。命中 block 引用計數 +1,但 token 不會重新計算。

邊界與失敗

• 非因果注意力層強制關 prefix caching:EngineCore._initialize_kv_caches 檢測到 non_causal=True 時把 enable_prefix_caching = False(core.py:265-269),否則雙向注意力的 prefix 命中會讀到錯誤的 KV。
• Prompt logprobs / pooling 跳過命中:Request.get_skip_reading_prefix_cache 返回 True 時 get_computed_blocks 直接返回空(kv_cache_manager.py:222-223),因為這些場景必須重新計算 prompt 的 KV。
• 全命中也要重算最後一個 token:max_cache_hit_length = request.num_tokens - 1(kv_cache_manager.py:231),否則沒有 logits 可採。
• xxhash 不是密碼學安全:prefix_caching_hash_algo = "xxhash" 速度更快,但理論上碰撞概率更高,多租戶場景可能洩漏前綴信息(cache.py:101-108)。
• 關閉路徑有專門 stub:KVCacheCoordinatorNoPrefixCache.find_longest_cache_hit 返回 ((), 0)(kv_cache_coordinator.py:416-424),讓 enable_caching=False 時不浪費 hash 計算。
• BlockPool LRU 淘汰:命中複用的 block 引用計數 +1,請求結束才 -1,長期沒人引用的 block 被 LRU 回收給新請求(kv_cache_manager.py:508)。

小結

前綴快取是 vLLM 高吞吐的關鍵機制之一:把 KV block 按「父 hash + token ids + extra_keys」鏈式哈希,跨請求複用相同前綴的計算結果。EngineCore 在初始化時根據 enable_prefix_caching 決定要不要構造 hasher,Scheduler.schedule 每次調 get_computed_blocks 查命中。複用的 block 在 /kv-cache/kv-cache-manager 裡管,BlockPool 細節見 /kv-cache/coordinator-blockpool,引擎核心的串聯見 /engine/engine-core。

對照官方資料:vLLM 文件 · README