[RFC]: PTA 层Matmul Checksum校验加速（V-ABFT 算子接入）-pytorch-AtomGit

已关闭

[RFC]: PTA 层Matmul Checksum校验加速（V-ABFT 算子接入） #4730

wuyouqi1创建于 17 天前关闭于 11 天前

17 天前创建

状态（Status）： Draft
作者（Authors）： @wuyouqi1
创建日期（Created）： 2026-09-10
更新日期（Updated）： 2026-09-10
相关 Issue/PR： Ascend/op-plugin MR !5853（算子接入）；Ascend/pytorch MR !38861（ASD 分发改造）

1. 概述

1.1 简介

本提案在 torch_npu（PTA）层接入基于 V-ABFT（方差估计自适应门限） 的 GEMM 容错校验算子 _npu_matmul_abft_verify（aclnn 层为 aclnnMatmulAbftVerify），替代 torch_npu.matmul_checksum 现有的纯 PyTorch 校验实现。对外接口 torch_npu.matmul_checksum(a, b, c) -> bool 的签名、返回值与 ASD 自动触发链保持完全不变；内部按"算子可用性 + 输入契约"两级分发：满足条件走 V-ABFT 算子路径，否则自动回退原纯 PyTorch 路径，stock CANN 环境零感知。

核心价值：将 matmul 静默数据损坏（SDC）校验从"经验阈值 + 行和重算"升级为"方差自适应阈值 + 分块校验和位流"，在实测中实现误检率 0、检出率 100%、大中规模 shape 开销从 ~600-880% 降至 ~16-47%。

1.2 动机

大模型长稳训练中，硬件级 SDC（静默数据损坏，如粒子翻转导致的计算错误）不报错、不崩溃，只会悄悄污染 checkpoint。ASD（Ascend Silent Detection）框架通过 NPU_ASD_CONFIG=enable:true,with_checksum:true 在梯度异常后对后续 matmul 附加校验，当前实现存在三个痛点：

误检：原实现用 5×累积舍入误差 经验阈值，且阈值由 C（待校验结果）自身统计推导——C 被污染时阈值跟着偏移。多模态模型特定激活分布下频繁误检，训练被迫中断。
性能：纯 PyTorch 实现组合十几个小 kernel 并重算一次 O(M·K) 矩阵乘，实测平均额外开销 BF16≈883% / FP32≈635%，只能抽样开启，无法全量保护。
精度格式：仅 BF16 做了检出优化；MOE gate 等大量使用 FP32 的场景存在保护缺口。

CANN 侧已交付 V-ABFT 算子（方差自适应阈值由 A、B 局部统计估计，不依赖 C；方阵规模下校验计算量 8a² vs 重算 2a³），本提案完成其 PTA 层接入。

1.3 目标

目标

误检率：高斯分布 + 常见模型训练数据分布下为 0（验收 1）
性能：大模型训练常见 shape 平均额外开销 < 45%（验收 2）
精度格式：FP32（MOE gate 场景）支持（验收 3）
兼容性：torch_npu.matmul_checksum 对外契约零变更；算子符号缺失环境自动回退，行为不劣于现状

非目标（边界）

不改动 ASD 框架触发链（asd.py hook / TCP 状态机 / NPU_ASD_CONFIG 语义）
不覆盖 3D batched matmul（沿用回退路径保护）
不涉及 FP16 公开支持（算子层支持，但无推荐 e_max，公开 API 维持 bf16/fp32）
不涉及图模式（torch.compile / torchair）下的校验

2. 用例分析

场景：LLM 训练中的 matmul SDC 检测

