状态（Status）： Reviewing
作者（Authors）： @yucaopanmu
创建日期（Created）： 2026-08-15
更新日期（Updated）： 2026-08-26
相关 Issue/PR： #5737

1. 概述

1.1 简介

本提案将新 CANN 算子 aclnnDropoutV3Grad 接入 op-plugin 中 native_dropout_backward 的 A5（Ascend950）路径，使 A5 上 dropout 反向的数值结果对齐 PyTorch/GPU（H20）。核心改动是：A5 上把 PyTorch 语义的缩放因子 scale 直接透传给新算子，内核以纯乘法链 gradX = gradY * mask * scale 完成计算，消除旧链路中"宿主侧由 scale 换算 p（1 - 1/scale）、算子内部再由 p 还原缩放因子（1/(1-p)）"两次浮点换算引入的精度误差。A5 上使用不支持新算子 aclnnDropoutV3Grad 的 CANN 版本则回退到之前的路径，A2/A3 代际调用路径与行为保持不变，对外 API 与语义不变。

1.2 动机

• 客户场景：字节客户预训练场景明确使用 dropout 算子，要求数值结果与 GPU 完全一致、性能符合预期。该算子精度直接影响字节预训练的商业验收面。
• 当前痛点：现有 NPU 通路中，dropout 反向（native_dropout_backward）在宿主编排层把 scale（PyTorch 语义 scale = 1/(1-p)）换算为 p = 1 - 1/scale 后下发给 aclnnDropoutDoMask，算子在内部再还原缩放因子。1 - 1/scale 与 1/(1-p) 两次浮点换算均存在舍入误差（如 p=0.3 时 1 - 1/1.428571... 无法精确表示），导致反向结果与 GPU 存在 ulp 级偏差；预训练多轮累积后偏差被放大，无法通过客户精度验收。
• 必要性：GPU（H20）侧 PyTorch 原生实现为单步乘法链 gradX = gradY * mask * scale，无中间换算。NPU 要对齐 GPU，必须消除链路中多余的浮点计算。
• 不做此提案的影响：字节预训练精度诉求无法满足，dropout 算子在 A5 上无法支撑客户商用验收。

1.3 目标

目标：

• A5 上 native_dropout_backward 反向数值结果对齐 PyTorch/GPU（H20），通过 UT 测试与 ATK 泛化精度验证。
• 性能不回退（纯乘法链与旧路径计算量等价；特殊值分支可免算子下发）。
• 兼容旧版本的 CANN，A5 上使用不支持新算子 aclnnDropoutV3Grad 的 CANN 版本则回退到之前的路径。
• A2/A3 调用路径与行为保持不变。
• PT 版本：2.7.1+；支持硬件：A5；交付时间：2026-08-27。

非目标（边界说明）：

• 前向 native_dropout 的精度/随机数序列对齐不在本提案范围内（前向路径不变）。
• A2/A3 代际不接入新路径，不改变其数值行为。
• 不涉及对外 API 新增或语义变更（torch.ops.aten.native_dropout_backward 接口不变）。

2. 用例分析

约束与限制：

• mask 必须为 UINT8 packed 位掩码（128bit 对齐、LSB-first，由前向 aclnnDropoutGenMaskV2 产出），与旧路径输入约定一致。
• scale 合法域为 0 或 [1, +inf)，与 PyTorch dropout 语义一致。
• 新路径仅在 GetSocVersion() >= Ascend950 且 CANN 支持新算子 aclnnDropoutV3Grad 时生效。

3. 方案设计

3.1 总体方案

整体思路：在 A5 上把 PyTorch 语义的 scale 原值直达新 CANN 算子 aclnnDropoutV3Grad，内核执行单步乘法链，从根因上消除旧链路两次浮点换算的精度损失。

核心处理流程（op_plugin/ops/opapi/NativeDropoutKernelNpuOpApi.cpp 的 native_dropout_backward）：

