Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

アテンションタイプ別 Manager 詳細

深度: [DEEP] 確信度: [VERIFIED] 最終更新: 2026-02-11

概要

SingleTypeKVCacheManager は1種類のアテンションタイプの KV キャッシュ管理ロジックを担当する抽象基底クラスである。アテンションタイプごとにサブクラスが存在し、ブロックの割り当て・解放・プレフィックスキャッシュ検索をそれぞれのセマンティクスに合わせて実装する。

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py

spec_manager_map [DEEP] [VERIFIED]

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py:1049

spec_manager_map = {
    FullAttentionSpec:            FullAttentionManager,
    MLAAttentionSpec:             FullAttentionManager,       # MLA も Full 扱い
    SlidingWindowSpec:            SlidingWindowManager,
    ChunkedLocalAttentionSpec:    ChunkedLocalAttentionManager,
    MambaSpec:                    MambaManager,
    CrossAttentionSpec:           CrossAttentionManager,
    SinkFullAttentionSpec:        SinkFullAttentionManager,
}

get_manager_for_kv_cache_spec(spec, **kwargs) ファクトリ関数がこのマップから Manager クラスをディスパッチする。

基底クラス: SingleTypeKVCacheManager [DEEP] [VERIFIED]

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py:24

状態管理

フィールド説明
req_to_blocksdefaultdict[str, list[KVCacheBlock]]リクエスト ID → 割り当て済みブロックリスト
num_cached_blockdict[str, int]リクエスト ID → キャッシュ登録済みブロック数。RUNNING リクエストのみ追跡
block_sizeint1 ブロックあたりのトークン数。DCP/PCP > 1 の場合は乗算される

コンストラクタ

def __init__(self, kv_cache_spec, block_pool, enable_caching,
             kv_cache_group_id, dcp_world_size=1, pcp_world_size=1):
    self.block_size = kv_cache_spec.block_size
    if dcp_world_size * pcp_world_size > 1:
        self.block_size *= dcp_world_size * pcp_world_size

DCP(Decode Context Parallelism)/ PCP(Prefill Context Parallelism)ではブロックサイズが並列度倍に拡大される。

get_num_blocks_to_allocate() [DEEP] [VERIFIED]

参照: L73

リクエストに必要な新規ブロック数を算出する。2 つのパスが存在:

get_num_blocks_to_allocate(request_id, num_tokens, new_computed_blocks, ...)
  │
  ├─ Fast-path: request_id in num_cached_block(RUNNING リクエスト)
  │  └─ max(num_required_blocks - num_req_blocks, 0)
  │     ※ Speculative Decoding のリジェクトで num_req_blocks > num_required_blocks もあり得る
  │
  └─ Slow-path: 新規リクエスト(プレフィックスキャッシュヒットあり)
     ├─ num_skipped_tokens = get_num_skipped_tokens(total_computed_tokens)
     ├─ num_skipped_blocks = num_skipped_tokens // block_size
     ├─ num_new_blocks = max(required - max(skipped, local_computed), 0)
     ├─ num_evictable_blocks = Σ(ref_cnt==0 かつ非null)
     │  ← touch() 時にキューから除去されるブロック分を加算
     └─ return num_new_blocks + num_evictable_blocks

Evictable blocks の加算理由: new_computed_blocks 内のブロックが空きキュー内(ref_cnt == 0)にある場合、touch() で空きキューから除去されるため、実質的に空きブロック数が減る。この分を事前に計上する。

allocate_new_computed_blocks() [DEEP] [VERIFIED]

参照: L137

プレフィックスキャッシュヒットしたブロックをリクエストに追加する。

allocate_new_computed_blocks(request_id, new_computed_blocks, ...)
  │
  ├─ RUNNING → assert len(new_computed_blocks) == 0 → return
  │
  └─ 新規リクエスト:
     ├─ num_skipped_blocks 計算
     ├─ スキップ分を new_computed_blocks から除去
     ├─ enable_caching → block_pool.touch(new_computed_blocks)
     ├─ req_blocks に null_block × num_skipped_blocks を追加
     ├─ req_blocks に new_computed_blocks を追加
     ├─ num_cached_block[request_id] = len(req_blocks)
     └─ external_computed_tokens > 0 → 追加ブロック割り当て

