vLLMにおけるKVCacheマネージャーの詳細解説

本稿では、大規模言語モデル推論フレームワーク「vLLM」におけるKVCacheマネージャーの内部構造と動作原理を解説する。vLLMのバージョン0.17.1を 기준으로、ブロック管理機構からメモリアロケーションまで、実装深处までを明らかにする。

検証環境はconda環境（Python 3.12）上に構築され、vllm==0.17.1で動作確認を行っている。

vLLMのスケジューリング機構を解説した以往の記事では、waitingキューとrunningキューの管理やトークンバジェットの割り当てについて紹介したが、各リクエストがGPUメモリ上のKVキャッシュブロックをどのように利用・解放しているかについては詳細を省略した。本稿では、このキーのとなるKVCacheManagerの役割と実装を掘り下げて解説する。

スケジューラーにおけるKVCacheManagerの利用箇所を以下に示す。runningキューでは、新しく生成されるトークンに対するブロック割り当てが可能かを判定する処理が行われる。

# 実行中リクエストのトークンに対するブロック割り当て確認
new_blocks = self.kv_cache_manager.allocate_slots(
    request_id,
    required_token_count,
    num_lookahead_tokens=self.lookahead_tokens_count,
)

waitingキューの処理では、既存の計算済みブロックの取得と、プレフィル処理時の新規ブロック割り当てが行われる。

# 計算済みブロックの取得（プリフィックスキャッシュ利用）
cached_blocks, newly_computed_count = (
    self.kv_cache_manager.get_computed_blocks(request_id)
)

# 新規トークンに対するブロック割り当て
new_blocks = self.kv_cache_manager.allocate_slots(
    request_id,
    token_count,
    num_new_computed_tokens=newly_computed_count,
    existing_cached_blocks=cached_blocks,
    num_lookahead_tokens=effective_lookahead,
    num_encoder_tokens=encoder_token_count,
    delay_cache_blocks=async_load_enabled,
)

KVCacheManagerの初期化処理では、複数の設定パラメータを受け取り、内部のコンポーネントを構成している。

# スケジューラーでの初期化処理
self.kv_cache_manager = KVCacheManager(
    kv_cache_config=kv_cache_config,
    max_model_len=self.max_model_len,
    enable_caching=self.cache_config.enable_prefix_caching,
    use_eagle=self.use_eagle,
    log_stats=self.log_stats,
    enable_kv_cache_events=self.enable_kv_cache_events,
    dcp_world_size=self.dcp_world_size,
    pcp_world_size=self.pcp_world_size,
    hash_block_size=self.block_size,
    metrics_collector=self.kv_metrics_collector,
)

初期化パラメータの主要な項目として、kv_cache_configはGPU上のブロック数設定を管理し、max_model_lenはモデルがサポートする最大シーケンス長を定義する。enable_prefix_cachingをtrueに設定すると、同一プレフィックスを持つリクエスト間でのKVキャッシュ再利用が有効化される。

ブロックアロケーションの基本概念

vLLMにおけるKVキャッシュブロック管理を理解する上で、 Three の重要な概念が存在する。

第一に、block_sizeは1ブロックに格納可能なトークン数を表す。通常は16トークン/ブロックに設定される。第二に、事前割当メモリはモデル初期化時に確保されるGPUメモリ領域であり、デバイスメモリ总量の割当比率（例：24GB×0.9=21.6GB）を基に計算される。この領域はモデル実行時のKVキャッシュ以外のメモリ使用量を考慮して動的に調整される。第三に、ブロック数は利用可能なGPUメモリをページサイズとレイヤー数で除算して算出される。

ページサイズの計算식은(K+V)×block_size×num_kv_heads×head_size×dtype_bytesとなる。num_layersはKVキャッシュグループ内の最大レイヤー数であり、vLLMはモデルの各レイヤーに対して形状が[2, num_blocks, block_size, num_kv_heads, head_size]の独立したKVキャッシュテンソルを割当てる。

