CANN ops-transformer 中 kv_quant_sparse_flash_attention 的 pytest 测试框架实战指南

CANN ops-transformer 中 kv_quant_sparse_flash_attention 的 pytest 测试框架实战指南

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

导读

本文围绕 CANN ops-transformer 仓库中attention/kv_quant_sparse_flash_attention算子的 pytest 测试框架展开，系统讲解其"CPU golden 复现 + TorchNPU 算子直调 + 精度对比"三层验证机制，以及single、batch_save、gen_excel_from_paramset、batch_exec四条主流程的使用方法、参数约束与批量用例管理。读完本文，你将掌握如何配置用例参数、组织 Excel 用例表、批量生成与回放.pt测试用例，并能结合源码理解该框架底层的数据流与精度判定逻辑，直接上手对该算子做功能验证。

框架定位与整体设计

为什么需要专门的测试框架

kv_quant_sparse_flash_attention是一个面向 Sparse Attention 场景的量化注意力算子，它在sparse_flash_attention基础上支持了 Per-Token-Head-Tile-128 量化输入。算子同时涉及多种 layout（BSND/TND/PA_BSND）、多种 KV 量化数据类型（hifloat8/float8_e4m3fn/int8）、PageAttention 块表以及离散稀疏索引，输入组合空间很大，单纯依赖手写脚本难以保证覆盖与可复现性。为此，仓库在 tests/pytest 目录下构建了一套基于 pytest 的完整测试框架。

三层验证机制

框架的核心验证思路可以概括为三条：

• CPU 侧：通过kv_quant_sparse_flash_attention_golden.py中的纯 PyTorch 实现复现算子数学流程，生成 golden 数据；
• NPU 侧：通过 TorchNPU 的torch_npu.npu_kv_quant_sparse_flash_attention接口进行算子直调，获取实际输出；
• 精度对比：通过result_compare_method.py对 CPU 与 NPU 结果做逐元素相对误差与命中率统计，输出Pass/Failed结论及fulfill_percent。

三者的调用关系在 single 测试主程序 中可以清晰看到：先由generate_input_tensors构造输入，再compute_cpu计算 golden，最后qsfa_run_npu拉起 NPU 并调用result_compare_method.check_result完成比对。

四条主流程

模式	作用	触发入口
single	基于 paramset 固定参数直接构造输入并拉起 NPU 单算子执行	test_kv_quant_sparse_flash_attention_single.py
batch_save	从 Excel 读取参数，生成包含 CPU golden 的.pt用例文件	batch/test_kv_quant_sparse_flash_attention_pt_save.py
gen_excel_from_paramset	从 paramset 生成 Excel 用例表	batch/gen_excel_from_paramset.py
batch_exec	从已有.pt文件批量回放执行 NPU 算子并对比精度	test_kv_quant_sparse_flash_attention_batch.py

其中single走"用例即代码"路线，batch_*系列走"Excel 驱动 + pt 中间产物"路线，两者共享同一套参数解析、golden 计算与精度对比逻辑，见 utils.py 与 kv_quant_sparse_flash_attention_golden.py。

当前实现范围与参数约束

框架将算子的参数约束下沉到 check_valid_param.py 统一拦截，保证非法组合在进入 NPU 执行前即被拒绝。以下是框架当前支持与强校验的约束全集（与算子侧 README.md 中的约束说明互相印证）：

参数	支持范围	校验逻辑
layout_query	BSND、TND	非 PA 场景要求layout_query == layout_kv
layout_kv	BSND、TND、PA_BSND	PA_BSND仅在开启 PageAttention 时使用
q_type	torch.float16、torch.bfloat16	其余类型直接抛ValueError
kv_dtype	hifloat8、float8_e4m3fn，兼容None（None走float8_e4m3fn默认生成路径）	见 paramset 注释：当前因 torchair 拦截，float8_e4m3fn只能传None
N1	1/2/4/8/16/32/48/64	注意：Atlas A3/A2 系列不支持 48
N2	仅1
sparse_mode	仅0、3	0为全部计算；3为 rightDownCausal 掩码
sparse_block_size	当前仅1（Token-wise 稀疏）
key_quant_mode/value_quant_mode	仅2（per_tile 量化）
tile_size	仅128
quant_scale_repo_mode	仅1（combine 模式，量化参数与数据混合存放）
attention_mode	0、2；取2时rope_head_dim必须为64	2表示 MLA-absorb 模式
block_size	仅在PA_BSND生效，且必须为 16 的倍数	非 PA 场景传值会被忽略
actual_seq_q/actual_seq_kv	若传入，长度必须等于B	未传时由generate_actual_seq自动生成

