news 2026/10/14 2:07:48

vllm-ascend Triton 算子深度解析：`_compute_slot_mapping_kernel` 的 KV Cache 槽位映射原理与实战

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
vllm-ascend Triton 算子深度解析：`_compute_slot_mapping_kernel` 的 KV Cache 槽位映射原理与实战
  • 人工智能
  • 大模型
  • 模型推理服务
  • Ascend
  • CANN

【免费下载链接】vllm-ascend

Community maintained hardware plugin for vLLM on Huawei Ascend

项目地址：https://gitcode.com/gh_mirrors/vl/vllm-ascend
点击查看免费下载

导读

本文围绕 vllm-ascend 在 Triton 上实现的 KV Cache 槽位映射算子_compute_slot_mapping_kernel（源码位于 vllm_ascend/ops/triton/compute_slot_mapping.py，配套文档为 vllm_ascend/ops/triton/compute_slot_mapping.md），完整讲解其「逻辑 token 位置 → 物理 KV Cache slot id」的映射算法、Context Parallel（CP）交错布局下的虚拟块换算、全部运行时与编译期参数含义，以及可直接运行的调用示例。读完本文，你将能独立理解 vllm-ascend 推理链路中 slot mapping 的生成过程，并具备在 NPU 上调试、定制该算子（如调整 tile 大小、块表窗口、CP 交错粒度）的能力。

1、算子功能：为什么 Attention 需要 slot mapping

在大模型推理过程中，KV Cache 通常采用block-based（分块）的物理存储方式：逻辑上属于同一序列的 token 会被分散到若干物理块中，块号由block_table维护。Attention 算子需要拿到每个 token 在物理 KV Cache 中的slot id（槽位号），才能正确地读写对应的 K/V 向量。

_compute_slot_mapping_kernel正是完成「逻辑位置 → 物理 slot id」映射计算的 Triton 算子，其核心行为如下：

  • 每个program_id(0) == req_idx的 program 处理一条请求（一个 sequence）；
  • 通过query_start_loc_ptr得到该请求在 batch 中的 token 起止区间[start_idx, end_idx)；
  • 对区间内每个 token：
    • 读取其逻辑序列位置pos；
    • 根据block_size将pos拆分为逻辑块号block_indices与块内偏移slot_offsets；
    • 通过block_table把逻辑块号转换为物理块号block_numbers；
    • 计算slot_id = block_numbers * block_size + slot_offsets，写回slot_mapping_ptr；
  • 当TOTAL_CP_WORLD_SIZE > 1（启用 Context Parallel）时，会按CP_KV_CACHE_INTERLEAVE_SIZE交错布局把虚拟块切分到各 rank：仅属于本 rank 的 token 写入真实 slot id，其余写入PAD_ID（详见第 3 节）；
  • 最后一个 program（req_idx == num_programs(0) - 1）专门用于把slot_mapping_ptr中[num_tokens, max_num_tokens)这段填充为PAD_ID，这是为了与 CUDA Graph 捕获（graph 兼容）下的固定形状输出保持一致。

从源码看（vllm_ascend/ops/triton/compute_slot_mapping.py），padding 分支的实现在req_idx == tl.num_programs(0) - 1时，以TILE_BLOCK_SIZE为步长循环，将[num_tokens, max_num_tokens)全部写入PAD_ID。该算子在 vllm-ascend 中的实际生产调用点位于 vllm_ascend/worker/block_table.py 的BlockTable.compute_slot_mapping（非 CP 路径），更完整的 v2/vLLM 主线适配内核见 vllm_ascend/ops/triton/v2/block_table/compute_slot_mappings.py。

1.1 从调用链理解算子定位

在推理每步执行前，vllm-ascend 的BlockTable会准备好query_start_loc（CSR 风格 indptr）、positions（各 token 的逻辑位置）以及 GPU 上的块表，然后启动本内核。非 CP 路径的启动方式（vllm_ascend/worker/block_table.py）展示了全部编译期常量的实际取值来源：