vLLMの全体的なメモリ割当フローは以下のようになる。入力プロンプトが到着すると、プレフィル段階でトークン数をブロックサイズで除算した数のブロックが事前割当される。ブロックサイズ16の場合、5トークンでも1ブロックが確保される。プレフィル段階ではプロンプトの全トークンのKV値がブロックに書き込まれ、各ブロックはblock_id（識別子）とref_cnt（参照回数）を持つ。デコード段階は增量方式で進行し、1トークンずつ（または少数ずつ）生成される。現在の最後のブロックに空きがあれば書き込みを継続し、満杯になった時点で新規ブロックを要求する。

GPU KVキャッシュメモリが不足した場合、スケジューラはFCFS（先入先出）を基本としたポリシーでプリエンプションを実行する。waitingキューのリクエストをCPUへ退避させたり、優先度に基づいて再スケジューリングを行ったりする。

複数入力プロンプトはvLLM内部では単一シーケンスとして連結されて処理される。これによりGPU利用率を最大化する。各シーケンスのposition idsは0から独立して開始され、アテンションマスクによって各トークングループが相互に隔離される。デコード段階ではslot_mappingテンソルを用いて、生成される各新トークンが物理KVキャッシュのどのスロットに対応するかを記録する。

KVCacheManagerの処理詳細

vLLMにはKVキャッシュとは別概念としてプレフィックスキャッシュが存在する。これは複数リクエスト間で同一プレフィックスが存在する場合に、当該部分のKV値を再計算せずに再利用するための機構である。重要な点として、共通部分是先頭から連続して一致している必要があり、中間部分が一致しているだけでは共有できない。

「你好、武汉について教えて？」と「你好、北京について教えて？」という2つのリクエストでは、「你好、」以降の共通プレフィックスが存在するためKVキャッシュ再利用が可能となる。一方、「旅游専門家として活躍している你好、北京について教えて？」のような場合、中間に異なるテキストが挿入されているため再利用はできない。

KVCacheManager（vllm/v1/core/kv_cache_manager.py）のコアロジックは以下の4つの主要メソッドで構成される。get_computed_blocksは指定されたリクエストに対するプレイフィックスキャッシュの検出を行い、allocate_slotsは新規トークンに対するブロック割当を実行する。freeはリクエスト完了時の全ブロック解放を担当し、cache_blocksはプレイフィックスキャッシュ用のブロック管理を行う。reset_prefix_cacheは全キャッシュブロックのリセットに使用される。

BlockPoolの生成過程

検証環境のモデルにはQwen/Qwen2-0.5B-Instructを使用し、32GBのGPUメモリと0.9の割当比率を前提とする。

BlockPool（vllm/v1/core/block_pool.py）では初期化時に全ブロックが同時に生成される。

self.blocks: list[KVCacheBlock] = [KVCacheBlock(idx) for idx in range(num_gpu_blocks)]

この例ではnum_gpu_blocksは131928となる。各KVCacheBlockオブジェクトはprevとnextポインタによる双方向リンクリスト構造を持ち、block_idはブロック識別子、ref_cntは参照回数を保持する。

ただし、この段階ではまだ物理GPUメモリの割当は行われていない。実際のメモリ割当はvllm/v1/core/kv_cache_utils.pyで実行される。

def get_kv_cache_config_from_groups(vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], available_memory: int):
    if len(kv_cache_groups) == 1 and isinstance(kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs):
        # 単一KVキャッシュグループの処理
        pass
    else:
        group_size = max(len(group.layer_names) for group in kv_cache_groups)  # 24レイヤー
        page_size = get_uniform_page_size([group.kv_cache_spec for group in kv_cache_groups])  # 8192
        assert group_size > 0, "group_size must be greater than 0"
        num_blocks = get_num_blocks(vllm_config, group_size, available_memory, page_size)  # 131928
        
        kv_cache_tensors = []
        for i in range(group_size):
            shared_by = []
            for j in range(len(kv_cache_groups)):
                if i < len(kv_cache_groups[j].layer_names):
                    shared_by.append(kv_cache_groups[j].layer_names[i])
            kv_cache_tensors.append(
                KVCacheTensor(size=page_size * num_blocks, shared_by=shared_by)
            )

