算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载导读本文基于 CANN ops-transformer 仓库中 StemIndexer 算子 pytest 测试框架说明完整讲解该算子如何在 CPU 侧实现 golden 参考实现、如何在 NPU 侧通过torch.ops.custom.npu_stem_indexer与npu_stem_indexer_metadata完成分核信息生成与算子调用以及如何以 pytest 组织 single / batch 两套白盒用例体系。读完本文你将掌握 StemIndexer 测试环境的搭建、STEM_INDEXER_MODE/STEM_INDEXER_CASE_ID等关键环境变量的用法、CSV 批量用例的生成与复跑流程并从源码级理解动态 TopK 预算、forced 索引、因果可见范围等核心校验逻辑。一、测试框架功能说明StemIndexer 的 pytest 测试框架位于experimental/attention/stem_indexer/tests/pytest围绕“CPU 出 golden、NPU 出结果、逐项比对”的核心思路实现算子功能验证CPU 侧按 stem_indexer_golden.py 中的stem_indexer_golden()实现参考算法直接使用 PyTorch 张量运算模拟选块与 TopK 排序作为结果比对的基准。NPU 侧先调用torch.ops.custom.npu_stem_indexer_metadata生成分核信息metadata再调用torch.ops.custom.npu_stem_indexer获取实际输出。这两个自定义算子接口由 torch_ops_extension 提供由 custom_ops.py 动态加载注册。执行模式支持eager直接调用算子与graph通过torch.compiletorchair编译为 aclgraph 图两种模式由环境变量STEM_INDEXER_MODE切换默认eager。结果比对比较sparse_seq_len与sparse_indices中sparse_seq_len范围内的有效前缀sparse_seq_len之后的尾部为未定义区域不做校验。从 test_stem_indexer_single.py 可以看到执行模式的解析逻辑为os.environ.get(STEM_INDEXER_MODE, eager)仅接受eager与graph两个合法值否则抛出ValueError。graph 模式对 FAIL 用例直接pytest.skip即图模式只验证正例。二、用例来源与白盒用例方案2.1 两个用例源白盒用例方案由两个文件共同维护test_stem_indexer_paramset.pysingle 模式的用例参数表用例直接维护在ENABLED_PARAMS中每条 case 前通过注释记录覆盖点和设计原因。csv/stem_indexer_generalized_cases.csvbatch 模式的用例表作为批量模式生成.pt文件的输入。single 与 batch 各自维护独立的用例集合互不干扰。2.2 TopK score 精度分支两种用例源均包含topk_score_precision字段用于切换 TopK score 的数据精度取值含义1uint32默认未显式配置时使用2uint16对应关系在 stem_indexer_golden.py 中定义为TOPK_SCORE_TYPE_BY_PRECISION {1: uint32, 2: uint16}。uint16 路径下 golden 会把排序分数先 round 到 bfloat16get_topk_sort_scores中执行scores.to(torch.bfloat16).to(torch.float32)以模拟 AICore 上低精度 TopK score 的数值行为。在 CSV 用例表中SI_WB_001~SI_WB_100固定使用 uint32 TopK score 路径topk_score_precision1SI_WB_101~SI_WB_150是前 50 条功能场景的镜像固定使用 uint16 TopK score 路径topk_score_precision2保持输入与属性完全不变、仅切换精度用于“精度分支配对”SI_WB_001_1、SI_WB_001_2及其 uint16 镜像SI_WB_101_1、SI_WB_101_2额外覆盖 batch8 的 64K/128K 等长 prefill 场景。2.3 CSV 用例字段说明以 stem_indexer_generalized_cases.csv 表头为准批量用例共包含以下字段字段说明case_id用例编号如SI_WB_001testcase_name用例名描述场景特征expected_result期望结果PASS/FAILbatch_sizebatch 大小q_heads/kv_headsquery / KV 头数q_seq_lens/kv_seq_lens按 batch 实际的 query / KV 长度数组num_prompt_tokens动态 TopK 预算基准长度数组causal是否因果true/falsealphaTPD 衰减系数stem_block_size/stem_stridestem 块大小 / 步长当前用例固定为 128 / 16initial_blocks/window_size强制保留的 sink / window 块数当前固定为 4 / 4topk_score_precisionTopK score 精度1uint322uint16qflat_dtype/kflat_dtype/vbias_dtype各输入张量数据类型如 BF16 / FP32special_setting特殊构造如qflatzeros、vbiasstrictly_ascending_by_k_blockdescription覆盖点与设计意图说明以SI_WB_001为例batch_size1, q_heads32, kv_heads8, q_seq_lens[3968], kv_seq_lens[32768], num_prompt_tokens[32768], causaltrue, alpha1设计意图是覆盖基础 prefill 链路与 32K token 的 KV/prompt 基准线。三、当前覆盖点全览框架的覆盖点设计高度贴近 StemIndexer 的实际计算分支以下逐条对应 CSV 与ENABLED_PARAMS中的真实用例空序列与尾块边界覆盖 q/kv 尾块以及 batch 内q_seq_lens0、kv_seq_lens0、二者同时为 0 的空序列。代表用例SI_WB_058batch 内混合 q0、kv0、短 KV 和 32K KV、SI_WB_061同一 batch 覆盖 q0、kv0、q/kv 同时为 0、SI_WB_0620/1/127/128/129 token 极小长度边界。sink/window 重叠关系initial_blocks4与window_size4在不同s2Valid下呈现完全重叠SI_WB_005s2Valid4、部分交集SI_WB_006s2Valid7、无交集SI_WB_007、短序列裁剪SI_WB_008s2Valid2四种形态。causal 与 non-causal 路径causal 下验证早期 query 块只能看到历史块SI_WB_009、S1/S2 不相等时的偏移SI_WB_040non-causal 下验证每个 q 块可见完整 KVSI_WB_010。TPD alpha 衰减alpha1无衰减SI_WB_011退化路径、alpha0.5普通衰减SI_WB_012、alpha0.01强衰减SI_WB_013覆盖 topkCurr clamp。动态 TopK 预算分支small/medium/large 三段分段边界为 55/56/159/160 blocksSI_WB_035~SI_WB_038逐边界覆盖SI_WB_048把四个边界放进同一 batch 混合验证。S2 方向baseN256整块SI_WB_016正好 256 块与尾块SI_WB_017257 块跨段归并。M 方向baseM64g4与g32下的完整 M 块与 M 尾块SI_WB_018~SI_WB_021、SI_WB_043。GQA 组合q_heads为 32/64、kv_heads为 2/4/8共覆盖 6 种合法组合SI_WB_022~SI_WB_026等其中q_heads64, kv_heads2g32是最大 GQA 组、风险最高。多 batch 变长与 prefill/decode 混合batch 覆盖 1 到 19、23 到 31、37、39 等非 2 次幂和较大 batch 场景decode 以q_seq_lens1的形态混入SI_WB_028、SI_WB_039、SI_WB_072、SI_WB_089等。长序列量级kv_seq_lens与num_prompt_tokens覆盖 32K/64K/128K/256K/1M token 基准线SI_WB_055为 decode 单 token 1M KV 上限SI_WB_091为 non-causal prompt1Mq_seq_lens以覆盖 query 尾块、decode 为主仅在张量规模可控时放大。OAMvbias与scoreScaleSI_WB_030构造 QK 全相等、仅由vbias决定排序验证DivRsAddVbias路径SI_WB_031验证scoreScale1/64的 DivRs 缩放方向与量级。num_prompt_tokens不整除stem_block_size如SI_WB_034131071128K 附近向上取整、SI_WB_086prompt 比 KV 多 1 token。短路 2 路径SI_WB_044覆盖curS2Len initial dynamicTopK window恰好相等时直接输出连续索引SI_WB_045覆盖刚越过该边界、进入正常 TopK 路径。四、Metadata 生成与输入约定4.1 Metadata 的容量与对齐metadata是 StemIndexer 主算子的前置输入。pytest 正例和普通单 case 脚本都会先调用stem_indexer_metadata按如下公式生成动态容量的 metadata16 B * kv_heads * (36 72) * 16然后向上对齐到 4096 个 INT32 元素。对应实现位于 stem_indexer_golden.pyMETADATA_HEADER_SIZE 16 METADATA_CORE_STRIDE 16 METADATA_AIC_CORE_NUM 36 METADATA_AIV_CORE_NUM 72 METADATA_ALIGN_SIZE 4096 def get_metadata_size(batch_size, kv_heads): max_section_num int(batch_size) * int(kv_heads) metadata_size METADATA_HEADER_SIZE max_section_num * ( METADATA_AIC_CORE_NUM METADATA_AIV_CORE_NUM ) * METADATA_CORE_STRIDE return ceil_div(metadata_size, METADATA_ALIGN_SIZE) * METADATA_ALIGN_SIZE即每个batch × kv_heads分片包含 AIC36 核与 AIV72 核两组核间同步信息每组以 16 个 INT32 为步长整个 metadata 以 4096 元素为单位向上对齐。在 test_stem_indexer_single.py 的build_metadata中可以看到metadata 由torch.ops.custom.npu_stem_indexer_metadata(q_seq_lens, kv_seq_lens, q_heads, kv_heads, causal..., stem_block_size..., dim_qkflat..., window_size...)实际生成其中dim_qkflat stem_stride * HEAD_DIMHEAD_DIM128。4.2 主算子的输入布局与约束主算子使用BNSD 布局qflat形状为[B, q_heads, max_qb, stride*HEAD_DIM]kflat形状为[B, kv_heads, max_kb, stride*HEAD_DIM]见 get_tensor_shapes。q_seq_lens、kv_seq_lens按 batch 实际长度传入num_prompt_tokens按 batch 传入动态 TopK 预算基准长度正例中保持num_prompt_tokens kv_seq_lens该输入缺省时由 OpHost 通过 TilingData 通知 Kernel 复用kv_seq_lens。metadata虽然在接口层声明为可选输入但当前主算子计算必须传入有效 Metadata缺省时会在 Tiling 阶段返回参数错误。测试用例不再单独维护额外 token 长度辅助字段qflat、kflat的 shape 由q_seq_lens、kv_seq_lens的最大值推导max_qb ceil_div(max(q_seq_lens), stem_block_size)max_kb同理。当前 case 表只保留可运行并可与 golden 比对的正例FAIL 用例仅在 single 模式的ENABLED_PARAMS中通过pytest.raises(Exception)断言异常见 test_stem_indexer_single.py。五、文件结构测试目录完整布局如下路径均相对于仓库根目录experimental/attention/stem_indexer/tests/pytest/ ├── test_run.sh # 执行脚本single/batch × eager/graph ├── test_stem_indexer_paramset.py # single 用例参数表ENABLED_PARAMS ├── stem_indexer_golden.py # CPU 侧 golden 实现 ├── result_compare_method.py # sparse 输出比较 ├── test_stem_indexer_single.py # single 主执行入口 ├── test_stem_indexer_batch.py # batch 主执行入口 ├── stem_indexer_aclgraph.py # single/batch 共用的 aclgraph(graph) 调用实现 ├── test_npu_stem_indexer.py # 参考 LI 写法的普通单 case 脚本 ├── pytest.ini # pytest 标记ci / graph ├── custom_ops.py # 自定义算子动态加载 ├── csv/stem_indexer_generalized_cases.csv # batch 用例表 ├── batch/stem_indexer_pt_save.py # 读取 CSV 并生成 pt ├── batch/stem_indexer_pt_loadprocess.py # 读取 pt 并调用算子 └── batch/replace_path.py # 替换 batch pytest 中的 pt 路径此外目录中还包含csv/stem_indexer_invalid_cases.csv、csv/stem_indexer_redline_cases.csv、csv/stem_indexer_stc_cases.csv等其他 CSV以及ttk/子目录中的 TTKTensor 测试套件用例含ttk_stem_indexer_generalized_cases.csv等和ut/下的 Host 侧单测tiling / infershape。pytest.ini定义了两个 markerciCI 测试与graph图模式编译测试。六、使用方法6.1 一键执行test_run.sh在当前 pytest 目录下执行bash test_run.sh single # single (eager) bash test_run.sh single_graph # single (graph)脚本内部test_run.sh支持 5 个参数single/single_graph/batch/batch_graph/help。eager 模式使用-m ci标记graph 模式使用-m graph标记并统一追加-W ignore::UserWarning -W ignore::DeprecationWarning屏蔽告警。批量模式先执行 pt 生成脚本、再执行 pytest最后输出result.csv。6.2 按 case_id 定向运行single 和 batch 模式都支持通过STEM_INDEXER_CASE_ID只运行指定用例多个 case_id 使用逗号分隔STEM_INDEXER_CASE_IDSI_WB_001,SI_WB_002 python3 -m pytest test_stem_indexer_single.py STEM_INDEXER_CASE_IDSI_WB_001,SI_WB_002 python3 -m pytest test_stem_indexer_batch.py过滤逻辑见 test_stem_indexer_single.py未设置该变量时运行全部用例设置后仅保留case_id命中的用例并打印匹配数量。batch 侧test_stem_indexer_batch.py则是先扫描 pt 目录再逐个加载 pt 中的case_id进行过滤。6.3 批量测试bash test_run.sh batch # eager 模式 bash test_run.sh batch_graph # graph 模式复跑已生成的.pt文件跳过生成阶段python3 -m pytest test_stem_indexer_batch.pybatch 模式流程与 QLIquantized lightning indexer保持一致1. 读取 csv/stem_indexer_generalized_cases.csv。 2. 生成每条 case 的 .pt 文件保存输入和 CPU golden。 3. pytest 逐个读取 .pt 文件运行时动态构造 metadata 并调用 NPU 算子。 4. 与 .pt 中保存的 golden 比对。 5. 生成 result.csv 记录批量执行结果。生成.pt时同样可用STEM_INDEXER_CASE_ID只选择指定用例STEM_INDEXER_CASE_IDSI_WB_001_1,SI_WB_101_1 \ python3 batch/stem_indexer_pt_save.py csv/stem_indexer_generalized_cases.csv pt_path生成脚本默认使用全部可用 CPU 核按 case 并行大用例并行时内存占用较高可用--workers限制进程数python3 batch/stem_indexer_pt_save.py csv/stem_indexer_generalized_cases.csv pt_path --workers 8--workers 0表示使用全部可用 CPU 核--workers 1表示按原方式串行生成。.pt文件和result.csv均为本地生成产物不需要提交。6.4 普通单 case 脚本不依赖 pytest 参数表的普通单 case 脚本可直接执行python3 test_npu_stem_indexer.py该脚本是参考 LILightningIndexer测试写法的最小示例便于快速复现单条路径。七、源码级原理解析7.1 golden 算法的关键环节stem_indexer_golden.py 的stem_indexer_golden()逐 batch、逐 kv_head、逐 q_block 计算核心逻辑可拆解为四个子函数causal 可见范围calc_causal_s2_valid计算s2_valid kv_block_num - q_block_num q_block_idx 1并 clamp 到[0, kv_block_num]decode 或 non-causal 时s2_valid kv_block_num。动态 TopK 预算calc_topk_budget先将prompt_len / stem_block_size向上取整得到 prompt block 数按 56small、 160medium、其余large三分支计算k_startmedium/large 分支使用rate与bias常数K_BLOCK_NUM_RATE_MEDIUM0.2、K_BLOCK_NUM_BIAS_MEDIUM30、K_BLOCK_NUM_RATE_LARGE0.1、K_BLOCK_NUM_BIAS_LARGE30再结合alpha在[k_start, k_start*alpha]之间线性插值并向下取整最终 clamp 到[1, k_start]与MAX_TOPK_COUNT256。注意插值乘加使用fma_float32模拟 AICore 标量表达式的浮点行为。forced 索引get_forced_indices将sink [0, min(initial_blocks, s2_valid))与window [max(s2_valid - window_size, 0), s2_valid)取并集保证这些块无条件保留。排序选块从[0, s2_valid)中剔除 forced 后按分数降序同分取小索引取前topk_select_num个再与 forced 合并、整体升序输出。sparse_indices初始化为全-1有效前缀按sparse_seq_len写入这与“尾部未定义区域不校验”的比对策略一致。7.2 比对方法有效前缀 TopK 边界容差result_compare_method.py 的assert_stem_indexer_result依次执行校验sparse_seq_len与sparse_indices的shape 完全一致校验sparse_seq_len逐元素相等对每个expected_seq_len 0的 (b, q_head, q_block) 行提取actual_row[:valid_len]与 golden 前缀做集合比对尾部区域只检查 padding 是否仍为-1bad_padding_count 0对动态 TopK 边界处的索引差异用explain_topk_mismatch计算原始分数get_row_scores重算matmul(k_group, q_vec) * score_scale vbias逐项给出相对误差并统计“超出相对误差 1e-3 的差异数占总有效 TopK 数比例”要求不超过MAX_BOUNDARY_RE_EXCEED_RATIO 5e-3输出结果行包含BadPadding / FailedRows / ExceedDiff / PctThd / PctRlt / Result六列统计方便 CI 日志快速定位。7.3 eager 与 graph 两种调用路径eagercall_stem_indexer将输入move_inputs_to_npu后直接调用torch.ops.custom.npu_stem_indexer(...)见 test_stem_indexer_single.py。graph由 stem_indexer_aclgraph.py 定义StemIndexerNetworktorch.nn.Module在forward中封装npu_stem_indexer调用再经torch.compiletorchair的CompilerConfig编译为 aclgraph 图执行。graph 模式只支持 PASS 用例。7.4 自定义算子的动态加载custom_ops.py 负责在测试进程启动时加载 NPU 自定义算子通过STEM_INDEXER_CUSTOM_OPS_PATH环境变量可选指向外部扩展目录默认回退到tests/pytest上级目录的torch_ops_extension加载顺序为共享库custom_ops_lib*.so通过torch.ops.load_library→ Python 包custom_ops/__init__.py→ converternpu_stem_indexer.py、npu_stem_indexer_metadata.py同时通过ASCEND_CUSTOM_OPP_PATH与LD_LIBRARY_PATH配置 vendor 算子包路径优先使用仓库build/_CPack_Packages/.../custom_transformer构建产物否则回退到固定开发机路径以torch.ops.custom.npu_stem_indexer与torch.ops.custom.npu_stem_indexer_metadata两个属性是否可访问作为加载成功的判据避免对单一开发 worktree 的硬编码。八、运行前提与注意事项环境需要安装torch_npu入口文件通过pytest.importorskip(torch_npu)做条件跳过以及可用的 torchair 环境graph 模式必需NPU 侧调用依赖已编译安装的 torch_ops_extension 扩展参考其 build_and_install.sh。test_run.sh必须在experimental/attention/stem_indexer/tests/pytest目录下执行脚本中的 CSV 与 pt 路径均为相对路径。当前 CSV 表内stem_block_size固定 128、stem_stride固定 16、initial_blocks与window_size固定 4如需覆盖其他取值需同步扩展ENABLED_PARAMS或 CSV。批量测试生成的.pt文件与result.csv均为本地产物无需提交仓库。参考资源测试框架入口experimental/attention/stem_indexer/tests/pytest/README.md算子设计文档experimental/attention/stem_indexer/docs/StemIndexer.mdOpHost 实现含 Tilingexperimental/attention/stem_indexer/op_host/stem_indexer_tiling.cpp、stem_indexer_def.cppKernel 实现experimental/attention/stem_indexer/op_kernel/stem_indexer.cpptorch 扩展与注册experimental/attention/stem_indexer/torch_ops_extension/custom_ops/csrc测试子目录tests/pytest、tests/ttk、tests/ut赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐ops-transformer 中 KvQuantSparseAttnSharedKv 算子 pytest 测试框架实战CPU Golden 生成、NPU 双模式调用与精度对比ops transformer 中 KvQuantSparseAttnSharedKv 算子 pytest 测试框架实战CPU Golden 生成、NPU 双算子库人工智能深度学习AscendCANN ops-transformerquant_lightning_indexer 算子 pytest 测试框架——CPU Golden、NPU 直调与批量用例实战CANN ops transformerquant_lightning_indexer 算子 pytest 测试框架——CPU Golden、NPU 直调与批算子库人工智能深度学习AscendCANN Compressor 算子 pytest 测试框架实战指南CPU golden 生成、NPU 精度对比与批量隔离执行CANN Compressor 算子 pytest 测试框架实战指南CPU golden 生成、NPU 精度对比与批量隔离执行 导读 本文围绕 CANN op算子库人工智能深度学习Ascend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考