fix: 13 算子 register_a1 签名升级（counter + platform，#19)

[Dashener2]

摘要

按 templates/operator/register.py 与 ops/embedding 的 _BACKENDS + _DEVICE_PROBE 范式，升级 13 个算子的 A1 注册签名（关联 #19 第 1 项）：

• 签名统一为 register_a1(dispatch_key="AutogradPrivateUse1", counter: dict | None = None, platform=None)
• 新增 _BACKENDS/_backend_for（platform 显式优先，dispatch_key 推断；覆盖 Autograd/PrivateUse1/AutogradPrivateUse1 → cambricon）与 _DEVICE_PROBE 注册守卫（MLU 可用性探测）
• 10 个可注册算子：命中计数支持 counter["n"]；test/op_level.py 同步传入 counter 并断言
• 3 个不可注册算子（empty/detach_/result_type）同步签名，保持按设计抛 NotImplementedError

自查

• 核心八件齐全：既有算子，八件已在库；本 PR 仅改 register.py / test/op_level.py
• 三层测试绿：13/13 example.py cambricon 三层全绿；真实 vLLM 框架级 tests/framework_level/test_a1_aten_ops.py 1 passed
• 性能对照：未改 kernel；example.py 内置 native 参照，无性能回归
• 黄金绿：goldendata 复用，script/check_accuracy.py 通过
• 包化回归绿：ops/test_packaging.py 顶部 import torch_npu，本环境仅 MLU 无法运行；已做等价验证（13 算子包态导入 + register_a1 可导入通过，见「说明」）
• 多平台变更时：本 PR 未做平台扩展，PLATFORM.md / MERGE.md 按 chore: PR#5 的 13 个算子按新库约定补两项——register 签名升级 + 多平台文档（平台扩展前） #19 约定暂欠
• 署名一致：改动处 Dashener2 2026-10-09-fix + # Dashener2 <dashen272@outlook.com>，与 commit 作者一致

说明

问题：#19 指出 13 算子仍为旧式单参注册签名 register_a1(dispatch_key="PrivateUse1")，缺 counter（op 层拦截计数）与 platform（第二平台显式选择），未对齐模板 / embedding 范式。

思路：

• 签名统一为 (dispatch_key="AutogradPrivateUse1", counter=None, platform=None)。
• _BACKENDS/_backend_for：platform 显式优先，否则按 dispatch_key 推断；保留第二平台扩展位。
• _DEVICE_PROBE 注册守卫：注册时探测 torch.mlu 可用性，跨平台误用即时报错，而非运行时静默错。
• 可注册算子命中计数写 counter["n"]；不可注册算子保持 NotImplementedError。

怎么用：

from ops.type_as import register_a1
register_a1(profile.dispatch_key, counter={"n": 0})   # op 层拦截验证
register_a1(platform="cambricon")                     # 显式平台
验证（v26.08.0 容器 / MLU）：
- 13/13 example.py cambricon 三层全绿；pytest tests/unit 24 passed
- counter / platform / 默认 key 实测通过
- 真实 vLLM A1 框架测试 1 passed
- 13 算子包态导入通过

[TruNcat3]

感谢补充真机验证 🙏 先说明总体结论：signature / counter / platform 这条主线是对的，13 个算子的真机三层和 vLLM 注入结果也很有价值。我提 Request Changes 不是否定这批修复，而是包态调用链还有一个可复现的运行时缺口，需要补上后再合。

1. 已确认完成的部分

这几点我已经对过 PR 和说明，方向正确：

1. 13 个算子都改成了新签名：

   register_a1(
       dispatch_key="AutogradPrivateUse1",
       counter: dict | None = None,
       platform=None,
   )

2. 10 个可注册算子接入了 counter["n"]，op 层测试也断言了计数。

3. empty / detach_ / result_type 没有被强行注册，仍然保留 NotImplementedError，符合原设计。

4. 增加了 platform/dispatch key 校验和 MLU 设备探测。

5. CI 三门已经通过；你们 MLU 环境的 13/13 三层、pytest tests/unit、真实 vLLM 注入结果也说明签名升级本身没有破坏现有脚本态用法。

下面的问题只针对 包态运行时路径。

2. 阻塞问题：包态不只是“能 import register_a1”，还要能执行到本算子 kernel

背景：仓库里现在有两种调用模式

这个仓库的算子需要同时支持两种模式。

模式 A：单算子脚本模式

例如 test/op_level.py 通常会先把算子目录插到 sys.path：

sys.path.insert(0, ops/clone)
import register

在这种模式下，下面这种裸导入可以工作：

from kernel.triton_level import clone_triton

因为 ops/clone 就在 sys.path 里，顶层 kernel 能解析到 ops/clone/kernel。

模式 B：多算子包态模式

也是 #14 包化要支持的场景：