TILE_BLOCK_SIZE = 1024 kernel_kwargs = { "KV_CACHE_BLOCK_SIZE": self.physical_block_size, "BLOCKS_PER_KV_BLOCK": self.blocks_per_phys_block, "TOTAL_CP_WORLD_SIZE": total_cp_world_size, "TOTAL_CP_RANK": total_cp_rank, "CP_KV_CACHE_INTERLEAVE_SIZE": self.cp_kv_cache_interleave_size, "PAD_ID": PAD_SLOT_ID, "TILE_BLOCK_SIZE": TILE_BLOCK_SIZE, "BLOCK_TABLE_WINDOW_SIZE": _next_power_of_2(cdiv(TILE_BLOCK_SIZE, self.block_size) + 1), "IS_CIRCULAR": self.is_circular, } _compute_slot_mapping_kernel(num_reqs + 1,), self.block_size, self.slot_mapping.gpu, **kernel_kwargs, )

从源码结构可以推断：当dcp_world_size > 1时走的是_compute_dcp_slot_mapping的 torch 实现（见 vllm_ascend/worker/block_table.py），其数学公式与本内核的 CP 分支完全一致；而 v2/vLLM 主线则通过AscendBlockTables.compute_slot_mappings启动_compute_slot_mappings_kernel（vllm_ascend/worker/v2/block_table.py），网格为(num_kv_cache_groups, num_reqs + 1)，说明该算子族同时服务于 v1 与 v2 两条推理路径。

2、非 CP 路径：基础映射公式

当TOTAL_CP_WORLD_SIZE == 1时，映射逻辑最简单（vllm_ascend/ops/triton/compute_slot_mapping.py）：

block_indices = pos // block_size # 逻辑块号 slot_offsets = pos - block_indices * block_size # 块内偏移

随后通过块表把逻辑块号翻译为物理块号并组装 slot id：

block_numbers = tl.gather(block_table_window, relative_block_indices, 0) slot_ids = block_numbers * block_size + slot_offsets

2.1 有界窗口 gather：避免整行加载

值得注意的实现细节是：内核并没有一次性加载整行block_table，而是采用「有界窗口 gather」策略（见 vllm_ascend/ops/triton/compute_slot_mapping.py）：

  1. 对每个TILE_BLOCK_SIZE大小的 token tile，先求出该 tile 内最小的有效逻辑块号block_idx_base（利用INT32_MAX哨兵值配合tl.min）；
  2. 只从块表行中加载从block_idx_base开始的BLOCK_TABLE_WINDOW_SIZE个连续条目；
  3. 用相对块号block_indices - block_idx_base执行tl.gather，得到物理块号。

这一设计的好处是：BLOCK_TABLE_WINDOW_SIZE是 2 的幂次编译期常量，便于 Triton 向量化；窗口大小只与TILE_BLOCK_SIZE / block_size有关，与请求长度无关，避免了长序列请求下加载整行块表带来的额外显存带宽开销。v2 内核文档 vllm_ascend/ops/triton/docs/v2/block_table/compute_slot_mappings.md 中对此有同样说明，并强调「长块表行不需要长 UB 分配」。

2.2 IS_CIRCULAR 分支