検証モデルのkv_cache_groups設定は以下の通り。

kv_cache_groups = [
    KVCacheGroupSpec(
        layer_names=[
            'model.layers.0.self_attn.attn',
            'model.layers.1.self_attn.attn',
            ...,
            'model.layers.23.self_attn.attn'
        ],
        kv_cache_spec=FullAttentionSpec(
            block_size=16,
            num_kv_heads=2,
            head_size=64,
            dtype=torch.bfloat16,
            page_size_padded=None,
            head_size_v=64,
            sliding_window=None,
            attention_chunk_size=None
        )
    )
]

返り値のKVCacheConfigは以下の構成となる。

kv_cache_config = KVCacheConfig(
    num_blocks=131928,
    kv_cache_tensors=[
        KVCacheTensor(size=1080754176, shared_by=['model.layers.0.self_attn.attn']),
        ...
    ],
    kv_cache_groups=[
        KVCacheGroupSpec(layer_names=[...], kv_cache_spec=FullAttentionSpec(...))
    ]
)
# 1080754176 = 8192 × 131928

物理GPUメモリの実際の割当はvllm/v1/worker/gpu_model_runner.pyのinitialize_kv_cache_tensorsメソッドで行われる。

def initialize_kv_cache_tensors(self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]):
    # ステップ1: KVキャッシュテンソルの物理メモリ割当
    kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
    
    # ステップ2: 各レイヤーが要求する形状へのリシェイプ
    kv_caches = self._reshape_kv_cache_tensors(kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes)
    
    return kv_caches

Allocate段階ではtorch.zerosを用いて各レイヤーのKVキャッシュテンソルを初期化する。

tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=self.device)
for layer_name in kv_cache_tensor.shared_by:
    kv_cache_raw_tensors[layer_name] = tensor

次元整形後の最終的なKVキャッシュ形状は(2, 131928, 16, 2, 64)となる。これは[K+V, num_blocks, block_size, num_kv_heads, head_size]に対応し、2はキーとバリューそれぞれを表す。

KVキャッシュメモリ管理は以下の2つの層で構成される。物理KVデータ格納層では、各アテンションレイヤー（またはKVキャッシュグループ）に対して事前割当されたKVCacheTensorが配置される。形状は(2, num_blocks, block_size, num_kv_heads, head_size)であり、これがKV値の実際の格納領域となる。全レイヤーは同一のnum_blocksを共有し、tensor間でメモリを共有する場合もある。

グローバルブロックメタデータ管理層では、BlockPoolにおいてKVCacheBlockオブジェクトのリストが生成される。num_gpu_blocks個のオブジェクトはblock_id、ref_cnt、block_hash、last_access_timeなどのメタデータのみを保持し、実際のKVデータは格納しない。この層のメモリ使用量は極めて小さい。

BlockPoolはメモリプールの概念として機能する。リクエストが新規ブロックを要求すると、free_block_queueから適切なブロックが取得される。メモリ不足で割り当てに失敗した場合、リクエストは待機状態となるか、プリエンプションの対象となる。リクエスト完了時に全ブロックが解放され、free_block_queueに返却される。

ブロック割当のロジック

新規リクエストに対するブロック割当処理の流れを解説する。KVCacheManagerは内部でブロックコーディネータ（kv_cache_coordinator.py）を使用しallocate_new_blocksメソッドを呼出す。

def allocate_new_blocks(
    self,
    request_id: str,
    num_tokens: int,
    num_tokens_main_model: int,
    num_encoder_tokens: int = 0,
) -> tuple[list[KVCacheBlock], ...]:
    return tuple(
        manager.allocate_new_blocks(
            request_id,
            num_encoder_tokens
            if isinstance(manager, CrossAttentionManager)
            else num_tokens,
            num_tokens_main_model,
        )
        for manager in self.single_type_managers
    )

