【RFC】HyperParallel Symmetric Memory（单边通信）特性设计

一、背景与动机

1.1 集合通信的三大痛点

随着大模型规模指数级增长，分布式训练中卡间通信开销成为性能瓶颈。传统集合通信（AllGather/ReduceScatter/AllReduce 等）存在三类固有缺陷：

1. 同步开销大——需要在通信前后执行强同步；MoE 训练中 AllToAll 通信负载不均衡时，强同步导致快慢卡间非必要等待；
2. 中间数据拷贝多——集合通信要求数据位于专用通信 buffer（HBM 上特定地址空间），上下游算子数据需在标准内存↔通信 buffer 之间来回拷贝；
3. 异构开销高——建链、初始化、下发等步骤需要 CPU 参与，动态 shape 场景下可能导致 device 侧断流。

1.2 单边通信的优势

单边通信由发起方单侧主动完成，接收方无需显式调用接收函数。发起方可直接读写远端内存，接收方仅需确保内存处于可访问状态。

核心优势：

• 无强同步屏障，生产者-消费者异步协作；
• 数据直达远端对称内存，消除中间拷贝；
• 计算单元（device 侧）直接发起通信，无需 CPU 中转。

通算融合算子的开发工作量往往达到数人月。在 Python 侧提供单边通信接口后，仅需天级工作量即可实现大范围、细粒度、性能较好的通算掩盖。

---

二、核心概念：对称内存

单边通信需在通信前初始化：在各 PE（Processing Element，对应一块 NPU 卡）的 HBM 上各自申请相同大小的内存，将首地址共享给所有 PE，形成对称内存（Symmetric Memory）。

• aclshmem_malloc(size) — 从本地对称堆分配；
• aclshmem_ptr(local_ptr, target_pe) — 将本地地址转换为 target_pe 上相同偏移的远端地址；
• 堆大小由环境变量 SYMMETRIC_MEMORY_HEAP_SIZE 控制（默认 1 GB）。

---

三、设计目标与约束

目标	说明
Python 级单边通信原语	用户无需编写 AscendC kernel，通过 Python API 完成 put/get/signal
平台无关接口	同一套 API 同时支持 PyTorch 和 MindSpore
信号与数据一体化	signal tensor 本身也由对称内存分配
支持高级融合算子	在原语之上提供 AllGather、AllToAll、AG+MatMul、MatMul+RS 融合实现

当前限制（v1.0.0）：

• rendezvous() 未实现（shmem v1.0.0 接口开发中）；
• MindSpore 侧 shmem_alltoall、fused_all_gather_matmul、fused_matmul_reduce_scatter 尚未实现；
• TLS 安全传输在初始化时显式关闭（测试环境）；
• 固定 bootstrap IP:PORT（单机多卡场景）。

---

四、Python API 设计

import hyper_parallel.core.symmetric_memory as symm

# 内存分配（从对称堆分配，自动 P2P 可访问）
tensor = symm.empty(shape, dtype)

# 原语：push（本地写入远端）
symm.shmem_put(target, target_offset, src, src_offset, size, target_rank)

# 原语：pull（从远端拉取到本地）
symm.shmem_get(target, target_offset, src, src_offset, size, target_rank)

# 原语：signal 原子操作（更新远端信号）
symm.shmem_signal_op(signal, signal_offset, signal_value, signal_op, target_rank)
# signal_op: 0=SET, 1=ADD

# 原语：等待 signal 满足条件（本地自旋）
symm.shmem_wait_for_signal(depend_tensor, signal, signal_offset, compare_value, compare_op)
# compare_op: 0=EQ, 1=GT, 2=LT

# 复合原语：put + signal（原子保证 fence 在 signal 更新之前）
symm.shmem_put_with_signal(target, target_offset, src, src_offset,
                            size, signal, signal_offset, signal_value, signal_op, target_rank)

# 高级集合通信（基于以上原语实现）
symm.shmem_allgather(output_tensor, input_tensor)
symm.shmem_alltoall(send_tensor_list, receive_tensor, receive_list)
symm.fused_all_gather_matmul(a, b, c, gather_out, signal, block_size)
symm.fused_matmul_reduce_scatter(x1, x2, symm_tensor, signal, reduce_op)

# 全局同步
symm.barrier()

关键设计决策：

1. 所有偏移/大小参数使用 tensor 而非标量：参数需下传到 AscendC device kernel，通过 GM 地址传递，标量须封装为 tensor；
2. depend_tensor 参数：不参与信号逻辑，但在 MindSpore 图模式下用于建立正确的算子执行顺序依赖；
3. signal tensor 也用对称内存分配：aclshmemx_signal_op 需要 aclshmem_ptr 将 signal 地址翻译为远端地址。