校验实现要点：check_valid_param以B/S1/S2/D/K均大于 0 为前置，PA 场景还要求block_num（若给定）大于 0。这些约束与算子侧 README.md 的参数说明一致，例如KV_N仅支持 1、KV_D为 656（nope 512 + rope2 128 + dequant_scale 44）、Q_D仅支持 576（q_nope+q_rope=512+64）。

paramset 用例配置

参数组合以字典形式集中定义在 kv_quant_sparse_flash_attention_paramset.py，当前内置四组典型场景：

• TND_PA：layout_query=TND+layout_kv=PA_BSND，B=22、N1=48、kv 为hifloat8；
• BSND_BSND：非 PA 场景，B=1、S2=4096、N1=64、kv 为None（即float8_e4m3fn默认路径）；
• BSND_PA：layout_query=BSND+layout_kv=PA_BSND，B=18、S1=3、block_size=16；
• TND_TND：纯 TND 非 PA，B=22、N1=48。

每组字典通过ENABLED_PARAMS汇总，框架再经combin_params对列表型参数做笛卡尔积展开，生成最终用例组合。每组参数还支持可选的range_query、range_key、range_query_rope、range_key_rope、range_dequant_scale数据生成区间覆盖。

环境配置

运行前需完成两件事：

1. 完成环境安装与环境变量配置，参见 环境部署；
2. 安装 TorchNPU 包，下载路径请参考 Attention 融合算子 Experimental 使用说明（需及时更换为最新版本）。

同时支持通过 custom 包方式调用算子，NPU 侧入口位于 batch/kv_quant_sparse_flash_attention_process.py，其中call_npu_eager直调torch_npu.npu_kv_quant_sparse_flash_attention，并保留call_npu_graph的图模式编译路径（默认走 eager）。

文件结构与职责

测试框架目录结构如下：

attention/kv_quant_sparse_flash_attention/tests/pytest/ ├── README.md ├── pytest.ini # 创建测试标记（ci / graph） ├── test_run.sh # 统一执行脚本 ├── check_valid_param.py # 参数约束拦截 ├── generate_tensor_data.py # 输入 tensor 随机生成 ├── kv_quant_sparse_flash_attention_golden.py # tensor 转换 / CPU 侧 golden 实现 ├── kv_quant_sparse_flash_attention_paramset.py # 单用例入参配置 ├── result_compare_method.py # 输出精度对比 ├── utils.py # 参数解析 / CPU、NPU 执行入口 ├── test_kv_quant_sparse_flash_attention_single.py # 单用例运行主程序 ├── test_kv_quant_sparse_flash_attention_batch.py # 从 pt 文件批量执行 NPU 测试 └── batch/ ├── kv_quant_sparse_flash_attention_process.py # npu 接口封装 ├── test_kv_quant_sparse_flash_attention_pt_save.py # 从 Excel 批量生成 pt ├── gen_excel_from_paramset.py # 从 paramset 生成 Excel └── excel/ ├── example.xlsx # 示例 Excel 用例文件 └── gen_example_xlsx.py # 生成示例 Excel 的脚本

各文件职责与源码对应关系：

• 参数层：paramset提供用例字典，check_valid_param做合法性校验，utils.convert_param_combination_to_cs_format将用例展开为算子调用所需的shape_input/dtype_input/range_input及属性集合；
• Golden 层：kv_quant_sparse_flash_attention_golden.py内实现 layout 转换（_n_trans_shape_to_bnsd、trans_bnsd_to_layout、trans_tnd_actseq）、hifloat8位级反量化（cvt_hifuint8_to_float）、PA 场景 kv_cache 构建（kv_concat_pa_preprocessing）以及核心的_t_increattention_bnsd逐 token 稀疏注意力计算；
• NPU 层：batch/kv_quant_sparse_flash_attention_process.py完成输入搬移（.npu()）、eager 直调与torch.npu.synchronize()同步；
• 比对层：result_compare_method.check_result输出Pass/Failed与fulfill_percent。

使用方法与命令详解