功能点	说明
触发方式	① ASD 自动：梯度 silent fault 触发后，hook 对每次 matmul 结果追加校验；② 手动：用户直接调用 `torch_npu.matmul_checksum(a, b, c)`
覆盖 shape	FFN up/down-proj（[B,4096]×[4096,11008] 等）、attention、decode 微批（M∈{1,4,8}）、MOE gate（[B,4096]×[4096,8]，FP32）
检出粒度	逐（行 × 256 列段）bit 级，压缩为 `comp_row` 位流；API 层归约为单 bool
精度	BF16 / FP32；FP16 算子层支持但无推荐阈值
性能指标	常见 shape 平均额外开销 < 45%（相对裸 matmul）
可靠性	误检率 0（校验不得误报中断训练）；bit 翻转检出率 ≥90%（实测 100%）
兼容性	算子符号缺失（stock CANN）自动回退纯 PyTorch 路径；调用方零改动
约束	2D 输入、a/b/w 同 dtype、c 恒 FP32、ND 格式、连续张量（桥接层负责转换）

3. 方案设计

3.1 总体方案

入口与返回契约不变，内部改为两级分发（torch_npu 侧改动收敛在 torch_npu/asd/checksum.py 单文件）：

torch.matmul hook（ASD 框架，不改）
  → torch_npu.matmul_checksum(a, b, result)          # 公开 API，懒加载导出，签名不变
    → _matmul_checksum(a, b, c)                       # 分发器
       ├─ _abft_path_applicable(a, b, c)              # 契约判断：2D + bf16/fp32 + shape一致 + 非退化
       ├─ _npu_matmul_abft_verify_available()          # 可用性探测（见 3.4）
       ├─ 可用 → _matmul_checksum_abft                 # V-ABFT 算子路径
       └─ 不可用/不适用 → _matmul_checksum_py          # 原纯 PyTorch 实现原样保留

算子路径语义桥接：

c.to(torch.float32)：真实 matmul 输出转型满足算子契约
checksum_weight = 全 1 向量（对应普通分块行校验和），按 (N, dtype, device) 缓存
e_max 按 dtype 取推荐值（bf16=0.001，fp32=2e-5）
a/b/c .contiguous()：容忍非连续输入（如转置权重）
位流→bool：comp_row 布局为 byte[seg*ceil(M/8)+rg]、bit=r%8 LSB-first（seg-major）；M%8≠0 时 padding bit 取值未定义，必须以缓存掩码屏蔽后判 !=255，否则误检

op-plugin 侧（MR !5853）：YAML 注册（custom 段，op_api: all_version，内部算子不 exposed）+ C++ 内核（fail-fast TORCH_CHECK + EXEC_NPU_CMD 运行期 dlsym）+ meta kernel（FakeTensor shape 推断）+ UT/FakeTensor 测试。

前置条件说明：算子符号 aclnnMatmulAbftVerify 当前由 CANN 9.2.0+ 的 vendors 算子包提供（libcust_opapi.so）；stock CANN 主线 libopapi.so 暂不含该符号——这正是回退机制存在的原因。

3.2 技术选型

备选方案	结论
op-plugin 注册 + Python 层探测回退（采纳）	wheel 构建不依赖符号存在（EXEC_NPU_CMD 仅字符串化符号名，运行期 dlsym）；获得 codegen 全链路（binding/schema/UT 基线）；探测失败 ⟺ 不可用，真实调用不吞异常
torch_npu 内直接 dlopen aclnn 符号	最终同样靠运行期 dlsym；缺 codegen 链路，与仓内惯例不一致
op-plugin 条件注册（有符号才注册）	TORCH_LIBRARY 无此机制；同 wheel 在不同 CANN 机器注册表不一致，难排查
C++ `DO_COMPATIBILITY` + 哨兵返回	降级表达式需可区分返回值；对返回位流的算子只能用空张量，与 M=0 合法空输出歧义，污染算子契约
每次调用 try/except	缺符号机器每次付 c10::Error 构造开销；且把算子真 bug 也吞掉静默降级

阈值算法选型：V-ABFT 方差自适应（由 A、B 局部统计估计，不依赖 C）替代原 5×error_total 经验公式（由 C 统计推导，C 被污染时阈值跟着偏）——这是误检归零的根因。

3.3 功能与性能设计

与原实现的行为差异（对直接调用方可见）