---

五、Push 模式与 Pull 模式

基于基础原语可实现两类单边通信模式：

5.1 Push 模式（主要采用）

发起方将数据主动写入到远端对称内存，并发送信号通知接收方。

1. 各 rank 分配 output_tensor（对称内存）+ signal（对称内存）
2. 发起方：shmem_put_with_signal(output_tensor, offset, src, ..., signal, ..., target_pe)
   → 数据写入远端 + aclshmem_fence() + 远端 signal ADD 1
3. 接收方：shmem_wait_for_signal(output_tensor, signal, 0, world_size, EQ)
   → 等待收到所有 rank 的信号后继续执行

优势：比 Pull 模式少一次本卡内存拷贝，性能更好。

5.2 Pull 模式

发送方先将数据拷贝到本卡对称内存并发出信号，接收方等待信号后主动拉取。

优势：数据消费顺序可控，有利于确定性计算。

5.3 典型接口设计示意

---

六、原语操作语义

shmem_put / shmem_get

• put：发起方分发数据到远端。Kernel 内部多 AIV core 并行分片，pingpong 双缓冲搬运（GM→UB→远端GM），最后 aclshmem_fence() 保证远端可见性；
• get：发起方从远端拉取数据到本地。数据流向相反，其余结构同 put。

shmem_signal_op

对 target_pe 的 signal 执行原子 SET 或 ADD 操作（底层调用 aclshmemx_signal_op），仅由 AIV core 0 执行，无需 SyncAll。

shmem_wait_for_signal

在本地 AICore 上自旋轮询 signal 地址（底层调用 aclshmem_signal_wait_until），直到满足条件（EQ/GT/LT）才继续。阻塞期间占用 AICore 资源。

shmem_put_with_signal（最重要的复合原语）

将"数据写入"与"信号通知"合并为一个 AscendC kernel：

数据搬运（所有 AIV core 并行） → aclshmem_fence() → SyncAll() → AIV0: signal_op

fence 必须在 signal 之前：这是 RDMA 语义的核心保证，确保接收方观察到 signal 变化时数据一定已完整到达。

---

七、架构与分层设计

用户 Python 代码
      │
      ▼
┌─────────────────────────────────────────────┐
│  hyper_parallel.core.symmetric_memory       │  ← 统一 API 层（平台无关）
│  core/symmetric_memory/__init__.py          │
└─────────────────┬───────────────────────────┘
                  │ get_platform().get_symmetric_memory_handler()
        ┌─────────┴──────────┐
        ▼                    ▼
┌───────────────┐   ┌────────────────────┐
│ Torch Handler │   │ MindSpore Handler  │  ← 平台适配层（Python）
└───────┬───────┘   └────────┬───────────┘
        │                    │
        ▼                    ▼
