ブロック表:論理 token から物理 KV block へのマッピング

源码版本v0.25.1

役割

PagedAttention の中心思想は KV cache を固定サイズの block に切り、各 token がどの block のどの slot にあるかを「ブロック表 (block table)」で lookup することです。vLLM v1 はブロック表を BlockTable クラスとして実装します:(max_num_reqs, max_num_blocks_per_req) の int32 テンソルで、各行が 1 つのリクエストに対応し、行の j 番目の要素がそのリクエストの j 番目の論理 block に対応する物理 block_id です。MultiGroupBlockTable は複数の KV cache group のブロック表を一緒にまとめ、ハイブリッドモデル (例: full attention + sliding window) で各 group ごとに 1 つずつ持たせます。BlockTable は同時に slot_mapping テンソルを保持します。各 query token を KV cache 内の物理 slot id に直接マッピングし、attention kernel はこれを使って KV を直接読み書きします。

ブロック表のライフサイクルはリクエストに紐付きます。InputBatch.add_request のとき block_table.add_row(request.block_ids, req_index) が manager から返された block_id リストをその行に書き込みます (gpu_input_batch.py:336-379)。各スケジュールステップのあと、新しく割り当てられた block_id は block_table.append_row(new_block_ids, req_index) で行末に追加されます (gpu_model_runner.py:1425-1428)。commit_block_table が numpy 配列を GPU にコピーし (block_table.py:166-167)、compute_slot_mapping が Triton kernel で各 token の slot id を計算します (block_table.py:141-164)。前方計算時、attention backend は common_attn_metadata.block_table_tensor (backend.py:420-421) からこのテンソルを読みます。

設計動機

なぜ block table を worker 側で改めて管理するのか?manager の req_to_blocks を直接使うのではだめなのか?

• numpy + GPU のダブルバッファ:block_table は CpuGpuBuffer で numpy / CPU / GPU の 3 つを同時に持ちます。CPU 上で書き (numpy)、GPU 上で読みます (int32 tensor)。これにより commit_block_table(num_reqs) は H2D copy 1 回で済みます (block_table.py:70-77)。CUDA graph キャプチャモードではサイズ固定で GPU テンソルをそのまま使い回します。
• kernel block_size ≠ manager block_size:一部の attention kernel (例: TRTLLM) の block_size は 16 ですが、KV cache manager の block_size は 32 です。map_to_kernel_blocks が 1 つの manager block を blocks_per_kv_block=2 個の kernel block に展開します (block_table.py:47-68)(block_table.py:173-201)。
• slot_mapping の kernel 化:compute_slot_mapping は Triton kernel _compute_slot_mapping_kernel を呼び出します。各 program が 1 つのリクエストを処理し、positions から block_index = pos // block_size を計算、block_table[req_idx, block_index] で block_number を引いて、最後に slot_id = block_number * block_size + offset を得ます (block_table.py:325-380)。numpy 実装より速く、CUDA graph 内でも実行できます。
• DCP / PCP のインターリーブ:TOTAL_CP_WORLD_SIZE / TOTAL_CP_RANK / CP_KV_CACHE_INTERLEAVE_SIZE の 3 つの constexpr により、kernel は decode context parallelism のもとで自 rank に属する slot だけに真の値を書き、それ以外は PAD_ID を書きます (block_table.py:368-379)。
• CUDA graph のパディング:num_reqs < num_reqs_padded のとき、_get_block_table はパディング行を NULL_BLOCK_ID で埋めます (gpu_model_runner.py:2279-2295)。attention kernel は NULL_BLOCK_ID を見たらその行をスキップします。
• 複数 group での builder キャッシュ共有:supports_update_block_table=True の builder は同じ spec の group が 2 度目に来たときに再構築せず、update_block_table だけを行います (gpu_model_runner.py:2477-2493)。メタデータ構築の重複を省きます。
• 行レベルのメンテナンス:append_row / clear_row / move_row / swap_row はブロック表を「リクエスト行ごとの動的配列」として扱い、condense (リクエストの並べ替え) や preempt (リクエストのスワップアウト) などのケースでインプレースに更新します。全体の再構築は不要です (block_table.py:102-139)。

主要ファイル