源码还支持环形缓冲 KV Cache（IS_CIRCULAR: tl.constexpr = False）。当IS_CIRCULAR为真时，逻辑块号固定为 0、块内偏移为pos % block_size，并且对pos < 0（无效位置）写入PAD_ID（[vllm_ascend/ops/triton/compute_slot_mapping.py](https://link.gitcode.com/i/5facefeec528edf8e21d67eb987e60f2#L53-L55, L86-L87)）。从 vllm_ascend/worker/block_table.py 可以看到，环形尾部缓存（circular tail cache）在 CP 下是不被允许的（dcp_world_size != 1时直接抛ValueError），且环形组的块表行数被固定为 1。

3、多 CP 域下的计算说明

当启用 Context Parallel（TOTAL_CP_WORLD_SIZE > 1）时，KV Cache 会以interleave（交错）方式分散存储到各 rank 上：序列中第i个 token 的 K/V 始终存放在dcp_rank == (i // CP_KV_CACHE_INTERLEAVE_SIZE) % TOTAL_CP_WORLD_SIZE的设备上。为保证各 rank 仅访问本地显存，本算子在 CP 模式下引入「虚拟块」概念并按以下步骤完成映射。

3.1 五步换算流程

步骤 1：虚拟块划分。将TOTAL_CP_WORLD_SIZE个 rank 的物理块在逻辑上合并为一个虚拟块：

virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE virtual_block_indices = pos // virtual_block_size # 虚拟块号 virtual_block_offsets = pos % virtual_block_size # 虚拟块内偏移

步骤 2：归属判定。通过 interleave 粒度判断当前 token 是否属于本 rank：

is_local = (virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK

is_local == False的 token 最终写入PAD_ID，不会访问本 rank 的block_table。

步骤 3：本地偏移换算。将虚拟偏移中属于本 rank 的部分抽取为本地物理偏移：

local_block_offsets = ( virtual_block_offsets // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE) ) * CP_KV_CACHE_INTERLEAVE_SIZE + (virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE)

该公式把每个虚拟块内本 rank 拥有的若干 interleave 段拼接成连续的本地偏移。

步骤 4：逻辑块号与块内偏移。基于本地偏移计算block_table索引：

block_indices = virtual_block_indices * BLOCKS_PER_KV_BLOCK + local_block_offsets // block_size slot_offsets = local_block_offsets % block_size

注意这里乘上了BLOCKS_PER_KV_BLOCK：当一个物理 KV 块被切分成多个逻辑 block（hybrid 块模式）时，虚拟块内也包含BLOCKS_PER_KV_BLOCK个逻辑块。

步骤 5：物理块号查询与 slot id 组装。通过block_table把逻辑块号转换为物理块号：

block_numbers = gather(block_table, block_indices) slot_ids = block_numbers * block_size + slot_offsets slot_ids = where(is_local, slot_ids, PAD_ID)

3.2 源码级验证

上述 CP 公式在 vllm_ascend/ops/triton/compute_slot_mapping.py 与 torch 参考实现 vllm_ascend/worker/block_table.py 中完全一致；v2 内核 vllm_ascend/ops/triton/v2/block_table/compute_slot_mappings.py 采用同样的虚拟块换算。测试用例 tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_compute_slot_mapping_v1.py 中提供了逐 token 的参考实现，可用于对照核验任意(cp_size, cp_rank, cp_interleave)组合下的映射结果；v2 内核还有专门的 CP 对拍测试（与 vLLM 上游参考内核逐元素比对，见 tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_compute_slot_mapping.py）。

3.3 一个直观的例子

假设TOTAL_CP_WORLD_SIZE = 2、CP_KV_CACHE_INTERLEAVE_SIZE = 128、block_size = 128，则 token 的 K/V 归属按 128 粒度在两个 rank 间交替：位置0..127属于 rank 0，128..255属于 rank 1，256..383又属于 rank 0，依此类推。对 rank 0 而言，位置0..127与256..383两段在虚拟块内通过local_block_offsets被拼接为连续的本地偏移0..255，从而映射到连续的物理 slot，这正是 attention 内核能够连续读写本地 K/V 的原因。

4、参数含义介绍

4.1 运行时参数（设备侧张量/标量）

参数名形状 / 类型含义
num_tokensint当前 batch 中实际的 token 数量。
max_num_tokensint预分配的最大 token 数（对应max_num_batched_tokens，通常为 CUDA Graph 捕获的固定批大小）。
query_start_loc_ptr[num_reqs + 1], int32每条请求在 batch 中的累加起始位置，类似 CSR 的indptr，req_idx的 token 区间为[query_start_loc[req_idx], query_start_loc[req_idx+1])。
positions_ptr[num_tokens], int64每个 token 在其所属请求中的逻辑序列位置。
block_table_ptr[max_num_reqs, max_num_blocks_per_req], int32（展平）块表，将「逻辑块号」映射为「物理块号」；按行展开存储，一行即一条请求。
block_table_strideintblock_table一行的元素个数，即max_num_blocks_per_req。
block_sizeintAttention 内核使用的逻辑块大小（一个逻辑块包含多少 token）。
slot_mapping_ptr[max_num_tokens], int32输出张量，写入每个 token 对应的物理 KV Cache slot id。

4.2 编译期常量参数（tl.constexpr）

参数名含义
KV_CACHE_BLOCK_SIZEKV Cache 物理分配块大小（物理 block 的大小）。
BLOCKS_PER_KV_BLOCK一个物理 KV 块包含多少个逻辑 block，满足KV_CACHE_BLOCK_SIZE = BLOCKS_PER_KV_BLOCK * block_size。hybrid 块模式下大于 1。
TOTAL_CP_WORLD_SIZEContext Parallel 通信域的总 rank 数；为1时走非 CP 路径。
TOTAL_CP_RANK当前设备在 CP 通信域中的 rank。
CP_KV_CACHE_INTERLEAVE_SIZECP 模式下 KV Cache 的交错大小，用于将虚拟块按 interleave 粒度切分到各 rank。
PAD_ID无效槽位的填充值（通常为PAD_SLOT_ID = -1，定义于vllm.v1.attention.backends.utils），用于 padding 区域与非本 rank 槽位。
TILE_BLOCK_SIZETriton 循环中的 tile 大小，控制单次迭代处理的 token 数；生产路径为 1024。
BLOCK_TABLE_WINDOW_SIZE一次性加载到寄存器的 block_table 窗口大小（需 ≥TILE_BLOCK_SIZE / block_size + 1，并取 2 的幂以便 Triton 向量化）。
IS_CIRCULAR是否为环形缓冲（circular）KV Cache 组；为真时逻辑块号恒为 0，负位置写入PAD_ID。

其中BLOCK_TABLE_WINDOW_SIZE的计算由工具函数_next_power_of_2（vllm_ascend/ops/triton/compute_slot_mapping.py）完成：

def _next_power_of_2(value: int) -> int: return 1 << (value - 1).bit_length()

+1的原因在于：一个 tile 可能从靠近块边界的位置开始，窗口需要多容纳一个块才能覆盖全部TILE_BLOCK_SIZE个 token 所需的块号。

4.3 启动网格

grid = (num_reqs + 1,)

前num_reqs个 program 分别处理一条请求，最后一个 program 负责把输出尾部 padding 为PAD_ID。这与 v2 内核的网格(num_kv_cache_groups, num_reqs + 1)在「请求轴多一个 padding program」的约定上保持一致。

5、算子使用示例

下面给出可直接运行的完整示例（对应 vllm_ascend/ops/triton/compute_slot_mapping.md 中的演示代码，补充了block_table_stride实参等必要细节）。假设：

  • block_size = 128，KV_CACHE_BLOCK_SIZE = 128，未启用 CP（TOTAL_CP_WORLD_SIZE = 1）；
  • 当前 batch 有 2 条请求，token 区间分别为[0, 5)与[5, 10)；
  • max_num_batched_tokens = 8192，需要把多余位置填充为PAD_ID。
import torch from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm_ascend.ops.triton.compute_slot_mapping import ( _compute_slot_mapping_kernel, _next_power_of_2, ) device = "npu" # ---- 1. 输入张量 ---- num_reqs = 2 num_tokens = 10 max_num_batched_tokens = 8192 block_size = 128 # 每条请求的 token 起止位置（CSR 风格 indptr） query_start_loc = torch.tensor([0, 5, 10], dtype=torch.int32, device=device) # 每个 token 在其请求中的逻辑位置 positions = torch.tensor( [0, 1, 2, 3, 4, # 请求 0 0, 1, 2, 3, 4], # 请求 1 dtype=torch.int64, device=device, ) # 块表：逻辑块号 -> 物理块号，形状 [max_num_reqs, max_num_blocks_per_req] max_num_reqs = 64 max_num_blocks_per_req = 320 block_table = torch.randint( 0, 320, (max_num_reqs, max_num_blocks_per_req), dtype=torch.int32, device=device ) # ---- 2. 输出张量 ---- slot_mapping = torch.zeros(max_num_batched_tokens, dtype=torch.int32, device=device) # ---- 3. 编译期常量 ---- KV_CACHE_BLOCK_SIZE = 128 BLOCKS_PER_KV_BLOCK = KV_CACHE_BLOCK_SIZE // block_size # = 1 TILE_BLOCK_SIZE = 1024 BLOCK_TABLE_WINDOW_SIZE = _next_power_of_2( cdiv(TILE_BLOCK_SIZE, block_size) + 1 ) # = _next_power_of_2(9) = 16 # ---- 4. 启动 kernel ---- grid = (num_reqs + 1,) _compute_slot_mapping_kernelgrid, # = max_num_blocks_per_req block_size, slot_mapping, KV_CACHE_BLOCK_SIZE=KV_CACHE_BLOCK_SIZE, BLOCKS_PER_KV_BLOCK=BLOCKS_PER_KV_BLOCK, TOTAL_CP_WORLD_SIZE=1, TOTAL_CP_RANK=0, CP_KV_CACHE_INTERLEAVE_SIZE=1, PAD_ID=PAD_SLOT_ID, # 通常为 -1 TILE_BLOCK_SIZE=TILE_BLOCK_SIZE, BLOCK_TABLE_WINDOW_SIZE=BLOCK_TABLE_WINDOW_SIZE, ) # slot_mapping[:num_tokens] 为各 token 的物理 slot id # slot_mapping[num_tokens:max_num_batched_tokens] 被填充为 PAD_ID print(slot_mapping[:num_tokens])

5.1 输出解读与验证

  • slot_mapping[0:5]对应请求 0 的 5 个 token，每个 token 的 slot id 等于block_table[0][pos // 128] * 128 + pos % 128；
  • slot_mapping[5:10]对应请求 1，物理块号取自块表的第 1 行；
  • slot_mapping[10:8192]全部为-1（PAD_SLOT_ID），这部分在 CUDA Graph 捕获场景下保证输出形状固定。

仓库测试 tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_compute_slot_mapping_v1.py 用同样的参考实现（逐 token 循环计算期望值）对多组(query_start_loc, positions, block_size, kv_cache_block_size)组合做了torch.testing.assert_close级别的对拍验证，其中还覆盖了 4 条请求、大序列的场景（tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_compute_slot_mapping_v1.py），可作为自行验证的参照。

5.2 启用 Context Parallel 的注意事项

当TOTAL_CP_WORLD_SIZE > 1时：

  • 物理块大小被虚拟化为virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE；
  • 仅is_local == True（属于本TOTAL_CP_RANK）的 token 才会写入真实 slot id，其余写PAD_ID；
  • block_indices与slot_offsets通过local_block_offsets重新换算，需要保证BLOCKS_PER_KV_BLOCK与CP_KV_CACHE_INTERLEAVE_SIZE配置正确——前者决定一个物理块被拆成几个逻辑块，后者决定交错切分粒度，两者共同决定本地块表索引的正确性。

从 vllm-ascend 的调度侧看（vllm_ascend/worker/block_table.py），每个 DCP rank 只保存max_model_len / dcp_world_size个 token 的 KV，因此max_num_blocks_per_req的计算也必须把dcp_world_size乘进block_size的分母中；这一点与内核的虚拟块换算互为印证。

6、与 v2 内核及 fused 变体的关系

compute_slot_mapping.py并非只有一个内核，它还包含面向多 KV Cache 组的变体，理解它们有助于把握该文件的全貌：

  • _compute_slot_mapping_fused_groups_kernel/_compute_slot_mapping_fused_groups_adaptive_kernel（vllm_ascend/ops/triton/compute_slot_mapping.py）：将多个 KV cache 组的 slot mapping 计算融合为一次启动，网格为(group_count, num_reqs * parallel_tiles + 1)；adaptive 变体还会根据请求长度动态选择小 tile（短请求）或大 tile（1024，长请求）。由compute_slot_mapping_fused_groups统一调度，生产入口为 vllm_ascend/worker/block_table.py，该路径要求所有组dcp_world_size == 1。
  • _pad_replayed_slots_kernel/pad_replayed_slot_mapping（vllm_ascend/ops/triton/compute_slot_mapping.py）：在 prefix cache 命中（replay）场景下，把可缓存组中「被重放」位置的 slot 覆盖为PAD_ID，避免共享块被重复写入不同值（docstring 中明确说明：写入本身不会错，但会把与共享者读取不同的值落进块里）。
  • v2 内核_compute_slot_mappings_kernel（vllm_ascend/ops/triton/v2/block_table/compute_slot_mappings.py）：服务 vLLM 主线的AscendBlockTables，支持多组、CP 交错以及slot_mapping_enabled（环形/状态类组直接输出PAD_ID），详细设计见 vllm_ascend/ops/triton/docs/v2/block_table/compute_slot_mappings.md。

6.1 约束与边界条件

综合内核实现、v2 文档与测试，使用该算子族时需注意以下约束：

  • positions转为 int32 后必须能容纳全部逻辑位置（超出会溢出）；
  • query_start_loc必须单调不减，且末尾值（总 token 数）不能超过slot_mapping输出的容量；
  • 有效位置的逻辑块号必须落在对应请求块表行内；
  • 空请求（start_idx == end_idx）的循环体执行 0 次，不会用哨兵值构成块表地址（vllm_ascend/ops/triton/docs/v2/block_table/compute_slot_mappings.md）；
  • 窗口加载带行 stride 掩码，保证永远不会越出分配的块表行边界。

总结

_compute_slot_mapping_kernel是 vllm-ascend 推理链路中承上启下的关键算子：它把调度器给出的逻辑 token 位置翻译成 Attention/写缓存算子可直接消费的物理 slot id，并通过有界窗口 gather、CP 虚拟块换算、tile 化 padding 等设计同时兼顾了性能、显存带宽与 CUDA Graph 兼容性。无论是排查 slot mapping 异常、适配新的 KV Cache 布局，还是优化 CP 场景下的显存访问，理解本文的映射公式与参数语义都是第一步。

  • 人工智能
  • 大模型
  • 模型推理服务
  • Ascend
  • CANN

【免费下载链接】vllm-ascend

Community maintained hardware plugin for vLLM on Huawei Ascend

项目地址：https://gitcode.com/gh_mirrors/vl/vllm-ascend
点击查看免费下载

相关推荐

上一篇：Sortable常见问题诊断：拖拽不生效解决方案
下一篇：Git-FTP 监控与日志分析终极指南：如何实时跟踪部署状态和快速排查问题

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

版权声明: 本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若内容造成侵权/违法违规/事实不符，请联系邮箱：809451989@qq.com进行投诉反馈，一经查实，立即删除！
网站建设 2026/10/14 2:07:43

Windows搭建Nginx流媒体服务器：RTMP与HLS直播配置实战

/* MD / 富文本中的 .toc（含博客园搬家等嵌套结构）；.toc-box 在侧栏，不受影响 */#content_views .toc,/* 编辑器常在目录前后插入空 p（:empty 仍占 20px），一并去掉避免顶空隙 */#content_views.markdown_views > p:empty:has(+ .toc),#content_views.markdown_views …

作者头像 李华
网站建设 2026/10/14 2:07:13

CDH 5.15.1离线部署完整指南：从介质制作到集群稳定运行

/* MD / 富文本中的 .toc（含博客园搬家等嵌套结构）；.toc-box 在侧栏，不受影响 */#content_views .toc,/* 编辑器常在目录前后插入空 p（:empty 仍占 20px），一并去掉避免顶空隙 */#content_views.markdown_views > p:empty:has(+ .toc),#content_views.markdown_views …

作者头像 李华