CANN ops-transformer 算子 PyTorch 接入实战：stem_indexer 稀疏注意力前处理算子的 TorchNPU 扩展构建与调用

CANN ops-transformer 算子 PyTorch 接入实战：stem_indexer 稀疏注意力前处理算子的 TorchNPU 扩展构建与调用

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer

本篇技术指南围绕 CANN ops-transformer 仓库中experimental/attention/stem_indexer/torch_ops_extension目录展开，系统讲解推理场景稀疏 Attention 前处理算子 StemIndexer 如何通过编译式NpuExtension桥接到 PyTorch，注册为torch.ops.custom.npu_stem_indexer（同时挂载到TorchNPU.npu_stem_indexer），并底层调用 ACLNN 接口aclnnStemIndexer/aclnnStemIndexerMetadata。读完本文，你将掌握该扩展的目录结构、构建安装流程、eager 单算子与 torch.compile 成图两种调用通路、接口与算子 IR 的对应关系，以及如何在算子 pytest 框架中集成验证。

一、StemIndexer 是什么：稀疏 Attention 的块级打分与动态选块

StemIndexer 是推理场景下稀疏 Attention 的前处理算子，承担块级打分与动态选块职责。对于每个 Query Block，算子基于qflat与kflat的相关性，并叠加 Value 量值偏置vbias（OAM，Output-Aware Metric）进行打分；随后按 Position-Decay 动态 TopK 预算选出关键 Key Block，输出sparse_indices与sparse_seq_len供下游 Block Sparse Attention 算子使用。完整的算子设计说明见 docs/StemIndexer.md。

固定参数为stem_block_size=128、stem_stride=16，因此代表向量聚合比例R = B_s/T_s = 8，扁平化特征维度D_f = T_s·D = 2048（原始 Head DimD = 128）。打分公式为：

S_{i,j} = <Q_i, K_j> / R^2 + V_j = <Q_i, K_j> / 64 + V_j

其中1/R^2 = 1/64对每组内部 64 个 Token Pair 贡献进行归一化。最终输出集合按Sink Block → 普通 TopK 结果 → Window Block顺序拼接，不执行额外去重。

本文聚焦的是该算子的PyTorch 接入层——即torch_ops_extension目录，它实现了算子从 Python 侧到 ACLNN C 接口的完整调用链。

二、接入层目录结构与职责划分

experimental/attention/stem_indexer/torch_ops_extension/ ├── setup.py # NpuExtension custom_ops.custom_ops_lib，编译 csrc/*.cpp ├── build_and_install.sh # 构建 wheel 并 pip 安装 ├── README.md └── custom_ops/ ├── __init__.py # 导入 custom_ops_lib(.so) + converter，并挂载到 TorchNPU ├── csrc/ │ ├── ops_def_registration.cpp # TORCH_LIBRARY(custom): npu_stem_indexer(_metadata) schema │ ├── ops_common.h / ops_common.cpp # EXEC_NPU_CMD_V1 / ConvertType 基础设施（拷贝） │ ├── npu_stem_indexer.cpp # NPU/Meta 前向实现 + TORCH_LIBRARY_IMPL │ └── npu_stem_indexer_metadata.cpp └── converter/ ├── __init__.py ├── npu_stem_indexer.py # torch.compile (GE) 成图 converter └── npu_stem_indexer_metadata.py

整体结构与参考实现quant_block_sparse_attn/torch_ops_extension对齐：采用编译式NpuExtension（custom_ops.custom_ops_lib）+TORCH_LIBRARY/TORCH_LIBRARY_IMPL+EXEC_NPU_CMD_V1的组合。其中csrc/ops_common.h、csrc/ops_common.cpp为该参考的整文件拷贝（自包含），提供 ACLNN 接口动态加载与参数类型转换的公共基础设施。

三、前置条件

接入层对运行环境有明确要求（与仓库 README.md 一致）：

• Linux 操作系统；
• Python 3.8+，GCC 9.4.0+；
• PyTorch >= 2.6.0 与匹配版本的TorchNPU；
• Ascend CANN Toolkit（环境变量ASCEND_HOME_PATH已设置）；
• 已部署aclnnStemIndexer/aclnnStemIndexerMetadata算子包，运行时可经ASCEND_CUSTOM_OPP_PATH/ASCEND_OPP_PATH检索到libcust_opapi.so。

