[Feature]: multicore shmem refactor

特性名称

Hyper-Parallel MultiCore 单边通信组件：CANN SHMEM Runtime、对称内存防护链与 Root PE 数据面

🚀 功能描述

背景

MegaMoE 等 MultiCore 融合 Kernel 需要 Root WORLD 对称内存上的 one-sided put/get/signal 语义，用于专家分发的数据面和跨 Rank 元数据交换。既有实现（逐算子 custom op kernel + 单文件 Python binding）在架构上存在四类结构性缺陷，无法通过增量修补达到可用状态：

1. host 侧阻塞语义：每次 put/get/signal 都发起 host 侧阻塞调用并自旋等待 device 完成，操作与周围 stream 串行化，从根本上排除通算交叠——而单边通信的核心收益场景正是"Tile/Expert/Shard 到达后立即被 AIV/AIC 消费"的流水，强同步 host 任务流与之直接冲突；
2. 无对称内存生命周期防护：free 后的 Tensor 及其视图仍可参与 RMA；且 CANN 堆的 best-fit 分配会确定性复用刚释放的同尺寸内存块，形成静默数据损坏——这是跨 Rank 对称堆特有的失效模式，单卡 allocator 语义无法覆盖；
3. 无运行中可观测性：泄漏只能靠 finalize 时点的确定性拒绝事后发现，定位需人工排查；
4. ABI 依赖扩散：aclshmem 头文件依赖散布在各 op 目录，CANN SHMEM 对部分接口明确禁止从 Custom Op 内调用，无白名单审计就无法判断"编译通过"是否等于"运行可用"，且无法开展无 NPU 依赖的 C++ UT。

方案概述

重建组件，核心决策：

分层结构（各层职责单一、依赖单向）：

层	职责	关键边界
cann/	CANN 类型、符号、枚举与调用隔离；按版本维护 Host/Device/CustomOp 可调用性白名单	不保存 Runtime 状态，收敛 CANN 升级的编译期语法变化
runtime/	进程唯一 Runtime：CANN 生命周期、冻结配置、设备身份、Root WORLD、Allocation Registry、最小 DFX 与分级日志，host mutex 全程串行化	不含 Torch/Python/消费者业务类型；Allocation ID 单调不复用
data_plane/	Root PE 坐标下的 MTE Put/Get/Signal/Wait/Barrier 与确定性字节切分（header-only 纯函数）	不拥有 Runtime 或 Allocation；消费者拥有 Tiling/Task/Event/流水
ops/all_gather/	内部参考消费者：910B UB ping-pong 分块 + 尾部重均衡的融合 AllGather	验证数据面，不形成统一 Kernel ABI
bindings/torch/ + Python 前端	引用计数生命周期、Tensor 转换、_api/_lifecycle/_runtime/_debug 四模块	Torch 只存在于生命周期与 Tensor 边界，不进入 C++ Runtime

stream-enqueue 通信语义：put/get/signal/wait_signal/barrier/all_gather 入队当前 NPU stream 后立即返回，完成由 stream 顺序保证；零字节输入为经校验的 no-op（numel>0 守卫位于 Runtime Resolve 层）。通信与计算可以在同一 stream 依赖图内流水化。

Torch 生命周期模型：acquire_shmem(root_group) 返回进程唯一 Runtime 的活动引用（Handle 只做引用，不承载功能方法）；首个 Handle 完成 Root 身份解析与 Bootstrap 初始化，末个 Handle 依次执行"NPU 静止 → ProcessGroup barrier → C++ Shutdown"。Root 身份用有序 global ranks 值语义判定——不同 ProcessGroup 对象只要成员、顺序、本地设备相同即共享 Runtime；析构/弱引用/atexit 不执行 collective 释放。

对称内存防护链（本特性新增的安全边界）：

• free 成功后底层 Storage 清零失效（绕过 torch resize 机制，直接 set_data_ptr+set_nbytes），Tensor 及其全部视图被所有 SHMEM 入口经活跃分配表校验拒绝——从构造上排除 use-after-free；
• 活跃 Allocation 丢失最后 Storage 引用时即时输出泄漏 ERROR（含 allocation_id/hex 基址/bytes/root_rank），并进入 debug_state 泄漏表；
• 设备身份初始化冻结，后续 set_device 漂移拒绝；跨生命周期失效 Tensor、重复 free、view 释放均确定性拒绝；
• empty 接口不设 device 参数，从构造上消除设备不匹配错误类。

首版能力边界（稳定范围）：Torch 唯一框架边界；一个进程一个本地 NPU 设备 + 一个 CANN Root WORLD（须覆盖完整 Torch WORLD）；数据面只验收节点内 MTE 路径；Bootstrap endpoint 仅用于初始化协调而非数据面。动态通信域、跨 SuperPod RoCE、Transport Selector/Topology Cache/Backend Factory 均明确不在首版。

现有替代方案

旧栈（本特性整体替换并删除）：五个逐算子 custom op kernel（get_mem/put_mem/put_mem_signal/signal_op/signal_wait_until）+ 单文件 Python binding。其局限在"背景"一节已述——补充两点替代性说明：