allocate_new_blocks() [VERIFIED]

参照: L208

def allocate_new_blocks(self, request_id, num_tokens, num_tokens_main_model):
    num_required_blocks = cdiv(num_tokens, self.block_size)
    num_new_blocks = num_required_blocks - len(req_blocks)
    if num_new_blocks <= 0:
        return []
    new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
    req_blocks.extend(new_blocks)
    return new_blocks  # 新規分のみ返す

cache_blocks() [VERIFIED]

参照: L235

def cache_blocks(self, request, num_tokens):
    num_full_blocks = num_tokens // self.block_size
    if num_cached_blocks >= num_full_blocks:
        return  # 既に登録済み
    block_pool.cache_full_blocks(request, req_blocks, num_cached_blocks,
                                  num_full_blocks, block_size, group_id)
    num_cached_block[request_id] = num_full_blocks

free() [VERIFIED]

参照: L261

def free(self, request_id):
    req_blocks = self.req_to_blocks.pop(request_id, [])
    ordered_blocks = reversed(req_blocks)  # 逆順でLRU最適化
    self.block_pool.free_blocks(ordered_blocks)
    self.num_cached_block.pop(request_id, None)

逆順の理由: チェーン末尾(最新トークン)がキューの先頭側に来ることで、次の割り当て時に最初に evict される。先頭ブロック(プレフィックス)は長く残り、共有確率が上がる。

remove_skipped_blocks() [VERIFIED]

参照: L343

def remove_skipped_blocks(self, request_id, total_computed_tokens):
    num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)
    if num_skipped_tokens <= 0:
        return  # Full Attention: 何もしない
    # 後方から走査して null_block に遭遇したら停止
    for i in range(num_skipped_blocks - 1, -1, -1):
        if blocks[i] == null_block:
            break  # 前回の呼び出しで既に解放済み
        removed_blocks.append(blocks[i])
        blocks[i] = null_block
    block_pool.free_blocks(removed_blocks)

get_num_skipped_tokens() [VERIFIED]

参照: L386

デフォルト実装は return 0(全トークンがアテンション対象)。サブクラスでオーバーライド。

FullAttentionManager [DEEP] [VERIFIED]

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py:400

標準的な全トークンアテンション。基底クラスの動作をそのまま継承し、2 つのメソッドのみ実装。

find_longest_cache_hit()

参照: L401

左→右にブロックハッシュを走査:
  キャッシュヒット → computed_blocks に追加
  キャッシュミス → break(チェーンが途切れたら以降は必ずミス)

EAGLE 使用時:
  最後のブロックを削除(hidden states 再計算が必要)

alignment_tokens でアライメント調整:
  Hybrid モデルで LCM ブロックサイズの倍数に切り詰め

Downward-closed 性質: Full Attention では blocks[0..n] がヒットするなら blocks[0..n-1] も必ずヒットする。この性質により、左→右の貪欲スキャンで最適解が得られる。

get_num_common_prefix_blocks()

参照: L450

def get_num_common_prefix_blocks(self, running_request_id):
    for block in blocks:
        if block.ref_cnt == len(self.req_to_blocks):
            num_common_blocks += 1
        else:
            break
    return num_common_blocks

原理: ref_cnt == 全リクエスト数 なら、そのブロックは全リクエストで共有されている → 共通プレフィックス。Cascade Attention で共通プレフィックスの再計算をスキップするために使用。

SlidingWindowManager [DEEP] [VERIFIED]

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py:461

Sliding Window Attention 用。ウィンドウ外のトークンの KV キャッシュを解放してメモリを節約する。

コンストラクタ

def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs):
    super().__init__(kv_cache_spec, **kwargs)
    self.sliding_window = kv_cache_spec.sliding_window