from ops.clone import register_a1

这时 Python 的模块名是：

ops.clone.register

如果 register.py 内部仍然裸写：

from kernel.triton_level import clone_triton

它找的是顶层 kernel，不是 ops.clone.kernel。在多算子同进程里，顶层 kernel 甚至可能已经被别的算子占用，这是 #13/#14 当初要根治的冲突来源。

可复现问题

在 PR 分支、仓库根目录下执行：

import torch
from ops.clone.register import clone_aten
clone_aten(torch.zeros(2))

当前结果：

ModuleNotFoundError: No module named 'kernel'

这说明包态不仅能导出 register_a1，而且已经拿到了 clone_aten；但一执行就失败在实现内部的顶层 kernel 导入。

为什么现有验证没有发现？

因为现有检查分成了几层，每层都没有覆盖到这个位置：

1. ops-packaging CI 只验证：

   import ops.<op>
   assert hasattr(module, "register_a1")

   它会触发 register.py 导入，但不会调用 clone_aten()，所以不会执行函数体里的：

   from kernel.triton_level import clone_triton

2. 13 个算子的 op/framework 测试主要还是单算子脚本模式，sys.path 已经插入了各自目录，所以裸导入能工作。

3. 真实 vLLM 注入大概率也是通过插件/脚本路径导入，这同样可能掩盖包态问题。

因此，这不是和你们验证结果矛盾，而是验证矩阵里还缺了一格：package mode + actual kernel execution。

建议的最小修复

10 个可注册算子的实现导入都改成“包态优先、脚本态 fallback”：

try:
    from .kernel.triton_level import clone_triton
except ImportError:
    from kernel.triton_level import clone_triton

需要覆盖这 10 个文件：

ops/_local_scalar_dense/register.py
ops/clone/register.py
ops/contiguous/register.py
ops/copy_/register.py
ops/detach/register.py
ops/dropout/register.py
ops/empty_like/register.py
ops/empty_strided/register.py
ops/item/register.py
ops/type_as/register.py

empty / detach_ / result_type 不注册，没有这个问题。

建议补的验证

至少在 MLU 真机上选一个代表算子，直接走包态执行路径。最小验证可以是直接调用实现函数，不注册全局 aten op，避免测试进程污染：

import torch
from ops.clone.register import clone_aten

x = torch.zeros(2, device="mlu:0")
y = clone_aten(x)
assert y.shape == x.shape
assert y.data_ptr() != x.data_ptr()

如果希望更接近消费方，也可以在独立子进程里做：

from ops.clone import register_a1

counter = {"n": 0}
lib = register_a1(
    "PrivateUse1",
    counter=counter,
    platform="cambricon",
)
x = torch.zeros(2, device="mlu:0")
y = x.clone()

assert counter["n"] > 0
assert y.data_ptr() != x.data_ptr()

这个测试的关键不是再证明 clone 数值正确，而是证明：从 ops.clone 包入口进来时，能找到并执行 ops.clone.kernel。

3. 关于 ops/test_packaging.py：MLU 环境应该怎么跑？

这里需要区分两个不同层级的“packaging test”。

层级 1：CPU 可跑的包入口扫描

CI 里的 ops-packaging job 做的是这一层：

import ops.<op>
assert hasattr(module, "register_a1")

它不依赖 NPU/MLU 设备，可以在 Ubuntu CPU 环境跑。你们的 PR 里这项已经 PASS，所以对 issue #19 的“包化导出”要求是足够的。

层级 2：ops/test_packaging.py 的 NPU 真机 E2E

当前 ops/test_packaging.py 不是纯 import 测试，而是一个 NPU 场景集成测试，包含：

• import torch_npu
• npu:0 上的 embedding 消费
• sdpa 大小 shape 路由
• NPU Triton/native 路径
• 多算子同进程真实数值验证

所以 MLU 机器跑不了它是预期行为，不需要强行改成 MLU 可跑。

建议的后续整理

以后可以把这个脚本拆成两个模式：

python3 ops/test_packaging.py --imports-only
python3 ops/test_packaging.py --npu-e2e

含义分别是：

• --imports-only：CPU/MLU 都能跑，CI 使用；
• --npu-e2e：保留现在完整的 NPU 多算子消费验证。

这个拆分建议后续单独做，不必塞进本 PR。你们当前 PR 只需要引用已经通过的 ops-packaging CI 即可。

4. MLU A1 key：我的建议和理由

你们问到：ops/embedding 里现在是 Autograd -> mlu、AutogradPrivateUse1 -> ascend，而这批算子把三个 key 都映射到 cambricon；如果希望 MLU 严格使用某个 key，应该怎么对齐。