• 仓库内无其他 one-sided 实现；HCCL 集合通信无法表达 target_pe 定向的 put/get/signal 语义与设备侧自发通信；
• CANN SymmetricMemory（torch_npu 侧）绑定其自身生命周期与分配器，不能嵌入 MultiCore Kernel 的执行资源组模型（多 Layer 共用一个执行资源组 Handle、collective 顺序一致约束）。

与DFX相关性DF

• 日志（Loggability）：HYPER_PARALLEL_SHMEM_LOG_LEVEL=0 开启分配/释放详细日志——hex 基址、op 类型（op=Empty/op=Free）、截断调用栈（剔除内部帧），用于泄漏点回溯；错误信息标准化为"xxx must be xx, but got xx"（契约违背）与"事实陈述: key=value"（状态断言）两种形态；
• 故障定位（Fault localization）：debug_state() 返回进程本地只读快照（_DebugState dict 子类，人类可读 __repr__：逐字段换行、地址 hex）——含活跃分配报表（allocation_id/base/bytes）、堆用量四元组（heap_size/allocated/max_allocated/remaining）、泄漏分配表；孤儿分配即时 ERROR；
• 故障注入（测试侧）：C++ UT 基于 fakes 注入 ACL/aclshmem 假实现，覆盖重初始化、注册表边界、设备漂移等故障路径，CI 无需 NPU。

提议的新API

Python 侧（hyper_parallel.core.multicore.shmem，内部 API，仅面向 MultiCore 子组件消费方，不向训练用户或模型实现提供稳定接口）：

# 生命周期（Handle 引用模型）
handle = shmem.acquire_shmem(root_group=None)   # 首个 Handle 初始化，末个 Handle 终态关闭
handle.closed; handle.close()                   # 幂等；非末个引用仅本地计数

# 对称分配（collective；无 device 参数，设备身份初始化冻结）
shmem.empty(*size, dtype=None, alignment=None)  # Root WORLD 对称堆分配
shmem.free(tensor)                              # collective 释放完整 Allocation；view 拒绝

# 单边通信（stream-enqueue：入队即返回，完成由 stream 顺序保证）
shmem.put(dst, src, pe) / shmem.get(dst, src, pe)      # MTE RMA
shmem.signal(signal_tensor, value, pe) / shmem.wait_signal(signal_tensor, value)
shmem.barrier()
shmem.all_gather(output, input)                 # 内部参考算子（910B）

# 诊断
shmem.debug_state()                             # 只读、进程本地、非 collective

CANN 侧构建产物：hyper_parallel_shmem_ops（custom op + all_gather kernel）与 hyper_parallel_shmem_torch（torch binding），统一由 ccsrc/CMakeLists.txt 两阶段构建编排；CANN SHMEM 依赖版本锁定并应用 Hyper patch。

是否影响现有API

• 删除：旧 Python 入口（lifecycle.py/_bindings.py）与五个逐算子 custom op 及其测试——均为内部实现，仓库内无外部消费者；
• 切换：MegaMoE 执行资源层（唯一消费方）迁移至新 API，常驻分配模式（执行资源组持有 Handle，一次性分配、模块关闭时按 collective 顺序释放）；
• 不新增对外稳定接口：acquire_shmem 等全部入口为 MultiCore 内部能力；Handle 不承载通信方法、不进入 Device 公共 ABI 的约束保持不变；
• 不修改：MegaMoE 对外 forward/backward 行为、训练接口与 checkpoint 兼容性。

补充信息

验证结果：

• C++ UT（fakes，无 NPU 依赖，CI 全绿）：config、分配注册表、Runtime 生命周期与重初始化、传输分块、host 适配层、device surface；
• NPU ST（单机 8 卡 torchrun）：生命周期/引用计数、device 漂移与失效 Tensor 拒绝、free 后防护链（含视图/重复 free/泄漏检测）、put/get/signal 正确性（含非默认 stream 顺序）、AllGather 字节边界矩阵（含零字节与未对齐尾部）；
• 端到端（qwen_moe_benchmark，8 卡 EP）：两次独立运行四层精度对比全过（初始参数逐位；前向 max err 1.9e-6/2.9e-6；梯度 99 tensor 占容差预算 0.6%；更新参数在预算内，2×lr 的 Adam 首步符号翻转为已知预期签名）；稳态完整 optimizer step（rank-max）缩短 21~26%，两次运行方向一致。

设计文档：总体设计原则、Runtime 代码级设计（v3.5 终态）、单边通信架构规划与实现评审文档随本特性在仓库内维护，可在评审时提供。

已知约束：AllGather kernel 面向 910B；对称堆大小由 HYPER_PARALLEL_SHMEM_HEAP_SIZE 决定，初始化后固定；moe_gating 类消费方接入须满足"各 Root Rank 执行资源/Allocation/Handle collective 顺序一致"。

配套 PR：pr_message_shmem_onesided.md（实现、测试与消费方切换随该 PR 上库）。

Thanks for contributing 🎉!

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

Start with pr_message_shmem_onesided.md for the proposed implementation, then inspect ccsrc/CMakeLists.txt and the existing lifecycle.py and _bindings.py entries being replaced. Run the referenced C++ fake-based unit tests and NPU system tests to understand the validation surface. Done means the new runtime, bindings, data plane, tests, and MegaMoE consumer migration satisfy the stated lifecycle, communication, and performance checks.