ops-transformer FlashAttentionScoreGrad 算子 Pytest 验证框架：CPU Golden、NPU 直调与精度对比实战

ops-transformer FlashAttentionScoreGrad 算子 Pytest 验证框架CPU Golden、NPU 直调与精度对比实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer在 CANN 的 ops-transformer 算子库中flash_attention_score_grad是用于训练场景的注意力反向算子其功能正确性的最终验证依赖一套位于 tests/pytest 目录下的 Pytest 测试框架。该框架的核心思路是CPU 侧以纯 PyTorch 复现算子数学语义生成 Golden 数据NPU 侧通过 TorchNPU 直调npu_fusion_attention_grad_v2接口获取实际输出再按相对误差占比与绝对偏差双重阈值做精度对比。读懂这套框架你就能掌握一个 NPU 算子端到端功能验证的完整工程方法如何组织用例、如何构造 FP8 量化与 Dropout 场景下的标杆数据、如何比对两侧结果并定位偏差。一、框架定位与测试体系flash_attention_score_grad算子实现了注意力的反向传播其正向计算公式为pseType1 时Y Dropout(Softmax(Mask((QK^T pse)/sqrt(d)), atten_mask), keep_prob)VpseType≠1 时Y Dropout(Softmax(Mask(QK^T/sqrt(d) pse), atten_mask), keep_prob)V对应的反向公式为dV P^TdY、dQ (dS * K)/sqrt(d)、dK (dS^T * Q)/sqrt(d)详见 算子 README。从源码结构看该算子的测试分为两层C 单元层tests/ut包含 aclnn 接口级测试test_aclnn_flash_attention_score_grad.cpp与 tiling 层测试验证 workspace 计算与 tiling 逻辑Pytest 精度层tests/pytest本文主体验证算子在真实 NPU 硬件上的数值精度。两者互补C UT 验证接口与切分逻辑Pytest 框架验证端到端数值行为。二、目录结构与文件职责Pytest 框架由 5 个 Python 文件组成见 tests/pytest/README.mdpytest/ - test_case.py # 测试用例集 - test_main.py # 执行主程序 - common.py # 公共定义文件 - cpu_impl.py # cpu实现 - npu_impl.py # npu实现各文件职责与关键实现对应关系如下文件职责关键实现common.py公共常量与上下文数据类CalculusContext数据类、日志工具logtest_case.py用例参数定义全局字典TestCases6 组用例cpu_impl.pyCPU Golden 数据生成attentionScoreWithGradBNSD/BSND/BSH/SBH 布局、run_unpadTND 布局、FP8 块量化、Philox 随机数npu_impl.pyNPU 算子直调PTAtest调用torch_npu.npu_fusion_attention_grad_v2test_main.py测试入口与精度比对checkResult、test_npu_flash_atten_grad其中 common.py 定义了框架级全局配置gtype torch.float32 # Golden 计算统一使用 fp32 PER_BLOCK_SIZE 128 # FP8 块量化的块大小S 方向 PER_VCORE_BLOCK_SIZE 64 # V 核分裂场景下的块大小 EPSILON 1e-8 # 量化 scale 计算的下限保护 DEVICE_ID 0 # 目标 NPU 设备号 pta_mode only_grad # 测试模式only_grad 或 auto_gradCalculusContext 是贯穿两侧的数据桥它是一个 dataclass持有输入用例字典input_case以及cu_seqlens_q_npu、q_npu、k_npu、dx_npu、x_max_npu、x_sum_npu、dscale_q_npu等字段。CPU 侧生成输入后写入这些字段NPU 侧直接从中取数上卡保证了两侧输入张量严格一致——这是精度对比有效性的前提。log函数common.py会同时打印到终端并追加写入run_log.txt便于事后回溯每个用例的比对结论。三、环境配置与运行方式tests/pytest/README.md 给出的前置要求与运行命令确认 TorchNPU 为最新版本NPU 侧直调接口npu_fusion_attention_grad_v2依赖 TorchNPU 提供source CANN 包环境变量支持 custom 包调用即自行编译的算子包。在 pytest 文件夹路径下执行pytest -s-s参数用于放行 print 输出使各用例的进度与比对日志直接可见。此外从 cpu_impl.py 与 npu_impl.py 的 import 语句可以确认运行环境还需要torch、torch_npu、numpy、einops、ml_dtypes等依赖。四、测试用例的组织方式用例全部集中在 test_case.py 的全局字典TestCases中每个用例是一个参数字典。当前仓库内置 6 组用例覆盖 GQA、MLA、FP8、FP32 与 TND 变长等场景用例名BN1/N2S1/S2D/D_Vdtypeout_dtypeinput_layout特点GQA_0112/1128/128177/177fp16fp16BNSD非对齐 Drope0MLA_0212/1128/128192/128fp16fp16BNSDMLA 形态D≠D_Vrope1Q/K 后 128 维做 RoPEFP8E5_0322/1125/125128/128fp8_e5m2fp16BSNDFP8 块量化FP8E4_0422/1129/125128/128fp8_e4m3fnbp16*SBHFP8 非对齐 SFP32_0598/280/80167/167fp32fp32BSNDpse_type3Alibi、atten_mask 为 uint8 的 B1SSTND_06416/16287/287128/128fp16fp16TND变长序列[114,132,14,287]、sparse_mode3、dropout0.9* 该用例的out_dtype字段在 test_case.py 中写作bp16而 npu_impl.py 的getOutDtype只映射fp16/bf16/fp32从源码结构看该字段取值需与映射函数保持一致。每个用例字典包含的关键字段可分四类形状类Bbatch、N1/N2Q 头数 / KV 头数其比值 G 即 GQA 分组数、S1/S2Q/KV 序列长、D/D_VK/Q 维与 V 维、TND 布局下另有actual_seq_qlen/actual_seq_kvlen变长列表数据类dtypefp16/bf16/fp32/fp8_e5m2/fp8_e4m3fn、out_dtype、seed、offset语义类sparse_mode稀疏/掩码模式、pre_tockens_input/next_tockens_input带 attention window 时的窗口边界、drop_out_possibilitykeep_prob、pse_type/pse_shape位置编码类型与布局、atten_mask_dtype/atten_mask_shape掩码类型与形状如SS/B1SS/BNSS、rope是否分离 RoPE 分量布局类input_layout支持BNSD、BSND、BSH、SBH、TND五种。新增用例只需在TestCases中追加一个字典test_main.py 的测试入口会自动遍历所有键值对执行无需改动其他文件。五、CPU 侧Golden 数据如何生成cpu_impl.py 是框架中篇幅最大的文件核心入口有两个由布局决定走哪一条见 test_main.pyfunc run_unpad if test_data[input_layout] TND else attentionScoreWithGrad5.1 输入生成DataGenDataGen 按[B, Nkv, G, S, D]的“分头展开”形状随机生成 Q/K/V/dX其中 K/V 的 G 维为 1供 GQA 广播并同步生成三类辅助数据位置编码 psepse_type in (0,1)时支持BN1S/1NSS/BNSS随机值与BNHS/1NHS的 Alibi 风格逐头斜率get_slopes按2^(-8/n)几何级数生成pse_type in (2,3)时由 get_all_alibi 生成 Alibi 偏置且pse_type3时距离项取平方根注意力掩码get_atten_mask 按sparse_mode生成——mode 0 用pre_tocken/next_tocken构造窗口带、mode 2/3 生成上/下三角因果掩码、mode 5/6 生成带prefix的掩码掩码统一以 uint8 的 0/1 表示FP8 量化 scaledtype为 fp8 时get_block_quant_scale 按 128 大小的块PER_BLOCK_SIZE对 S 方向分块取块内绝对值最大值计算scale FP8_MAX / max_vale5m2 的 FP8_MAX 为 57344e4m3fn 为 448并支持 V 核 64 元素分裂PER_VCORE_BLOCK_SIZE与尾块处理。5.2 前向/反向数学复现Golden 的数学核心是 tforward 与 tbackward它们用 fp32gtype逐张量复现算子语义与算子 README 中的公式一一对应# pseType1 时先 add 再 mul否则先 mul 再 add与算子语义严格对齐 if pse_type 1: qk torch.matmul(q, k.permute(0, 1, 2, 4, 3)).add(pse).mul(scale) else: qk torch.matmul(q, k.permute(0, 1, 2, 4, 3)).mul(scale).add(pse) qk qk.masked_fill(atten_mask.bool(), valuetorch.tensor(-40000.0)) softmax_res, x_max, x_sum tsoftmax(qk) drop_res softmax_res * drop_mask * (1.0 / keep_prob) y torch.matmul(drop_res, v)反向部分按 FlashAttention 反向的推导实现tsoftmax_graddrop_res softmax_res.mul(drop_mask).mul(1.0 / keep_prob) dv torch.matmul(drop_res.permute(0, 1, 2, 4, 3), dx) dp torch.matmul(dx, v.permute(0, 1, 2, 4, 3)) dp_drop dp * drop_mask * (1.0 / keep_prob) softmax_grad_res tsoftmax_grad(dp_drop, softmax_res) * scale # (dp - (dp*P).sum) * P dq torch.matmul(softmax_grad_res, k) dk torch.matmul(softmax_grad_res.permute(0, 1, 2, 4, 3), q) dk torch.sum(dk, dim2, keepdimTrue) # GQA组内求和回 Nkv dv torch.sum(dv, dim2, keepdimTrue)注意tsoftmax会额外返回x_max/x_sum——这正是反向算子必需的中间量对应算子接口的softmax_max/softmax_sumattentionScoreWithGrad 将它们广播为[B, Nq, S, 8]后写入ctx.x_max_npu/ctx.x_sum_npu供 NPU 侧消费。FP8 用例还完整模拟了量化-反量化链路block_quant_or_dequant_single 按 128 块乘 scale 模拟量化chunked_matmul_with_quant 等分块 matmul 函数模拟硬件按块反量化累加的行为使 Golden 与 NPU 硬件的数值行为处于同一“模拟层次”。5.3 Dropout 掩码与硬件同源随机数训练算子的 Dropout 掩码必须与硬件生成的完全一致否则对比无意义。cpu_impl.py 在 CPU 上纯 Python 实现了Philox4_32随机数生成器philox_random/philox4_32/philox4_bumpkeyget_drop_mask 用用例中的seed拆成高/低 32 位作为 key、offset/16作为 counter 起点逐字节比较philox_uint8 keep_prob*255生成 0/1 掩码。这与 NPU 侧seed/offset参数见 5.4 节直调代码构成同一随机流保证两侧掩码逐位相同。所有内置用例drop_out_possibility1时掩码退化为标量 1不做丢弃TND_06 用 0.9 覆盖真实丢弃路径。5.4 输出整形CPU 侧最后按用例的input_layout用einops.rearrange把[B, Nkv, G, S, D]的中间表示重排为算子实际输入布局BNSD 合并 GQA 头、BSH/SBH 合并头维、BSND/TND 保留头维TND 布局还通过 unpad_input 把 padding 掉的变长序列压回紧凑[S1, N, D]并计算cu_seqlens前缀和写入上下文。RoPE 场景rope1则将 Q/K 的[:128]与[128:]分量拆开分别作为query/query_rope输入。六、NPU 侧TorchNPU 直调实现npu_impl.py 的 PTAtest 从CalculusContext取出 CPU 生成的全部输入搬到npu:{DEVICE_ID}后直调 TorchNPU 融合接口。当前 common.py 固定pta_mode only_grad走如下调用npu_impl.pydq, dk, dv, *_ torch_npu.npu_fusion_attention_grad_v2( q, k, v, dx, Nq, psepse1, padding_maskNone, atten_maskatten_masks, softmax_maxx_max.float().to(device), softmax_sumx_sum.float().to(device), softmax_inNone, attention_inout.to(pttype).to(device), scale_valuescalar, keep_probkeep_prob, input_layoutinput_layout, pre_tokenspre_tockens, next_tokensnext_tockens, seedseed, offsetoffset, numelsB * Nq * S * Skv, inner_precise0, prefixprefix, sparse_modectx.input_case[sparse_mode], pse_typepse_type)几个值得注意的实现细节正向结果作为反向输入attention_in传的是 CPU 侧tforward得到的输出outsoftmax_max/softmax_sum传 CPU 侧的x_max/x_sum——即反向算子的全部输入都来自 Golden 前向隔离了正向误差使精度对比只检验反向路径随机流对齐seed/offset与 CPU 侧 Philox 生成器使用同一组参数numels按B*Nq*S*SkvTND 下为实际变长序列乘积之和确定掩码数量FP8 专用参数npu_impl.py 在 fp8 场景额外传入d_scale_q/d_scale_k/d_scale_v/d_scale_dy/d_scale_o即 CPU 侧dscale_* 1/scale布局固定重排为 BNSD与out_dtype0 表示 fp161 表示 bf16见 getOutDtypeZeroOrOne输入张量先转换为torch.float8_e4m3fn/torch.float8_e5m2再上卡auto_grad 备选路径npu_impl.py 还实现了auto_grad模式——对 q/k/v 设置requires_gradTrue调用正向npu_fusion_attentionrope 场景用npu_fusion_attention_v2后npu_out.backward(dx)由框架自动求导。该模式在代码注释中明确说明FP8 场景暂不支持且当前配置下默认不启用。七、精度比对标准比对逻辑集中在 test_main.py 的checkResult函数中逐输出张量dq/dk/dvrope 场景另有dq_rope/dk_rope执行Shape 校验NPU 结果与 Golden 形状不一致直接判失败逐元素阈值偏差门限取max(|a|, |b|) * 0.005与2.5e-5的较大者ratio_threshold 0.005test_main.py统计超出门限的元素个数与占比绝对偏差护栏若max(diff) 0.1无论占比多少都判为用例执行失败error: ... 用例执行失败否则仅记录info级别的占比统计。每个用例执行完毕后终端与run_log.txt中会留下diff_max/diff_sum及占比结论例如GQA_01 dq diff_max : 0.00048828 GQA_01 dq diff_sum : 0.001332主测试函数 test_npu_flash_atten_grad 的完整流程即遍历TestCases→ 构造CalculusContext→ 按布局选择run_unpad或attentionScoreWithGrad生成 Golden 与全部 NPU 输入 → 调用PTAtest获取 NPU 输出 → 逐张量checkResult。八、框架设计要点小结从这套框架的实现中可以提炼出 NPU 训练类算子精度验证的几个通用工程实践两侧共享同一个上下文对象CalculusContextCPU 侧生产数据、NPU 侧消费数据从机制上杜绝“两侧输入不一致”这一最常见的对比失真来源Golden 与硬件行为同层模拟FP8 场景不直接对比“理想数学值”而是模拟块量化/反量化、按块反量化累加使 Golden 落在与硬件相同的数值行为层级上随机性显式可控Dropout 掩码用 Philox4_32 在 CPU 复现硬件随机流seed/offset成为用例参数而非隐式状态比对阈值双层设计相对占比阈值容纳浮点噪声绝对偏差上限0.1拦截量级错误用例即数据所有场景差异布局、精度、稀疏模式、变长序列、RoPE都参数化在TestCases字典中扩展新场景零代码侵入。如果你需要验证该算子在某个新布局、新精度或新稀疏模式下的行为最短路径就是在 test_case.py 追加一个用例字典 → 在配置好 TorchNPU 与 CANN 环境变量的机器上进入 tests/pytest 执行pytest -s→ 查看run_log.txt中该用例的diff_max、diff_sum与偏差占比结论。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考