KVCacheManager: リクエスト単位の KV キャッシュライフサイクル

源码版本v0.25.1

役割

KVCacheManager はスケジューラ (scheduler) から見える KV キャッシュ (KV cache) インターフェースです。「このリクエストに何個の block を割り当てる」「このリクエストが完了したので block を返却する」「この prompt がどれだけプレフィックスキャッシュ (prefix caching) にヒットしたか確認する」といった動作を allocate_slots / free / get_computed_blocks の 3 つの高レベル API にまとめます。内部では異なる KV cache タイプ (full attention / sliding window / Mamba 状態など) の協調はすべて KVCacheCoordinator に委譲し、物理 block プールの管理は BlockPool に任せます。スケジューラは KVCacheBlocks という不変な結果 (kv_cache_manager.py:29-50) を受け取るだけでよく、下が単一タイプか混合モデルかを知る必要はありません。

構築時、KVCacheManager.__init__ は get_kv_cache_coordinator を通じて coordinator を組み立て、ついでに block_pool / num_kv_cache_groups / kv_cache_config を自身に保持します (kv_cache_manager.py:147-183)。watermark (watermark_blocks = int(watermark * num_blocks)) はすでにスケジュールされたリクエスト用のバッファで、WAITING / PREEMPTED 状態にのみ効き、新規リクエストが入ってきたときに走行中のリクエストを押し出さないようにします (kv_cache_manager.py:164-167)(kv_cache_manager.py:367-374)。manager 全体は per-request 状態をほとんど持たず、すべての req_to_blocks は SingleTypeKVCacheManager の中にあります。

設計動機

なぜ manager はこれほど薄く見え、本当にロジックはすべて coordinator にあるのか?

• 責務の階層化:KVCacheManager は「リクエスト向け」の意味論 (allocate_slots / free / get_computed_blocks) を担い、KVCacheCoordinator は「KV cache group 向け」の協調 (混合モデル時に複数の group をどう一緒に割り当てるか) を担い、BlockPool は「物理 block 向け」のプール化を担います。3 層はそれぞれの役割を持ち、manager はそれらを糊付けします。
• KVCacheBlocks が安定インターフェース:blocks[i][j] は i 番目の KV cache group の j 番目の block を表し、get_block_ids() は tuple[list[int], ...] に変換して attention backend に渡します (kv_cache_manager.py:73-88)。__add__ により 2 回の allocate_slots の結果を結合でき (<SrcLink path="vllm/v1/core/kv_cache_manager.py" lines="52-59" label="kv_cache_manager.py"/>)、new_empty は「block なし」のプレースホルダーを返します。
• 遅延キャッシュコミット:allocate_slots の delay_cache_blocks=True ブランチは P/D 非同期転送時に cache_blocks をスキップし (kv_cache_manager.py:448-451)、リモート KV が本当に到着してから hash を付与します。事前にキャッシュすると誤ヒットを招くためです。
• watermark は待機中リクエストにのみ効く:has_scheduled_reqs and request.status in (WAITING, PREEMPTED) のときのみ watermark を加算し、すでに実行中のリクエストの増分割り当てでは予約は不要です (kv_cache_manager.py:367-374)。
• 全シーケンスのアドミッションチェック:full_sequence_must_fit=True のときはまず prompt 全体で何個の block が必要か見積もり、超過すれば None を返します。chunked prefill がリクエストを入れたあとに行き詰まるのを防ぐためです (kv_cache_manager.py:376-391)。
• reserved_blocks で非同期 KV connector を確保:すでに in-flight な prefill シーケンスが特定の block に依存しているとき、新しく入ってきた非同期 connector load にそれらを奪わせないようにします。required_blocks > available_blocks のときは割り当てを拒否します (kv_cache_manager.py:420-426)。

主要ファイル

