【RFC】分布式算子分发架构演进与三阶段接口迁移

0. RFC 摘要

本 RFC 基于 OpDispatcher 重构方案，
用于指导分布式算子从旧的 infer_layout(layouts, extra_args) 模式迁移到
preprocess / infer_layout(cache_values) / get_expand_impl 三阶段模式。

本次工作不是单纯修改函数签名，而是对算子分发职责进行重新划分：

• OpDispatcher 负责稳定的分发骨架、缓存查询、算子执行和结果包装；
• DistributedOp 子类负责各算子的参数规范化、Layout 推导和可选的本地实现扩展；
• 平台差异在 preprocess 中完成接口归一化，后续阶段只处理统一语义；
• 通过不可变缓存键和最小化 cache_values 保证缓存正确性和命中率；
• 通过分阶段迁移控制存量算子的兼容性风险。

关联交付：

• PR 508：OpDispatcher 重构：完成三阶段架构骨架、缓存键重构和 Linear/Sort 试点迁移；
• PR 736：分布式算子迁移：阶段性迁移 11 个算子，集中处理 aclop 参数标准化并更新对应测试；
• 本 RFC：统筹剩余分布式算子迁移、旧 suffix 流程收敛及测试适配。

PR 508 采用新旧流程并存的渐进式迁移方式，没有在该 PR 中一次性删除全部旧流程。本 RFC 承接后续迁移和收敛工作。

范围

• 分布式算子三阶段接口及职责边界；
• PyTorch、MindSpore Primitive 和 MindSpore mint 接口的参数归一化；
• Layout 推导缓存、算子展开实现和 DTensor 输出包装；
• 存量算子迁移、UT/ST 适配和旧流程清理。

非目标

• 不新增或改变用户可见的分布式算子 API；
• 不改变算子的数学语义和既有切分规则；
• 不在本 RFC 中重新设计 DTensor、DeviceMesh 或底层通信实现。

---

1. 背景与问题

HyperParallel 通过 PyTorch __torch_function__ 和 MindSpore __fallback__ 接管 DTensor 算子调用，
再由 OpDispatcher 完成分布式算子查找、Layout 推导、本地计算和结果包装。旧流程将算子差异集中在
OpDispatcher 内，随着算子类型增加，暴露出以下架构问题：

问题	根因	架构影响
extra_args 丢失参数名和调用语义	位置参数、关键字参数只保留 value，缺少统一参数绑定	可能产生错误缓存命中和错误 Layout 推导
参数处理存在副作用	参数解析函数原地修改可变的 LayoutCacheKey	状态难追踪，缓存正确性和可测试性下降
suffix 分支硬编码	WithShape、Reshape、Slice 等差异集中在分发器	新增特殊算子需要修改核心流程，违反开闭原则
infer_layout 输入和返回格式不统一	不同算子自行约定 layouts、extra_args 和额外返回值	接口语义模糊，调用方存在多套处理分支
缓存键包含无关参数	分发器无法判断哪些值真正影响 Layout	缓存命中率下降，重复执行 Layout 推导
PyTorch/MindSpore 调用约定不同	Primitive、functional 和 keyword-only 参数规则不同	跨平台实现容易出现参数路由不一致

1.1 新旧接口对比

旧模式（存量，逐步迁移）

class XxxDistributedOp(DistributedOp):
    def infer_layout(self, layouts: tuple, extra_args: Optional[tuple] = None) -> tuple:
        ...
    def get_expand_impl(self, func, infer_result, layouts: tuple,
                        extra_args: Optional[tuple] = None) -> Optional[callable]:
        ...

新模式（推荐，已迁移算子参考）

class XxxDistributedOp(DistributedOp):
    def preprocess(self, args: tuple, kwargs: dict) -> tuple:
        """返回 (local_args, local_kwargs, cache_values)"""
        ...
    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """基于 cache_values 推导输出 Layout"""
        ...
    def get_expand_impl(self, func: Optional[callable], infer_result: tuple,
                        cache_values: list) -> Optional[callable]:
        """返回 expand 闭包，或 None"""
        ...

1.2 三阶段分发流程

