CANN ops-transformer Quest Prefill Metadata 算子实验指南KV Cache 预填充阶段稀疏注意力元数据构建与性能验证【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读本篇文章围绕 CANN ops-transformer 仓库中experimental/select_attention_operators子项目的第二个实验Experiment 2 — Quest Prefill Metadata展开深入讲解如何在 Ascend 910B 系列 NPU 上通过quest_prefill_metadataAscendC 算子从已完成预填充prefill的 KV Cache 中构建 QUEST 稀疏注意力所需的 maxblock / minblock 元数据。读完本文你将掌握该算子的输入输出约定、构建与运行方式、功能正确性测试方法、基准测试方法论与实测性能数据以及底层 AscendC 内核的实现原理。实验背景为什么需要 Prefill 阶段的元数据在 LLM 解码阶段QUESTQuery-Aware Sparsity类稀疏注意力方法通过为每个 KV Cache 块维护一个低成本的元数据向量仅根据当前解码 query 与元数据之间的近似匹配选出最重要的 top-k 个 KV 块参与注意力计算从而显著降低解码阶段的访存开销。本实验中的quest_prefill_metadata算子承担元数据初始化这一环节在 KV Cache 已经完成 prefill 之后、正式进入逐 token 解码之前一次性扫描全部 K Cache 块为每个 KV 块计算两个 D 维元数据向量maxblock对块内 BLOCK_SIZE 个 token 沿 token 维做逐通道per-channel最大值归约minblock同样的归约操作取最小值。每 BLOCK_SIZE128个 KV 块产生的 128 个元数据向量被打包成一个元数据块metadata block写入maxblocks/minblocks的对应位置。该元数据随后被同目录下的quest_block_select_paged系列算子消费用于在解码时预测稀疏掩码。关于算子间的整体协作关系可参见 select_attention_operators 总览 README 与 select_attn_ops 算子库 README。实验定位与场景参数本实验的目的是同时验证quest_prefill_metadata内核的功能正确性与 Python 参考实现逐元素对比与性能延迟与有效带宽。实验默认场景参数如下参数取值说明BLOCK_SIZE128固定Paged attention 中常用的每块 token 数同时是元数据块与 KV Cache 块的粒度D128固定Head 维度LLM 常用值N8固定KV head 数量与 Qwen3-8B 的配置一致B扫描大 batch批量大小作为扫参维度MKBPR扫参每个请求最大 KV Cache 块数决定序列长度MMBPR由 MKBPR 推导MMBPR cdiv(MKBPR, BLOCK_SIZE)即每个请求最大元数据块数其中cdiv(x, y) (x y - 1) // y。需要特别注意的是元数据块的覆盖粒度是BLOCK_SIZE * BLOCK_SIZE 16384个 token一个元数据块中恰好收纳 128 个 KV 块的元数据向量这构成了 QUEST 两级分块的层级结构详见下文源码级实现原理。环境准备与算子构建实验位于experimental/select_attention_operators子项目内运行前需要完成 CANN 环境初始化与算子 Python 包构建# 1. 创建 conda 环境项目验证过的版本组合见 requirements.txt conda create -n sa python3.11.10 -y conda activate sa pip install -r experimental/select_attention_operators/requirements.txt # 2. 激活 CANN 环境并指定卡型示例为 Ascend910B4可按实际卡型修改 source experimental/select_attention_operators/scripts/init_cann.sh Ascend910B4 # 3. 编译全部 AscendC 算子并构建 Python 接口 bash experimental/select_attention_operators/scripts/build_kernels.sh构建过程通过 compile.sh 调用bisheng编译器以--cce-soc-version$SOC_VERSION --cce-soc-core-typeVecCore将quest_prefill_metadata.cpp编译为共享库并配合torch_interface.cpp中基于PYBIND11_MODULE的绑定向 Python 暴露select_attn_ops.quest_prefill_metadata接口。构建完成后可先用一个简单的 Python 脚本验证算子是否已正确加载并返回结果import torch import torch_npu from select_attn_ops import quest_prefill_metadata # 构造最小输入B1, N8, BLOCK_SIZE128, D128 device torch.device(npu:0) B, N, BLOCK_SIZE, D 1, 8, 128, 128 MKBPR 128 # 每请求 KV 块数 MMBPR 1 # 每请求元数据块数 num_kv_blocks B * MKBPR num_meta_blocks B * MMBPR k_cache torch.randn(num_kv_blocks, BLOCK_SIZE, N, D, dtypetorch.float16, devicedevice) block_tables torch.randperm(num_kv_blocks, devicedevice).reshape(B, MKBPR).to(torch.int32) seq_lens torch.tensor([MKBPR * BLOCK_SIZE] * B, dtypetorch.int32, devicedevice) metadata_block_tables torch.zeros(B, MMBPR, dtypetorch.int32, devicedevice) maxblocks torch.zeros(num_meta_blocks, BLOCK_SIZE, N, D, dtypetorch.float16, devicedevice) minblocks torch.zeros(num_meta_blocks, BLOCK_SIZE, N, D, dtypetorch.float16, devicedevice) quest_prefill_metadata(k_cache, block_tables, seq_lens, metadata_block_tables, maxblocks, minblocks) print(maxblocks.shape, minblocks.shape) # 输出会写入 maxblocks / minblocks接口层在 torch_interface.cpp 中会对输入做严格校验D与BLOCK_SIZE必须为 128B、N、BLOCK_SIZE、D各维度在k_cache、maxblocks、minblocks之间必须一致否则抛出TORCH_CHECK错误。功能正确性验证pytest 测试套件实验目录下的test_quest_prefill_metadata.py将参考实现对比封装为可参数化的 pytest 测试。运行方式与对应覆盖范围如下pytest -k basic -v # 仅运行 basic 扫参组 pytest . # 运行全部测试从 test_quest_prefill_metadata.py 的源码结构看测试套件按四个维度组织Test 1 – Basic functionalitybatch_size ∈ {1, 2}、num_kv_heads ∈ {4, 8}、mkbpr ∈ {1, 64, 126, 128}dtype 覆盖 fp16 与 bf16Test 2 – Edge casesKV head 数量覆盖到非对称值{1, 2, 4, 7, 8, 9, 16, 21, 32, 33}mkbpr覆盖块边界附近的{63, 64, 65, 126, 127, 128, 129, 150, 255, 256, 257}用于验证对齐边界与尾块处理Test 3 – Large sequencebatch_size ∈ {1, 2, 4, 8}mkbpr扩展到512验证超长序列下多层元数据块多meta_blk的正确性Test 4 – Large batchbatch_size ∈ {16, 20, 24, 32}逼近真实服务场景的高并发批处理。每个参数组合通过内部工作函数test_prefill_kernel执行以下步骤用 gen_data_quest_prefill_metadata.py 中的gen_quest_prefill_inputs生成伪随机输入固定 SEED42保证失败可复现用 Python 参考实现ref_quest_prefill_metadata计算期望的 maxblocks / minblocks调用 AscendC 内核quest_prefill_metadata计算实际输出通过compare_tensors以宽松的 fp16 容差rtol1e-2, atol1e-3断言两者逐元素一致任何不匹配即判测试失败。其中same_seq_len_all_reqsssar参数控制批内各请求序列长度是否一致为True时所有请求取最大序列长度为False时每个请求随机取[0, MKBPR*BLOCK_SIZE]区间内的长度从而同时覆盖变长序列 请求内 KV 块不满的稀疏页表场景。单场景深入对比dump 全部输出张量如果需要针对某个具体输入做深入调试可以跳过 pytest 参数化直接以脚本方式运行测试文件python test_quest_prefill_metadata.py脚本的__main__入口会执行一次手动的冒烟测试batch_size20, num_kv_heads8, block_size128, head_dim128, mkbpr128, ssarFalse并将参考实现与内核输出的 maxblocks / minblocks 全部打印出来逐项对比输出 PASSED 或 FAILED 及详细差异信息。这适合定位某个特定形状下的数值偏差或调试内核行为。基准测试延迟与有效带宽运行方式与测什么python benchmark_quest_prefill_metadata.pybenchmark_quest_prefill_metadata.py 会针对N8、B ∈ {10, 20, 24, 32}、MKBPR ∈ {63, 80, 94, 128, 150, 200, 256}的笛卡尔积组合默认SAME_SEQ_LEN_ALL_REQSTrue依次测量延迟μs通过torch.npu.Event计时每个配置执行 1 次 warmup 6 次重复取平均有效带宽TB/s基于bytes_moved_prefill计算的全局内存读写字节总量K Cache 读取 block 表读取 maxblocks/minblocks 写入除以耗时正确性每个配置点同时运行参考实现与内核输出Outputs_equal列yes/no作为可信度标记。默认DTYPE在脚本中为torch.bfloat16文档中给出的 fp16 数据通过将DTYPE修改为torch.float16复现。测试在x86 宿主机 910B4 卡上完成。fp16 实测结果文档给出的 fp16 基准数据如下BLOCK_SIZE128, D128, SAME_SEQ_LEN_ALL_REQSTrue其中Ref_*列为 Python 参考实现Our_*列为 AscendC 内核 DTYPEtorch.float16 BLOCK_SIZE128 D128 SAME_SEQ_LEN_ALL_REQSTrue N B Seq_len Outputs_equal Ref_Latency_[usec] Our_Latency_[usec] Ref_BW_[TB/sec] Our_BW_[TB/sec] ---------------------------------------------------------------------------------------------------------- 8 10 8064 yes 93027.76 309.27 0.002 0.551 8 10 10240 yes 117024.04 384.37 0.002 0.559 8 10 12032 yes 136078.70 484.75 0.002 0.519 8 10 16384 yes 182469.81 602.44 0.002 0.566 8 10 19200 yes 219475.81 718.18 0.002 0.562 8 10 25600 yes 329046.24 948.14 0.002 0.564 8 10 32768 yes 417126.71 1321.51 0.002 0.516 8 20 8064 yes 208446.04 664.72 0.002 0.513 8 20 10240 yes 270382.06 835.40 0.002 0.515 8 20 12032 yes 314411.11 889.26 0.002 0.566 8 20 16384 yes 446804.24 1325.83 0.002 0.514 8 20 19200 yes 492977.95 1572.19 0.002 0.514 8 20 25600 yes 671196.25 1890.16 0.002 0.566 8 20 32768 yes 835718.99 2653.32 0.002 0.514 8 24 8064 yes 258454.59 756.71 0.002 0.540 8 24 10240 yes 329535.64 943.38 0.002 0.547 8 24 12032 yes 376479.13 1213.10 0.002 0.498 8 24 16384 yes 518973.06 1635.08 0.002 0.500 8 24 19200 yes 609663.98 1785.78 0.002 0.543 8 24 25600 yes 807009.20 2584.06 0.002 0.497 8 24 32768 yes 1023406.74 2992.35 0.002 0.547 8 32 8064 yes 349012.08 1079.05 0.002 0.505 8 32 10240 yes 407759.97 1233.08 0.002 0.558 8 32 12032 yes 515659.63 1438.41 0.002 0.560 8 32 16384 yes 681750.24 2144.14 0.002 0.509 8 32 19200 yes 809321.04 2323.78 0.002 0.556 8 32 25600 yes 1036969.16 3398.52 0.002 0.504 8 32 32768 yes 1353555.26 4333.40 0.002 0.503 bf16 实测结果bfloat16 下文档给出的基准数据如下 DTYPEtorch.bfloat16 BLOCK_SIZE128 D128 SAME_SEQ_LEN_ALL_REQSTrue N B Seq_len Outputs_equal Ref_Latency_[usec] Our_Latency_[usec] Ref_BW_[TB/sec] Our_BW_[TB/sec] ---------------------------------------------------------------------------------------------------------- 8 10 8064 yes 91797.47 341.11 0.002 0.500 8 10 10240 yes 119940.88 416.60 0.002 0.516 8 10 12032 yes 142493.83 444.33 0.002 0.566 8 10 16384 yes 190612.51 659.39 0.002 0.517 8 10 19200 yes 228080.26 788.18 0.002 0.512 8 10 25600 yes 331777.49 1042.17 0.002 0.513 8 10 32768 yes 422492.27 1208.98 0.002 0.564 8 20 8064 yes 213834.41 610.95 0.002 0.558 8 20 10240 yes 268580.89 767.21 0.002 0.560 8 20 12032 yes 320251.59 980.86 0.002 0.513 8 20 16384 yes 425229.09 1209.94 0.002 0.563 8 20 19200 yes 501996.54 1439.80 0.002 0.561 8 20 25600 yes 677030.97 2083.51 0.002 0.513 8 20 32768 yes 853946.94 2409.99 0.002 0.566 8 24 8064 yes 259979.37 827.37 0.002 0.494 8 24 10240 yes 311260.29 1038.10 0.002 0.497 8 24 12032 yes 370500.20 1109.80 0.002 0.544 8 24 16384 yes 523368.08 1489.21 0.002 0.549 8 24 19200 yes 616558.68 1952.58 0.002 0.496 8 24 25600 yes 808925.46 2347.28 0.002 0.547 8 24 32768 yes 1033353.92 3306.62 0.002 0.495 8 32 8064 yes 354466.72 989.11 0.002 0.551 8 32 10240 yes 443984.86 1354.58 0.002 0.508 8 32 12032 yes 523036.50 1585.94 0.002 0.508 8 32 16384 yes 704769.45 1947.90 0.002 0.560 8 32 19200 yes 805073.97 2558.09 0.002 0.505 8 32 25600 yes 1086389.57 3075.26 0.002 0.556 8 32 32768 yes 1375007.49 3912.25 0.002 0.557 性能结论fp16 下内核有效带宽达到至少 0.55 TB/sec而 910B4 的标称全局内存带宽为0.80 TB/sec即约69% 的带宽利用率对比 Python 参考实现仅约 0.002 TB/sec内核加速效果显著例如B10, Seq_len8064时延迟从约 93 毫秒降至约 0.31 毫秒。bf16 下有相近的性能水平约 0.49~0.57 TB/sec。从表格看各序列长度下Outputs_equal均为 yes说明在全部基准配置点上内核与参考实现输出一致。需要说明的是上述数据来源于该实验 README 在 x86 宿主 910B4 环境下的记录实际数字会随 CANN 版本、固件、batch 组成及SAME_SEQ_LEN_ALL_REQS取值变化建议在目标环境上自行复测。源码级实现原理参考实现算法的语义基准ref_quest_prefill_metadata.py 以最直白的方式定义了算子语义对每个请求r、每个 KV headn由seq_lens[r]计算实际 KV 块数⌈seq_lens[r]/BLOCK_SIZE⌉与元数据块数⌈num_kv_blocks_in_request/BLOCK_SIZE⌉逐元数据块、逐 KV 块取出 K Cache 切片k_cache[kv_block_id, :, n, :]沿 token 维dim 0分别做 reduce-max 与 reduce-min得到maxblock[meta_blk_id, blk, n, :]与minblock[meta_blk_id, blk, n, :]对每个元数据块中未使用的行尾块填充置零保证不存在的 KV 块对应的元数据为 0。其中_calculate_tokens_to_reduce处理尾块只有当当前 KV 块是最后一个元数据块中的最后一个 KV 块时才只对seq_len - 已归约 token 数个有效 token 做归约其余情况均按整块BLOCK_SIZE处理。这与内核中的 tail check 逻辑完全对应。数据生成器gen_data_quest_prefill_metadata.py 中的gen_quest_prefill_inputs以固定随机种子SEED42生成k_cache(num_kv_blocks, BLOCK_SIZE, N, D)torch.randn * 1.5放大数值范围以充分考验 max/min 归约seq_lensssarTrue时全部为MKBPR*BLOCK_SIZE否则每个请求在[0, max_seq_len]内随机block_tables与metadata_block_tables通过torch.randperm生成随机置换的块索引表int32模拟 vLLM 式 PagedAttention 的非连续块布局maxblocks/minblocks先以torch.empty分配输出占位再统一contiguous()化。这种随机置换块表 随机/等长序列长度的生成方式可以覆盖连续与非连续物理块、整块与尾块、满块与空块等真实 PagedAttention 场景下的边界条件。AscendC 内核向量核上的两级分块归约quest_prefill_metadata.cpp 实现了向量核VecCore内核核心设计可以总结为以下几点1. 并行模型1 个核处理 1 个 (batch, head) 组合num_batch_heads B * N按GetBlockIdx()在 AI Core 间轮转分配实现 batch 维与 KV head 维的天然并行每个核内循环处理该 (r, h) 对应的全部 KV 块因此每个 KV 块只会被读取一次。2. 两级分块结构外层循环遍历meta_blk每个元数据块覆盖BLOCK_SIZE * BLOCK_SIZE 16384个 token内层循环遍历该元数据块收纳的BLOCK_SIZE个 KV 块对每个 KV 块先通过DataCopy只拷贝k_cache[kv_block_id, :, h, :]这一 head 切片利用srcStride跳过其他 head再做归约。3. 对数归约ReduceTokenDim对(BLOCK_SIZE, D)的 UB 局部张量采用逐轮两两合并的对数归约Max/Min(vec[0], vec[0], vec[len/2])奇数向量时先搬运到前半段再继续直到len DisMaxtrue走MaxisMaxfalse走Min从而用同一套循环同时产出 maxblock 与 minblock归约前先将 K 块复制到工作缓冲保证k_block_lt原值不被破坏避免重复加载 GM。4. 缓冲与流水设计输入队列k_block_in_q_使用双缓冲DOUBLEBUFFER2配合MTE2_V事件同步实现 GM→UB 拷贝与向量计算的流水重叠max_out_q_/min_out_q_在归约前用 ±65504fp16 极值初始化尾块行用 0 填充最终通过带dstStride跳过其他 head的DataCopy将BLOCK_SIZE*D的元数据行写回maxblocks/minblocks的指定meta_blk_id位置。5. 接口与约束内核启动参数为B, N, BLOCK_SIZE, D, MKBPR, MMBPR由 torch_interface.cpp 从张量 shape 自动推导并做TORCH_CHECK校验D 128、BLOCK_SIZE 128等启动时blockDims min(B * N, NUM_CORES)即任务数超过物理核数时自动截断。内核源码注释还特别说明了一个值得注意的工程事实bf16 输入同样能通过测试尽管内核内部的 min/max 处理实际发生在 fp16half路径上这一点被记录在算子 README 的 TODO 项中Dobule check why bf16 passes even though the whole min/max processing happens in fp16从源码结构看这得益于 bf16 与 fp16 同为 16 位宽、且归约结果在测试容差内一致。内存占用视角为什么这是带宽受限算子从bytes_moved_prefill的计算逻辑可以看出该算子的全局内存流量由三部分组成K Cache 读取num_effective_kv_blocks * BLOCK_SIZE * N * D * 2字节、块表读取每有效 KV 块 4 字节、以及 maxblocks/minblocks 双份写入。整个内核几乎没有计算密集操作只有逐元素 Max/Min因此性能瓶颈在全局内存带宽这也解释了实测 69% 带宽利用率的达成路径单次读取 双缓冲流水 对数归约把向量单元的空闲时间压到了最低。在解码流水线中的位置与后续实验quest_prefill_metadata只是 QUEST 稀疏注意力流程的第一步。按照 select_attention_operators 总览 README 中记录的当前最佳实践已在 vllm-ascend 中验证使用quest_prefill_metadata()在prefill 后创建元数据并在每解码 128 个 token即 BLOCK_SIZE 的整数倍时更新元数据使用quest_block_select_paged_in_out_w()在解码阶段根据当前 token 的 query 向量预测重要的 KV 块索引。同目录下的后续实验3_quest_block_select_paged、4_quest_block_select_paged_w分别实现了基于元数据的稀疏掩码预测器以及额外融合 sink window 特征的版本——它们都以本实验产出的 maxblocks / minblocks 作为输入因此本实验的正确性验证是整条稀疏注意力链路的基石。本实验目录中的 conftest.py 会把每个实验子目录注入sys.path使得各实验可以共享gen_data_*、ref_*等辅助模块在子项目根目录执行pytest -v experiments即可一次性运行全部实验的测试。总结quest_prefill_metadata实验给出了一个完整的参考实现 → 数据生成 → 功能测试 → 基准测试的算子验证范式功能上通过四级 pytest 参数化基础、边界、长序列、大 batch与随机页表/变长序列生成全面覆盖了 PagedAttention 场景下可能出现的对齐与尾块问题性能上在 910B4 上以约 0.5~0.57 TB/sec 的有效带宽、69% 的标称带宽利用率将原本毫秒级的元数据构建压缩到亚毫秒级为解码阶段每秒数十次元数据刷新提供了可行基础实现上向量核 (batch, head) 并行 双缓冲流水 对数归约的设计是一份可以直接复用的 AscendC 带宽密集型算子模板。对于希望在 Ascend NPU 上落地 QUEST 类稀疏注意力、或需要理解 PagedAttention 元数据管理细节的开发者建议按本文的步骤先运行测试与基准再结合源码逐段对照理解两级分块与尾块处理逻辑。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/18 22:32:53 · 阅读 1.2k CANN ops-transformer Quest Prefill Metadata 算子实验指南：KV Cache 预填充阶段稀疏注意力元数据构建与性能验证
CANN ops-transformer Quest Prefill Metadata 算子实验指南KV Cache 预填充阶段稀疏注意力元数据构建与性能验证【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读本篇文章围绕 CANN ops-transformer 仓库中experimental/select_attention_operators子项目的第二个实验Experiment 2 — Quest Prefill Metadata展开深入讲解如何在 Ascend 910B 系列 NPU 上通过quest_prefill_metadataAscendC 算子从已完成预填充prefill的 KV Cache 中构建 QUEST 稀疏注意力所需的 maxblock / minblock 元数据。读完本文你将掌握该算子的输入输出约定、构建与运行方式、功能正确性测试方法、基准测试方法论与实测性能数据以及底层 AscendC 内核的实现原理。实验背景为什么需要 Prefill 阶段的元数据在 LLM 解码阶段QUESTQuery-Aware Sparsity类稀疏注意力方法通过为每个 KV Cache 块维护一个低成本的元数据向量仅根据当前解码 query 与元数据之间的近似匹配选出最重要的 top-k 个 KV 块参与注意力计算从而显著降低解码阶段的访存开销。本实验中的quest_prefill_metadata算子承担元数据初始化这一环节在 KV Cache 已经完成 prefill 之后、正式进入逐 token 解码之前一次性扫描全部 K Cache 块为每个 KV 块计算两个 D 维元数据向量maxblock对块内 BLOCK_SIZE 个 token 沿 token 维做逐通道per-channel最大值归约minblock同样的归约操作取最小值。每 BLOCK_SIZE128个 KV 块产生的 128 个元数据向量被打包成一个元数据块metadata block写入maxblocks/minblocks的对应位置。该元数据随后被同目录下的quest_block_select_paged系列算子消费用于在解码时预测稀疏掩码。关于算子间的整体协作关系可参见 select_attention_operators 总览 README 与 select_attn_ops 算子库 README。实验定位与场景参数本实验的目的是同时验证quest_prefill_metadata内核的功能正确性与 Python 参考实现逐元素对比与性能延迟与有效带宽。实验默认场景参数如下参数取值说明BLOCK_SIZE128固定Paged attention 中常用的每块 token 数同时是元数据块与 KV Cache 块的粒度D128固定Head 维度LLM 常用值N8固定KV head 数量与 Qwen3-8B 的配置一致B扫描大 batch批量大小作为扫参维度MKBPR扫参每个请求最大 KV Cache 块数决定序列长度MMBPR由 MKBPR 推导MMBPR cdiv(MKBPR, BLOCK_SIZE)即每个请求最大元数据块数其中cdiv(x, y) (x y - 1) // y。需要特别注意的是元数据块的覆盖粒度是BLOCK_SIZE * BLOCK_SIZE 16384个 token一个元数据块中恰好收纳 128 个 KV 块的元数据向量这构成了 QUEST 两级分块的层级结构详见下文源码级实现原理。环境准备与算子构建实验位于experimental/select_attention_operators子项目内运行前需要完成 CANN 环境初始化与算子 Python 包构建# 1. 创建 conda 环境项目验证过的版本组合见 requirements.txt conda create -n sa python3.11.10 -y conda activate sa pip install -r experimental/select_attention_operators/requirements.txt # 2. 激活 CANN 环境并指定卡型示例为 Ascend910B4可按实际卡型修改 source experimental/select_attention_operators/scripts/init_cann.sh Ascend910B4 # 3. 编译全部 AscendC 算子并构建 Python 接口 bash experimental/select_attention_operators/scripts/build_kernels.sh构建过程通过 compile.sh 调用bisheng编译器以--cce-soc-version$SOC_VERSION --cce-soc-core-typeVecCore将quest_prefill_metadata.cpp编译为共享库并配合torch_interface.cpp中基于PYBIND11_MODULE的绑定向 Python 暴露select_attn_ops.quest_prefill_metadata接口。构建完成后可先用一个简单的 Python 脚本验证算子是否已正确加载并返回结果import torch import torch_npu from select_attn_ops import quest_prefill_metadata # 构造最小输入B1, N8, BLOCK_SIZE128, D128 device torch.device(npu:0) B, N, BLOCK_SIZE, D 1, 8, 128, 128 MKBPR 128 # 每请求 KV 块数 MMBPR 1 # 每请求元数据块数 num_kv_blocks B * MKBPR num_meta_blocks B * MMBPR k_cache torch.randn(num_kv_blocks, BLOCK_SIZE, N, D, dtypetorch.float16, devicedevice) block_tables torch.randperm(num_kv_blocks, devicedevice).reshape(B, MKBPR).to(torch.int32) seq_lens torch.tensor([MKBPR * BLOCK_SIZE] * B, dtypetorch.int32, devicedevice) metadata_block_tables torch.zeros(B, MMBPR, dtypetorch.int32, devicedevice) maxblocks torch.zeros(num_meta_blocks, BLOCK_SIZE, N, D, dtypetorch.float16, devicedevice) minblocks torch.zeros(num_meta_blocks, BLOCK_SIZE, N, D, dtypetorch.float16, devicedevice) quest_prefill_metadata(k_cache, block_tables, seq_lens, metadata_block_tables, maxblocks, minblocks) print(maxblocks.shape, minblocks.shape) # 输出会写入 maxblocks / minblocks接口层在 torch_interface.cpp 中会对输入做严格校验D与BLOCK_SIZE必须为 128B、N、BLOCK_SIZE、D各维度在k_cache、maxblocks、minblocks之间必须一致否则抛出TORCH_CHECK错误。功能正确性验证pytest 测试套件实验目录下的test_quest_prefill_metadata.py将参考实现对比封装为可参数化的 pytest 测试。运行方式与对应覆盖范围如下pytest -k basic -v # 仅运行 basic 扫参组 pytest . # 运行全部测试从 test_quest_prefill_metadata.py 的源码结构看测试套件按四个维度组织Test 1 – Basic functionalitybatch_size ∈ {1, 2}、num_kv_heads ∈ {4, 8}、mkbpr ∈ {1, 64, 126, 128}dtype 覆盖 fp16 与 bf16Test 2 – Edge casesKV head 数量覆盖到非对称值{1, 2, 4, 7, 8, 9, 16, 21, 32, 33}mkbpr覆盖块边界附近的{63, 64, 65, 126, 127, 128, 129, 150, 255, 256, 257}用于验证对齐边界与尾块处理Test 3 – Large sequencebatch_size ∈ {1, 2, 4, 8}mkbpr扩展到512验证超长序列下多层元数据块多meta_blk的正确性Test 4 – Large batchbatch_size ∈ {16, 20, 24, 32}逼近真实服务场景的高并发批处理。每个参数组合通过内部工作函数test_prefill_kernel执行以下步骤用 gen_data_quest_prefill_metadata.py 中的gen_quest_prefill_inputs生成伪随机输入固定 SEED42保证失败可复现用 Python 参考实现ref_quest_prefill_metadata计算期望的 maxblocks / minblocks调用 AscendC 内核quest_prefill_metadata计算实际输出通过compare_tensors以宽松的 fp16 容差rtol1e-2, atol1e-3断言两者逐元素一致任何不匹配即判测试失败。其中same_seq_len_all_reqsssar参数控制批内各请求序列长度是否一致为True时所有请求取最大序列长度为False时每个请求随机取[0, MKBPR*BLOCK_SIZE]区间内的长度从而同时覆盖变长序列 请求内 KV 块不满的稀疏页表场景。单场景深入对比dump 全部输出张量如果需要针对某个具体输入做深入调试可以跳过 pytest 参数化直接以脚本方式运行测试文件python test_quest_prefill_metadata.py脚本的__main__入口会执行一次手动的冒烟测试batch_size20, num_kv_heads8, block_size128, head_dim128, mkbpr128, ssarFalse并将参考实现与内核输出的 maxblocks / minblocks 全部打印出来逐项对比输出 PASSED 或 FAILED 及详细差异信息。这适合定位某个特定形状下的数值偏差或调试内核行为。基准测试延迟与有效带宽运行方式与测什么python benchmark_quest_prefill_metadata.pybenchmark_quest_prefill_metadata.py 会针对N8、B ∈ {10, 20, 24, 32}、MKBPR ∈ {63, 80, 94, 128, 150, 200, 256}的笛卡尔积组合默认SAME_SEQ_LEN_ALL_REQSTrue依次测量延迟μs通过torch.npu.Event计时每个配置执行 1 次 warmup 6 次重复取平均有效带宽TB/s基于bytes_moved_prefill计算的全局内存读写字节总量K Cache 读取 block 表读取 maxblocks/minblocks 写入除以耗时正确性每个配置点同时运行参考实现与内核输出Outputs_equal列yes/no作为可信度标记。默认DTYPE在脚本中为torch.bfloat16文档中给出的 fp16 数据通过将DTYPE修改为torch.float16复现。测试在x86 宿主机 910B4 卡上完成。fp16 实测结果文档给出的 fp16 基准数据如下BLOCK_SIZE128, D128, SAME_SEQ_LEN_ALL_REQSTrue其中Ref_*列为 Python 参考实现Our_*列为 AscendC 内核 DTYPEtorch.float16 BLOCK_SIZE128 D128 SAME_SEQ_LEN_ALL_REQSTrue N B Seq_len Outputs_equal Ref_Latency_[usec] Our_Latency_[usec] Ref_BW_[TB/sec] Our_BW_[TB/sec] ---------------------------------------------------------------------------------------------------------- 8 10 8064 yes 93027.76 309.27 0.002 0.551 8 10 10240 yes 117024.04 384.37 0.002 0.559 8 10 12032 yes 136078.70 484.75 0.002 0.519 8 10 16384 yes 182469.81 602.44 0.002 0.566 8 10 19200 yes 219475.81 718.18 0.002 0.562 8 10 25600 yes 329046.24 948.14 0.002 0.564 8 10 32768 yes 417126.71 1321.51 0.002 0.516 8 20 8064 yes 208446.04 664.72 0.002 0.513 8 20 10240 yes 270382.06 835.40 0.002 0.515 8 20 12032 yes 314411.11 889.26 0.002 0.566 8 20 16384 yes 446804.24 1325.83 0.002 0.514 8 20 19200 yes 492977.95 1572.19 0.002 0.514 8 20 25600 yes 671196.25 1890.16 0.002 0.566 8 20 32768 yes 835718.99 2653.32 0.002 0.514 8 24 8064 yes 258454.59 756.71 0.002 0.540 8 24 10240 yes 329535.64 943.38 0.002 0.547 8 24 12032 yes 376479.13 1213.10 0.002 0.498 8 24 16384 yes 518973.06 1635.08 0.002 0.500 8 24 19200 yes 609663.98 1785.78 0.002 0.543 8 24 25600 yes 807009.20 2584.06 0.002 0.497 8 24 32768 yes 1023406.74 2992.35 0.002 0.547 8 32 8064 yes 349012.08 1079.05 0.002 0.505 8 32 10240 yes 407759.97 1233.08 0.002 0.558 8 32 12032 yes 515659.63 1438.41 0.002 0.560 8 32 16384 yes 681750.24 2144.14 0.002 0.509 8 32 19200 yes 809321.04 2323.78 0.002 0.556 8 32 25600 yes 1036969.16 3398.52 0.002 0.504 8 32 32768 yes 1353555.26 4333.40 0.002 0.503 bf16 实测结果bfloat16 下文档给出的基准数据如下 DTYPEtorch.bfloat16 BLOCK_SIZE128 D128 SAME_SEQ_LEN_ALL_REQSTrue N B Seq_len Outputs_equal Ref_Latency_[usec] Our_Latency_[usec] Ref_BW_[TB/sec] Our_BW_[TB/sec] ---------------------------------------------------------------------------------------------------------- 8 10 8064 yes 91797.47 341.11 0.002 0.500 8 10 10240 yes 119940.88 416.60 0.002 0.516 8 10 12032 yes 142493.83 444.33 0.002 0.566 8 10 16384 yes 190612.51 659.39 0.002 0.517 8 10 19200 yes 228080.26 788.18 0.002 0.512 8 10 25600 yes 331777.49 1042.17 0.002 0.513 8 10 32768 yes 422492.27 1208.98 0.002 0.564 8 20 8064 yes 213834.41 610.95 0.002 0.558 8 20 10240 yes 268580.89 767.21 0.002 0.560 8 20 12032 yes 320251.59 980.86 0.002 0.513 8 20 16384 yes 425229.09 1209.94 0.002 0.563 8 20 19200 yes 501996.54 1439.80 0.002 0.561 8 20 25600 yes 677030.97 2083.51 0.002 0.513 8 20 32768 yes 853946.94 2409.99 0.002 0.566 8 24 8064 yes 259979.37 827.37 0.002 0.494 8 24 10240 yes 311260.29 1038.10 0.002 0.497 8 24 12032 yes 370500.20 1109.80 0.002 0.544 8 24 16384 yes 523368.08 1489.21 0.002 0.549 8 24 19200 yes 616558.68 1952.58 0.002 0.496 8 24 25600 yes 808925.46 2347.28 0.002 0.547 8 24 32768 yes 1033353.92 3306.62 0.002 0.495 8 32 8064 yes 354466.72 989.11 0.002 0.551 8 32 10240 yes 443984.86 1354.58 0.002 0.508 8 32 12032 yes 523036.50 1585.94 0.002 0.508 8 32 16384 yes 704769.45 1947.90 0.002 0.560 8 32 19200 yes 805073.97 2558.09 0.002 0.505 8 32 25600 yes 1086389.57 3075.26 0.002 0.556 8 32 32768 yes 1375007.49 3912.25 0.002 0.557 性能结论fp16 下内核有效带宽达到至少 0.55 TB/sec而 910B4 的标称全局内存带宽为0.80 TB/sec即约69% 的带宽利用率对比 Python 参考实现仅约 0.002 TB/sec内核加速效果显著例如B10, Seq_len8064时延迟从约 93 毫秒降至约 0.31 毫秒。bf16 下有相近的性能水平约 0.49~0.57 TB/sec。从表格看各序列长度下Outputs_equal均为 yes说明在全部基准配置点上内核与参考实现输出一致。需要说明的是上述数据来源于该实验 README 在 x86 宿主 910B4 环境下的记录实际数字会随 CANN 版本、固件、batch 组成及SAME_SEQ_LEN_ALL_REQS取值变化建议在目标环境上自行复测。源码级实现原理参考实现算法的语义基准ref_quest_prefill_metadata.py 以最直白的方式定义了算子语义对每个请求r、每个 KV headn由seq_lens[r]计算实际 KV 块数⌈seq_lens[r]/BLOCK_SIZE⌉与元数据块数⌈num_kv_blocks_in_request/BLOCK_SIZE⌉逐元数据块、逐 KV 块取出 K Cache 切片k_cache[kv_block_id, :, n, :]沿 token 维dim 0分别做 reduce-max 与 reduce-min得到maxblock[meta_blk_id, blk, n, :]与minblock[meta_blk_id, blk, n, :]对每个元数据块中未使用的行尾块填充置零保证不存在的 KV 块对应的元数据为 0。其中_calculate_tokens_to_reduce处理尾块只有当当前 KV 块是最后一个元数据块中的最后一个 KV 块时才只对seq_len - 已归约 token 数个有效 token 做归约其余情况均按整块BLOCK_SIZE处理。这与内核中的 tail check 逻辑完全对应。数据生成器gen_data_quest_prefill_metadata.py 中的gen_quest_prefill_inputs以固定随机种子SEED42生成k_cache(num_kv_blocks, BLOCK_SIZE, N, D)torch.randn * 1.5放大数值范围以充分考验 max/min 归约seq_lensssarTrue时全部为MKBPR*BLOCK_SIZE否则每个请求在[0, max_seq_len]内随机block_tables与metadata_block_tables通过torch.randperm生成随机置换的块索引表int32模拟 vLLM 式 PagedAttention 的非连续块布局maxblocks/minblocks先以torch.empty分配输出占位再统一contiguous()化。这种随机置换块表 随机/等长序列长度的生成方式可以覆盖连续与非连续物理块、整块与尾块、满块与空块等真实 PagedAttention 场景下的边界条件。AscendC 内核向量核上的两级分块归约quest_prefill_metadata.cpp 实现了向量核VecCore内核核心设计可以总结为以下几点1. 并行模型1 个核处理 1 个 (batch, head) 组合num_batch_heads B * N按GetBlockIdx()在 AI Core 间轮转分配实现 batch 维与 KV head 维的天然并行每个核内循环处理该 (r, h) 对应的全部 KV 块因此每个 KV 块只会被读取一次。2. 两级分块结构外层循环遍历meta_blk每个元数据块覆盖BLOCK_SIZE * BLOCK_SIZE 16384个 token内层循环遍历该元数据块收纳的BLOCK_SIZE个 KV 块对每个 KV 块先通过DataCopy只拷贝k_cache[kv_block_id, :, h, :]这一 head 切片利用srcStride跳过其他 head再做归约。3. 对数归约ReduceTokenDim对(BLOCK_SIZE, D)的 UB 局部张量采用逐轮两两合并的对数归约Max/Min(vec[0], vec[0], vec[len/2])奇数向量时先搬运到前半段再继续直到len DisMaxtrue走MaxisMaxfalse走Min从而用同一套循环同时产出 maxblock 与 minblock归约前先将 K 块复制到工作缓冲保证k_block_lt原值不被破坏避免重复加载 GM。4. 缓冲与流水设计输入队列k_block_in_q_使用双缓冲DOUBLEBUFFER2配合MTE2_V事件同步实现 GM→UB 拷贝与向量计算的流水重叠max_out_q_/min_out_q_在归约前用 ±65504fp16 极值初始化尾块行用 0 填充最终通过带dstStride跳过其他 head的DataCopy将BLOCK_SIZE*D的元数据行写回maxblocks/minblocks的指定meta_blk_id位置。5. 接口与约束内核启动参数为B, N, BLOCK_SIZE, D, MKBPR, MMBPR由 torch_interface.cpp 从张量 shape 自动推导并做TORCH_CHECK校验D 128、BLOCK_SIZE 128等启动时blockDims min(B * N, NUM_CORES)即任务数超过物理核数时自动截断。内核源码注释还特别说明了一个值得注意的工程事实bf16 输入同样能通过测试尽管内核内部的 min/max 处理实际发生在 fp16half路径上这一点被记录在算子 README 的 TODO 项中Dobule check why bf16 passes even though the whole min/max processing happens in fp16从源码结构看这得益于 bf16 与 fp16 同为 16 位宽、且归约结果在测试容差内一致。内存占用视角为什么这是带宽受限算子从bytes_moved_prefill的计算逻辑可以看出该算子的全局内存流量由三部分组成K Cache 读取num_effective_kv_blocks * BLOCK_SIZE * N * D * 2字节、块表读取每有效 KV 块 4 字节、以及 maxblocks/minblocks 双份写入。整个内核几乎没有计算密集操作只有逐元素 Max/Min因此性能瓶颈在全局内存带宽这也解释了实测 69% 带宽利用率的达成路径单次读取 双缓冲流水 对数归约把向量单元的空闲时间压到了最低。在解码流水线中的位置与后续实验quest_prefill_metadata只是 QUEST 稀疏注意力流程的第一步。按照 select_attention_operators 总览 README 中记录的当前最佳实践已在 vllm-ascend 中验证使用quest_prefill_metadata()在prefill 后创建元数据并在每解码 128 个 token即 BLOCK_SIZE 的整数倍时更新元数据使用quest_block_select_paged_in_out_w()在解码阶段根据当前 token 的 query 向量预测重要的 KV 块索引。同目录下的后续实验3_quest_block_select_paged、4_quest_block_select_paged_w分别实现了基于元数据的稀疏掩码预测器以及额外融合 sink window 特征的版本——它们都以本实验产出的 maxblocks / minblocks 作为输入因此本实验的正确性验证是整条稀疏注意力链路的基石。本实验目录中的 conftest.py 会把每个实验子目录注入sys.path使得各实验可以共享gen_data_*、ref_*等辅助模块在子项目根目录执行pytest -v experiments即可一次性运行全部实验的测试。总结quest_prefill_metadata实验给出了一个完整的参考实现 → 数据生成 → 功能测试 → 基准测试的算子验证范式功能上通过四级 pytest 参数化基础、边界、长序列、大 batch与随机页表/变长序列生成全面覆盖了 PagedAttention 场景下可能出现的对齐与尾块问题性能上在 910B4 上以约 0.5~0.57 TB/sec 的有效带宽、69% 的标称带宽利用率将原本毫秒级的元数据构建压缩到亚毫秒级为解码阶段每秒数十次元数据刷新提供了可行基础实现上向量核 (batch, head) 并行 双缓冲流水 对数归约的设计是一份可以直接复用的 AscendC 带宽密集型算子模板。对于希望在 Ascend NPU 上落地 QUEST 类稀疏注意力、或需要理解 PagedAttention 元数据管理细节的开发者建议按本文的步骤先运行测试与基准再结合源码逐段对照理解两级分块与尾块处理逻辑。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考