维度	原纯 PyTorch 路径	V-ABFT 算子路径
阈值	经验公式（C 统计推导）	方差自适应（A/B 统计推导，不依赖 C）
粒度	整行一个 bool	行×列段 bit 位流（API 层归约）
副作用	疑似时重跑 matmul 覆盖 c（"自愈"）	纯只读校验
二次确认	有（重算比对均值）	无（自适应阈值替代，实测误检 0）
实测开销	BF16≈883% / FP32≈635%	FP32 M≥1024 时 16-34%

性能实测（Atlas A2 / 910B2，CANN 9.2.0 + vendors 包，中位数计时）：

M（K=N=4096）	BF16 开销	FP32 开销
128	286%	119%
512	185%	57%
1024	90%	34% ✅
2048	65%	16% ✅
4096	51%	23% ✅

瓶颈定位：算子本体存在 ~0.2-0.3ms 近固定开销（executor/tiling/16MiB workspace/launch + B 全量读取），M 从 128→4096 几乎不变。验收 shape 集（M≤64 decode 微批，matmul 本身仅 ~0.1ms）在该固定开销下结构性不可达（平均 BF16 309.5% / FP32 186.2%）；FP32 在 prefill 规模（M≥1024）已达标。

3.4 安全隐私与DFX设计

可靠性：误检率 0（三种分布 × 两 dtype × 10 shape × 每组 1000 次，全部 0/1000）；bit 翻转注入检出率 100%（bit1-3，bit0 符号位翻转被注入过滤器拒绝）
兼容性：
- 对外 API torch_npu.matmul_checksum 契约零变更（签名/返回/懒加载导出/ASD hook 全不动）
- 算子以内部算子形式注册（_npu_matmul_abft_verify 前缀下划线，不进 torch_npu.__all__、不暴露公开别名），避免公开 API 兼容性承诺
- 可用性探测一次微型真实调用后缓存（探测输入已知合法，失败 ⟺ 不可用），stock CANN 自动回退
可测试性：算子注册 meta kernel，FakeTensor 模式可做 shape/dtype 推断测试（无 CANN/设备依赖）；真机 UT 按 CANN 版本门控（≥9.2.0）
可维护性：torch_npu 侧改动收敛单文件；comp_row 位流布局/padding 语义是唯一与算子版本耦合的知识点，已用掩码缓存封装并有钉死实验脚本
安全：纯只读校验路径，不再改写用户张量 c（原路径"自愈"副作用移除，见 §5 风险）

3.5 编程与调用设计

3.5.1 编程模型基本设计

开发环境：Atlas A2/A3 系列 + CANN ≥ 9.2.0（含 vendors 算子包，提供 aclnnMatmulAbftVerify 符号）；PyTorch 2.1+ / torch_npu 配套版本
开发约束：无符号环境（stock CANN）自动回退纯 PyTorch 路径，功能可用但性能/误检指标以算子路径为准；FP16 无推荐 e_max，公开 API 不支持
可验收设计：
- 验收 1（误检率=0）：gaussian/uniform/scaled 三分布 × bf16/fp32 × LLM+MOE gate shape，每组 1000 次
- 验收 2（开销<45%）：常见 shape 集中位数计时，分 dtype 判定
- 检出率（参考项）：bit-flip 故障注入，按 bit 分组统计

3.5.2 接口定义与设计

公开 API 契约不变，详见官方文档 torch_npu.matmul_checksum。内部算子定义如下（不对外暴露，仅供算子开发者参考）。

3.5.2.1 `torch_npu._npu_matmul_abft_verify`（内部算子）

接口描述：V-ABFT GEMM 容错校验。接收矩阵 A、B 与预计算的 matmul 结果 C，做分块行校验和校验，输出行级检错位流。
接口原型：torch_npu._npu_matmul_abft_verify(a, b, c, checksum_weight, *, e_max=0.001) -> Tensor
输入/输出参数：