算子本体的硬件适配情况（见 docs/StemIndexer.md）：当前仅适配Ascend 950PR / Ascend 950DT，A3、A2、Atlas 200I/500 A2 及其他训练/推理系列产品均不支持。在 op_host/stem_indexer_def.cpp 中，算子定义仅执行AICore().AddConfig("ascend950", aicoreConfig)，代码注释亦明确“当前仅适配 A5（ascend950），A2/A3 暂未适配，不注册”。

四、构建与安装：wheel 与就地编译两种方式

进入experimental/attention/stem_indexer/torch_ops_extension目录后，提供两种构建方式：

# 方式一：构建 wheel 并安装（推荐） bash build_and_install.sh # 方式二：就地编译（.so 生成在 custom_ops/custom_ops_lib*.so） python3 setup.py build_ext --inplace

build_and_install.sh 的内部流程为：清理历史build目录 →python3 setup.py build bdist_wheel构建 wheel → 在dist目录执行pip3 install *.whl -I强制安装。

从 setup.py 可以看到编译细节：

• 使用torch_npu.utils.cpp_extension.NpuExtension，扩展模块命名为custom_ops.custom_ops_lib，源文件为custom_ops/csrc/*.cpp（glob 收集）；
• 通过torch_npu安装根目录定位 ACL 头文件：include/third_party/acl/inc，并作为-I编译参数传入；
• 支持USE_NINJA=1环境变量启用 ninja 加速（默认关闭）；
• package_data显式包含*.py与*.so，保证安装后 Python 侧from . import custom_ops_lib可加载编译产物。

五、Python 侧注册与挂载机制

custom_ops/__init__.py是整个接入层的入口，其职责有三：

1. 导入编译产物custom_ops_lib（.so），触发TORCH_LIBRARY(custom, m)中的 schema 注册；
2. 导入converter子模块，注册 torch.compile 成图所需的 GE converter；
3. 遍历torch.ops.custom命名空间下所有算子，逐个setattr(torch_npu, op_name, custom_op_func)挂载到torch_npu，从而支持TorchNPU.npu_stem_indexer(...)调用形式；若torch.ops.custom不存在则输出 warning 并降级为仅torch.ops.custom.xxx调用。

schema 定义位于 csrc/ops_def_registration.cpp：

TORCH_LIBRARY(custom, m) { m.def("npu_stem_indexer(Tensor qflat, Tensor kflat, Tensor vbias, Tensor q_seq_lens, " "Tensor kv_seq_lens, *, " "Tensor? num_prompt_tokens=None, Tensor? metadata=None, " "bool causal=True, int stem_block_size=128, int stem_stride=16, float alpha=1.0, " "int initial_blocks=4, int window_size=4, float k_block_num_rate_medium=0.2, " "int k_block_num_bias_medium=30, float k_block_num_rate_large=0.1, " "int k_block_num_bias_large=30, int topk_score_precision=1) -> (Tensor, Tensor)"); m.def("npu_stem_indexer_metadata(Tensor q_seq_lens, Tensor kv_seq_lens, int q_heads, " "int kv_heads, *, bool causal=True, int stem_block_size=128, int dim_qkflat=2048, " "int window_size=4) -> Tensor"); }

schema 设计遵循“必选张量 + 必选标量在前，*之后为可选张量与带默认值的属性”的约定，张量 dtype 与属性默认值对齐算子原型 op_host/stem_indexer_def.cpp 中Input/Attr的声明（causal默认true、stem_block_size=128、stem_stride=16、alpha=1.0、initial_blocks=4、window_size=4、k_block_num_rate_medium=0.2、k_block_num_bias_medium=30、k_block_num_rate_large=0.1、k_block_num_bias_large=30、topk_score_precision=1）。

六、eager 调用：单算子前向实现

NPU/Meta 前向实现位于 csrc/npu_stem_indexer.cpp，实现要点如下：

• 入参校验：对qflat/kflat断言 4D、vbias断言 3D；qflat/kflat必须为 BFloat16，vbias必须为 Float32；kflat的 Kb 维与vbias一致，qflat/kflat最后一维一致；topk_score_precision仅允许 1（UINT32）或 2（UINT16）；
• 输出推导：sparse_indices为 INT32[B, q_heads, Qb, Kb]，sparse_seq_len为 INT32[B, q_heads, Qb]；
• ACLNN 调用：通过EXEC_NPU_CMD_V1(aclnnStemIndexer, ...)下发，实参顺序严格按算子 IR 声明顺序（输入 → 属性 → 输出），与 Python schema 的形参顺序不同——schema 为便于调用将可选张量前置，C++ 实现内部已按 IR 顺序重排；
• 设备注册：TORCH_LIBRARY_IMPL(custom, PrivateUse1, m)注册 NPU 实现，TORCH_LIBRARY_IMPL(custom, Meta, m)注册 Meta 实现（仅做 shape/dtype 推导，不真正计算）。

Metadata 接口 csrc/npu_stem_indexer_metadata.cpp 用于生成分核调度信息，其输出容量按16 + B·kv_heads·(36 + 72)·16计算并向上对齐到 4096 个 INT32 元素（常量在源码中以STEM_INDEXER_METADATA_HEADER_SIZE=16、STEM_INDEXER_METADATA_AIC_CORE_NUM=36、STEM_INDEXER_METADATA_AIV_CORE_NUM=72、STEM_INDEXER_METADATA_CORE_STRIDE=16、STEM_INDEXER_METADATA_ALIGN_SIZE=4096定义），与文档中M = AlignUp(16 + C_max(36+72)×16, 4096)、C_max = B·N_k的容量公式一致。

eager 单算子调用示例：

import torch, torch_npu import custom_ops # 注册 torch.ops.custom.npu_stem_indexer(_metadata) 并挂载到 torch_npu metadata = torch.ops.custom.npu_stem_indexer_metadata( q_seq_lens, kv_seq_lens, q_heads, kv_heads, causal=True, stem_block_size=128, dim_qkflat=2048, window_size=4, ) sparse_indices, sparse_seq_len = torch.ops.custom.npu_stem_indexer( qflat, kflat, vbias, q_seq_lens, kv_seq_lens, causal=True, stem_block_size=128, stem_stride=16, alpha=1.0, initial_blocks=4, window_size=4, k_block_num_rate_medium=0.2, k_block_num_bias_medium=30, k_block_num_rate_large=0.1, k_block_num_bias_large=30, topk_score_precision=1, num_prompt_tokens=num_prompt_tokens, metadata=metadata, ) # 也可：TorchNPU.npu_stem_indexer(...)

七、torch.compile 成图：GE converter

converter目录下的 Python 文件为 torch.compile（GE 图模式）场景提供成图转换能力：

• converter/npu_stem_indexer.py 通过@register_fx_node_ge_converter(torch.ops.custom.npu_stem_indexer.default)注册，将 FX 节点转换为torchair.ge.custom_op("StemIndexer", ...)，输入为qflat/kflat/vbias/q_seq_lens/kv_seq_lens/num_prompt_tokens/metadata，属性逐一映射（causal→attr.Bool、整型属性→attr.Int、浮点属性→attr.Float），输出为["sparse_indices", "sparse_seq_len"]；
• converter/npu_stem_indexer_metadata.py 同理注册npu_stem_indexer_metadata，生成 GE 节点StemIndexerMetadata。

converter 中meta_outputs形参为固定写法，若写错会影响 GE 节点的输出 dtype 与 shape 推导。由于接口“支持图模式”（见 docs/StemIndexer.md 约束说明），converter 是图模式下可用的前提。

八、接口与 IR 对齐：张量布局与可选张量语义

接入层与算子原型的对应关系（摘自 torch_ops_extension/README.md）：

• 入参/属性/输出与算子原型 op_host/stem_indexer_def.cpp 一一对应；
• EXEC_NPU_CMD_V1实参按 IR 声明顺序传入（输入→属性→输出）；Python schema 为便于调用将必选张量前置，C++ 实现内部已按 IR 顺序重排；
• qflat/kflat为 BF16[B, N, Qb/Kb, stem_stride*D]；vbias为 FP32[B, kv_heads, Kb]；其余q_seq_lens/kv_seq_lens/num_prompt_tokens/metadata为 INT32，metadata长度按16 + B * kv_heads * (36 + 72) * 16计算并向上对齐到 4096 个 INT32 元素；输出sparse_indices/sparse_seq_len为 INT32；
• num_prompt_tokens/metadata为可选张量（c10::optional<at::Tensor>），两条调用通路（eager 与 graph）均将缺省状态传递给 OpHost。未提供num_prompt_tokens时，TilingData 记录复用标志，Kernel 使用kv_seq_lens；未提供metadata时，OpHost Tiling 按当前计算要求返回明确的参数错误。

关于可选张量的缺省行为，测试框架 tests/pytest/README.md 也有印证：正例中保持num_prompt_tokens >= kv_seq_lens，缺省时由 OpHost 通过 TilingData 通知 Kernel 复用kv_seq_lens；metadata虽然在接口层声明为可选，但当前主算子计算必须传入有效 Metadata，缺省时会在 Tiling 阶段返回参数错误。

九、完整参数说明与约束

以下参数表完整继承自 docs/StemIndexer.md，维度符号：$B$ 为 Batch，$N_q$/$N_k$ 为 Query/KV Head 数，$Q_{\max}$/$K_{\max}$ 为最大 Query/KV Block 数，$D_f = T_s·D = 2048$。

关键约束速览（详见 docs/StemIndexer.md）：

• 支持图模式；仅适配 Ascend 950PR/950DT；
• $B$ 取值范围 $[1,65536]$；$N_q$ 仅支持 32 或 64，$N_k$ 仅支持 2、4 或 8，且 $N_q \bmod N_k = 0$；
• metadata必须由与主算子相同的q_seq_lens、kv_seq_lens、$N_q$、$N_k$、causal、stem_block_size、$D_f$ 和window_size生成；
• Sink Block 和 Window Block 不参与普通 TopK 候选，最终按 Sink → TopK → Window 顺序拼接，不执行去重；
• sparse_indices有效前缀元素为 Key Block 逻辑索引，第 $b$ 个 Batch 取值 $[0, K_b-1]$；调用者不应依赖有效索引的排列顺序。

十、动态 TopK 预算算法

动态 TopK 预算（Position-Decay）的计算逻辑是选块行为的核心，公式如下（$B_s=128$）：

1. Prompt Token 数转 Block 数：$P_b = \lceil L_b^p / B_s \rceil$；
2. 初始预算分档：

K_s = P_b , P_b < 56 K_s = floor(0.2·P_b + 30), 56 ≤ P_b < 160 K_s = floor(0.1·P_b + 30), P_b ≥ 160

3. 衰减终点 $K_e = \alpha·K_s$；
4. 对第 $i$ 个 Query Block，按 Right-down 对齐位置 $p_i = i + K_b - Q_b$ 线性插值并 clamp 到 $[1, K_s]$，实际普通 TopK 数量再限制为不超过 256。

alpha=1.0表示不衰减；值越小序列后部选择的 Key Block 越少。该分段边界（55/56/159/160 blocks）在 tests/pytest/README.md 中被列为显式测试覆盖点。

十一、Ascend 950PR/950DT 完整调用示例

以下示例完整继承自 docs/StemIndexer.md，展示了从构造输入、生成 metadata 到执行选块的完整流程：

import torch import torch_npu import custom_ops batch_size = 4 q_heads = 64 kv_heads = 8 d = 128 stem_block_size = 128 stem_stride = 16 q_block_num = 16 k_block_num = 64 q_seq_len = q_block_num * stem_block_size kv_seq_len = k_block_num * stem_block_size prompt_token_num = kv_seq_len torch.manual_seed(0) qflat = torch.randn( batch_size, q_heads, q_block_num, stem_stride * d, dtype=torch.bfloat16, ).npu() kflat = torch.randn( batch_size, kv_heads, k_block_num, stem_stride * d, dtype=torch.bfloat16, ).npu() vbias = torch.randn( batch_size, kv_heads, k_block_num, dtype=torch.float32, ).npu() q_seq_lens = torch.full((batch_size,), q_seq_len, dtype=torch.int32).npu() kv_seq_lens = torch.full((batch_size,), kv_seq_len, dtype=torch.int32).npu() num_prompt_tokens = torch.full((batch_size,), prompt_token_num, dtype=torch.int32).npu() # 1. 生成与主算子输入匹配的分核调度 Metadata。 metadata = torch.ops.custom.npu_stem_indexer_metadata( q_seq_lens, kv_seq_lens, q_heads, kv_heads, causal=True, stem_block_size=stem_block_size, dim_qkflat=stem_stride * d, window_size=4, ) # 2. 执行块级打分与动态选块。 sparse_indices, sparse_seq_len = torch.ops.custom.npu_stem_indexer( qflat, kflat, vbias, q_seq_lens, kv_seq_lens, num_prompt_tokens=num_prompt_tokens, metadata=metadata, causal=True, stem_block_size=stem_block_size, stem_stride=stem_stride, alpha=1.0, initial_blocks=4, window_size=4, k_block_num_rate_medium=0.2, k_block_num_bias_medium=30, k_block_num_rate_large=0.1, k_block_num_bias_large=30, topk_score_precision=1, )

十二、与 pytest 测试框架集成

扩展可通过环境变量STEM_INDEXER_CUSTOM_OPS_PATH（或默认相对路径<stem_indexer>/torch_ops_extension）被测试检索，既支持 globcustom_ops_lib*.so（torch.ops.load_library），也支持 execcustom_ops/__init__.py。

pytest 测试框架 的验证思路为：CPU 侧按设计实现 golden；NPU 侧先经npu_stem_indexer_metadata生成分核信息，再经npu_stem_indexer获取实际结果，支持 eager 与 graph（torch.compile+torchair编译为 aclgraph）两种模式（STEM_INDEXER_MODE切换，默认 eager）；结果比对仅比较sparse_seq_len与sparse_indices有效前缀，尾部未定义区域不校验。

执行方式（在tests/pytest目录下）：

bash test_run.sh single # single (eager) bash test_run.sh single_graph # single (graph) bash test_run.sh batch # batch (eager) bash test_run.sh batch_graph # batch (graph)

指定用例运行（STEM_INDEXER_CASE_ID支持逗号分隔多个 case_id）：

STEM_INDEXER_CASE_ID=SI_WB_001,SI_WB_002 python3 -m pytest test_stem_indexer_single.py STEM_INDEXER_CASE_ID=SI_WB_001,SI_WB_002 python3 -m pytest test_stem_indexer_batch.py

测试覆盖点包括：q/kv 尾块与空序列边界、Sink/Window 的重叠与裁剪、causal 与 non-causal 路径、TPDalpha无衰减/普通衰减/强衰减、动态 TopK small/medium/large 分段边界（55/56/159/160 blocks）、GQA 组合（q_heads32/64 ×kv_heads2/4/8 共 6 种合法组合）、多 batch 变长与 prefill/decode 混合、32K/64K/128K/256K/1M token 长序列量级、OAMvbias影响选块与scoreScale=1/64路径，以及topk_score_precision的 uint32（SI_WB_001~SI_WB_100）与 uint16（SI_WB_101~SI_WB_150镜像）两条路径。

十三、底层基础设施：EXEC_NPU_CMD_V1 与动态库加载

csrc/ops_common.h（自包含拷贝）提供了接入层运行时的关键支撑，理解它能更好地把握整条调用链：

• ACLNN 接口动态加载：GetOpApiFuncAddr按ASCEND_CUSTOM_OPP_PATH→ASCEND_OPP_PATH/vendors（读取config.ini的load_priority）→ 各 feature 库（libopapi_math.so/libopapi_nn.so/libopapi_cv.so/libopapi_transformer.so/libopapi_legacy.so）→libopapi.so的顺序查找aclnnStemIndexer等接口符号；自定义算子优先从libcust_opapi.so解析；
• 类型转换：ConvertType将at::Tensor/at::Scalar/at::IntArrayRef/c10::optional<at::Tensor>等 PyTorch 类型转换为 ACL 的aclTensor/aclScalar/aclIntArray，dtype 经ConvertScalarTypeToAclDataType映射（如BFloat16→ACL_BF16、Float→ACL_FLOAT、Int→ACL_INT32）；
• 执行宏：EXEC_NPU_CMD_V1(aclnn_api, ...)依次完成xxxGetWorkspaceSize查询 workspace、按需分配 workspace Tensor、在当前 NPU 流（c10_npu::getCurrentNPUStream()）上异步执行 ACLNN 接口，并在执行后释放转换产物。

正是这套基础设施，使得npu_stem_indexer/npu_stem_indexer_metadata两个自定义算子能够以“与原生算子一致”的方式在 NPU 流上排队执行。

十四、小结

torch_ops_extension是 StemIndexer 算子从 Python 侧触达 NPU 的唯一入口，其“schema 定义 → NPU/Meta 双实现 → converter 成图”三层结构清晰可复用：ops_def_registration.cpp声明接口契约，npu_stem_indexer.cpp/npu_stem_indexer_metadata.cpp实现 eager 通路，converter/支撑 torch.compile 图模式，ops_common屏蔽了 ACLNN 动态加载与类型转换的复杂度。如需在自有业务中接入 StemIndexer，推荐路径是：按本文第四节完成构建安装 → 按第十一节组织输入并先调用 metadata 接口 → 再调用主算子；若需验证正确性，可直接复用 pytest 框架 的 golden 比对流程。

相关参考：

• 算子完整设计：docs/StemIndexer.md
• 算子原型定义：op_host/stem_indexer_def.cpp
• 接入层 README：torch_ops_extension/README.md
• pytest 验证框架：tests/pytest/README.md

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer