塊表:邏輯 token 到物理 KV 塊的映射

源码版本v0.25.1

職責

PagedAttention 的核心思想是把 KV cache 切成固定大小的 block,每個 token 在哪個 block 的哪個 slot 由"塊表 (block table)"查出來。vLLM v1 把塊表實現成 BlockTable 類:一個 (max_num_reqs, max_num_blocks_per_req) 的 int32 張量,每行對應一條請求,行裡第 j 個元素是該請求第 j 個邏輯 block 對應的物理 block_id。MultiGroupBlockTable 把多個 KV cache group 的塊表打包在一起,讓 hybrid 模型(例如 full attention + sliding window)各持一份。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 三份,在 CPU 上寫(numpy),在 GPU 上讀(int32 tensor);這樣 commit_block_table(num_reqs) 只需一次 H2D copy(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 把一個 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 處理一條請求,從 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 三個 constexpr 讓 kernel 在 decode context parallelism 下,只把屬於本 rank 的 slot 寫真實值、其餘寫 PAD_ID(block_table.py:368-379)。
• CUDA graph padding:num_reqs < num_reqs_padded 時,_get_block_table 把 padding 行用 NULL_BLOCK_ID 填充(gpu_model_runner.py:2279-2295),attention kernel 看到 NULL_BLOCK_ID 就跳過那些行。
• 多 group 共享 builder 快取:supports_update_block_table=True 的 builder 在第二次同樣 spec 的 group 上不重建,只 update_block_table(gpu_model_runner.py:2477-2493),省去重複的 metadata 構建。
• 行級維護:append_row / clear_row / move_row / swap_row 把塊表當成"按請求行的動態陣列",在 condense(請求重排)、preempt(請求換出)等場景下原地更新,不需要整體重建(block_table.py:102-139)。

關鍵檔案

• BlockTable.__init__:18-100 — 雙緩衝構造,處理 kernel_block_size != block_size 的 hybrid 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 — 返回 GPU 張量切片給 attention backend。
• 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 padding。
• 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) 把新塊加進行尾,對 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] 清零再把計數清零(block_table.py:124-128)。
• CUDA graph padding: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