参数名称	输入/输出	类型	描述	取值范围
a	输入	Tensor	矩阵乘输入 A，[M, K]，ND，连续	fp16/bf16/fp32
b	输入	Tensor	矩阵乘输入 B，[K, N]，ND，连续，与 a 同 dtype	fp16/bf16/fp32
c	输入	Tensor	预计算结果 C=A×B，[M, N]，ND，连续	恒 float32
checksum_weight	输入	Tensor	行校验和编码向量 r，[N]，与 a 同 dtype；全 1 = 普通分块行校验和	同 a dtype
e_max	输入	float	误差阈值系数；bf16 推荐 0.001（默认），fp32 推荐 0.00002；调大可消误报，调小增检出	≥ 0
comp_row	输出	Tensor(uint8)	行检错位流 [ceil(M/8)×ceil(N/256)]，每 bit 一行一段结果，1 正确 0 错误；布局 seg-major、LSB-first，padding bit 未定义	—

异常处理：维度非 2D、shape 不一致、dtype 组合非法、e_max<0、非连续、非 ND 格式 → RuntimeError（fail-fast）
约束说明：不支持图模式；确定性实现
变更说明：新增（op-plugin MR !5853）；曾以 npu_matmul_abft_verify 公开名评审，按社区意见改为内部算子（_ 前缀 + 不 exposed），对外统一走 torch_npu.matmul_checksum
调用参考代码：

import torch
import torch_npu

a = torch.randn(64, 128, dtype=torch.bfloat16, device="npu")
b = torch.randn(128, 256, dtype=torch.bfloat16, device="npu")
c = torch.matmul(a, b)                      # 待校验的 matmul 结果

# 推荐用法：公开 API（自动分发 + 自动回退）
flag = torch_npu.matmul_checksum(a, b, c)  # bool 标量张量，True = 检出异常

3.5.3 编程手册设计

在既有 docs/zh/custom_APIs/torch_npu/torch_npu-matmul_checksum.md 中补充：V-ABFT 路径生效条件（CANN ≥ 9.2.0 + vendors 包）、回退行为说明、性能适用规模建议（M≥512 走算子路径收益显著）、NPU_ASD_CONFIG 联动的 dtype 支持说明（bf16/fp32）。内部算子不出公开文档。

4. 测试设计

层级	用例	覆盖
算子 UT（op-plugin 仓，21 例，CANN≥9.2.0 门控）	`test_npu_matmul_abft_verify.py`	三 dtype 基础、e_max 默认/显式、ceil 边界（M%8、N%256）、大 shape（1024³）、检错（破坏 C 前 8 行 → 字节 0 全 0）、输出 shape/dtype、11 个负向参数校验（与内核 TORCH_CHECK 一一对应）
Meta/FakeTensor（2 例，无 CANN 依赖）	`test_fake_tensor.py`	shape/dtype/device 推断，多 dtype、非整除 shape
分发功能（自测套件，20+ 例）	`test_dispatch_abft.py`	双路径分发、探测缓存、强制回退功能不变、3D/dtype 不匹配回退、非连续 b 内部连续化、CPU 张量 TypeError、缓存生效
验收 1 误检率	`test_false_positive.py`	gaussian/uniform/scaled × bf16/fp32 × 6 LLM shape + 4 MOE gate shape，每组 1000 次 → 全部 0/1000
检出率（参考）	`test_detection_rate.py`	bit-flip 注入（bit1-3 × 两 dtype）→ 100% 检出
验收 2 性能	`test_performance.py` + `scaling_bench.py`/`profile_checksum.py`	中位数计时；开销随 M 的 scaling；分段开销分解（cast/算子/位流转换）

断言设计要点：利用位流语义做与 bit 序无关的断言——正确 C → 完整字节全 255；破坏行 → 对应字节全 0；padding 字节仅断言有效位（取值未定义）；参考 C 由 float64 金标准矩阵乘生成，避免参考值自身舍入误差造成假失败。

5. 缺点和风险