我的建议是：不要把 embedding 当前的 selector 作为这批算子的通用规则。embedding 是多平台算子，而且它的 Autograd -> mlu 是一条历史兼容/特殊映射；同一个 embedding 代码里的 kernel/platform.py、README、PLATFORM 报告又写的是 MLU 使用 AutogradPrivateUse1。这本身有不一致，我们会在 embedding 侧单独清理，但不建议把这条特殊映射继续泛化。

对这批当前只有 Cambricon 一个平台的算子，建议这样理解三个 key：

key	含义	建议
PrivateUse1	设备后端 key，也是旧 profile/旧测试常用值	可保留为兼容 key
AutogradPrivateUse1	PrivateUse1 对应的 autograd key	建议作为 canonical A1 key
Autograd	更泛的 autograd key，覆盖面更大	不建议默认映射，除非某个算子有真机证据说明必须使用

当前单平台环境下，三个 key 都可能命中；你们实测到这一点是合理的。风险主要出现在多后端同进程：generic Autograd 的覆盖范围更大，可能影响 NPU/CUDA 等其他后端的 autograd 分发。等到第二平台接入时，这个风险会更明显。

所以我的建议是：

canonical: AutogradPrivateUse1
compatibility: PrivateUse1
avoid by default: Autograd

如果某个算子确实因为 torch_mlu/functorch/dispatch 行为必须使用 generic Autograd，可以保留，但希望在 PR 里写清楚：

1. 为什么要用 generic key；
2. 哪个真机用例证明 AutogradPrivateUse1 不满足；
3. 是否会影响多后端同进程。

这一点是设计建议，不是本次最主要的阻塞；阻塞点仍是第 2 节的包态执行 import。

5. 合并前 checklist

为了方便后续复审，建议按这个小清单收尾：

• 10 个可注册算子的实现导入改成包态优先、脚本态 fallback
• MLU 真机增加至少一个包态实际执行用例
• 保持 3 个不可注册算子的 NotImplementedError
• key 映射建议收敛为 AutogradPrivateUse1 + PrivateUse1；如保留 Autograd，补充必要性说明
• CI 继续保持 ops-packaging / unit-tests / docs-and-syntax 全绿

其余工作做得很好；补上包态执行这一块后，我应该可以直接复审通过。

[Dashener2]

@TruNcat3 感谢 review，缺口已按建议补上（新 commit e283944）：

1. 包态执行缺口：10 个可注册算子的实现导入改为「包态优先、脚本态 fallback」
   （try: from .kernel.<mod> import <fn> / except ImportError: from kernel...），
   修复 ops.<op>.register 包态下顶层 kernel 遮蔽的 ModuleNotFoundError。
   你给的 repro（from ops.clone.register import clone_aten）已通过。
2. 包态实际执行用例：ops/clone/test/op_level.py 增加子进程用例，经
   ops.clone.register → ops.clone.kernel 执行并断言存储独立（代表算子，避免进程污染）。
3. key 收敛：_backend_for 收敛为 canonical AutogradPrivateUse1 + 兼容
   PrivateUse1，移除 generic Autograd（无真机必要性，且多后端同进程风险更大）。
4. ops/test_packaging.py：按你的说明理解——它不是纯 import 测试，MLU 跑不了属预期；
   本 PR 引用已通过的 ops-packaging CI，脚本拆 --imports-only/--npu-e2e 后续单独做。
5. 3 个不可注册算子的 NotImplementedError 保持不动。

本地（v26.08.0 容器 / MLU）复验：13/13 example.py 三层全绿、pytest tests/unit 24 passed、
真实 vLLM A1 注入 1 passed。麻烦再看一下，谢谢！

[TruNcat3]

感谢补齐，复查通过 ✅

确认点：

1. 10 个可注册算子的实现导入都改为包态优先、脚本态 fallback；静态检查到 10 处 from .kernel...，顶层裸导入不再作为首选路径。
2. ops/clone/test/op_level.py 新增子进程包态执行用例，从 ops.clone.register 调到 clone_aten，并实际执行到 ops.clone.kernel，同时断言输出 shape 和 storage independence。这个用例与脚本态 A1 测试、CI 包导出扫描互补，覆盖了此前缺口。
3. key 映射收敛为 canonical AutogradPrivateUse1 + 兼容 PrivateUse1，移除 generic Autograd，符合多后端同进程的保守边界。
4. 3 个不可注册算子仍保持 NotImplementedError，没有为了统一签名强行注册。
5. CI 三门全绿；你们 MLU 环境复验 13/13 三层、unit tests 和真实 vLLM 注入也都通过。

一个非阻塞 follow-up：后续可以再把包态用例增强为 from ops.clone import register_a1 → 注册 → x.clone() → counter 断言的一体化子进程用例；当前 ops.clone.register 直调实现已足够覆盖本次 review 指出的 kernel import 缺口。

LGTM。