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

人工智能大模型模型推理服务AscendCANN【免费下载链接】vllm-ascendCommunity 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 ParallelCP交错布局下的虚拟块换算、全部运行时与编译期参数含义以及可直接运行的调用示例。读完本文你将能独立理解 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 节最后一个 programreq_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.pypadding 分支的实现在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_locCSR 风格 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_kernelvllm_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.pyblock_indices pos // block_size # 逻辑块号 slot_offsets pos - block_indices * block_size # 块内偏移随后通过块表把逻辑块号翻译为物理块号并组装 slot idblock_numbers tl.gather(block_table_window, relative_block_indices, 0) slot_ids block_numbers * block_size slot_offsets2.1 有界窗口 gather避免整行加载值得注意的实现细节是内核并没有一次性加载整行block_table而是采用「有界窗口 gather」策略见 vllm_ascend/ops/triton/compute_slot_mapping.py对每个TILE_BLOCK_SIZE大小的 token tile先求出该 tile 内最小的有效逻辑块号block_idx_base利用INT32_MAX哨兵值配合tl.min只从块表行中加载从block_idx_base开始的BLOCK_TABLE_WINDOW_SIZE个连续条目用相对块号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 CacheIS_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 ParallelTOTAL_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 是否属于本 rankis_local (virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE) % TOTAL_CP_WORLD_SIZE TOTAL_CP_RANKis_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 块被切分成多个逻辑 blockhybrid 块模式时虚拟块内也包含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 0128..255属于 rank 1256..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 的indptrreq_idx的 token 区间为[query_start_loc[req_idx], query_start_loc[req_idx1])。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是否为环形缓冲circularKV Cache 组为真时逻辑块号恒为 0负位置写入PAD_ID。其中BLOCK_TABLE_WINDOW_SIZE的计算由工具函数_next_power_of_2vllm_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 128KV_CACHE_BLOCK_SIZE 128未启用 CPTOTAL_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], dtypetorch.int32, devicedevice) # 每个 token 在其请求中的逻辑位置 positions torch.tensor( [0, 1, 2, 3, 4, # 请求 0 0, 1, 2, 3, 4], # 请求 1 dtypetorch.int64, devicedevice, ) # 块表逻辑块号 - 物理块号形状 [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), dtypetorch.int32, devicedevice ) # ---- 2. 输出张量 ---- slot_mapping torch.zeros(max_num_batched_tokens, dtypetorch.int32, devicedevice) # ---- 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_SIZEKV_CACHE_BLOCK_SIZE, BLOCKS_PER_KV_BLOCKBLOCKS_PER_KV_BLOCK, TOTAL_CP_WORLD_SIZE1, TOTAL_CP_RANK0, CP_KV_CACHE_INTERLEAVE_SIZE1, PAD_IDPAD_SLOT_ID, # 通常为 -1 TILE_BLOCK_SIZETILE_BLOCK_SIZE, BLOCK_TABLE_WINDOW_SIZEBLOCK_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 % 128slot_mapping[5:10]对应请求 1物理块号取自块表的第 1 行slot_mapping[10:8192]全部为-1PAD_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_IDblock_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_kernelvllm_ascend/ops/triton/compute_slot_mapping.py将多个 KV cache 组的 slot mapping 计算融合为一次启动网格为(group_count, num_reqs * parallel_tiles 1)adaptive 变体还会根据请求长度动态选择小 tile短请求或大 tile1024长请求。由compute_slot_mapping_fused_groups统一调度生产入口为 vllm_ascend/worker/block_table.py该路径要求所有组dcp_world_size 1。_pad_replayed_slots_kernel/pad_replayed_slot_mappingvllm_ascend/ops/triton/compute_slot_mapping.py在 prefix cache 命中replay场景下把可缓存组中「被重放」位置的 slot 覆盖为PAD_ID避免共享块被重复写入不同值docstring 中明确说明写入本身不会错但会把与共享者读取不同的值落进块里。v2 内核_compute_slot_mappings_kernelvllm_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 场景下的显存访问理解本文的映射公式与参数语义都是第一步。赞分享人工智能大模型模型推理服务AscendCANN【免费下载链接】vllm-ascendCommunity maintained hardware plugin for vLLM on Huawei Ascend项目地址https://gitcode.com/gh_mirrors/vl/vllm-ascend点击查看免费下载相关推荐vllm-ascend Compressor 算子深度解析KV Cache 压缩前处理与稀疏注意力加速实践vllm ascend Compressor 算子深度解析KV Cache 压缩前处理与稀疏注意力加速实践 本文基于 vllm ascend 开源仓库中 cs人工智能大模型模型推理服务AscendCANNvllm-ascend SparseAttentionScore 算子深度解析面向 Paged KV Cache 的 Top-K 稀疏注意力设计与实现vllm ascend SparseAttentionScore 算子深度解析面向 Paged KV Cache 的 Top K 稀疏注意力设计与实现 导读人工智能大模型模型推理服务AscendCANNvLLM-Ascend SparseAttnSharedkvMetadata 算子深度解析共享 KV 稀疏注意力分核调度元数据生成原理与实战vLLM Ascend SparseAttnSharedkvMetadata 算子深度解析共享 KV 稀疏注意力分核调度元数据生成原理与实战 本篇技术指南围绕人工智能大模型模型推理服务AscendCANN上一篇Sortable常见问题诊断拖拽不生效解决方案下一篇Git-FTP 监控与日志分析终极指南如何实时跟踪部署状态和快速排查问题创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考