风险	说明	应对
小 M 性能未达标	算子 ~0.2-0.3ms 近固定开销使 decode 微批（M≤64）shape 平均开销 BF16 309%/FP32 186%，超 45% 阈值，属结构性不可达	三个方向：① CANN 降固定开销（executor/tiling/workspace）② 算子 API 扩展支持 B^r 及统计量预计算（训练中 B 多为固定权重）③ 验收口径按 prefill 规模（M≥512/1024）或绝对开销重定。需社区决策（见 §7）
"自愈"行为差异	原 py 路径疑似时会重跑 matmul 覆盖 c（副作用自愈）；算子路径纯只读。依赖该行为的上层代码切路径后此行为消失	ASD hook（结果仅 logical_or_ 累积）不受影响；文档明示差异；如需保留可在分发层按需补
comp_row 布局版本耦合	位流布局/padding 语义属算子实现细节，若 CANN 侧变更会破坏掩码逻辑	布局已实测钉死并有探针脚本；建议 CANN 侧将布局写入接口契约
upstream 门禁冲突	upstream 已合入 bf16-only 校验门禁（e96f560923），与 FP32 支持目标矛盾	MR !38861 按特性意图移除该限制，设备校验保留；合入说明中已注明分歧
符号依赖	算子符号在 vendors 包而非主线 libopapi.so，环境差异化	探测回退机制保证 stock CANN 行为不劣于现状；建议推动符号进主线
环境回退导致指标口径漂移	无符号环境走 py 路径，性能/误检指标不代表算子路径	测试报告须注明实际生效路径（探测结果可查询）

6. 现有技术

经典 ABFT（Huang & Abraham, 1984）：算法级容错，用校验编码（行/列校验和）在线检测矩阵计算错误。V-ABFT 是其在 GEMM 场景的工程化演进：分块校验 + 由输入统计推导的方差自适应阈值，替代固定经验阈值。
重算比对（本提案前实现）：重算期望校验和比对，O(M·K) 起步且阈值依赖待校验结果自身，误检/性能双输。
硬件检错（ECC/RAS/内存_scrub）：覆盖存储与传输位翻转，但不覆盖计算过程中的静默错误；ASD + checksum 与其互补。
GPU 侧同类实践：训练长稳场景的在线 checksum（如各云厂商的 silent data corruption detection 方案）多为抽样重算或 loss 异常监测，bit 级 matmul 校验和方案在本领域属较深度的算子级集成。

7. 未解决问题

验收 2 口径：小 M（decode 微批）在算子固定开销下结构性不可达——按 prefill 规模重定 shape 集、改绝对开销指标、还是等 CANN 降开销？需社区决策。
B^r 预计算接口：是否推动 CANN 扩展算子 API 支持权重侧统计量预计算（训练场景 B 固定，可省去全量读 B 的大头开销）？
FP32 自动注入：ASD hook（asd.py）的 dtype 门控从 bf16 集合扩至 {bf16, fp32} 的收尾改动是否随本 MR 合入（当前 FP32 需手动调用 API 保护）。
算子符号进主线：aclnnMatmulAbftVerify 是否/何时从 vendors 包进入主线 libopapi.so（影响默认环境的算子路径命中率）。
FP16 公开支持：算子层已支持 fp16，是否补充推荐 e_max 后放开公开 API。

附录

参考资料链接
- op-plugin MR !5853：算子注册 + C++ 内核 + meta + UT（https://gitcode.com/Ascend/op-plugin/pulls/5853）
- pytorch MR !38861：ASD checksum 分发改造（https://gitcode.com/Ascend/pytorch/pulls/38861）
- V-ABFT 自适应阈值算法与推导（CANN 侧交付文档）
术语表
- SDC：Silent Data Corruption，静默数据损坏
- ASD：Ascend Silent Detection，昇腾静默故障检测框架（NPU_ASD_CONFIG）
- V-ABFT：Variance-based ABFT，方差估计自适应门限的算法级容错校验
- comp_row：行方向检错位流输出，uint8，每 bit 一行一段检错结果
文档更新计划
- torch_npu.matmul_checksum API 文档：补充 V-ABFT 路径生效条件与回退说明（随 MR !38861）
- NPU_ASD_CONFIG 环境变量文档：checksum 联动 dtype 支持说明更新（bf16 → bf16/fp32）