self.single_type_managersはget_manager_for_kv_cache_specによって取得され、アテンションの種類に応じた適切なマネージャーインスタンスが選択される。

spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
    FullAttentionSpec: FullAttentionManager,
    MLAAttentionSpec: FullAttentionManager,
    SlidingWindowSpec: SlidingWindowManager,
    ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,
    MambaSpec: MambaManager,
    CrossAttentionSpec: CrossAttentionManager,
    SinkFullAttentionSpec: SinkFullAttentionManager,
}

これらのマネージャーはSingleTypeKVCacheManagerを継承しており、allocate_new_blocksとallocate_new_computed_blocksの2つの主要メソッドを実装している。

allocate_new_blocksは単純に必要ブロック数を計算し、block poolから要求数量のブロックを取得する。計算式はcdiv(num_tokens, block_size)（天井除算）であり、10トークンでblock_size=16の場合は1ブロックが割当られる。

allocate_new_computed_blocksはプレイフィックスキャッシュから既に計算済みのブロックを取得する処理を担当する。

具体的な割当例を以下に示す。4レイヤーのモデルで50トークンの入力を扱う場合、block_size=16なので4ブロックが必要となる。まずプレイフィックスキャッシュを確認し、前2ブロック（block_id: 10, 21）がキャッシュされている場合を考える。未割当部分はnullで初期化され、allocate_new_blocksで残り2ブロック（block_id: 45, 67）が取得される。最終的なブロックテーブルは[10, 21, 45, 67]となる。

デコード段階では、各レイヤーのKVキャッシュテンソルがこのブロックテーブルに基づいてアクセスされる。

layer_0: kv_cache_layer0[:, [10,21,45,67], ...]
layer_1: kv_cache_layer1[:, [10,21,45,67], ...]
...
layer_3: kv_cache_layer3[:, [10,21,45,67], ...]

ブロックが満杯になって新規トークンを処理する場合は、新しいブロックID（例：89）が追加され、ブロックテーブルに拡張される。

プレイフィックスキャッシュの詳細

プレイフィックスキャッシュは完了済みリクエストのKVキャッシュブロックを保持し、新リクエストが同一プレフィックスを持つ場合にそれらを再利用するための機構である。vLLMはブロック単位のキャッシュ管理を採用しており、SGLangのトークンレベルとは異なる設計となっている。

ブロックのハッシュ計算はチェーンハッシュ方式を採用している。各ブロックは双方向リンクリストで接続され、前後のブロックとの関係性を保持する。

# ブロック0（トークン0〜15）のハッシュ計算
hash0 = hash(NONE_HASH + tuple(tokens[0:16]) + extra_keys)

# ブロック1（トークン16〜31）のハッシュ計算
hash1 = hash(hash0 + tuple(tokens[16:32]) + extra_keys)

重要な点として、ブロックは完全に満杯状態（block_size個のトークンを格納）の場合にのみハッシュ値が計算され、キャッシュに登録される。部分的に填充されたブロックはハッシュを持たない。

スケジューラーがwaitingキューを処理する際、get_computed_blocksメソッドが呼ばれて計算済みブロックと新規計算トークン数を取得する。

def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]:
    if not self.enable_caching or request.skip_reading_prefix_cache:
        return self.empty_kv_cache_blocks, 0

    computed_blocks, num_new_computed_tokens = (
        self.coordinator.find_longest_cache_hit(
            request.block_hashes,  # 事前計算されたブロックハッシュリスト
            request.num_tokens - 1
        )
    )
    
    return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens

find_longest_cache_hitメソッドは全KVCacheBlockを走査し、リクエストのブロックハッシュと一致するものを探索する。一致したキャッシュブロックが見つかれば、当該ブロックが再利用対象となる。

