基于 CANNBot-DSL 的 Ascend NPU 复杂算子示例集合。

CANNBot-DSL

基于 CANNBot-DSL 的 Ascend NPU 复杂算子样例集合，配套完整的文档站与精度测试。

📖 概述 · 💻 软硬件配套 · ⚡️ 快速上手 · 📦 算子列表 · 📚 文档导航 · 📂 目录结构 · 🔥 更新日志 · 📜 许可证

🔥 最新动态

2026-10-30

📦 发布 CANNBot-DSL 0.8.0 wheel 包：master 分支配套 cannbot-dsl==0.8.0，可直接通过 pip 安装。

2026-09-30

✅ 支持 DeepSeek V4.1：本仓算子集已覆盖相关算子，包括 indexer_prologue_qw、indexer_prologue_k、attn_prologue、attn_epilogue、quant_lightning_indexer、quant_sparse_lightning_indexer、mixed_quant_sparse_flash_mla。

🚀 新增 17 个算子样例，并更新既有算子适配最新 CANNBot-DSL。其中，Sparse Flash Attention 在典型场景下最高加速比达到 1.63 倍；Qwen Sparse Attention 在长 KV decode 场景下平均加速比达到 1.23 倍。更多性能数据与测试口径请参见对应样例 README，其余算子的详细信息请参见下方的算子列表。

📌 概述

CANNBot 是 CANN 社区的 Infra 智能体层，用 Agent 完成 AscendC/PyPTO/TileLang/Triton 等各类语言的算子开发、模型迁移与推理优化，并延伸至图模式、Runtime 等更多 CANN 开发场景。

本仓（cannbot-dsl）是其 DSL 仓，提供 Agent 亲和的编程范式；仓群还包括 cannbot、cannbot-skills、cannbot-knowledge、cann-bench、cannbot-sentry 等仓库，结构如下。

CANNBot 仓群结构

💻 软硬件配套说明

维度	支持的版本
昇腾产品	Ascend 950PR / Ascend 950DT（NPU ARCH 3510），见硬件兼容性查询
CPU 架构	x86_64 / aarch64
操作系统	CANN 支持的 Linux 发行版，见硬件兼容性查询
Python	3.10、3.11、3.12
CANN	9.2.1，见 CANN 安装部署
CANNBot-DSL	0.8.0
torch	2.7.1 / 2.9.0 / 2.10.0 / 2.11.0 / 2.12.0

上表口径截至 2026-10-29。CANN 软件包请以 CANN 安装部署页面提供的产品和架构匹配版本为准。

⚡️ 快速上手

1. 安装 CANN

版本要求见软硬件配套说明。CANN 软件包为 Linux 版本，分 x86_64 与 aarch64 两种架构。请从 CANN 安装部署页面选择 CANN 9.2.1 以及与目标产品、操作系统架构匹配的 Toolkit 和算子包。

${install_path} 为安装路径，默认 /usr/local/Ascend：

bash ./Ascend-cann_9.2.1_linux-x86_64.run --install --force --install-path=${install_path}
bash ./Ascend-cann-950-ops_9.2.1_linux-x86_64.run --install --force --install-path=${install_path}

source ${install_path}/ascend-toolkit/set_env.sh

2. 安装 CANNBot-DSL

CANNBot-DSL 0.8.0 已发布到 pip 源。发行包名为 cannbot-dsl，导入名为 cannbotdsl：

python -m pip install cannbot-dsl==0.8.0

安装完成后验证版本：

python -c 'import cannbotdsl; print(cannbotdsl.__version__)'

3. 运行一个算子

以 flash_kda 为例。该样例依赖 CANNBot-DSL 0.8.0，先生成 metadata，再执行 FlashKDA：

import math

import torch
import torch_npu

from flash_kda.flash_kda import flash_kda
from flash_kda.flash_kda_metadata import flash_kda_metadata