输入 grad_output, mask, scale
  │
  ├─ SocVersion >= Ascend950 && check_aclnn_kernel_available("aclnnDropoutV3Grad")？（否 → 旧路径 aclnnDropoutDoMask，A2/A3 行为不变）
  │
  ├─ TORCH_CHECK：scale 必须为 0 或 >= 1，否则报错（与旧路径校验一致）
  ├─ mask.numel() == 0 → 直接返回与 mask 同形状的空 tensor（不下发）
  ├─ p = (scale == 0) ? 1 : 1 - 1/scale（仅用于分支判定，不参与计算）
  │     ├─ p == 0（scale==1，无丢弃）→ 返回 grad_output.clone()
  │     └─ p == 1（scale==0 或 +inf，全部丢弃）→ 返回 zeros
  └─ 其余情形 → EXEC_NPU_CMD(aclnnDropoutV3Grad, grad_output, mask, scale, result)
                 内核语义：gradX = gradY * mask * scale（纯乘法链，与 GPU 一致）

设计要点：

1. 精度对齐：scale 原值透传，内核一次乘法完成缩放，与 PyTorch CUDA 反向实现的计算链路完全一致。
2. 分支前置：p==0、p==1、空 mask 三类退化场景在宿主侧短路，减少无效算子下发，同时保证极端入参下语义与 GPU 一致（GPU 侧同样短路）。注意 p = 1 - 1/scale 在宿主侧仅用于分支判定，p==0/p==1 的判定对舍入误差不敏感（scale==1 时 1-1/1 精确为 0；scale→+inf 时结果精确趋向 1）。
3. 版本兼容：通过 check_aclnn_kernel_available 判定，对旧版本的 CANN 回退到之前的路径。
4. 代际隔离：通过 GetSocVersion() 判定，A2/A3 完整保留 aclnnDropoutDoMask 路径。

3.2 技术选型

3.3 功能与性能设计

功能设计：

• A5 分支新增 SOC 版本判定、新算子可用性检查、参数校验、退化场景短路与 aclnnDropoutV3Grad 下发（约 +31 行）。
• 对外功能与输入输出契约不变：输入 grad_output、mask、scale，输出与 grad_output 同 shape/dtype 的 tensor。
• 数据模型无变更：mask 沿用 UINT8 packed 位掩码格式。

性能设计：

• 新内核为单次乘加链，与旧 aclnnDropoutDoMask 计算量等价，理论无性能回退。
• scale==1/scale==0/空 mask 场景由宿主侧短路，省去算子下发与 kernel 启动开销，性能略有收益。

影响范围： 仅 native_dropout_backward 一个接口、A5 一个平台；native_dropout 前向与 dropout_backward 不受影响。

3.4 安全隐私与DFX设计

• 兼容性：A2/A3 行为逐位不变（路径隔离）；A5 上输出 shape/dtype 契约不变；参数校验语义（TORCH_CHECK 报错文本）与旧路径一致。
• 可维护性：双路径以 SOC 版本与是否支持新算子清晰隔离，代码注释说明分支意图；新增路径与旧路径结构对称，降低维护成本。
• 可测试性：新增 7 个 A5 单测用例（见第 4 节），覆盖正常值、边界值、非法入参与空 mask；旧用例补充设备注解保证代际覆盖。
• 可靠性：新算子可用性检查机制保证 CANN 版本不匹配时自动回退，不会因新算子缺失导致功能不可用。
• 安全隐私：不涉及任何数据采集、存储或传输，无安全隐私影响。

3.5 编程与调用设计

3.5.1 编程模型基本设计

• 开发环境：A5（Ascend950）服务器、PT 2.7.1+、配套 CANN 版本（提供 aclnnDropoutV3Grad）。
• 开发约束：新算子仅在 A5 及后续代际生效；A2/A3 上通过 SOC 分支自动走旧路径。
• 可验收设计：UT 全部通过 + ATK 泛化精度验证（A5 vs GPU(H20)）通过。

3.5.2 接口定义与设计

3.5.2.1 aclnnDropoutV3Grad（新增 CANN 算子接口）

• 接口描述：计算 dropout 反向，内核语义 gradX = gradY * mask * scale，纯乘法链，对齐 GPU 精度。
• 接口原型：aclnnStatus aclnnDropoutV3Grad(const aclTensor* gradOutput, const aclTensor* mask, double scale, aclTensor* result)（以 CANN 发布接口文档为准）。
• 输入/输出参数：

• 返回参数：