DTensorBase.__torch_function__ / __fallback__
  → OpDispatcher.dispatch()
  → distribute_op.preprocess(args, kwargs)          ← 新：参数规范化
      → distribute_op.infer_layout(cache_values)    ← 新：接收 cache_values
      → distribute_op.get_expand_impl(func, ...)    ← 新：接收 cache_values
  → op_impl(*local_args, **local_kwargs)            ← 无 extra_info 展开
  → DTensor.from_local()

extra_info 不显式传给 op_impl。infer_result = (output_layouts, extra_info) 中的 extra_info
仅用于 get_expand_impl 闭包内部捕获，三阶段分发流程始终以
op_impl(*local_args, **local_kwargs) 调用本地实现。

---

2. 架构目标与质量属性

本次重构以正确性为首要约束，同时改善可扩展性、可维护性、跨平台一致性、性能和可测试性。

质量属性	质量属性场景	设计响应与验收方式
正确性	同一算子使用位置参数、关键字参数或混合参数调用	preprocess 先完成参数规范化；语义等价调用生成一致的 cache_values，语义不同调用不得错误命中缓存
可扩展性	新增具有特殊参数或本地执行逻辑的算子	只新增/复用 DistributedOp 策略及 YAML 注册，不修改 OpDispatcher 主流程
可维护性	修改某个算子的参数解析、Layout 规则或展开实现	变更分别限定在 preprocess、infer_layout、get_expand_impl，避免霰弹式修改
跨平台一致性	同一算子同时支持 PyTorch、MindSpore Primitive 和 mint	_normalize_*_args 统一语义，preprocess 按接口类型路由 local_args/local_kwargs
性能	相同 Layout 和切分相关参数重复调用	使用不可变缓存键并只缓存影响 Layout 的最小信息；重复调用复用 infer_result/op_impl
可测试性	验证参数处理、Layout 数学和展开实现	三个阶段可独立进行 UT；ST 只验证成功路径的分布式数值结果
可演进性	大量存量算子无法一次性迁移	先引入兼容分支和试点算子，再分批迁移；全部迁移后删除 legacy/suffix 流程

---

3. 4+1 架构视图

本 RFC 使用 4+1 视图组织架构信息。UML 图用于表达稳定的组件关系和一次算子调用的动态协作，
具体迁移代码仍以后续章节的接口契约为准。

3.1 逻辑视图

逻辑上将系统划分为接管层、分发编排层、算子策略层、缓存层和本地执行层：

classDiagram
    direction LR

    class DTensorBase {
        +__torch_function__()
        +__fallback__()
    }

    class OpDispatcher {
        +dispatch(op_call, args, kwargs)
        -_dispatch_layout_infer(op_name, op_call, args, kwargs)
        -_lookup_or_infer_layout(...)
    }

    class LayoutCacheManager {
        +distributed_op(op_name)
        +get_layout_cache()
    }

    class DistributedOp {
        +preprocess(args, kwargs)
        +infer_layout(cache_values)
        +get_expand_impl(func, infer_result, cache_values)
        +wrap_output(py_output, output_layouts)
    }

    class XxxDistributedOp

    DTensorBase --> OpDispatcher : intercept and dispatch
    OpDispatcher --> LayoutCacheManager : resolve and cache
    LayoutCacheManager --> DistributedOp : registered strategy
    DistributedOp <|-- XxxDistributedOp : operator-specific policy

职责约束：

• preprocess 是唯一接触原始 args/kwargs 的阶段；
• infer_layout 只做 Layout 数学和输入约束校验，不执行算子计算；
• get_expand_impl 只提供可选的本地实现包装，不承担输入校验；
• OpDispatcher 不感知具体算子的 shape、tuple 展开或参数调整规则。

3.2 开发视图