get_num_skipped_tokens()

参照: L556

def get_num_skipped_tokens(self, num_computed_tokens):
    return max(0, num_computed_tokens - self.sliding_window + 1)
例: sliding_window=4, num_computed_tokens=7

Tokens: [0 1 2 3 4 5 6 | 7]
                          ↑ 次に計算するトークン
                  [4 5 6 7]  ← sliding window(サイズ4)
        [0 1 2 3]            ← skipped(4トークン)

find_longest_cache_hit()

参照: L466

右→左にブロックハッシュを走査:
  キャッシュヒット → computed_blocks[i] にセット、連続カウント++
  キャッシュミス → 連続カウント = 0(リセット)

  連続カウント >= sliding_window_contiguous_blocks:
    末尾をトリミングして break

sliding_window_contiguous_blocks = ceil((window - 1) / block_size)

Right-to-left の理由: Sliding Window は最新のトークン付近のブロックが重要。右端から連続ヒットを探すことで、ウィンドウ内の有用なキャッシュを効率的に発見する。

初期値: computed_blocksnull_block で埋められ、ヒットした位置のみ実ブロックで置換される。

制約事項:

  • DCP/PCP 非対応 (assert dcp_world_size == 1)
  • EAGLE 使用時は sliding_window_contiguous_blocks += 1

get_num_common_prefix_blocks()

参照: L584

常に 0 を返す。プレフィックスブロックは全て null_block に置換されているため、Cascade Attention は使用不可。

ChunkedLocalAttentionManager [DEEP] [VERIFIED]

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py:594

チャンク境界でアテンションが分割されるモデル用。各チャンク内のトークンのみが互いにアテンションする。

コンストラクタ

def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs):
    super().__init__(kv_cache_spec, **kwargs)
    self.attention_chunk_size = kv_cache_spec.attention_chunk_size

get_num_skipped_tokens()

参照: L691

