vLLM 架构拆解

日本語

简体中文
English
Español
Français
Deutsch
繁體中文

日本語

简体中文
English
Español
Français
Deutsch
繁體中文

Appearance

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 予算を減らして割り当てます。

設計動機 ​

  • token 単位ではなく block 単位の再利用: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 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 を防ぎます。

主要ファイル ​

データフロー ​

リクエストが来ると Scheduler は先に get_computed_blocks を呼び、このステップは coordinator 内部で find_longest_cache_hit に進み、block_hashes リストを順に BlockPool で検索して同じ hash の block があるか調べます:

python
# 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 を繋ぎます:

python
# 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)、さもなければ双方向注意のプレフィックスヒットが誤った 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