• 异常处理：空 mask、scale==1、scale==0 等退化场景由宿主侧短路，不下发该接口；非法 scale 由宿主侧 TORCH_CHECK 抛出 RuntimeError。
• 约束说明：仅 A5（Ascend950）及后续代际调用；输入 mask 必须为 packed UINT8 格式。
• 变更说明：新增接口；对宿主侧 torch.ops.aten.native_dropout_backward 的调用方无感知。
• 调用参考代码：

if (c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950 &&
 	       check_aclnn_kernel_available("aclnnDropoutV3Grad")) {
  // ... 参数校验与退化场景短路 ...
  at::Tensor result = at_npu::native::OpPreparation::apply_tensor_without_format(grad_output);
  EXEC_NPU_CMD(aclnnDropoutV3Grad, grad_output, mask, scale, result);
  return result;
}

3.5.3 编程手册设计

本提案不新增对外 API，torch.ops.aten.native_dropout_backward 调用方式与语义不变，无需新增《编程手册》。建议在已有 native_dropout_backward 相关文档中补充说明：A5 平台反向计算采用 aclnnDropoutV3Grad，数值对齐 GPU。

4. 测试设计

4.1 单元测试（UT）

已在 test/test_base_ops/test_native_dropout_backward.py中添加用例：

测试要点：

1. mask 由测试侧构造 UINT8 packed 位掩码（_packed_bit_mask），并在 CPU 侧展开为 bool 掩码（_expand_bit_mask）计算 golden，保证 golden 语义与 GPU 一致。
2. bf16 用例比对语句为 assertRtolEqual(output_cpu.float().numpy(), output_npu.cpu().float().numpy(), 0.004)，比对前两侧转 float32，避免 bf16 直接比对引入的比对误差。

4.2 泛化精度验证（ATK）

• A5 上运行 ATK 泛化精度验证：同一 dropout 反向用例分别在 NPU 与 GPU（H20）上执行，比对输出数值，要求与 GPU 一致，与算子侧对齐的精度标准为双千双万。
• 覆盖 fp32/fp16/bf16 多种 dtype、不同 p/scale 取值与多种 shape。

4.3 集成/端到端测试

• 字节预训练场景端到端验证：dropout 反向接入真实训练任务，验证数值一致性与训练吞吐。

5. 缺点和风险

实现成本：宿主侧约 +31 行代码、单测约 +100 行，改动较小且集中在 2 个文件。

6. 现有技术

• PyTorch CUDA 实现：native_dropout_backward 在 GPU 上直接以 gradX = gradY * mask * scale 单步乘法链计算，无中间换算——本提案的精度对齐目标即与其计算链一致。
• CANN Dropout 接口演进：CANN 侧已有 aclnnDropoutDoMask（A2/A3 在用）与 aclnnDropoutV3/aclnnDropoutV3Grad 代际接口，本提案复用 CANN 既有演进成果，宿主侧仅做接入。
• 差异点：本提案相较其他算子精度对齐方案，最小化改动面——仅替换 A5 后端调用，不动前向、不动语义、不动 A2/A3。

7. 未解决问题

• 远期是否将 A2/A3 统一迁移至 V3 路径（需评估存量验收影响与收益）。
• aclnnDropoutV3Grad 支持的 dtype 范围是否需进一步扩展（如 int8 量化场景），待 CANN 侧确认。
• 文档更新计划的具体落点（docs/zh 中是否补充 A5 精度说明），待评审确认。

附录

• 参考资料链接
  • 需求描述：A5 代际 Dropout 反向算子数值对齐 PyTorch/GPU（字节预训练场景，PT 2.7.1+）
  • PyTorch dropout 参考：torch/_refs/nn/functional/dropout.py 与 ATen native_dropout_backward CUDA kernel
• 术语表
  • A5 / Ascend950：第五代昇腾 AI 处理器代际标识，c10_npu::SocVersion::Ascend950
  • packed bit mask：UINT8 位掩码，按 128bit 对齐、LSB-first 打包存储，由 aclnnDropoutGenMaskV2 生成
  • scale：PyTorch dropout 反向缩放因子，语义 1/(1-p)
  • ATK：Ascend 测试套件（泛化精度验证，NPU 与 GPU 输出数值对比）
• 文档更新计划
  • 随 RFC 合入后，在 docs/zh 相关算子文档中补充 A5 反向精度说明（无对外 API 变更，不新增手册）

欢迎加入社区，感谢您对社区的贡献 🎉!