• KVCacheBlocks dataclass:29-50 — スケジューラ側の不変な結果、blocks[i][j] = i 番目の group の j 番目の block。
• KVCacheBlocks.__add__:52-59 — 2 回の割り当て結果の結合、増分割り当てでよく使われます。
• KVCacheBlocks.get_block_ids:73-88 — tuple[list[int], ...] に変換して attention backend に渡します。
• KVCacheManager.__init__:114-183 — coordinator の組み立て、watermark_blocks の計算、empty_kv_cache_blocks をキャッシュして GC オーバーヘッドを回避。
• get_computed_blocks:206-246 — coordinator.find_longest_cache_hit を呼んでプレフィックスヒットを探し、(KVCacheBlocks, num_new_computed_tokens) を返します。
• allocate_slots:248-343 — 3 段構成:スキップされた block の解放 → prefix token の処理 → 新しい token の block 割り当て。失敗時は None を返します。
• full_sequence_must_fit:376-391 — アドミッションチェック:prompt 全体の block 数 + watermark が空き block 数以下である必要があります。
• required_blocks gate:410-426 — available = free - reserved、required > available なら割り当て拒否。
• allocate_new + cache_blocks:428-464 — coordinator に新 block 割り当てを委譲 + cache_blocks 呼び出しで hash をコミット。
• free:466-474 — coordinator.free(request_id) を呼び出し、tail block から先に淘汰されるよう逆順で解放。
• pop_blocks_for_free:495-506 — per-request の会計を持ち出すが BlockPool には返さない、P/D 転移で使います。
• reset_prefix_cache:516-530 — RLHF で重みが変わったあとにキャッシュ全体を無効化。

データフロー

スケジューラは各ステップで各リクエストに対し、まず get_computed_blocks でプレフィックスヒットを確認し、次に allocate_slots で新しい token の block を割り当てます。以下は allocate_slots が watermark とアドミッションチェックを済ませたあとに実際に処理を行う部分です:

if (
    new_computed_block_list is not self.empty_kv_cache_blocks.blocks
    or num_external_computed_tokens > 0
):
    # Append the new computed blocks to the request blocks until now to
    # avoid the case where the new blocks cannot be allocated.
    self.coordinator.allocate_new_computed_blocks(
        request_id=request.request_id,
        new_computed_blocks=new_computed_block_list,
        num_local_computed_tokens=num_local_computed_tokens,
        num_external_computed_tokens=num_external_computed_tokens,
    )

new_blocks = self.coordinator.allocate_new_blocks(
    request.request_id,
    num_tokens_need_slot,
    num_tokens_main_model,
    num_encoder_tokens,
)

# P/D: delay caching blocks if we have to recv from
# remote. Update state for locally cached blocks.
if not self.enable_caching or delay_cache_blocks:
    return self.create_kv_cache_blocks(new_blocks)

(kv_cache_manager.py:428-451) まずプレフィックスヒットした block を吸収し (ref_cnt を増やし)、次に coordinator に新しい token の物理 block を割り当てさせ、最後に遅延キャッシュしない場合だけ cache_blocks を呼んで埋まったばかりの block に hash を付与して prefix cache に登録します (詳細は Coordinator と BlockPool)。返された KVCacheBlocks は get_block_ids で block_table に変換され、ブロック表 に詰められて attention kernel に渡されます。

境界と失敗

• num_new_tokens == 0 かつ外部 token なし:ValueError を送出します。呼び出し側は新しい token を計算するか、外部で計算済みの token を保存するかのどちらかを保証する必要があります (kv_cache_manager.py:346-350)。
• get_computed_blocks が全ヒット時に 1 token 足りない:max_cache_hit_length = request.num_tokens - 1。全ヒット時も最後の token を再計算して logits を得る必要があるため、キャッシュヒット長は prompt_length - 1 で止まります (kv_cache_manager.py:225-231)。
• prefix cache をスキップするリクエスト:request.skip_reading_prefix_cache=True または enable_caching=False のとき get_computed_blocks は空を返します (kv_cache_manager.py:222-223)。prompt logprobs が必要なリクエストなどが該当します。
• reset_prefix_cache の拒否:num_used_blocks != 1 (null block のみ) でなければまだ解放されていない block があるため、reset は失敗して False を返します (kv_cache_manager.py:516-530)。
• pop_blocks_for_free はプールに返さない:呼び出し側が自分で逆順に block_pool.free_blocks を呼ぶ必要があり、そうしないと block がリークします (kv_cache_manager.py:495-506)。
• take_events が group メタ情報を付与:BlockStored イベントは group_idx しか持たないため、manager が kv_cache_event_metadata で kind と sliding_window を補います (kv_cache_manager.py:566-589)。

まとめ

KVCacheManager はスケジューラから見える KV キャッシュのファサードで、実際に処理を行う KVCacheCoordinator + BlockPool は Coordinator と BlockPool で分解します。割り当てられた block id は最終的に ブロック表 に載り、attention kernel が物理 KV cache を位置特定するのに使われます。モデル重み自体は DefaultModelLoader の経路でロードされ、このレイヤーとは結合していません。

公式資料: vLLM 文档 · README