┌────────────────┐  ┌─────────────────────┐
│libaclshmem_    │  │aclshmem_ms.so       │  ← C++ 平台绑定层
│torch.so        │  │(MindSpore Custom Op)│
│(TorchScript    │  │                     │
│ CustomClass)   │  │                     │
└───────┬────────┘  └──────────┬──────────┘
        └──────────┬───────────┘
                   ▼
        ┌──────────────────────┐
        │  AscendC Kernel 层   │  ← bisheng 编译，ops/*/kernel/*.cpp
        │  (ShmemKernel 命名空间)│
        └──────────┬───────────┘
                   ▼
        ┌──────────────────────┐
        │  CANN SHMEM 库       │  ← libhpshmem.so
        │  aclshmem_ptr()      │
        │  aclshmemx_signal_op │
        │  aclshmem_signal_wait│
        └──────────────────────┘

层	职责	关键文件
统一 API 层	平台无关入口，模块加载时选 Handler	core/symmetric_memory/__init__.py
平台适配层	持有初始化状态、Stream 池，转发调用	platform/{torch,mindspore}/symmetric_memory/symmetric_memory.py
C++ 绑定层	提取 tensor 地址、获取 stream、调用 kernel	ccsrc/symmetric_memory/platform/{torch,mindspore}/
AscendC Kernel 层	多 AIV core 数据搬运、信号原子操作	ccsrc/symmetric_memory/ops/*/kernel/*_kernel.cpp
shmem 底层	P2P 地址映射、RDMA 传输	第三方库 libhpshmem.so

---

八、C++ 内核实现要点

分核控制

单边通信在 NPU 上占用 AIV 核心进行数据搬运，需控制核数以打满卡间带宽同时不过多占用 AIV 资源。通过 block_dim 参数控制：PyTorch 侧默认 block_dim=1，MindSpore PyNative 为 4。

Kernel 内通过 aiv_idx_（本 AIV core 索引）和 aiv_num_（总数）均匀切分数据，各 core 独立处理自己的分片，最后 SyncAll() 汇聚。

put_mem_kernel 核心流程

Init():
  从 GM 读取参数（target_offset, src_offset, size）到寄存器

Process():
  size_per_core = size / aiv_num_            # 均匀分片
  remote_ptr = aclshmem_ptr(target_ptr, target_pe)  # 获取远端地址

  # Pingpong 双缓冲搬运（ping/pong 交替，MTE2 读时 MTE3 写，吞吐翻倍）
  for each block:
      WaitFlag(MTE3_MTE2)                    # 等待 buf 可用
      aclshmemi_copy_gm2ub(buf, src, size)   # 本地 GM→UB
      SetFlag(MTE2_MTE3) / WaitFlag(MTE2_MTE3)
      aclshmemi_copy_ub2gm(remote, buf, size)  # UB→远端 GM
      SetFlag(MTE3_MTE2)

  aclshmem_fence()    # 内存屏障：确保远端可见
  SyncAll()           # 多 AIV core 汇聚

put_mem_signal_kernel 核心顺序

数据搬运（同 put_mem）→ aclshmem_fence() → SyncAll() → AIV0: aclshmemx_signal_op()

fence 必须在 signal 之前，保证接收方看到 signal 时数据已完整写入。

signal_wait_until_kernel

调用 aclshmem_signal_wait_until(signal_addr, ACLSHMEM_CMP_EQ/GT/LT, compare_value) 在 AICore 上自旋等待。depend_tensor 作为输入 tensor 用于建立框架内的执行顺序依赖。

---

九、平台绑定层

PyTorch（TorchScript CustomClass）

libaclshmem_torch.so 注册两个 CustomClass：

• Manager：attr_init() 完成 shmem 初始化（禁用 TLS + aclshmemx_init_attr()）；malloc() 调用 aclshmem_malloc() 后用 at_npu::native::from_blob 零拷贝包装为 PyTorch tensor；
• Ops：每个方法提取 tensor data_ptr()，获取当前 NPU stream，调用对应 ShmemKernel::aclshmem_* 函数。

Python 侧初始化：延迟到首次调用 empty() 时触发，此时 distributed group 已建立。

MindSpore（PluggableAllocator + PyboostRunner）

两个层面：

1. 内存分配器（symmetric_memory_allocator.cc）：实现 PluggableAllocator（C 函数 Alloc/Free），通过 ms.runtime.MemPool 注入到 MindSpore 内存管理系统。首次 Alloc() 时读取 RANK_ID/RANK_SIZE/MS_SCHED_HOST/MS_SCHED_PORT 完成 shmem 初始化。

   

2. Custom Op（c_api/*/）：继承 PyboostRunner，在 LaunchKernel() 中调用 ShmemKernel::aclshmem_*。通过 MS_CUSTOM_OPS_EXTENSION_MODULE 宏和 op YAML 描述注册到 MindSpore。

---

十、高级算法实现要点

AllGather（Push 模式）

每个 rank 向所有其他 rank 的 output_tensor[rank_id*size] 处 push 自己的数据，每发完一个用 signal ADD 1 通知；接收方等待 signal 达到 world_size。

最多并发 16 路通信（comm_streams 池大小）。

AllToAll（Push 模式）

先用 dist.all_to_all_single 交换各 rank 的接收偏移（元数据），再各自向所有 rank push 对应 segment，signal 累加到 world_size 后等待完成。

fused_all_gather_matmul（通算流水）

计算公式：gather_out = AllGather(a)；c = ReduceScatter(gather_out @ b)

流水线设计：

• comm_streams：分块 push（每 rank 的数据按 block_size 分块，逐块向所有 rank 发送）；
• compute_streams[rank]：为每个 rank 分配独立 stream，等待各块 signal 到达后立即执行该块的 MatMul——通信与计算深度重叠。

fused_matmul_reduce_scatter（通算流水）

计算公式：output = ReduceScatter(x1 @ x2)

采用 (rank_id + rank) % world_size 轮转分配计算任务，避免所有 rank 同时争抢同一目标的网络带宽；接收方边等待边累加局部结果。

---

十一、初始化流程

PyTorch

dist.init_process_group()        # 分布式初始化
torch.npu.set_device(rank)
import hyper_parallel.core.symmetric_memory as symm
# 模块加载：torch.ops.load_library("libaclshmem_torch.so")

# 首次 symm.empty() 触发：
_manager.attr_init(rank_id, world_size, heap_size, "tcp://127.0.0.1:8662")
# → aclshmemx_set_conf_store_tls(false) + aclshmemx_init_attr()
# → 所有 rank 完成 P2P 握手，建立对称堆

