【RFC】基于 DFunction 接入自定义算子的端到端流程

背景

在分布式训练场景中，存在若干需要用户自定义正反向计算逻辑的场景，使得现有算子级并行框架难以直接处理：

1. 自定义 autograd 函数与分布式并行框架的割裂：用户通过 torch.autograd.Function 或 MindSpore 等效类自定义正反向时，并行框架对正向插入的 layout 推导、DTensor 包装等额外操作对用户不透明，导致自定义反向无法感知并正确处理分布式张量。
2. 缺少跨平台统一抽象：用户在 PyTorch 和 MindSpore 后端上编写的自定义 autograd 函数无法复用，需为不同后端分别适配。
3. 无法与 DTensor dispatch 系统衔接：现有用户自定义函数无法接入 DistributedOp 的 layout 推导和 DTensor 输入/输出包装流程，导致多卡场景下无法直接使用分布式张量作为输入输出。

为解决上述问题，我们引入 DFunction 机制，允许用户以 local tensor 视角编写自定义分布式 autograd 函数，框架自动完成 DTensor 的提取与包装。

本次提交基于 DFunction 机制，成功接入了 npu_dense_lightning_indexer_softmax_lse 接口，打通了从用户 API 到底层 CustomOpBuilder 的完整链路。

设计方案

DFunction 是 platform.Function 的子类（在 PyTorch 上继承 torch.autograd.Function，在 MindSpore 上继承对应的 _Function）。用户子类实现 forward / backward 静态方法，操作的是 local tensor；当输入包含 DTensor 时，apply() 自动路由到 _OP_DISPATCHER.dispatch() 完成 layout 推导和 DTensor 包装，用户感知不到多卡和单卡的区别。

• 单卡路径：DFunction.apply(local_x, local_y) → super().apply() → 平台 autograd 机制 → forward(ctx, local_x, local_y)
• 多卡路径：DFunction.apply(dtensor_x, dtensor_y) → _OP_DISPATCHER.dispatch() → 提取 local tensors → forward(ctx, local_x, local_y) → DTensor.from_local(output, mesh, placements)

layout 推导逻辑由用户配套实现的 DistributedOp 子类提供，通过 _op_name 字符串与 DFunction 子类关联。

用户接口到底层 CustomOpBuilder 的实现链路

本章节详细梳理 npu_dense_lightning_indexer_softmax_lse 接口从用户调用到底层 C++ kernel 的完整实现链路。

目录结构

hyper_parallel/
├── custom_ops/                              # 平台无关的用户接口层
│   └── __init__.py                          # 公开 API 入口
│
├── platform/
│   ├── platform.py                          # Platform 基类，定义 custom_ops 抽象属性
│   ├── torch/
│   │   ├── platform.py                      # TorchPlatform.custom_ops 属性实现
│   │   └── custom_ops/
│   │       ├── __init__.py
│   │       └── custom_ops.py                # TorchCustomOps (PyTorch 路径抛异常)
│   │
│   └── mindspore/
│       ├── platform.py                      # MindSporePlatform.custom_ops 属性实现
│       └── custom_ops/
│           ├── __init__.py
│           ├── custom_ops.py                # MindSporeCustomOps (调用 DFunction)
│           ├── custom_op_impl.py            # DFunction 子类 + CustomOpBuilder 加载
│           ├── module.cc                    # pybind11 模块入口
│           └── dense_lightning_indexer_softmax_lse.cc  # Ascend C++ kernel 封装
│
└── core/shard/
    ├── dfunction.py                         # DFunction 基类
    ├── _op_dispatch.py                      # OpDispatcher.dispatch() 实现
    └── ops/
        ├── parallel_ops.py                  # DistributedOp 基类
        ├── parallel_ops_register.py         # 分布式算子注册表
        ├── parallel_npu_dense_lightning_indexer_softmax_lse.py  # Layout 推导实现
        └── yaml/
            └── npu_dense_lightning_indexer_softmax_lse_ops.yaml  # YAML 配置

类与接口设计

1. 用户接口层 (hyper_parallel/custom_ops/__init__.py)

def npu_dense_lightning_indexer_softmax_lse(
        query_index, key_index, weights, *,
        actual_seq_qlen=None, actual_seq_klen=None,
        layout='BSND', sparse_mode=3, ...) -> Tuple:
    """平台无关的公开 API，委托给 platform.custom_ops"""
    return _platform.custom_ops.npu_dense_lightning_indexer_softmax_lse(...)

职责：提供统一的用户入口，参数校验，文档字符串。

2. Platform 抽象层

基类 (hyper_parallel/platform/platform.py):

class Platform(ABC):
    @property
    @abstractmethod
    def custom_ops(self):
        """子类必须实现，返回平台特定的 custom ops 类"""
        raise NotImplementedError

