/assign @xuejinghui
torch.autograd.Function 与 custom_op + register_autograd 的机制差异
核心结论
这不是 Torch 版本依赖问题，而是两种机制在 torch.compile 下的本质差异：
torch.autograd.Function的forward函数体会被 Dynamo 追踪执行；custom_op + register_autograd的forward会被当作不透明的图节点，Dynamo 不会深入其实现函数体。
因此，autograd.Function 在 eager 模式下可以正常运行，并不意味着它在 torch.compile 下也一定可用。只要 forward 或 backward 内部包含 Dynamo 无法理解的 Python/C++ 扩展调用，就可能在编译阶段失败。
eager 与 torch.compile 的对比
| 机制 | eager | torch.compile(inductor) |
|---|---|---|
torch.autograd.Function |
forward / backward 直接执行，通常可以正常运行。 |
Dynamo 遇到 apply() 时通过高阶算子（HOP）处理，并重新追踪 forward 和 backward 的函数体。因此，函数体内的代码必须全部能够被 Dynamo 理解。 |
custom_op + register_autograd |
通过 dispatcher 执行，通常可以正常运行。 | forward 是带有 schema 和 fake/meta 实现的 dispatcher 节点。Dynamo 将其作为图节点处理，不深入追踪实现函数体；backward 同样可以调用已注册的 dispatcher op，形成黑盒节点。 |
为什么 mHC 开发者的 autograd.Function 在 compile 下不可用
相关实现的 forward 大致如下：
class MhcPreSinkhornFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, x, phi, alpha, bias, ...):
op_module = mhc_pre_sinkhorn_op_builder.load()
hin, h_post, h_res, ... = op_module.mhc_pre_sinkhorn(...)
...
在 eager 模式下，这段代码没有问题：OpBuilder.load() 加载扩展模块，然后通过 pybind 直接调用算子。
但在 torch.compile 下，Dynamo 会重新追踪 MhcPreSinkhornFunction.forward 的函数体。当它遇到下面这个 pybind 调用时：
op_module.mhc_pre_sinkhorn(...)
Dynamo 不知道这个 builtin 函数如何追踪，于是报错：
Dynamo does not know how to trace the builtin
mhc_pre_sinkhorn.pybind11_detail_function_record...
这正是 issue #4289 所体现的问题。根因不是 autograd.Function 在 eager 下不能工作，而是其函数体在 compile 下会被追踪，而函数体中存在无法被 Dynamo 解释的 pybind 直调。
为什么 compressor 的实现方式可以工作
compressor 将 pybind 调用藏在 custom_op 的实现函数内部：
@torch.library.custom_op(
"cann_ops_transformer::_compressor_forward",
...,
)
def _compressor_forward(x, wkv, wgate, ...):
op_module = compressor_op_builder.load()
... = op_module.compressor(...)
同时为该 op 注册 fake/meta 实现，用于编译阶段的形状推导。
对 Dynamo 来说，模型侧调用的是：
torch.ops.cann_ops_transformer._compressor_forward(...)
它是一个已经注册的 dispatcher op，具有明确的 schema 和 fake 实现。因此：
- Dynamo 将它作为图节点记录下来；
- 编译阶段使用 fake/meta 实现进行形状推导；
- Dynamo 不会进入
custom_op实现函数体，也不会看到其中的OpBuilder.load()或 pybind 调用； - 编译后的运行时通过 dispatcher 调用
PrivateUse1实现，再由实现层调用 pybind 算子。
调用链可以概括为：
模型代码
-> torch.ops.cann_ops_transformer._compressor_forward
-> dispatcher 图节点
-> PrivateUse1 实现
-> pybind 扩展
这里的关键是 forward 被黑盒化为 dispatcher 节点，而不是 pybind 调用本身有什么特殊之处。
背景
DeepSeek-V4 mHC 需要在训练的 forward/backward 中调用以下 dispatcher 算子：
cann_ops_transformer::mhc_pre_sinkhorncann_ops_transformer::mhc_post当前
cann_ops_transformerPython wrapper 在输入需要梯度时，会切换到 Pythontorch.autograd.Function，并在其forward内调用OpBuilder.load()加载后的 pybind 扩展。该调用路径无法被 PyTorch Dynamo/AOTAutograd 作为图节点追踪，开启torch.compile(backend="inductor")时会失败，典型报错为：我们已验证直接调用 dispatcher 算子可以被 Inductor 入图；但当前 forward 算子没有官方的 dispatcher-level Autograd 注册，因此业务侧需要额外编写
torch.autograd.Functionbridge，手动保存 forward 中间量并调用 backward 算子。环境信息
/data/x00617812/CANN/cann-9.2.0/python/site-packages/cann_ops_transformer重现步骤
下面的探针使用公开的 dispatcher schema，且不会预先加载任何
OpBuilder：import torch import torch_npu import cann_ops_transformer def fn(x, phi, alpha, bias): outputs = torch.ops.cann_ops_transformer.mhc_pre_sinkhorn( x, phi, alpha, bias, 4, 20, 1e-6, 1e-6, True ) h_in, h_post, h_res = outputs[:3] return torch.ops.cann_ops_transformer.mhc_post( x, h_res.reshape(1, 128, 4, 4), h_in, h_post ) compiled = torch.compile(fn, backend="inductor") compiled(x, phi, alpha, bias)在当前软件包中，直接 forward 路径可以注册并执行，但 AOTAutograd 没有这些 forward dispatcher 算子的导数规则，因此模型需要使用自定义 bridge。
交付需求
参考同仓
cann_ops_transformer/ops/compressor.py的官方交付范式（torch.library.custom_op+register_fake+register_autograd(setup_context=...)+ 公开 wrapper 统一走 dispatcher），为mhc_pre_sinkhorn和mhc_post补齐以下三项交付：一、官方 Autograd 注册（对齐
compressor.py:189-226）mhc_pre_sinkhorn注册 autograd：setup_context保存前向输入x, phi, alpha, bias以及out_flag=True返回的 5 个中间量（h_pre, hc_before_norm, inv_rms, sum_out, norm_out，即 forward 8 元组输出的后 5 个），并保存hc_eps；backward调用torch.ops.cann_ops_transformer.mhc_pre_sinkhorn_backward，对 9 个 forward 输入返回 4 个张量梯度 + 5 个 None。mhc_post注册 autograd：save_for_backward(x, h_res, h_out, h_post)；backward调用torch.ops.cann_ops_transformer.mhc_post_backward返回 4 个梯度。注册完成后，直接调用
torch.ops.cann_ops_transformer.mhc_pre_sinkhorn/mhc_post即可参与 eager Autograd，并可通过torch.compile(backend="inductor")追踪，不再需要业务侧 bridge。二、公开 wrapper 去掉双轨分支（对齐
compressor.py:248-264）mhc_pre_sinkhorn()/mhc_post()公开接口删除needs_grad分支中切换到 Pythontorch.autograd.Function（MhcPreSinkhornFunction.apply/MhcPostFunction.apply）+OpBuilder.load()pybind 直调的路径，改为统一调用 dispatcher：def mhc_pre_sinkhorn(x, phi, alpha, bias, hc_mult, num_iters, hc_eps=1e-6, norm_eps=1e-6): out_flag = x.requires_grad or phi.requires_grad or alpha.requires_grad or bias.requires_grad return torch.ops.cann_ops_transformer.mhc_pre_sinkhorn( x, phi, alpha, bias, hc_mult, num_iters, hc_eps, norm_eps, out_flag)[:3]由 wrapper 计算
out_flag传入 forward 算子，始终走torch.ops.cann_ops_transformer.*。编译场景下不再存在不可追踪的 pybindautograd.Function路径。三、补齐导出（对齐
ops/__init__.py中其它算子全量导出）在
ops/__init__.py中导出mhc_post_backward：from .mhc_post_backward import mhc_post_backward与
mhc_pre_sinkhorn_backward等 backward 算子保持一致，业务侧无需再显式 import 该模块。关联 issue
mhc_post_backward提供 TorchAir/torch-extension 支持。验证记录
在上述环境中，使用围绕 dispatcher 算子编写的本地
torch.autograd.Functionbridge 进行的最小 NPU 探针，已通过 eager 和编译模式下的 forward/backward 验证。4 个可训练输入均成功获得梯度。该 bridge 只是针对当前缺少官方 Autograd 注册问题的临时下游规避方案。