vLLM 架构拆解

日本語

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

日本語

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

Appearance

KVCacheCoordinator + BlockPool:混合モデルの協調とブロックプール ​

源码版本v0.25.1

役割 ​

KVCacheCoordinator は vLLM v1 が複数の KV cache タイプ (full attention / sliding window / Mamba 状態 / MLA など) の上に設けた協調レイヤーです。各 KV cache group に 1 つの SingleTypeKVCacheManager (例: FullAttentionManager / SlidingWindowManager / RSWAManager) を割り当て、coordinator は「N 個の token の block を割り当てる」「最長プレフィックスヒットを探す」「リクエストを解放する」「埋まった block をキャッシュする」といった動作を group ごとに分割して対応する single-type manager に振ります。BlockPool は共有の物理 block プールで、すべての group が同じ KVCacheBlock 群を使います。ただし hash key に group_id を付けて区別するため、同じ block を複数の group (例: full attention と sliding window の同じ token) が同時にキャッシュできます。

構築時 KVCacheCoordinator.__init__ は BlockPool を直接作り、get_manager_for_kv_cache_spec で各 group に single-type manager を起動します (kv_cache_coordinator.py:91-120)。具体的な戦略は get_kv_cache_coordinator ファクトリが選びます。キャッシュしない場合は KVCacheCoordinatorNoPrefixCache、単一 group なら UnitaryKVCacheCoordinator、複数 group なら HybridKVCacheCoordinator になります (kv_cache_coordinator.py:782-822)。BlockPool 自身は FreeKVCacheBlockQueue (双方向リスト) で LRU 淘汰を行い、cached_block_hash_to_block ハッシュテーブルでプレフィックスキャッシュ (prefix caching) のインデックスを管理します (block_pool.py:144-197)。

設計動機 ​

なぜ coordinator を NoPrefix / Unitary / Hybrid の 3 種に分けるのか?

  • 空実行のフォールバック:KVCacheCoordinatorNoPrefixCache は任意の group 数 (0 を含む) をサポートし、プレフィックスキャッシュ関連のインターフェースを一切実装しません。find_longest_cache_hit は (空 blocks, 0) を返し、get_num_common_prefix_blocks はすべて 0 を返します (kv_cache_coordinator.py:413-424)。これにより prefix caching 無効時にはキャッシュパスを一切通りません。
  • 単一 group の高速パス:UnitaryKVCacheCoordinator は hash_block_size == block_size かつ group が 1 つだけであることを前提とし (kv_cache_coordinator.py:471-476)、find_longest_cache_hit は single manager を直接呼び、マージをせず済ませます (kv_cache_coordinator.py:480-496)。hybrid の spec-group マージのオーバーヘッドを避けるためです。
  • Hybrid は spec 単位で集約:HybridKVCacheCoordinator.verify_and_split_kv_cache_groups は同じ spec の複数 group を 1 つの SpecGroup にまとめ、一緒に cache hit 探索をします (kv_cache_coordinator.py:560-588)。さらに full attention を先頭に並べ、後続 group にタイトな上界を与えます (kv_cache_coordinator.py:590-595)。
  • 共有 BlockPool:すべての group が num_gpu_blocks 個の物理 block を共有します。異なる group のキャッシュ hash は make_block_hash_with_group_id(block_hash, group_id) で連結し、get_cached_block は group_id リストごとに検索します (block_pool.py:199-224)。BlockHashToBlockMap は 1 つの hash が複数 group の block に対応できます。
  • LRU + tail block 優先:FreeKVCacheBlockQueue は双方向リストで O(1) の途中削除を実現し、free 時には逆順に戻すことで tail block が先に淘汰されるようにします (kv_cache_utils.py:179-197)。free_blocks は hash のない block を hash 付きより先に淘汰します (block_pool.py:622-635)。
  • null_block のプレースホルダー:block_id=0 は null_block で、remove_skipped_blocks が sliding window 外や preempt された block の置き換えに使います。ref_cnt は管理されず、解放時にはスキップするよう注意が必要です (block_pool.py:188-192)。

主要ファイル ​

データフロー ​

BlockPool.get_new_blocks が物理割り当ての真のエントリです。free queue の先頭から N 個の block をポップし、それぞれの古い hash を evict して ref_cnt を 1 に設定します:

python
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
    if num_blocks > self.get_num_free_blocks():
        raise ValueError(f"Cannot get {num_blocks} free blocks from the pool")

    ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)

    if self.enable_caching:
        for block in ret:
            self._maybe_evict_cached_block(block)
            assert block.ref_cnt == 0
            block.ref_cnt += 1
            if self.metrics_collector:
                self.metrics_collector.on_block_allocated(block)
    else:
        for block in ret:
            assert block.ref_cnt == 0
            block.ref_cnt += 1
            if self.metrics_collector:
                self.metrics_collector.on_block_allocated(block)
    return ret

(block_pool.py:542-572) これは SingleTypeKVCacheManager.allocate_new_blocks に呼ばれ (single_type_kv_cache_manager.py:279-311)、後者は新規 block を req_to_blocks[request_id] に append します。解放は BlockPool.free_blocks で行い、まず ref_cnt を -1 し、ref_cnt が 0 になった block は「hash あり / なし」で分けて free queue の末尾に戻します (block_pool.py:614-635)。リクエストの prompt が 1 つの block を埋めると、cache_full_blocks が block の内容から hash を計算して cached_block_hash_to_block に書き込みます。次に同じプレフィックスを持つリクエストが来ると get_cached_block で直接ヒットし、touch で ref_cnt を +1 すれば再計算は不要です (block_pool.py:199-224)(block_pool.py:597-612)。

境界と失敗 ​

  • 空き block 数の超過:get_new_blocks は事前チェックせず直接 raise ValueError します (block_pool.py:553-554)。呼び出し側 (KVCacheManager.allocate_slots) は get_num_blocks_to_allocate で事前に計算しておく必要があります (kv_cache_coordinator.py:130-185)。
  • Hybrid は DCP / PCP 非対応:assert dcp_world_size == 1 and pcp_world_size == 1 (kv_cache_coordinator.py:556-557)。
  • Unitary は hash_block_size == block_size を強制:assert not enable_caching or hash_block_size == self.block_size (kv_cache_coordinator.py:471-473)。
  • Hybrid は最低 2 つの attention group:assert len(self.attention_groups) > 1 (kv_cache_coordinator.py:586-588)。
  • block にすでに hash があるのに set しようとした:set_block_hash は assert で重複を防ぎます (kv_cache_utils.py:153-157)。
  • reset_prefix_cache のクリア拒否:null block 以外にまだ block が使われていれば logger.warning + False を返します (block_pool.py:665-672)。
  • evict_blocks が外部 block_id を受け取る:KV connector が報告した block id が範囲外なら即座に assert します (block_pool.py:647-652)。
  • retention_interval のバリデーション:_validate_prefix_cache_retention_interval は正の整数で scheduler_block_size の倍数であることを要求します (kv_cache_coordinator.py:31-58)。

まとめ ​

Coordinator + BlockPool が KVCacheManager の下の物理層を構成します。coordinator は戦略に従って single-type managers に振り分け、BlockPool は物理 block プールとプレフィックスキャッシュのインデックスを管理します。block 割り当ての結果は最終的に block_id リストになり、ブロック表 が attention kernel が直接消費できるテンソルに包み込みます。

公式資料: vLLM 文档 · README