MindSpore 实现 (hyper_parallel/platform/mindspore/platform.py):

class MindSporePlatform(Platform):
    Function = _Function  # MindSpore autograd 基类
    _custom_ops_cls = None

    @property
    def custom_ops(self):
        if MindSporePlatform._custom_ops_cls is None:
            from hyper_parallel.platform.mindspore.custom_ops.custom_ops import MindSporeCustomOps
            MindSporePlatform._custom_ops_cls = MindSporeCustomOps
        return MindSporePlatform._custom_ops_cls

PyTorch 实现 (hyper_parallel/platform/torch/platform.py):

class TorchPlatform(Platform):
    Function = torch.autograd.Function  # PyTorch autograd 基类
    _custom_ops_cls = None

    @property
    def custom_ops(self):
        if self._custom_ops_cls is None:
            from hyper_parallel.platform.torch.custom_ops.custom_ops import TorchCustomOps
            self._custom_ops_cls = TorchCustomOps
        return self._custom_ops_cls

3. 平台特定 CustomOps 类

MindSpore (hyper_parallel/platform/mindspore/custom_ops/custom_ops.py):

class MindSporeCustomOps:
    @staticmethod
    def npu_dense_lightning_indexer_softmax_lse(...):
        from hyper_parallel.platform.mindspore.custom_ops.custom_op_impl import (
            NpuDenseLightningIndexerSoftmaxLseDFunction,
        )
        return NpuDenseLightningIndexerSoftmaxLseDFunction.apply(...)

PyTorch (hyper_parallel/platform/torch/custom_ops/custom_ops.py):

class TorchCustomOps:
    @staticmethod
    def npu_dense_lightning_indexer_softmax_lse(*args, **kwargs):
        raise RuntimeError(
            "npu_dense_lightning_indexer_softmax_lse is not supported via HyperParallel "
            "on the PyTorch platform. Please call torch_npu.npu_dense_lightning_indexer_softmax_lse directly."
        )

4. DFunction 子类 (hyper_parallel/platform/mindspore/custom_ops/custom_op_impl.py)

# 1. 加载 C++ 自定义算子
_custom_ops = ms.ops.CustomOpBuilder(
    "custom_ops",
    [
        "module.cc",
        "dense_lightning_indexer_softmax_lse.cc",
        ...
    ],
    backend="Ascend",
).load()

# 2. DFunction 子类
class NpuDenseLightningIndexerSoftmaxLseDFunction(DFunction):
    _op_name = "npu_dense_lightning_indexer_softmax_lse"

    @staticmethod
    def forward(ctx, query_index, key_index, weights, ...):
        return _custom_ops.npu_dense_lightning_indexer_softmax_lse(...)

    @staticmethod
    def backward(ctx, *grad_outputs):
        return (None,) * 9  # 无梯度需求

5. DFunction 基类 (hyper_parallel/core/shard/dfunction.py)

class DFunction(platform.Function):
    _op_name: str = None

    @classmethod
    def apply(cls, *args, **kwargs):
        has_dtensor = any(isinstance(a, DTensor) for a in args)
        if has_dtensor:
            if cls._op_name is None:
                raise ValueError("DTensor inputs require '_op_name'")
            return _OP_DISPATCHER.dispatch(cls._get_local_callable(), args, kwargs)
        return super().apply(*args, **kwargs)

    @classmethod
    def _get_local_callable(cls) -> _LocalCallable:
        """返回带 op_name 的 callable，用于 dispatcher 查找 DistributedOp"""
        ...

6. OpDispatcher (hyper_parallel/core/shard/_op_dispatch.py)

class OpDispatcher:
    def dispatch(self, op_call: callable, args: tuple, kwargs: dict) -> object:
        op_name = platform.get_op_name(op_call)

        # 自动注册通过 DistributedOp 程序化注册的算子
        if op_name not in self.layout_infer_ops and get_distributed_op(op_name) is not None:
            self.layout_infer_ops[op_name] = {}

        return self._dispatch_layout_infer(op_name, op_call, args, kwargs)

7. DistributedOp 子类 (hyper_parallel/core/shard/ops/parallel_npu_dense_lightning_indexer_softmax_lse.py)

class NpuDenseLightningIndexerSoftmaxLseDistributedOp(DistributedOp):
    def __init__(self):
        super().__init__("npu_dense_lightning_indexer_softmax_lse")

    def preprocess(self, args, kwargs) -> tuple:
        """提取 local tensors，构建 cache_values"""
        ...

    def infer_layout(self, cache_values) -> Tuple[tuple, None]:
        """推导输出 layout：BSND → (B, N2index, S1)，TND → (N2index, T1)"""
        ...

    def get_expand_impl(self, func, infer_result, cache_values, ...) -> Optional[Callable]:
        """Context Parallel 场景下的 key 切片 / seq_len 调整"""
        ...

