块表:逻辑 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。