def get_num_skipped_tokens(self, num_computed_tokens):
    return (num_computed_tokens // self.attention_chunk_size) * self.attention_chunk_size
例1: chunk_size=8, computed=13 → skipped=8  (チャンク[0,7]全体)
例2: chunk_size=8, computed=8  → skipped=8  (チャンク[0,7]全体)
例3: chunk_size=8, computed=7  → skipped=0  (まだチャンク内)

find_longest_cache_hit()

参照: L599

1. local_attention_start_idx = (max_length // chunk_size) * chunk_size
2. computed_blocks = [null_block] × (start_idx // block_size)  ← ウィンドウ外
3. start_idx から max_num_blocks まで左→右スキャン:
   ヒット → append、ミス → break

ウィンドウ外のブロックは null_block でパディングし、ウィンドウ内のみ FullAttention と同様の左→右スキャンを行う。

制約事項:

  • EAGLE 非対応 (assert use_eagle is False)
  • DCP/PCP 非対応
  • 異なるブロックサイズの混在非対応

MambaManager [DEEP] [VERIFIED]

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py:744

Mamba(State Space Model / 線形アテンション)用の Manager。Transformer ベースのアテンションとは根本的に異なり、KV キャッシュではなく「状態」を管理する。

2つのキャッシュモード

モード説明
none(デフォルト)基底クラスの動作 + speculative blocks 追加
align最後の状態ブロックのみ追跡、null_block パディング、speculative blocks 再利用

追加状態(align モード)

フィールド説明
last_state_block_idxdict[str, int]前ステップで割り当てた状態ブロックのインデックス
_allocated_block_reqsset[str]ブロック割り当て済みリクエストの集合
num_speculative_blocksintSpeculative Decoding 用の余分なブロック数

get_num_skipped_tokens()

参照: L967

def get_num_skipped_tokens(self, num_computed_tokens):
    return num_computed_tokens - 1  # 最後の状態のみ必要

Mamba は状態の累積的な更新なので、最後のトークンの状態さえあれば以前のトークンの状態は不要。

find_longest_cache_hit()

参照: L756

右→左に走査:
  最初のヒットで即座に break
  ヒット位置の前を null_block で埋める

最後の状態のみ必要なため、最右のヒット 1 つで十分。

get_num_blocks_to_allocate()(align モード) [DEEP] [VERIFIED]

参照: L832

align モード:
  ├─ 既存リクエスト(_allocated_block_reqs に存在):
  │  └─ 最大 1 ブロック追加(speculative blocks 再利用のため)
  │
  └─ 新規リクエスト:
     └─ 1 + num_speculative_blocks ブロック

allocate_new_blocks()(align モード) [DEEP] [VERIFIED]

参照: L885

align モードの割り当ては複雑:

1. num_tokens を main model 分に制限(lookahead 除外)
2. last_state_block_idx を記録:
   - 既存: prev_len - 1 - num_speculative_blocks
   - 新規(キャッシュヒット有): prev_len - 1
3. null_block でスキップ位置をパディング
4. 既存リクエスト: speculative blocks をスキップ位置に移動して再利用
5. 残りの新規ブロックを割り当て

remove_skipped_blocks()(align モード)

参照: L804

基底クラスの remove_skipped_blocks() に加え、last_state_block_idx のブロックも解放する:

if last_state_block_idx < cdiv(num_computed_tokens, block_size) - 1:
    block_pool.free_blocks([blocks[last_state_block_idx]])
    blocks[last_state_block_idx] = null_block

2 ステップ前のブロックが不要になるタイミングで解放する。

free()

参照: L961

def free(self, request_id):
    if self.mamba_cache_mode == "align":
        self._allocated_block_reqs.discard(request_id)
        self.last_state_block_idx.pop(request_id, None)
    super().free(request_id)

CrossAttentionManager [DEEP] [VERIFIED]

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py:976

エンコーダ-デコーダモデル(Whisper 等)のクロスアテンション用。エンコーダ出力はリクエスト固有(異なる音声/画像入力)のため、プレフィックスキャッシュの恩恵がない。

制約

メソッド動作
allocate_new_computed_blocks()assert len(new_computed_blocks) == 0
cache_blocks()raise ValueError
find_longest_cache_hit()raise NotImplementedError
get_num_common_prefix_blocks()return 0

エンコーダブロックはリクエスト開始時に num_encoder_tokens に基づいて静的に割り当てられ、デコード中は変化しない。

SinkFullAttentionManager [DEEP] [VERIFIED]

参照: target/vllm/vllm/v1/core/single_type_kv_cache_manager.py:1025

StreamingLLM のための Attention Sink 実装。FullAttentionManager を継承し、初期化時に先頭の sink ブロックを事前確保する。

コンストラクタ

class SinkFullAttentionManager(FullAttentionManager):
    def __init__(self, kv_cache_spec: SinkFullAttentionSpec, ...):
        super().__init__(...)
        sink_len = kv_cache_spec.sink_len
        assert sink_len > 0 and sink_len % self.block_size == 0
        num_sink_block = sink_len // self.block_size
        self.sink_blocks = self.block_pool.free_block_queue.popleft_n(num_sink_block)

特徴:

  • sink_lenblock_size の倍数でなければならない
  • sink ブロックは初期化時に popleft_n() で確保され、以降解放されない
  • FullAttentionManager の find_longest_cache_hit()get_num_common_prefix_blocks() をそのまま使用

各 Manager の比較表 [VERIFIED]

Managerスキップ計算キャッシュ検索CascadeDCP/PCPEAGLE
FullAttention0(全トークン)左→右ref_cnt 基準対応対応
SlidingWindowmax(0, n-w+1)右→左(連続)非対応非対応対応
ChunkedLocal(n//c)*cnull_pad + 左→右非対応非対応非対応
Mamban - 1右→左(単一)非対応非対応-
CrossAttention0非対応非対応--
SinkFullAttention0左→右ref_cnt 基準対応対応

関連ドキュメント