8. YAML 配置 (hyper_parallel/core/shard/ops/yaml/npu_dense_lightning_indexer_softmax_lse_ops.yaml)

npu_dense_lightning_indexer_softmax_lse:
  dist_op_name: _npu_dense_lightning_indexer_softmax_lse_dist_op
  distributed_op_class: NpuDenseLightningIndexerSoftmaxLseDistributedOp
  distributed_op_file: parallel_npu_dense_lightning_indexer_softmax_lse

9. C++ Kernel 封装 (hyper_parallel/platform/mindspore/custom_ops/dense_lightning_indexer_softmax_lse.cc)

class DenseLightningIndexerSoftmaxLseRunner : public ms::pynative::PyboostRunner {
  void RunAclnnDirect() {
    // 调用 Ascend ACLNN API: aclnnDenseLightningIndexerSoftmaxLse
    ...
  }
};

调用链路图

用户调用
    │
    ▼
hyper_parallel.custom_ops.npu_dense_lightning_indexer_softmax_lse()
    │
    ▼
_platform.custom_ops.npu_dense_lightning_indexer_softmax_lse()
    │
    ├── PyTorch: TorchCustomOps → RuntimeError (提示使用 torch_npu)
    │
    └── MindSpore: MindSporeCustomOps
            │
            ▼
        NpuDenseLightningIndexerSoftmaxLseDFunction.apply()
            │
            ├── 无 DTensor 输入 → super().apply() → forward() → _custom_ops.npu_dense_lightning_indexer_softmax_lse()
            │
            └── 有 DTensor 输入 → _OP_DISPATCHER.dispatch()
                    │
                    ▼
                查找 DistributedOp (通过 op_name)
                    │
                    ▼
                NpuDenseLightningIndexerSoftmaxLseDistributedOp
                    │
                    ├── preprocess(): 提取 local tensors
                    ├── infer_layout(): 推导输出 layout
                    └── get_expand_impl(): CP 调整 (可选)
                    │
                    ▼
                调用 local callable → forward() → _custom_ops.npu_dense_lightning_indexer_softmax_lse()
                    │
                    ▼
                DTensor.from_local() 包装输出

对外 API

from hyper_parallel import custom_ops

# 单卡调用 (plain tensor)
softmax_max, softmax_sum = custom_ops.npu_dense_lightning_indexer_softmax_lse(
    query_index, key_index, weights, layout='BSND')

# 多卡调用 (DTensor)
from hyper_parallel import init_device_mesh
from hyper_parallel.core.dtensor.dtensor import distribute_tensor
from hyper_parallel.core.dtensor.placement_types import Shard, Replicate

mesh = init_device_mesh("npu", (2, 4), mesh_dim_names=("dp", "tp"))
query_dist = distribute_tensor(query, mesh, (Shard(0), Replicate()))
key_dist = distribute_tensor(key, mesh, (Replicate(), Replicate()))
weights_dist = distribute_tensor(weights, mesh, (Shard(0), Replicate()))

softmax_max_dist, softmax_sum_dist = custom_ops.npu_dense_lightning_indexer_softmax_lse(
    query_dist, key_dist, weights_dist, layout='BSND')

使用约束

1. DFunction 子类必须设置 _op_name，且与注册的 DistributedOp 实例的 op_name 完全一致；
2. forward 和 backward 内部必须操作 local tensor，不得递归调用 DFunction.apply；
3. 非 Tensor 位置参数需通过 preprocess 方法处理，或改用 kwargs 传递；
4. 当前仅支持动态图，暂不支持静态图；
5. PyTorch 平台需直接调用 torch_npu.npu_dense_lightning_indexer_softmax_lse，DTensor dispatch 通过 YAML 注册的 DistributedOp 自动处理；

测试设计

1. 单卡精度验证：MindSpore vs torch_npu (PTA) bit-exact 比较
2. 多卡分布式验证：DTensor 输入，验证输出 layout 正确、值与单卡对齐
3. Context Parallel 场景：BSND+CP / TND+CP 的 key 切片和 seq_len 调整正确性

测试文件：tests/mindspore/st/custom_ops/test_npu_dense_lightning_indexer_softmax_lse.py

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 99
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/99

Start with hyper_parallel/core/shard/dfunction.py and trace dispatch through _op_dispatch.py and the named DistributedOp implementation. Then run tests/mindspore/st/custom_ops/test_npu_dense_lightning_indexer_softmax_lse.py to understand the single-card, DTensor, and Context Parallel cases. Done means the DFunction path, layout handling, custom kernel integration, and listed validation scenarios work as described.