• BlockTable.__init__:18-100 — ダブルバッファ構築、kernel_block_size != block_size のハイブリッド block モードを処理。
• BlockTable.append_row/add_row:102-122 — 行末に追加 / 行全体を block_id リストで書き込み。
• BlockTable.clear/move/swap_row:124-139 — 行レベルのメンテナンス、condense / preempt の並べ替えを支える。
• BlockTable.compute_slot_mapping:141-164 — Triton kernel を呼び出し positions を slot_mapping テンソルに変換。
• BlockTable.commit_block_table/clear:166-171 — H2D コピー / ゼロクリア。
• BlockTable.map_to_kernel_blocks:173-201 — manager block_id リストを kernel block_id リストに展開。
• BlockTable.get_device_tensor:203-213 — attention backend に GPU テンソルのスライスを返します。
• MultiGroupBlockTable:223-322 — 複数 group のラッパー、add_row / append_row / compute_slot_mapping などを各 group の BlockTable にブロードキャスト。
• _compute_slot_mapping_kernel:325-380 — Triton kernel:per-req program、slot_id = block_table[req, pos // block_size] * block_size + offset を計算、DCP 時は TOTAL_CP_RANK でフィルタ。
• InputBatch.block_table init:172-180 — MultiGroupBlockTable の構築ポイント。
• InputBatch.add_request:336-379 — 新規リクエスト入場時に block_table.add_row(request.block_ids, req_index) を実行。
• _get_block_table:2279-2298 — 各ステップで block_table_tensor を用意、NULL_BLOCK_ID のパディングを補完。
• append_row on new blocks:1425-1428 — 新規割り当ての block_id を既存行の末尾に追加。
• commit_block_table:1933-1933 — スケジュール完了後に numpy を H2D で GPU に書き込み。
• CommonAttentionMetadata:2378-2396 — block_table_tensor を common metadata に保持。
• update_block_table:2477-2493 — 同じ spec の group が構築済みの metadata を再利用。
• block_table_tensor field:420-421 — CommonAttentionMetadata の block_table フィールド。
• FlashAttentionMetadata.block_table:234-234 — Flash backend が持つ block_table tensor。

データフロー

各スケジュールサイクルで、gpu_model_runner._prepare_inputs が SchedulerOutput を input batch の状態に変換します。既存リクエストごとに block_table.append_row(new_block_ids, req_index) で新 block を行末に追加し、condense 後の空スロットには move_row / swap_row で整理します (gpu_input_batch.py:629-758)。その後 commit_block_table(num_reqs) が numpy 配列を H2D で GPU にコピーし、compute_slot_mapping が Triton kernel で slot_mapping を計算します:

@triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"])
def _compute_slot_mapping_kernel(
    num_tokens,
    max_num_tokens,
    query_start_loc_ptr,  # [num_reqs + 1], int32
    positions_ptr,  # [num_tokens], int64
    block_table_ptr,  # [max_num_reqs, max_num_blocks_per_req], int32 (flat)
    block_table_stride,  # max_num_blocks_per_req
    block_size,
    slot_mapping_ptr,  # [max_num_tokens], int64
    TOTAL_CP_WORLD_SIZE: tl.constexpr,
    TOTAL_CP_RANK: tl.constexpr,
    CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
    PAD_ID: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    req_idx = tl.program_id(0)

    if req_idx == tl.num_programs(0) - 1:
        # Pad remaining slots for CUDA graph compatibility.
        for i in range(num_tokens, max_num_tokens, BLOCK_SIZE):
            offsets = i + tl.arange(0, BLOCK_SIZE)
            tl.store(
                slot_mapping_ptr + offsets,
                PAD_ID,
                mask=offsets < max_num_tokens,
            )
        return

(block_table.py:325-352) 最後の program は padding slot に PAD_ID を書き込む担当で、CUDA graph 再利用時に残データが attention を汚染しないようにします。その後 _get_block_table(kv_cache_gid) がこの group の GPU テンソルを取り出し (gpu_model_runner.py:2279-2295)、CommonAttentionMetadata.block_table_tensor に詰めます (gpu_model_runner.py:2378-2396)。Flash attention / Triton / rocM などの backend はそれぞれ attn_metadata.block_table を読み (flash_attn.py:234-234)(triton_attn.py:208-234)、slot_mapping と組み合わせて KV cache の読み書きを完了します。block_id は KVCacheManager から Coordinator と BlockPool を経て割り当てられたものです。

境界と失敗

• kernel_block_size が block_size を割り切らない:block_size % kernel_block_size != 0 なら直接 raise ValueError します (block_table.py:58-62)。
• 空リストの append:append_row([]) はそのまま return し、num_blocks_per_row を更新しません (block_table.py:107-108)。
• clear_row は num_blocks_per_row を 0 にする前にゼロクリア:clear_row は block_table.np[row, :num_blocks] をゼロクリアしてからカウンタを 0 にします (block_table.py:124-128)。
• CUDA graph のパディング:num_reqs:num_reqs_padded の行は NULL_BLOCK_ID で埋め (gpu_model_runner.py:2292-2294)、kernel は NULL_BLOCK_ID を見たらスキップします。
• EncoderOnlyAttentionSpec はゼロテンソルを使用:encoder-only の KV cache には block_table がなく、_get_block_table は (num_reqs_padded, 1) の全ゼロテンソルをプレースホルダーとして返します (gpu_model_runner.py:2282-2287)。
• PCP / DCP の未初期化:get_pcp_group() / get_dcp_group() はテストで未初期化のことがあり、try/except で world_size=1 にフォールバックします (block_table.py:86-99)。
• Mamba は block_table_tensor を再利用:Mamba / GDN / Linear attention バックエンドは mamba_get_block_table_tensor で block_table を state_indices として使い、cache mode で gather する必要があります (utils.py:925-965)。
• max_num_blocks を TRTLLM に合わせる:MultiGroupBlockTable.__init__ は max_num_blocks を 128 // block_size の倍数に揃えます。一部の attention backend (TRTLLM) が要求するためです (block_table.py:260-265)。

まとめ

BlockTable は KVCacheManager / Coordinator と BlockPool が割り当てた block_id リストを GPU テンソルにまとめ、slot_mapping Triton kernel で各 token の物理 slot を計算し、CommonAttentionMetadata に詰めてすべての attention backend に渡します。モデル重みのロード (DefaultModelLoader 参照) とは結合しておらず、KV cache 管理層とのみ結合しています。

公式資料: vLLM 文档 · README