所有模式统一通过 test_run.sh 入口执行，该脚本内部使用set -euo pipefail严格模式，并对每个模式调用对应的 pytest 目标（-m ci标记）或 Python 脚本。

命令格式

bash test_run.sh <模式> [-E excel_path] [-S sheet] [-P path] [-O output_path]

参数选项说明：

选项	说明	适用模式
-E excel_path	指定 Excel 文件路径，默认./excel/example.xlsx	batch_save
-S sheet	指定 Excel Sheet 页名，默认Sheet1	batch_save / gen_excel_from_paramset
-P path	指定路径（不同模式含义不同，见下表）	single / batch_save / batch_exec / gen_excel_from_paramset

-P在各模式下的含义与默认值：

模式	-P参数含义	默认值
single	paramset 文件名	kv_quant_sparse_flash_attention_paramset
batch_save	pt 文件保存路径	./pt_files/
batch_exec	pt 文件执行路径（目录或单个文件）	./pt_files/
gen_excel_from_paramset	paramset 文件名	kv_quant_sparse_flash_attention_paramset

实现细节：脚本通过PARAMSET_FILE、EXCEL_PATH、EXCEL_SHEET、PT_FILES_PATH等环境变量把选项传给测试模块，-E/-S/-P三个选项顺序可任意，可省略使用默认值。

single：单用例算子调测

手动配置kv_quant_sparse_flash_attention_paramset.py的参数，或使用-P指定其他 paramset 文件，即可完成"CPU golden + NPU + 精度对比"的完整闭环：

bash test_run.sh single # 使用默认 paramset bash test_run.sh single -P my_paramset # 使用指定的 paramset 文件

对应 pytest 入口为test_kv_quant_sparse_flash_attention_single.py，其内部用ThreadPoolExecutor(max_workers=1)串行执行每个用例，任一用例精度Failed时通过pytest.fail报出用例名与fulfill_percent；执行期间每用例结果会追加写入result.xlsx。

batch_save：从 Excel 批量生成 pt 用例

从 Excel 读取参数，生成包含 CPU golden 的.pt用例文件，供后续batch_exec回放：

bash test_run.sh batch_save # 使用默认 Excel 和 Sheet bash test_run.sh batch_save -E ./test.xlsx # 指定 Excel 文件 bash test_run.sh batch_save -E ./test.xlsx -S Sheet1 # 指定 Excel 和 Sheet bash test_run.sh batch_save -E ./test.xlsx -S Sheet1 -P ./output_pt/ # 指定全部参数 bash test_run.sh batch_save -S Sheet1 -E ./test.xlsx # 参数顺序可任意

该模式入口为 batch/test_kv_quant_sparse_flash_attention_pt_save.py，每个用例以Testcase_Name命名的.pt文件保存到目标目录；保存路径通过sed动态写入脚本中的PT_SAVE_PATH。

gen_excel_from_paramset：从 paramset 生成 Excel

当手头只有 paramset 字典、需要产出可供批量编排的 Excel 用例表时使用：

bash test_run.sh gen_excel_from_paramset # 使用默认 paramset bash test_run.sh gen_excel_from_paramset -P my_paramset # 指定 paramset 文件 bash test_run.sh gen_excel_from_paramset -P my_paramset -E ./output/example.xlsx # 指定输出路径 bash test_run.sh gen_excel_from_paramset -P my_paramset -E ./output/example.xlsx -S decode # 指定 Sheet 名

batch_exec：从 pt 文件批量回放执行 NPU 测试

对batch_save生成的.pt用例做批量回放，执行 NPU 算子并与文件内的 CPU golden 对比精度：

bash test_run.sh batch_exec # 执行默认目录下所有 pt 文件 bash test_run.sh batch_exec -P ./pt_files/test.pt # 执行单个 pt 文件 bash test_run.sh batch_exec -P ./custom_pt_dir/ # 执行指定目录下所有 pt 文件

对应入口为 test_kv_quant_sparse_flash_attention_batch.py，其从PT_FILES_PATH收集.pt文件后逐个加载params/input/cpu_output，走与 single 相同的qsfa_run_npu比对链路。

Excel 用例表格式

batch_save读取的 Excel 列名必须与批量框架字段严格一致。以下是一行可直接参考的示例用例（来自 README.md）：

