背景与目标

DeepSeek V3.2 的 DSA 需要接入 CANN 提供的 Indexer、Sparse Attention、Sparse KL，以及 dense warm-up 所需的 LSE/KL 算子。接入的约束是：V3.2 能明确选择 CANN 路径，同时不改变已有 Omni 用户的调用方式、默认行为和算子实现。

兼容性问题：当前接入的DSA相关算子的CANN版本是9.2.0；hp上的omni在pangu上有使用，CANN版本是9.1.0；

总体方案

保留现有 dsa_indexer、dsa_sparse_attention、dsa_kl_loss 三个对外接口，在接口上增加仅限关键字的 backend 参数，默认值为 omni：

dsa_indexer(..., backend="omni")          # 默认保持原 Omni 行为
dsa_sparse_attention(..., backend="omni")
dsa_kl_loss(..., backend="omni")

dsa_indexer(..., backend="cann")          # 显式选择 CANN
dsa_sparse_attention(..., backend="cann")
dsa_kl_loss(..., backend="cann")

Omni 分支继续调用原有 torch.ops.custom.* 算子，不重写、不替换 Omni 的算子实现。CANN 分支调用 cann_ops_transformer 的 Torch API。后端依赖在选中相应分支时加载，普通 Omni 调用方不需要主动传参或改造调用代码。DeepSeek V3.2 在普通前向和 CP 路径中显式传 backend="cann"；其他沿用默认值的调用仍走 Omni。

dense warm-up 在 Omni 路径中没有对应算子，因此提供独立的 CANN dense 接口，而不是给 Omni 接口强行增加一个不存在的实现：

• dsa_dense_indexer_lse(index_query, index_key, head_weights, *, query_lengths=None, key_lengths=None, layout="TND")
• dsa_dense_indexer_kl_loss(index_query, index_key, head_weights, query, key, softmax_max, softmax_sum, indexer_lse, *, scale, query_rope=None, key_rope=None, query_lengths=None, key_lengths=None, layout="TND", loss_coefficient=1.0)

后续：等pangu升级CANN版本之后，日落掉onni算子实现，统一改为CANN算子实现。

V3.2 算子调用与数据流

Sparse 阶段：

1. Indexer 投影生成 query/key/weights，由 CANN lightning_indexer 产生因果 TopK 索引。
2. CANN sparse_flash_attention 计算稀疏注意力输出，并返回 softmax maximum/sum。
3. 训练启用辅助 KL 时，将 FA 统计量传给 CANN Sparse KL，取得 indexer 投影需要的梯度和 loss。

Dense warm-up 阶段：

1. 对每个 packed document 调用 FA，得到输出及其 softmax maximum/sum。
2. 通过 dense LSE 算子得到 indexer LSE，再将 FA 统计量和 indexer LSE 传给融合 Dense KL。
3. Python 不构造完整的 teacher 概率矩阵；KL 梯度由融合算子返回并路由到 indexer 投影。

CP： CP wrapper 继续负责已有的 gather/通信和分片边界处理；CANN adapter 接收通信完成后的本地 Tensor，不接管分布式通信。packed TND 的空 query 分段由当前 dense 路径按非空 document 调用 FA 处理。

Layout 与版本适配

公开接口中的 TND 长度使用 cumulative document ends，不含前导 0；CANN package adapter 在需要时转换成 cu_seqlens。BSND 则传每个 batch 的有效长度。调用方不需要根据 CANN API 的 metadata 格式修改模型接口。

Dense LSE 根据设备自动选择版本：当前实现对 A2/A3 使用独立的轻量 V1 Torch binding，对 950/910_95 使用 package V2；不要求调用方设置版本参数。910B 上 dense LSE/KL 使用 V1 binding，不走该平台不支持的 package dense native kernel。V1 binding 只负责调用已安装的 ACLNN，不修改 CANN 源码。

CANN 9.1/9.2 的 Python 包、ACLNN 符号和 native OPP 必须作为一个匹配组合验证。接口核对材料显示 9.2 增加带可选 sinks 的 GradKLLoss V2，同时保留 V1；当前上层接口没有 sinks 参数，因此版本差异不应泄漏给 Omni 调用方或模型层。当前 NPU 结果是 CANN 9.1.0-beta.1 + 从 ops-transformer master 构建的 package + 匹配 native OPP 的验证，不代表原生 9.1 Python wheel 或 CANN 9.2 已通过；CANN 9.2 正式包/OPP 及 950 实机仍需单独验证。仅能 import Python API 不代表对应 kernel 已部署或可执行。

Omni 兼容性边界

• backend 默认是 omni，未显式选择 CANN 的既有调用保持原路径。
• Omni 算子的输入/输出约定与实现不因 CANN 接入而改写。
• V3.2 显式选择 CANN；不要求所有模型或所有 DSA 调用一起迁移。
• Dense API 是新增的 CANN-only 能力，不伪装成原有 Omni 接口。
• 依赖缺失、SoC/OPP 不匹配应在 CANN 分支暴露，不应通过静默切回 Omni 改变用户选择。

已完成的 NPU 验证

910B3，Python 3.10.20、Torch 2.6.0+cpu、torch_npu 2.6.0.post6、CANN 9.1.0-beta.1，使用文档指定的 matching native OPP 和标准 cann_ops_transformer package：

• tests/torch/components/functional/test_dsa_ops.py -k 'TND and not dense'：4 passed，覆盖 TND Indexer（FP16/BF16）、Sparse Attention 和 Sparse KL。
• tests/torch/components/functional/test_dsa_ops.py -k 'TND and dense'：3 passed，覆盖 TND Dense LSE（FP16/BF16）和 Dense KL。
• tests/torch/context_parallel/test_deepseek_v32_dsa_cp.py：1 passed，双卡覆盖 packed TND、dense/sparse 分支、辅助 KL 开关，以及 CP 与单卡参考的输出和梯度对照。

验收建议

1. Omni 默认调用的回归测试通过，并确认没有引入对 cann_ops_transformer 的无条件依赖。
2. V3.2 普通 sparse/dense 和 CP sparse/dense 按预期选择 CANN 分支。
3. BSND 与 TND 的边界、TopK padding、dense KL 梯度及 CP packed boundaries 都覆盖数值对照。
4. 在目标 SoC 上使用匹配的 Python 包、CANN ACLNN 和 native OPP 验证；将 CANN 9.2 与 950 的结论标为待实机验证，直到完成对应测试。