模块	位置	职责
分发编排	hyper_parallel/core/shard/_op_dispatch.py	接管后路由、缓存查询、调用和输出包装
算子协议	hyper_parallel/core/shard/ops/parallel_ops.py	定义 DistributedOp 三阶段接口
算子策略	hyper_parallel/core/shard/ops/parallel_*.py	参数规范化、Layout 规则和展开实现
声明式注册	hyper_parallel/core/shard/ops/yaml/*.yaml	算子名到 DistributedOp 策略的映射
单元测试	tests/ut/core/shard/ops/test_parallel_*.py	分阶段验证参数、Layout、错误路径和闭包逻辑
系统测试	tests/{torch,mindspore}/.../shard/ops/	验证真实后端上的分布式数值结果

平台无关的三阶段协议位于 core/shard，不得直接依赖 PyTorch 或 MindSpore 实现；平台接口差异通过
get_platform() 和每个算子的参数规范化逻辑收敛。

3.3 进程视图

一次 DTensor 算子调用按以下顺序执行。每个 rank 在本地执行相同的分发流程；如算子需要通信，
通信逻辑由选中的本地实现或展开闭包通过平台抽象发起。

sequenceDiagram
    actor User
    participant DTensorBase
    participant Dispatcher as OpDispatcher
    participant DistOp as DistributedOp
    participant Cache as LayoutCacheManager
    participant LocalOp as Local op_impl

    User->>DTensorBase: call operator(DTensor, ...)
    DTensorBase->>Dispatcher: dispatch(op_call, args, kwargs)
    Dispatcher->>DistOp: preprocess(args, kwargs)
    DistOp-->>Dispatcher: local_args, local_kwargs, cache_values
    Dispatcher->>Cache: lookup(op_name, cache_key)

    alt cache miss
        Dispatcher->>DistOp: infer_layout(cache_values)
        DistOp-->>Dispatcher: output_layouts, extra_info
        Dispatcher->>DistOp: get_expand_impl(func, infer_result, cache_values)
        DistOp-->>Dispatcher: op_impl or None
        Dispatcher->>Cache: save(infer_result, op_impl)
    else cache hit
        Cache-->>Dispatcher: infer_result, op_impl
    end

    Dispatcher->>LocalOp: op_impl(*local_args, **local_kwargs)
    LocalOp-->>Dispatcher: local output
    Dispatcher->>DistOp: wrap_output(local output, output_layouts)
    DistOp-->>User: DTensor output

3.4 部署视图

部署路径	接管入口	参数接口特点	统一后的执行路径
PyTorch rank 进程	DTensorBase.__torch_function__	支持位置参数、关键字参数和 keyword-only 参数	OpDispatcher → DistributedOp → platform/local op
MindSpore rank 进程	DTensorBase.__fallback__	Primitive 多为 positional；mint 可能包含 keyword-only 参数	OpDispatcher → DistributedOp → platform/local op

同一个 DistributedOp 可以服务多个平台入口，但必须在 preprocess 中将平台调用约定归一化。
DeviceMesh 描述 rank 拓扑和 Layout，实际通信由平台后端在各 rank 进程中执行。

3.5 场景视图（+1）

场景	关键路径	预期结果
等价参数调用	sort(x, 1) 与 sort(x, dim=1) 分别进入 preprocess	生成一致的规范参数和缓存键
特殊 shape 算子	reshape/slice 在 preprocess 中保存必要的全局信息	infer_layout 得到输出 Layout，闭包使用局部 shape/索引执行
多输出算子	sort/topk 返回多个本地 Tensor	wrap_output 按输出 Layout 一一包装为 DTensor
Partial 输入	infer_layout 收到包含 Partial 的 Layout	按算子规则传播 Partial 或在执行前明确报错
跨平台调用	PyTorch function、MindSpore Primitive、mint 调用同一算子	参数路由符合各自接口约束，Layout 语义和数值结果一致

---

4. 架构模式与关键决策

4.1 使用的架构模式

模式	在本方案中的体现	目的
Dispatcher	OpDispatcher 提供统一接管和路由入口	隔离框架接管机制与分布式算子实现
Strategy / 扩展点	注册表按算子名选择 DistributedOp 子类	新增算子策略时保持主流程稳定
Pipeline	preprocess → infer_layout → get_expand_impl → execute → wrap_output	明确阶段职责和数据契约
Registry	YAML 和分布式算子注册表建立算子名到策略的映射	避免在分发器中硬编码具体算子
Immutable Value Object	LayoutCacheKey 使用不可变 tuple 和预计算 hash	避免缓存键被意外修改并降低重复计算开销
Cache-Aside	分发器先查询缓存，未命中时再推导并写回	减少重复 Layout 推导和闭包构建

4.2 方案权衡

决策点	备选方案	选择	原因与代价
特殊算子逻辑归属	在 OpDispatcher 增加 suffix 分支 / 下沉到算子策略	下沉到 DistributedOp	提升可扩展性；代价是每个算子必须遵循更严格的接口契约
参数表示	沿用 layouts + extra_args / 使用规范化 cache_values	cache_values	保留参数语义并最小化缓存输入；代价是迁移时必须逐算子确认哪些值影响 Layout
迁移方式	一次性替换 / 新旧流程兼容迁移	分阶段兼容迁移	降低大量算子同时回归的风险；代价是过渡期需要维护两条路径
缓存键	可变 list、动态计算 hash / 不可变 tuple、预计算 hash	不可变键	消除副作用并稳定哈希；代价是构建时必须一次性收集完整语义
额外执行参数	由分发器展开 extra_info / 闭包捕获	get_expand_impl 闭包捕获	保持主调用签名统一；代价是闭包实现需清晰区分构建期与运行期逻辑

---

5. 设计与交付过程

5.1 分阶段交付

阶段	主要工作	交付物 / 验证
现状分析	梳理 DTensor 接管链路、suffix 分支、缓存和参数语义问题	op_dispatch_refactor_v2.md 现状架构及问题清单
架构设计	定义质量属性、三阶段职责、缓存协议和兼容策略	方案评审、接口契约、UML/4+1 视图
骨架与试点	引入不可变缓存键、新分发骨架，迁移 Linear/Sort	PR 508 及相关 UT
分批迁移	按算子族迁移 preprocess/infer_layout/get_expand_impl	PR 736 等阶段性交付、UT、已有 ST 回归
架构收敛	删除 suffix 和 legacy 分支，统一注册与调用路径	全仓搜索无旧接口残留
验收交付	运行跨平台 UT/ST、静态检查和性能对比	门禁通过，无功能和性能回退

5.2 风险与缓解

风险	缓解措施
平台参数约定不同导致本地调用失败	为每个多平台算子建立 _normalize_*_args 和 preprocess 路由 UT
cache_values 缺少影响 Layout 的参数	明确每个算子的缓存契约，并验证等价调用和非等价调用的缓存键
Partial、StridedShard 等边界语义回归	在 infer_layout UT 覆盖 Partial、负维度、多 mesh 轴映射和错误路径
批量迁移导致回归定位困难	以算子族分批迁移，每批独立执行相关 UT/ST 后再合入
Layout 对象缓存污染测试结果	setUp/tearDown 同时清理 _LAYOUT_CACHE、_DEVICE_MESH_MAP 和通信组状态
新流程引入运行时开销	缓存 Layout 推导与闭包；对比迁移前后的热点调用性能，确认无显著回退

5.3 验收标准

• 目标算子均实现三阶段接口，不再回退旧 infer_layout(layouts, extra_args) 流程；
• YAML 中不再使用 infer_layout_suffix，OpDispatcher 中删除对应 legacy handler；
• 语义等价的 args/kwargs 调用生成一致的 cache_values 和缓存键；
• infer_layout 返回统一的 (output_layouts, extra_info)，校验和错误路径具备 UT；
• PyTorch、MindSpore Primitive、MindSpore mint 的参数路由具备针对性 UT；
• 相关 UT 全量通过，已有多卡 ST 无数值回归；
• 静态检查和门禁通过，热点算子分发性能无显著回退；
• 迁移规范、参考实现和实际主干代码保持一致。

---

6. 迁移范围与涉及文件

6.1 算子实现文件

已迁移：

• hyper_parallel/core/shard/ops/parallel_matmul.py — LinearDistributedOp ✅
• hyper_parallel/core/shard/ops/parallel_sort.py — SortDistributedOp ✅

待迁移（共 38 个文件）：

parallel_conv3d.py, parallel_repeat.py, parallel_pad.py, parallel_gather.py,
parallel_activation_with_axis.py, parallel_elementwise.py, parallel_repeat_interleave.py,
parallel_squeeze.py, parallel_embedding.py, parallel_flatten.py, parallel_new_ones.py,
parallel_tuple_elementwise.py, parallel_cumsum.py, parallel_slice.py,
parallel_expand_dims.py, parallel_concat.py, parallel_multinomial.py, parallel_split.py,
parallel_one_hot_ext.py, parallel_nonzero.py, parallel_slice_ext.py, parallel_expand.py,
parallel_argsort.py, parallel_unbind.py, parallel_isin.py, parallel_reshape.py,
parallel_atleast_1d.py, parallel_masked_scatter.py, parallel_norm.py,
parallel_argmax_with_value_ops.py, parallel_topk.py, parallel_reduce.py,
parallel_transpose.py, parallel_outer.py, parallel_scatter.py,
parallel_npu_flash_attention_score.py, parallel_ms_flash_attention_score.py,
parallel_scaled_dot_product_attention.py

6.2 UT 测试文件

测试文件	对应算子
tests/ut/core/shard/ops/test_parallel_linear.py	Linear
tests/ut/core/shard/ops/test_parallel_sort.py	Sort
tests/ut/core/shard/ops/test_parallel_*.py	其他算子（迁移时同步适配）

---

7. 迁移步骤详解

7.1 preprocess 实现规范

职责：参数规范化 + to_local + 构建 cache_values。不含任何校验逻辑（校验全部放 infer_layout）。

_normalize_*_args 输出规范

默认规则：全部以 args 返回，kwargs 为空。 MindSpore Primitive 不接受 kwargs，必须全部为 positional args。

# ✅ 正确：全 args + 空 kwargs（MindSpore Primitive 兼容）
def _normalize_linear_args(x, weight, bias=None):
    return (x, weight, bias), {}

# ❌ 错误：bias 放 kwargs，MindSpore Dense 算子无法接收
def _normalize_linear_args(x, weight, bias=None):
    return (x, weight), {'bias': bias}

例外：Python 接口 * 后的 keyword-only 参数必须放 kwargs。 MindSpore mint.* functional_overload 接口行为相同。

# torch.sort(input, dim=-1, descending=False, *, stable=False)
# stable 在 * 后 → keyword-only → 必须放 kwargs
# MindSpore mint.sort 为 functional_overload → 行为同 kwargs 接口
def _normalize_sort_args(x, dim=-1, descending=False, stable=False):
    return (x, dim, descending), {'stable': stable}  # ✅ x/dim/descending 为位置参数，stable 为 kwargs

跨平台接口路由

当同一分布式算子类服务于多个平台或接口类型时，在 preprocess 中按 self.op_name 路由 local_args / local_kwargs。

# MindSpore Primitive（如 SortExt）不接受 kwargs → 全 args
# PyTorch 函数（sort）/ MindSpore functional_overload（Sort）→ keyword-only 参数放 kwargs
_MS_PRIMITIVE_OP_NAMES = frozenset({'SortExt'})

if self.op_name in self._MS_PRIMITIVE_OP_NAMES:
    local_args = (tensor.to_local(), dim, descending, stable)
    local_kwargs = {}
else:
    local_args = (tensor.to_local(), dim, descending)
    local_kwargs = {'stable': stable}

cache_values 构建

cache_values 是 infer_layout 和 get_expand_impl 的唯一输入。约定：Layout 对象在前，标量参数（dim、int、bool 等）在后。

cache_values = [input_tensor.layout, dim]            # 单输入 + 1 个标量
cache_values = [x_layout, w_layout, bias_layout]     # 多输入（无 layout 用 None 占位）

Suffix 处理

同一分布式算子类支持多种 infer_layout_suffix 时，在 preprocess 中统一处理差异：

• WithShape：将 shape 信息包含在 cache_values
• WithTupleExpand：将 tuple/list 参数展开后放入 cache_values

---

7.2 infer_layout 实现规范

新签名：infer_layout(self, cache_values: list) -> Tuple[tuple, None]

实现要点：

1. 必须手动调用 _check_partial_inputs()

旧流程的 base class infer_layout 会自动调用 partial 校验；新流程子类完全覆盖了基类方法，必须在开头手动调用：

def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
    layout = cache_values[0]
    self._check_partial_inputs([layout])          # ← 单输入
    # self._check_partial_inputs([x_layout, w_layout])  # 多输入
    ...

若算子允许 Partial 输入（如 MatMul 的 partial accumulation），设置 self._allow_partial_inputs = True。

2. 所有校验集中在此处

类型校验、范围校验、layout 兼容性（mesh_shape 一致性）、sharding 约束，全部在 infer_layout 中完成，不分散到 preprocess 或 get_expand_impl。

3. 检查分片状态推荐用 alias_tensor_map

alias_tensor_map 返回字符串，"None" 表示 Replicate，同时支持 StridedShard 的 tuple 映射，是更通用的选择：

alias_map = layout.alias_tensor_map
if alias_map[dim] != "None":     # ✅ 推荐：该维度被分片
    raise ValueError(...)

# tensor_map 用整数 -1 表示 Replicate，但 StridedShard tuple 场景需额外处理：
mapping = layout.tensor_map[dim]
is_sharded = any(m != -1 for m in mapping) if isinstance(mapping, (list, tuple)) else mapping != -1

4. 错误信息使用 self.op_name，不硬编码

raise ValueError(
    f"For {self.op_name}, sorting along a sharded dimension "
    f"(dim {dim} mapped to {mapping}) is not supported."
)

5. docstring 使用 Rules 格式

def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
    """
    Infer output layouts for Xxx operator.

    Rules:
        1. Input must not have Partial status.
        2. <算子特有约束>
        3. <输出 layout 规则>

    Args:
        cache_values (list): [input_layout, dim, ...]

    Returns:
        tuple: ((output_layout,), None)

    Raises:
        ValueError: If ...
    """

---

7.3 get_expand_impl 实现规范

新签名：get_expand_impl(self, func: Optional[callable], infer_result: tuple, cache_values: list) -> Optional[callable]

实现要点：

1. 卫语句（guard clause）优先

先判断不需要特殊处理的情况并提前返回 None，再定义闭包后立即返回：

def get_expand_impl(self, func, infer_result, cache_values):
    # 卫语句：无 contract dim 分片 或 无 bias → 不需要额外处理
    if not_needed:
        return None

    # 计算运行时参数
    scaling_factor = compute_scaling(infer_result, cache_values)

    def expand_impl(x, w, bias):
        return func(x, w, bias / scaling_factor)   # scaling_factor 通过闭包捕获

    return expand_impl                              # 定义后立即返回

2. 额外参数通过闭包捕获，不经 extra_info 传递

infer_result[1]（即 extra_info）可在 get_expand_impl 内读取，但三阶段分发流程不会将其展开传给
op_impl。需要的运行时信息应在 get_expand_impl 中计算并通过闭包捕获。

---

8. 迁移示例

8.1 Sort 算子（已迁移，完整示例）

from typing import Tuple, Optional
from .parallel_ops import DistributedOp


# stable 是 torch.sort 中 * 后的 keyword-only 参数，必须放 kwargs
# x、dim、descending 是普通位置参数，放 args
def _normalize_sort_args(x, dim=-1, descending=False, stable=False):
    return (x, dim, descending), {'stable': stable}


class SortDistributedOp(DistributedOp):
    """Distributed implementation for Sort operator."""

    # MindSpore SortExt Primitive 不接受 kwargs；Sort（mint.sort）为 functional_overload，
    # 行为同 PyTorch torch.sort，keyword-only 参数走 kwargs
    _MS_PRIMITIVE_OP_NAMES = frozenset({'SortExt'})

    def preprocess(self, args: tuple, kwargs: dict) -> tuple:
        args, kwargs = _normalize_sort_args(*args, **kwargs)
        input_tensor = args[0]
        dim, descending, stable = args[1], args[2], kwargs['stable']

        if self.op_name in self._MS_PRIMITIVE_OP_NAMES:
            # MindSpore Primitive：全部 positional
            local_args = (input_tensor.to_local(), dim, descending, stable)
            local_kwargs = {}
        else:
            # PyTorch torch.sort / MindSpore mint.sort：stable 为 keyword-only
            local_args = (input_tensor.to_local(), dim, descending)
            local_kwargs = {'stable': stable}

        cache_values = [input_tensor.layout, dim]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layouts for Sort operator.

        Rules:
            1. Input must not have Partial status.
            2. dim must be an integer within the valid range [-ndim, ndim-1].
            3. The sort dimension must not be sharded (including StridedShard multi-axis mappings).
            4. Output values and indices layouts are identical to the input layout.

        Args:
            cache_values (list): [input_layout, dim]

        Returns:
            tuple: ((values_layout, indices_layout), None)

        Raises:
            ValueError: If input has Partial status, dim is out of range, or the sort
                dimension is sharded.
        """
        layout, dim = cache_values[0], cache_values[1]
        self._check_partial_inputs([layout])

        if not isinstance(dim, int):
            raise ValueError(
                f"For {self.op_name}, dimension should be int, but got {type(dim)}"
            )

        in_tensor_map = layout.tensor_map
        ndim = len(in_tensor_map)
        if dim < -ndim or dim >= ndim:
            raise ValueError(
                f"For {self.op_name}, dimension out of range "
                f"(expected [{-ndim}, {ndim - 1}], but got {dim})"
            )
        if dim < 0:
            dim += ndim

        mapping = in_tensor_map[dim]
        is_sharded = (any(m != -1 for m in mapping) if isinstance(mapping, (list, tuple))
                      else mapping != -1)
        if is_sharded:
            raise ValueError(
                f"For {self.op_name}, sorting along a sharded dimension "
                f"(dim {dim} mapped to {mapping}) is not supported."
            )

        return ((layout, layout), None)

    # Sort 不需要 expand_impl，无需覆盖 get_expand_impl（基类返回 None）

8.2 Linear 算子（已迁移，完整示例）

from typing import Callable, List, Optional, Tuple
from .parallel_ops import DistributedOp


# bias 为普通位置参数，放 args；local_kwargs 为空（兼容 MindSpore Dense Primitive）
def _normalize_linear_args(x, weight, bias=None):
    return (x, weight, bias), {}


class LinearDistributedOp(DistributedOp):
    """Distributed implementation for Linear operator."""

    def preprocess(self, args: tuple, kwargs: dict) -> tuple:
        args, _ = _normalize_linear_args(*args, **kwargs)
        x_tensor, w_tensor, bias = args[0], args[1], args[2]
        local_args = (
            x_tensor.to_local(),
            w_tensor.to_local(),
            bias.to_local() if hasattr(bias, '_layout') else bias,
        )
        local_kwargs = {}
        cache_values = [
            x_tensor.layout,
            w_tensor.layout,
            bias.layout if hasattr(bias, '_layout') else None,
        ]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """
        Infer output layout for Linear operator.

        Rules:
            1. Input x and weight must not have Partial status.
            2. Weight must be 2D [out_features, in_features].
            3. x and weight must share the same mesh_shape.
            4. Output dim inherits sharding from weight's output dim (Shard(0)).
            5. If the contracting dim (in_features) is sharded, output carries Partial('sum').
            6. If bias is present and its output dim sharding differs from weight's, raise ValueError.

        Args:
            cache_values (list): [x_layout, w_layout, bias_layout_or_None]

        Returns:
            tuple: ((output_layout,), None)

        Raises:
            ValueError: If any rule above is violated.
        """
        x_layout, w_layout, bias_layout = cache_values[0], cache_values[1], cache_values[2]
        self._check_partial_inputs([x_layout, w_layout])

        x_map = x_layout.alias_tensor_map
        w_map = w_layout.alias_tensor_map

        if len(w_map) != 2:
            raise ValueError(
                f"For {self.op_name}, weight should be 2D [out_features, in_features], "
                f"but got {len(w_map)}D"
            )
        if x_layout.mesh.mesh_shape != w_layout.mesh.mesh_shape:
            raise ValueError(
                f"For {self.op_name}, x and weight must have the same mesh_shape, "
                f"but got x={x_layout.mesh.mesh_shape}, weight={w_layout.mesh.mesh_shape}"
            )

        # 输出维度推导、Partial 标记、bias 校验 …（省略具体推导，参考已迁移实现）
        ...
        return (output_layout,), None

    def get_expand_impl(self, func: Optional[Callable], infer_result: tuple,
                        cache_values: list) -> Optional[Callable]:
        x_layout = cache_values[0]
        bias_layout = cache_values[2]
        x_map = x_layout.alias_tensor_map
        x_contract_dim = len(x_map) - 1

        # 卫语句：无 contract dim 分片 或 无 bias → 不需要额外处理
        if x_map[x_contract_dim] == "None" or not bias_layout:
            return None

        output_layout = infer_result[0][0]
        scaling_factor = 1
        if isinstance(x_map[x_contract_dim], tuple):
            for axis in x_map[x_contract_dim]:
                scaling_factor *= output_layout.mesh.get_device_num_along_axis(axis)
        else:
            scaling_factor *= output_layout.mesh.get_device_num_along_axis(x_map[x_contract_dim])

        def expand_impl(x: object, w: object, bias: object) -> object:
            """Pre-scale bias to counteract the AllReduce accumulation over shards."""
            return func(x, w, bias / scaling_factor)   # scaling_factor 通过闭包捕获

        return expand_impl

---

9. UT 测试适配

9.1 必须清除 _LAYOUT_CACHE（关键）

根因：_build_layout 以 (device_mesh.to_hash(), placements, ndim) 为 key 缓存 Layout 对象。to_hash() 是内容哈希（基于 mesh_shape + dim_names），与 Python 对象 id 无关。

若某个测试调用 layout.set_partial_by_dev_axis() 修改了缓存中的 Layout 对象（如 test_*_partial_input_raises_error），后续使用相同 mesh 配置的测试会拿到已被污染的 Layout，导致误判。

解决方案：setUp/tearDown 中必须同时清理 _LAYOUT_CACHE：

from hyper_parallel.core.dtensor.dtensor import _build_layout, _LAYOUT_CACHE
from hyper_parallel.core.dtensor.device_mesh import _DEVICE_MESH_MAP
from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS

class TestParallelXxx(unittest.TestCase):
    def setUp(self) -> None:
        EXISTING_COMM_GROUPS.clear()
        _DEVICE_MESH_MAP.clear()
        _LAYOUT_CACHE.clear()   # ← 防止跨测试 Layout 状态污染

    def tearDown(self) -> None:
        EXISTING_COMM_GROUPS.clear()
        _DEVICE_MESH_MAP.clear()
        _LAYOUT_CACHE.clear()

9.2 测试调用方式适配

新方式（统一使用 cache_values）：

cache_values = [x_layout, w_layout, dim]
output_layouts, extra_info = op.infer_layout(cache_values)
output_layout = output_layouts[0]

get_expand_impl 调用方式：

# 旧：op.get_expand_impl(None, output_layout, (x_layout,), extra_args)
# 新：
impl = op.get_expand_impl(None, (output_layouts, None), cache_values)

9.3 UT 不使用 @arg_mark

UT 测试是平台无关的本地单元测试，不加 @arg_mark 装饰器（@arg_mark 仅用于 ST 集成测试）。

9.4 preprocess 测试

def test_sort_preprocess_stable_routing(self, mock_platform):
    """验证 stable 参数按算子类型正确路由到 local_args 或 local_kwargs。"""
    mock_tensor = MagicMock()
    mock_tensor.layout = x_layout
    mock_tensor.to_local.return_value = MagicMock()

    # SortExt Primitive：stable 在 positional args
    op_ext = SortDistributedOp("SortExt")
    local_args, local_kwargs, _ = op_ext.preprocess((mock_tensor,), {})
    assert local_kwargs == {}
    assert len(local_args) == 4  # tensor, dim, descending, stable

    # sort / Sort functional：stable 在 kwargs
    op_fn = SortDistributedOp("sort")
    local_args, local_kwargs, _ = op_fn.preprocess((mock_tensor,), {})
    assert local_kwargs == {'stable': False}
    assert len(local_args) == 3  # tensor, dim, descending

---

10. 参考实现

文件	说明
hyper_parallel/core/shard/ops/parallel_matmul.py	LinearDistributedOp（新模式完整示例）
hyper_parallel/core/shard/ops/parallel_sort.py	SortDistributedOp（新模式完整示例）
tests/ut/core/shard/ops/test_parallel_linear.py	Linear UT 测试（新调用方式）
tests/ut/core/shard/ops/test_parallel_sort.py	Sort UT 测试（新调用方式）
hyper_parallel/core/shard/_op_dispatch.py	_dispatch_layout_infer（三阶段分发路径）