Testcase_Prefix	Testcase_Number	layout_query	layout_kv	q_type	kv_dtype	B	S1	S2	N1	N2	D	K	scale_value	key_quant_mode	value_quant_mode	sparse_block_size	tile_size	rope_head_dim	sparse_mode	attention_mode	quant_scale_repo_mode	block_size	block_num	actual_seq_q	actual_seq_kv
tnd_sample	1	TND	TND	torch.bfloat16	hifloat8	2	8	8	16	1	512	4	0.04166666666666666	2	2	1	128	64	3	2	1	256	[5,8]	[6,8]

字段解析要点（与 utils.py 的load_excel_test_cases/_parse_excel_cell_value对应）：

• q_type支持字符串torch.bfloat16/torch.float16，由STR_MAP_DICT映射为 torch 类型；
• kv_dtype为空、None或float8_e4m3fn时统一走float8_e4m3fn默认生成路径；
• actual_seq_q/actual_seq_kv支持[5,8]这类列表字符串，经ast.literal_eval解析；TND 场景下语义为前缀和；
• block_num为空时，PA 场景由代码按ceil(actual_seq_kv/block_size)逐 batch 累加自动计算；
• 数值单元格统一做归一化（如整型浮点转为int），并支持range_query等数据生成区间列；
• 缺少必需列时框架会pytest.skip并提示缺失列名。

核心实现深度解析

CPU golden 的计算链路

golden 主流程位于 kv_quant_sparse_flash_attention_golden.py 的compute_golden，可分为五个阶段：

1. 输入归一化：全部 tensor 转为 CPU float32 参与计算，同时保存原始 dtype 的 raw bytes 供 kv_cache 拼接使用；
2. 参数装配（_prepare_fa_param）：执行 q/k 拼接（q 由 q_nope 与 q_rope 沿 D 维拼接；k 由 k_nope、k_rope 与 dequant_scale 拼接，对应算子侧Q_D=576、KV_D=656的构造）、TND 前缀和转逐 batch 长度、layout 到 BNSD 的转换；
3. 稀疏索引生成（_generate_sparse_indices）：按sparse_mode计算每个 (batch, head, token) 的有效阈值——sparse_mode=0时阈值为act_seqlen_kv，sparse_mode=3（rightDownCausal）时阈值为act_seqlen_kv - act_seqlen_q + s + 1，并用torch.randperm随机选取不超过sparse_blockcount的 block 索引，尾部用-1填充无效值；
4. kv_cache 构建（_generate_block_table_and_cache）：PA 场景按block_table将 KV 数据写入[block_num, block_size, N, D]的 cache，并做 0 轴非连续（key_stride）支持；非 PA 场景走kv_concat_nopa_preprocessing直接拼接key_cache；
5. 核心计算（_t_increattention_bnsd）：逐 batch、逐 KV head、逐 query token 循环，经gatherKV按稀疏索引收集 K/V 片段，执行Q@K^T、乘以scale_value、softmax、再与 V 相乘，最后把结果还原回原始 layout。

其中对hifloat8的位级反量化（cvt_hifuint8_to_float）按 8bit 的 4-bit 点段 + 指数/尾数字段手工解码，并处理 0、NaN、±inf、±32768/32768 溢出等特殊值；MLA 场景（attention_mode=2）下 V 直接复用 K（v_tensor = k_tensor）。

NPU 直调与两种执行路径

batch/kv_quant_sparse_flash_attention_process.py 提供两条调用路径：

• eager 路径（默认）：call_npu_eager将query_cache/key_cache/value_cache/sparse_indices/block_table等输入.npu()后直接调用torch_npu.npu_kv_quant_sparse_flash_attention，并透传scale_value、sparse_block_size、key_quant_mode、value_quant_mode、layout_query、layout_kv、sparse_mode、attention_mode、quant_scale_repo_mode、tile_size、rope_head_dim、pre_tokens、next_tokens等属性；
• graph 路径（保留）：call_npu_graph通过torchair的CompilerConfig（reduce-overhead模式 +_aclnn_static_shape_kernel静态 shape 编译）对封装好的Network模块做torch.compile，用于图模式编译验证，与算子 README 中"该接口支持图模式"的说明呼应。

精度对比与判定规则

result_compare_method.py 的check_result实现以下判定逻辑：