注意：attr_init 是集合操作，所有 rank 必须同时调用，通常在 dist.barrier() 之后。

MindSpore

ms.communication.init()          # MindSpore 分布式初始化
import hyper_parallel.core.symmetric_memory as symm

# 首次 symm.empty() 创建 MemPool → 首次内存分配触发 Alloc()：
# 读取 RANK_ID, RANK_SIZE, MS_SCHED_HOST/PORT（端口 +13 避免冲突）
# aclshmemx_init_attr() 完成初始化

---

十二、已知限制与后续工作

限制	状态
rendezvous() 未实现	shmem v1.0.0 开发中
MindSpore 侧高级算子未实现	后续实现
PyTorch block_dim=1	只用 1 个 AICore block，待优化为 4
TLS 安全传输关闭	生产环境需评估

---

十四、快速上手示例

示例 1：Put + Signal 生产者-消费者

import torch
import torch.distributed as dist
import hyper_parallel.core.symmetric_memory as symm

dist.init_process_group()
rank = dist.get_rank()
torch.npu.set_device(rank)

data_buf = symm.empty(4096, dtype=torch.float16)
signal   = symm.empty(1, dtype=torch.int32)
signal.zero_()

dist.barrier()

if rank == 0:  # 生产者
    local_data = torch.randn(4096, dtype=torch.float16, device='npu')
    symm.shmem_put_with_signal(
        target=data_buf,
        target_offset=torch.tensor([0], dtype=torch.int64, device='npu'),
        src=local_data,
        src_offset=torch.tensor([0], dtype=torch.int64, device='npu'),
        size=torch.tensor([4096], dtype=torch.int64, device='npu'),
        signal=signal,
        signal_offset=torch.tensor([0], dtype=torch.int64, device='npu'),
        signal_value=torch.tensor([1], dtype=torch.int32, device='npu'),
        signal_op=0,    # SET
        target_rank=1
    )

elif rank == 1:  # 消费者
    symm.shmem_wait_for_signal(
        depend_tensor=data_buf,
        signal=signal,
        signal_offset=torch.tensor([0], dtype=torch.int64, device='npu'),
        compare_value=torch.tensor([1], dtype=torch.int32, device='npu'),
        compare_op=0    # EQ
    )
    result = data_buf.sum()  # 数据已就绪

示例 2：融合 AllGather + MatMul

m_local, k, n = 512, 1024, 2048
m_total = m_local * world_size

a = torch.randn(m_local, k, dtype=torch.float16, device='npu')
b = torch.randn(k, n, dtype=torch.float16, device='npu')

gather_out = symm.empty((m_total, k), dtype=torch.float16)
c          = torch.zeros(m_total, n, dtype=torch.float16, device='npu')
signal     = symm.empty(world_size * 4, dtype=torch.int32)
signal.zero_()

dist.barrier()

gather_result, matmul_result = symm.fused_all_gather_matmul(
    a, b, c, gather_out, signal, block_size=128
)
# gather_result: (m_total, k)  matmul_result: (m_total, n)

---

十五、代码结构速查

ccsrc/symmetric_memory/
├── ops/
│   ├── put_mem/kernel/put_mem_kernel.cpp        # AscendC put kernel
│   ├── get_mem/kernel/get_mem_kernel.cpp        # AscendC get kernel
│   ├── put_mem_signal/kernel/                   # AscendC put+signal kernel
│   ├── signal_op/kernel/                        # AscendC signal 原子操作
│   └── signal_wait_until/kernel/               # AscendC 自旋等待
└── platform/
    ├── torch/torch_bindings.cpp                 # TorchScript CustomClass
    └── mindspore/
        ├── c_api/allocator/                     # PluggableAllocator
        └── c_api/{put_mem,get_mem,...}/         # MindSpore Custom Op

platform/torch/symmetric_memory/symmetric_memory.py   # Torch Handler（含高级算法）
platform/mindspore/symmetric_memory/symmetric_memory.py  # MS Handler
core/symmetric_memory/__init__.py                     # 统一 API 入口

---

附录：底层 shmem API 速查

API	说明
aclshmemx_init_attr()	初始化 shmem，建立所有 PE 的对称堆（集合操作）
aclshmem_malloc(size)	从对称堆分配，返回本地 VA
aclshmem_ptr(local_ptr, pe)	返回 pe 上同等偏移的远端地址
aclshmemi_copy_gm2ub/ub2gm	GM↔UB 搬运（MTE2/MTE3 引擎）
aclshmem_fence()	保证之前所有 put 对远端可见
aclshmemx_signal_op(sig, val, op, pe)	对 pe 的 signal 原子 SET/ADD
aclshmem_signal_wait_until(sig, cmp, val)	本地自旋等待 signal 满足条件