特性名称

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。

方案概述

重建组件，核心决策：

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

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 🎉!