• 使用np.isclose逐元素判定，默认rtol=0.005、atol=0.000025；当 NPU 输出为bfloat16时自动放宽为rtol=0.0078125、atol=0.0001；float8_e4m3fn/float8_e5m2输出走原始字节比对；
• fulfill_percent = (总元素数 - 不达标元素数) / 总元素数 * 100%，需大于等于(1 - 0.005) * 100% = 99.5%才判定Pass；
• 另设max_diff_hd=10作为最大相对误差上限，超过即判Failed；
• 输出内容包含元素级对照表（ExpectOut/RealOut/FpDiff/RateDiff）、最大相对误差定位（Max-RE line）以及Rtol/Atol/PctThd/PctRlt/Result汇总行；
• 对 CPU golden 中的 inf/nan 溢出元素做统计与展示，equal_nan=True保证 NaN 位置不误判。

结果产物与用例流转

• result.xlsx：记录每个用例的关键信息（layout、dtype、shape、量化模式、数据生成区间等）、执行状态（Pass/Failed）与fulfill_percent。该文件由 utils.save_result 维护，列结构与 example.xlsx 对齐，可直接回灌用于后续批量生成 pt；
• ./pt_files/*.pt：batch_save流程生成的中间测试用例，每个文件包含Testcase_Name、params、input、cpu_output四要素，是batch_exec回放与跨环境复现的载体。

推荐的标准流转链路为：paramset维护用例 →gen_excel_from_paramset产出 Excel → 人工/脚本增删改 Excel 行 →batch_save生成带 golden 的 pt 用例 →batch_exec批量回放验证；单条用例快速排查则直接用single。

与算子实现及单元测试的衔接

• 算子功能定义、shape 推导与 tiling 实现分别位于 op_host/kv_quant_sparse_flash_attention_def.cpp、op_host/kv_quant_sparse_flash_attention_infershape.cpp 与 op_host/kv_quant_sparse_flash_attention_tiling.cpp，算子支持的平台范围（Ascend 950PR/950DT、Atlas A3、Atlas A2）与sparse_block_size、Q_N、key 连续性的平台差异约束详见 算子 README；
• 内核侧实现位于 op_kernel 下，按 arch22 / arch35 分目录维护 kernel、cube（MLA）与 vector 服务代码；
• 与 pytest 框架互补的还有 tests/ut 下的 C++ 单元测试（test_aclnn_kv_quant_sparse_flash_attention.cpp、各 arch 的 tiling 测试与 infershape 测试），用于在 Python 框架之外对 host 侧逻辑做编译期与单测验证；
• 若希望不依赖 pytest 框架快速验证单算子调用，可参考 examples/test_npu_kv_quant_sparse_flash_attention.py 的直调示例。

常见问题与排查建议

• 用例被 skip：Excel 文件不存在或缺少必需列时框架会pytest.skip并打印提示，先检查列名是否与上文清单一致、-E/-S是否指向正确文件；
• TND 场景前缀和校验：actual_seq_q/actual_seq_kv为 TND 前缀和，convert_param_combination_to_cs_format会校验T1 <= B*S1、T2 <= B*S2，且后一个元素必须不小于前一个，否则 golden 侧trans_tnd_actseq会报负长度错误；
• float8_e4m3fn传 None：由于当前 torchair 拦截原因，kv_type传float8_e4m3fn时只能传None，paramset 中有明确注释；
• PA 场景 block_size：必须为 16 的倍数且不小于等于 0，block_num未给时框架按ceil(act_seq_kv / block_size)自动累计，block_table 不足时会报错退出；
• 精度 Failed 定位：查看result.xlsx中对应用例的fulfill_percent，再用batch_exec -P ./pt_files/用例名.pt单点回放，结合打印输出的 Max-RE line 定位最大相对误差元素位置。

总结

kv_quant_sparse_flash_attention的 pytest 测试框架以"CPU golden + TorchNPU 直调 + 精度对比"为支柱，通过single与batch_*两套流程覆盖了从单点调试到批量回归的完整测试生命周期。参数约束、Excel 用例驱动、pt 中间产物、结果回写等设计使其既适合算子开发阶段的快速验证，也适合上线前的规模化回归。结合 test_run.sh、utils.py、kv_quant_sparse_flash_attention_golden.py 与 result_compare_method.py 等源码，开发者可以按需扩展 paramset 场景、补充 Excel 用例行，或调整精度阈值，将该框架复用到后续算子版本的功能验证中。

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