公式ドキュメントの具体例で動作を確認する。

時刻1: キャッシュは空のリクエスト0（A-O、19トークン）が到着。Block poolから4ブロックが割当られる。3ブロックは満杯状態でキャッシュに登録される（SHA256を使用）。第4ブロックは3トークンのみ填充される。

時刻2: デコード処理開始。第4ブロックが満杯になり、4つのキャッシュハッシュが建立される。新規ブロックが要求される。

時刻3: リクエスト1（A-J、k-n、14トークン）が到着。先頭10トークンはリクエスト0と同一。最初の2ブロック（8トークン）のみがキャッシュヒットとなる。第3ブロックは4トークン中2トークンのみ一致するためヒットしない。

キャッシュヒット判定の内部ロジックを以下に示す。

@classmethod
def find_longest_cache_hit(cls, block_hashes, max_length, ...):
    computed_blocks = tuple([] for _ in range(len(kv_cache_group_ids)))
    max_num_blocks = max_length // block_size
    
    for block_hash in itertools.islice(block_hashes, max_num_blocks):
        # キャッシュプールから一致するブロック探索
        if cached_block := block_pool.get_cached_block(block_hash, kv_cache_group_ids):
            # 一致ブロック発見
            for computed, cached in zip(computed_blocks, cached_block):
                computed.append(cached)
        else:
            # 不一致 발견時点で終了
            break
    
    if use_eagle and computed_blocks[0]:
        for computed in computed_blocks:
            computed.pop()
    
    return computed_blocks

時刻4: リクエスト0が完了し解放される。ブロック2、3、4が逆順で空きキューに追加される。ただしブロック2と3は依然としてキャッシュされている。ブロック0と1はリクエスト1で使用中のため空きキューには追加されない。

時刻5: リクエスト1が完了し解放される。

時刻6: リクエスト2（29トークン、先頭12トークンはリクエスト0と同一）が到着。空きキューは7-8-9-4-3-2-6-5-1-0の順だが、キャッシュヒットしたブロック（0、1、2）は割り当て前にキューから除去され、7-8-9-4-3-6-5となる。結果として割り当てられるブロックは0（キャッシュ）、1（キャッシュ）、2（キャッシュ）、7、8、9、4、3となる。

主要コンポーネントとファイル構成

KVキャッシュ管理に関わる主要なコードファイルは以下の通りである。

kv_cache_manager.pyはブロック割り当てのインターフェースを提供し、コアロジックはkv_cache_coordinator.pyに実装されている。プレイフィックスキャッシュの利用と新規トークンのブロック割り当てを管理する。block_pool.pyは計算されたnum_blocksに基づいてブロックプールを生成する。KVCacheBlockオブジェクトのリスト管理を担当する。gpu_model_runner.pyはモデルの各レイヤーに対する物理GPUメモリの割当を実行し、グローバルなKVCacheBlockによるメモリ管理を行う。

BlockPoolは容量num_blocksのメモリプールとして機能し、各ブロックはblock_size個のトークンを格納できる。初期化時に全レイヤーに対して(K+V, num_blocks, block_size, num_kv_heads, head_size)サイズのテンソルが割当われる。全物理ブロックはこの連続GPUメモリを共有し、block_idによるページングインデックスでアクセスされる。

新規プロンプト入力時に必要ブロック数が計算され（len(tokens) // block_size）、トークンがプールに配置される。滿杯状態のブロックにはハッシュ値が計算され、block hashに登録される。新規入力時にハッシュ計算が行われ、block hashに一致があればプレイフィックスキャッシュが再利用される。具体的には、ブロックテーブル[1, 2]が割り当てられた場合、計算は(K+V, [2, 3], block_size, num_kv_heads, head_size)に基づいてKVキャッシュにアクセスされる。

リクエストが完了すると、使用済みブロックがblock poolに返却される。キャッシュは保持されたままとなるため、同一プレフィックスを持つ後続リクエストは即座にキャッシュを再利用できる。