vLLM Code Reading (Section 6): KV Cache 中的 BlockTable 与 Slot Mapping
vLLM 的 KV cache 管理有两个核心映射层:BlockTable 解决 request 到 KV block id 的映射,slot_mapping 解决本次 forward 中每个 token 到物理 KV slot 的映射。前者是请求级、长期维护的数据结构;后者是 step 级、临时生成的 GPU 工作区。
1. KV cache group 的含义
KV cache group 表示一批具有相同 KV cache 管理规格的层。这里的 group 不是 tensor parallel group,也不是 worker group,而是 KVCacheConfig.kv_cache_groups 里的 KVCacheGroupSpec。一个 group 里保存两类信息:layer_names 表示哪些 layer 属于这个组,kv_cache_spec 表示这些 layer 使用什么 KV cache 规格。
普通 decoder-only full attention 模型通常只有一个 group。所有 attention layer 的 KV cache 规格一致,例如相同的 block_size、num_kv_heads、head_size、dtype 和 full attention 管理方式,因此可以共用一套 block table。这个共用不是说所有层共享同一份 K/V 数据,而是说它们使用同一套 request 到 block id 的映射。
混合模型通常会出现多个 group。比如 full attention 和 sliding window attention 混合时,full attention layer 和 sliding window layer 对历史 token 的保留规则不同,因此可能被放到不同 KV cache group。再比如含有 Mamba、MLA 或 cross attention 的模型,也可能因为 cache layout、block size 或管理策略不同而拥有多个 group。
flowchart TD
A["KVCacheConfig.kv_cache_groups"] --> B["group 0: FullAttentionSpec"]
A --> C["group 1: SlidingWindowSpec"]
A --> D["group 2: MambaSpec / MLA Spec"]
B --> B1["layer_names: decoder.layers.*.self_attn"]
C --> C1["layer_names: sliding window layers"]
D --> D1["layer_names: hybrid state layers"]
2. MultiGroupBlockTable 的结构
MultiGroupBlockTable 是 KV cache group 到 BlockTable 的容器。它里面不是每层一张表,而是每个 KV cache group 一张表。假设模型只有一个 full attention group,那么 MultiGroupBlockTable 里通常只有一张 BlockTable;假设模型有 full attention 和 sliding window 两类 group,那么里面通常有两张 BlockTable。
一张 BlockTable 内部主要有三块状态。block_table 是二维 int32 表,形状是 [max_num_reqs, max_num_blocks_per_req],用于记录每个 request row 拥有哪些物理 block id。num_blocks_per_row 是 CPU 侧数组,用于记录每个 row 当前已经写入多少 block。slot_mapping 是一维 int64 buffer,形状是 [max_num_batched_tokens],用于记录当前 step 中每个 scheduled token 要访问的物理 KV slot。
1 | |
3. Row 是 Worker 本地的 request 槽位
row 表示 InputBatch 里的 request index。它不是 token index,也不是 layer index,也不是 block id。Worker 侧维护一个 req_id_to_index 字典,把稳定的 req_id 映射到当前 persistent batch 里的 row。
这种设计的意义是让 batch 内部状态保持紧凑。请求完成或被移除后,vLLM 可以移动、交换或压缩 row,从而避免 batch 表里出现太多空洞。由于 row 可能变化,所以稳定身份是 req_id,而 row 是每个 Worker 当前 step 使用的局部执行槽位。
在 BlockTable 里,row 的语义可以写成下面这个公式:
1 | |
例如 row 0 对应 request A,block_table[0][0] = [10, 25, 31],表示 request A 在 group 0 中的第 0、1、2 个逻辑块分别映射到物理 block 10、25、31。
4. add_row 与 append_row 的分工
add_row 的语义是覆盖一整行。它会先把 num_blocks_per_row[row_idx] 清零,然后调用 append_row 从 row 的开头写入 block ids。因此,新请求进入 batch,或者 request 从 preemption 恢复并需要重建完整 block 列表时,适合走 add_row。
append_row 的语义是在已有 block 后继续追加。已有 request 在新的 scheduler step 中可能又获得了新的 KV block,这时不能清空旧 block,因为 prompt 和历史 decode token 的 KV cache 还要继续被 attention 读取。因此 append_row 会读取 start = self.num_blocks_per_row[row_idx],然后把新增 block ids 写到 [start, start + num_blocks) 这个区间。
1 | |
这也解释了为什么 append_row 里必须保留 start。如果所有路径都只是新建 request,那么 start 确实总是 0;但持续 decode 和 chunked prefill 都会追加 block,旧 block 不能丢。
5. 真正的 block 分配发生在 EngineCore
Worker 侧更新 BlockTable 不等于分配 block。真正决定“这个 request 需要哪些新 KV block”的地方在 EngineCore 的 Scheduler 和 KVCacheManager。Scheduler 根据 token budget、KV cache 空间、prefix cache、preemption 等条件调用 KVCacheManager.allocate_slots(),得到本 step 新增的 block ids。
这些 block ids 会通过 SchedulerOutput 下发给 Worker。新请求使用 scheduled_new_reqs[].block_ids 携带完整 block 列表,老请求使用 scheduled_cached_reqs.new_block_ids 携带增量 block 列表。Worker 收到后只是在本地 InputBatch.block_table 中执行 add_row 或 append_row,让本 rank 的模型执行能够按照 EngineCore 的分配结果访问 KV cache。
sequenceDiagram
participant S as Scheduler
participant K as KVCacheManager
participant O as SchedulerOutput
participant W as Worker InputBatch
participant B as BlockTable
S->>K: allocate_slots(request, num_new_tokens)
K-->>S: KVCacheBlocks
S->>O: block_ids / new_block_ids
O->>W: execute_model(scheduler_output)
W->>B: add_row(...) or append_row(...)
6. commit_block_table 只提交 block_table
commit_block_table() 的作用是把 CPU 侧更新过的 block table 拷贝到 GPU。Worker 更新 row 时写的是 block_table.np,也就是 CPU/Numpy 视图;但后面的 slot mapping Triton kernel 需要在 GPU 上读取 block id。因此在 _prepare_inputs() 开始阶段,vLLM 会提前调用 commit_block_table(num_reqs)。
这个提前调用是一个性能优化。block_table.copy_to_gpu(num_reqs) 发起的是 non-blocking copy,后面 CPU 还要准备 req_indices、query_start_loc、positions、采样元数据等内容。把 copy 放在前面,可以让 H2D 拷贝和 CPU 侧输入准备重叠。
需要强调的是,commit 不是提交给 EngineCore,而是 Worker 内部的 CPU staging buffer 到 GPU tensor 的同步。EngineCore 已经在 Scheduler 阶段完成了 block 分配,Worker 这里做的是执行前的数据准备。
7. slot_mapping 是 token 到物理 KV slot 的映射
slot_mapping 解决的是 token 级写入位置问题。模型 forward 看到的是本 step 的 scheduled tokens,它们在输入张量里是连续排列的;但 KV cache 在显存里是 paged/block layout,物理位置由 block id 和 block 内 offset 共同决定。因此 attention backend 写 K/V cache 时,需要知道第 i 个 scheduled token 应该写到哪个物理 slot。
计算公式可以简化成:
1 | |
例如 block_size = 16,request A 的 block ids 是 [10, 25],某个 token 的 absolute position 是 17。那么 block_idx = 1,block_id = 25,offset = 1,最终 slot_id = 25 * 16 + 1 = 401。后续 kernel 就可以用 slot_mapping[token_idx] = 401 把这个 token 的 K/V 写到对应 KV cache slot。
8. slot_mapping 在 GPU 上直接计算
slot_mapping 在正常 GPU runner 路径里不是 CPU 计算后拷贝到 GPU。它虽然使用 CpuGpuBuffer 分配了 CPU/GPU 两份 buffer,但主路径只使用 slot_mapping.gpu。真正的计算发生在 BlockTable.compute_slot_mapping() 里,Triton kernel 读取 block_table.gpu、query_start_loc.gpu 和 positions.gpu,然后直接写入 slot_mapping.gpu。
这条路径和 block_table 正好相反。block_table 是 CPU 更新、GPU 消费,所以需要 copy_to_gpu;slot_mapping 是 GPU 生成、GPU 消费,所以不需要 copy_to_gpu。它的 CPU buffer 主要是通用封装带来的,不是推理主路径的计算位置。
flowchart LR
A["block_table.np"] --> B["commit_block_table"]
B --> C["block_table.gpu"]
D["positions.gpu"] --> F["Triton slot mapping kernel"]
E["query_start_loc.gpu"] --> F
C --> F
F --> G["slot_mapping.gpu"]
G --> H["reshape_and_cache / attention backend"]
9. _get_slot_mappings 只是取结果和分发
_get_slot_mappings() 不是 kernel 调用点,也不是 block 分配点。它发生在 compute_slot_mapping() 之后,作用是把每个 KV cache group 已经生成好的 slot_mapping.gpu 取出来,并转换成两个消费者需要的格式。
第一个格式是 slot_mappings_by_gid,也就是 gid -> slot_mapping。这个结构给 attention metadata 构建使用,因为 metadata 需要按 KV cache group 理解 block table 和 slot mapping。第二个格式是 slot_mappings_by_layer,也就是 layer_name -> slot_mapping。这个结构给 forward context 使用,因为模型 forward 过程中 attention layer 通常按 layer name 取自己的运行时上下文。
这也再次说明 group 内每层不是各有一张 BlockTable。同一个 group 里的所有 layer 会指向同一个 slot_mapping tensor。它们共享的是 token 到 slot 的映射,但每个 layer 仍然有自己的 KV cache tensor。
10. Kernel 调用位置与异步语义
slot mapping kernel 的调用点在 _prepare_inputs()。GpuModelRunner 先准备当前 step 的 request 范围、position、sequence length 等元数据,然后调用 self.input_batch.block_table.compute_slot_mapping(...)。MultiGroupBlockTable 会遍历每个 BlockTable,每个 KV cache group 调一次 _compute_slot_mapping_kernel。
Triton kernel launch 是 CUDA 层面的异步 enqueue。Python 调用返回时,不代表 GPU 已经完成计算,而是代表 kernel 已经被放入当前 CUDA stream。后续 attention 或 reshape_and_cache kernel 在同一个 stream 上排队,因此 stream 顺序会保证它们读取 slot_mapping.gpu 之前,slot mapping kernel 已经完成写入。
这不是 Python 的 async/await,也不是返回 future 的异步 API,而是 CUDA runtime 的异步执行模型。CPU 线程可以继续 enqueue 后续 GPU 工作,GPU 按 stream 顺序执行。
11. 小 batch 下为什么仍然值得单独调用 kernel
单流 decode 时,slot mapping kernel 可能只为一个 token 工作。比如 num_reqs = 1 且 total_num_scheduled_tokens = 1,vLLM 仍然会 launch 一次 Triton kernel。这从单流 latency 看确实有固定 overhead,但这个设计服务的是 continuous batching 和吞吐优先的执行模型。
在大 batch decode 中,一个 step 往往对应许多 active request,每个 request 通常生成一个 token,slot mapping 的输入规模就不再是 1。在 prefill 或 chunked prefill 中,一个 step 可能包含几千个 scheduled tokens,这时提前用 GPU kernel 一次性生成 slot mapping 很自然。
更重要的是,这个 kernel 是按 KV cache group 调用,不是按 layer 调用。一个 group 内几十层共用同一个 slot_mapping。如果不提前生成映射,每个 attention layer 或 cache 写入 kernel 都可能重复执行 position -> block_id -> slot_id 这套计算。vLLM 用一次轻量 kernel 换掉后续多层重复计算,整体上更适合高吞吐场景。
12. BlockTable 和 slot_mapping 的大小
BlockTable 是元数据结构,大小通常远小于真正的 KV cache。单张表中最大的部分是 block_table,形状为 [max_num_reqs, max_num_blocks_per_req],元素类型是 int32,并且 CPU 和 GPU 各有一份。max_num_blocks_per_req 大致等于 ceil(max_model_len / block_size),因此它会随上下文长度线性增长。
以 max_model_len = 200000、block_size = 16、max_num_reqs = 256 为例,max_num_blocks_per_req = 12500。单侧 block_table 大小是 256 * 12500 * 4 bytes = 12.8 MB,CPU 和 GPU 两份合计约 25.6 MB。相比之下,真正的 KV cache 会按 layer、KV head、head dim 和 dtype 存储 K/V 数据,通常是 GB 级别。
slot_mapping 更小,因为它只按当前最大 batch token 数分配。它的大小是 max_num_batched_tokens * 8 bytes。如果 max_num_batched_tokens = 8192,单侧只有 64 KB;即使 max_num_batched_tokens = 131072,单侧也只有 1 MB。因此 slot_mapping.gpu 更像是每个 group 的运行时 scratch buffer,而不是长期占用大量显存的 cache。
13. 复用边界
slot_mapping.gpu 的内存可以跨 step 复用,但内容不能跨 step 复用。每个 step 的 active request、row 排列、scheduled token 数量、positions 和新增 block ids 都可能变化,因此 slot mapping 的值必须重新计算。vLLM 通过预分配 buffer 避免每 step 重新申请显存,通过覆盖写入保证每 step 的映射正确。
BlockTable 的内容则是跨 step 维护的。一个 request 持续 decode 时,它的 row 会保留已有 block ids,并在需要新 block 时追加。request 完成、被移除或 batch 被压缩时,row 可能被清空、移动或交换。也就是说,BlockTable 是 persistent batch 的一部分,而 slot_mapping 是一次 forward 的派生数据。
14. 总结
vLLM 在 KV cache 访问上使用了两级映射。第一级是 BlockTable,它把 request row 和序列 block index 映射到物理 block id;第二级是 slot_mapping,它把当前 step 的 token index 映射到物理 KV slot。第一级由 EngineCore 分配 block 后在 Worker 侧长期维护,第二级由 Worker 在每次 forward 前用 GPU kernel 临时生成。
这套设计的关键收益是职责清晰。EngineCore 统一做 block 分配,Worker 只维护本地执行所需的表;CPU 负责 request 级状态更新,GPU 负责 token 级 slot 映射;group 内所有 layer 共用同一套映射,避免按层重复计算。代价是小 batch decode 时会多一次轻量 kernel launch,但在 continuous batching、chunked prefill 和多层 attention 写 cache 的整体场景下,这个开销是可控且容易复用的。