B, NQK, NV, S, D = 1, 3, 3, 8192, 128
q = torch.randn(B, NQK, S, D, device="npu", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn(B, NV, S, D, device="npu", dtype=torch.bfloat16)
g = torch.randn_like(v)
beta = torch.randn(B, NV, S, device="npu", dtype=torch.bfloat16)
initial_state = torch.zeros(B, NV, D, D, device="npu", dtype=torch.float32)
A_log = torch.randn(NV, device="npu", dtype=torch.float32)
dt_bias = torch.randn(NV, D, device="npu", dtype=torch.float32)

metadata = flash_kda_metadata(q, v, initial_state, "BNSD")
out, final_state = flash_kda(
    q, k, v, g, beta, 1 / math.sqrt(D), initial_state,
    A_log, dt_bias, -5.0, "BNSD", metadata,
)

print("out", tuple(out.shape))
print("final_state", tuple(final_state.shape))

样例以目录为单位组织，运行前请在仓库根目录将 samples 加入 PYTHONPATH：

export PYTHONPATH="$PWD/samples:$PYTHONPATH"
python3 your_script.py

运行成功时打印：

out (1, 3, 8192, 128)
final_state (1, 3, 128, 128)

各算子的接口参数、数据类型约束与返回值说明，见算子列表中对应样例的 README。

4. 运行精度测试

python3 -m pytest test/flash_kda/test_flash_kda.py -v

📦 算子列表

本仓算子持续合入中，当前已开源算子如下：

样例	接口	说明
matmul/matmul	`matmul()`	A16W16 矩阵乘，`C[M,N] = A[M,K] × B[N,K]^T`
matmul/matmul	`matmul_streamk()`	同目录的 Stream-K（DPSK）实现，DP + SK 混合调度，适合小 M/N、大 K
matmul/batch_matmul	`batch_matmul()`	批量矩阵乘，batch 维右对齐广播，支持 rank 2~6
matmul/fused_matmul	`fused_matmul()`	矩阵乘融合逐列 Add / Mul 和 GELU / ReLU，支持 FP16 / BF16 输入及 FP32 输出
matmul/dual_level_quant_matmul	`npu_dual_level_quant_matmul()`	MXFP4 双层 Scale 矩阵乘，使用 NZ 权重，支持 FP16 / BF16 输出和 Bias
matmul/quant_matmul	`npu_quant_matmul()`	MXFP8 / MXFP4 全量化矩阵乘（`quant_batch_matmul_mx.py`）
matmul/quant_matmul	`npu_quant_matmul()`	per-tensor（TT）量化矩阵乘，支持 HiFloat8 / INT8 / FP8（`quant_batch_matmul_hif8_tt.py`）
matmul/quant_matmul	`matmul_mix_quant()`	MXA8W4：MXFP8 激活 × MXFP4 权重（`quant_batch_matmul_mxa8w4.py`）
grouped_matmul	`group_matmul()`	分组矩阵乘，支持 M 轴 / K 轴分组
flash_attn	`flash_attn()`、`flash_attn_metadata()`	Flash Attention，支持 GQA、变长序列与分页 KV Cache
flash_attn_fp8_fullquant	`flash_attn_fp8_fullquant()`	FP8 全量化 Attention，支持 GQA、分页 KV Cache 与 causal mask
sparse_flash_attention	`sparse_flash_attention()`	Sparse Flash Attention（SFA），按 `sparse_indices` 选取参与计算的 KV token
qwen_sparse_attn	`qwen_sparse_attn()`	固定 `block_size=128` 的分块稀疏注意力，含 AICPU metadata 规划
flash_mla_with_kvcache	`flash_mla_with_kvcache()`、`flash_mla_with_kvcache_metadata()`	Flash MLA 推理，Q 为 TND、KV 走分页 cache
quant_block_sparse_attn	`quant_block_sparse_attn()`、`quant_block_sparse_attn_metadata()`	FP8 量化块稀疏注意力，面向分页 prefill
mixed_quant_sparse_flash_mla	`mixed_quant_sparse_flash_mla()`	混合量化稀疏 MLA：BF16 Q + FP8 原始 KV + 可选 FP4 压缩 KV
mixed_quant_sparse_flash_mla_metadata	`mixed_quant_sparse_flash_mla_metadata()`	MQSMLA 的 AICPU 分核调度，输出整行 / FlashDecode 分片计划
flash_kda	`flash_kda()`、`flash_kda_metadata()`	Kimi Delta Attention prefill 融合算子，含 AICPU 调度 metadata
fused_recurrent_kda_replayssm	`fused_recurrent_kda_replayssm()`	独立 ReplaySSM verify，读取 FP32 checkpoint，写出 FP16 U/K 和 FP32 decay；支持 BSND/BNSD/TND 和 8/24/28/32 控核
attn_prologue	`attn_prologue()`	MXFP8 的 QA/KV 投影、RMSNorm、QR 量化、QB 投影、RoPE 与 KV cache 写回
attn_epilogue	`attn_epilogue()`	Attention 输出投影：inverse RoPE、两次 MXFP8 激活量化与 WoA/WoB 两级矩阵乘；BF16 输入/输出，E4M3FN 权重与 E8M0 scale
indexer_prologue_qw	`indexer_prologue_qw()`	MXFP8 Q GEMM、尾部 RoPE、MXFP4 量化，以及 BF16 W GEMM
indexer_prologue_k	`indexer_prologue_k()`	Indexer K 路前处理：BF16 投影、RMSNorm、RoPE、MXFP4 量化与分页 cache 写入
qsa_indexer	`qsa_indexer()`、`qsa_indexer_metadata()`	压缩 Key 稀疏索引，选高分压缩块并展开为 token 索引
stem_indexer	`stem_indexer()`、`stem_indexer_metadata()`	块级特征打分 + 动态 TopK 选择 Key Block
quant_lightning_indexer	`quant_lightning_indexer()`	MXFP4 Lightning Indexer（QLI），遍历上下文选 TopK token
quant_lightning_indexer_metadata	`quant_lightning_indexer_metadata()`	QLI 配套的 AICPU 调度算子
quant_sparse_lightning_indexer	`quant_sparse_lightning_indexer()`	MXFP4 Sparse Lightning Indexer（QSLI），仅访问候选 Key 块
quant_sparse_lightning_indexer_metadata	`quant_sparse_lightning_indexer_metadata()`	QSLI 配套的 AICPU 调度算子
kv_compress_epilog	`kv_compress_epilog()`	KV Cache 量化压缩与按槽位原地更新
reduce_scatter	`AivReduceScatterSoleMesh`	HIF8 多卡 ReduceScatter，支持 sum/max/min 与 UBMem/URMA
dynamic_mx_quant_sr	`dynamic_mx_quant_sr()`	block32 MXFP8 随机舍入量化，FP32/BF16 → E4M3FN/E5M2；依赖 CANNBot-DSL 0.8.0，环境要求见样例文档
rms_norm	`rms_norm()`	RMSNorm，`y = x · rstd · γ`
engram_gate	`engram_gate()`	Engram 残差门：双路 RMS、加权点积与 signed-sqrt sigmoid 门控
pointnet_sa	`pointnet_sa()`	PointNet++ SA 层的 shared MLP + max-pool
voxel_conv	`voxel_conv()`	VoxelNet Convolutional Middle Layers
fused_norm_rope_attn_rope_cast_fwd	`fused_norm_rope_attn_rope_cast_fwd()`	Q RMSNorm、Q RoPE、Sparse Attention、O 共轭 RoPE 与 MXFP8 量化（含 SF 打包）的融合算子，其中 Q RMSNorm 为可选融合

📚 文档导航

接口文档的入口是 API 参考；算子样例、测试与仓库结构不另设站点页面，直接阅读仓内文件。

内容	说明	入口
API 总览	按语言、数据结构、操作、执行、编译发布和 AI CPU 组织的 API 参考	docs/zh/api/index.md
语言 API	装饰器、控制流、VF、编译期表达式和 Kernel 启动	docs/zh/api/language/index.md
数据结构	数据类型、Tensor、布局、片上存储、Buffer、RingBuffer 和数组	docs/zh/api/data-structures/index.md
操作 API	数据搬运、标量、矢量、矩阵、原子和张量操作	docs/zh/api/compute/index.md
设备执行	设备执行上下文、流水类型、调试和已公开的同步接口	docs/zh/api/execution/index.md
编译与发布	Host 参数规格、JIT/AOT 编译、Native 发布和 ACLNN 工程接入	docs/zh/api/compilation/index.md
AI CPU API	AI CPU Kernel 的声明、参数、编译、验证和启动	docs/zh/api/aicpu/index.md
算子样例	各样例目录下的 README：接口签名、数据类型约束与运行方式	samples/
精度测试	与样例一一对应的 pytest 用例	test/
Native 打包	把 `samples/` 下的算子编译打包为 Native wheel	net/native_package/README.md
更新日志	版本与算子变更记录	CHANGELOG.md

📂 目录结构

cannbot-dsl/
├── docs/          # 文档站源码（VitePress）
├── samples/       # 算子样例及各自的 README
├── test/          # 与样例对应的测试代码
├── figures/       # 样例 README 使用的性能对比图
├── net/           # Native 算子编译与 wheel 打包
└── scripts/       # CI 与合规检查脚本