mori通信库分析(一)——对称内存传输数据过程

前言

 在高性能计算领域，对称内存（Symmetric Memory）是简化多进程/多GPU通信的关键抽象。AMD开源的MORI库（基于IBGDA）借助对称堆（Symmetric Heap）和虚拟内存管理（VMM），实现了高效的GPU间RDMA数据传输。本文简单分析设备端执行shmem_put时的数据发送路径。

主机端初始化：建立对称映射

 注册对称对象, RegisterSymmMemObj：

// 1. 分配CPU侧对象
SymmMemObj* cpuMemObj = new SymmMemObj();
cpuMemObj->localPtr = localPtr;

// 2. 通过Allgather交换本地基地址
cpuMemObj->peerPtrs = calloc(worldSize, sizeof(uintptr_t));
bootNet.Allgather(&localPtr, cpuMemObj->peerPtrs, sizeof(uintptr_t));
// 现在 peerPtrs[pe] 是远端进程 pe 的本地堆基地址（虚拟地址）

// 3. 处理P2P传输（同节点）
cpuMemObj->p2pPeerPtrs = calloc(worldSize, sizeof(uintptr_t));
// 交换IPC句柄并打开，得到可直接访问的映射地址
hipIpcGetMemHandle(&handle, localPtr);
bootNet.Allgather(&handle, ipcMemHandles, sizeof(hipIpcMemHandle_t));
for each peer i:
    if (CanUseP2P(i))
        hipIpcOpenMemHandle(&cpuMemObj->p2pPeerPtrs[i], ipcMemHandles[i], ...);
// 对于非RDMA传输（P2P/SDMA），覆盖peerPtrs为本地可访问地址
if (transportType != RDMA) cpuMemObj->peerPtrs[i] = cpuMemObj->p2pPeerPtrs[i];

// 4. RDMA注册（仅当有RDMA peer）
if (anyRdmaPeer) {
    RdmaMemoryRegion mr = rdmaContext->RegisterRdmaMemoryRegion(localPtr, size);
    cpuMemObj->lkey = mr.lkey;
    cpuMemObj->peerRkeys[rank] = mr.rkey;
}
bootNet.Allgather(&cpuMemObj->peerRkeys[rank], cpuMemObj->peerRkeys, sizeof(uint32_t));

 将cpuMemObj及其指针数组通过hipMemcpy拷贝到GPU，并通过globalGpuStates->heapObj供设备端内核访问。ConfigureHeapInfoForGpu

 关键点：peerPtrs在RDMA模式下存储的是远端虚拟地址，在P2P模式下存储的是本进程映射的本地地址，但设备端代码统一通过peerPtrspe + offset计算目标地址，无需区分底层传输。

设备端发送流程

PutNbi APIs

PutNbi APIs

// ============================================================================
// PutNbi APIs - Address-based only
// ============================================================================
__device__ __attribute__((visibility("default"))) int mori_shmem_putmem_nbi_thread(
    void* dest, const void* source, size_t bytes, int pe, int qpId) {
  mori::shmem::ShmemPutMemNbiThread(dest, source, bytes, pe, qpId);
  return 0;
}

// ============================================================================
// PutNbi APIs - Warp Scope (Address-based only)
// ============================================================================
__device__ __attribute__((visibility("default"))) int mori_shmem_putmem_nbi_warp(void* dest,
                                                                                 const void* source,
                                                                                 size_t bytes,
                                                                                 int pe, int qpId) {
  mori::shmem::ShmemPutMemNbiWarp(dest, source, bytes, pe, qpId);
  return 0;
}

// ============================================================================
// PutNbi APIs - Block Scope (Address-based only)
// ============================================================================
__device__ __attribute__((visibility("default"))) int mori_shmem_putmem_nbi_block(
    void* dest, const void* source, size_t bytes, int pe, int qpId) {
  mori::shmem::ShmemPutMemNbiBlock(dest, source, bytes, pe, qpId);
  return 0;
}

OpenSHMEM Style PutNbi APIs

#define DEFINE_SHMEM_PUT_MEM_NBI_ADDR_API_TEMPLATE(Scope)                                      \
  inline __device__ void ShmemPutMemNbi##Scope(void* dest, const void* source, size_t bytes,   \
                                               int pe, int qpId = 0) {                         \
    DISPATCH_TRANSPORT_TYPE(ShmemPutMemNbi##Scope##Kernel, pe, dest, source, bytes, pe, qpId); \
  }

DEFINE_SHMEM_PUT_MEM_NBI_ADDR_API_TEMPLATE(Thread)
DEFINE_SHMEM_PUT_MEM_NBI_ADDR_API_TEMPLATE(Warp)
DEFINE_SHMEM_PUT_MEM_NBI_ADDR_API_TEMPLATE(Block)

DISPATCH_TRANSPORT_TYPE

#define DISPATCH_TRANSPORT_TYPE(func, pe, ...)                                    \
  GpuStates* globalGpuStates = GetGlobalGpuStatesPtr();                           \
  application::TransportType transportType = globalGpuStates->transportTypes[pe]; \
  if (transportType == application::TransportType::RDMA) {                        \
    func<application::TransportType::RDMA>(__VA_ARGS__);                          \
  } else if (transportType == application::TransportType::P2P) {                  \
    func<application::TransportType::P2P>(__VA_ARGS__);                           \
  } else if (transportType == application::TransportType::SDMA) {                 \
    func<application::TransportType::SDMA>(__VA_ARGS__);                          \
  } else {                                                                        \
    assert(false);                                                                \
  }

ShmemPutMemNbiThread调用的函数为ShmemPutMemNbiThreadKernel。DISPATCH_TRANSPORT_TYPE根据配置的传输类型，分发到对应函数。

以ShmemPutMemNbiThreadKernel<application::TransportType::RDMA>为例。

template <>
inline __device__ void ShmemPutMemNbiThreadKernel<application::TransportType::RDMA>(
    const application::SymmMemObjPtr dest, size_t destOffset,
    const application::SymmMemObjPtr source, size_t sourceOffset, size_t bytes, int pe, int qpId) {
  bool need_turn{true};
  uint64_t turns = __ballot(need_turn);
  while (turns) {
    uint8_t lane = __ffsll((unsigned long long)turns) - 1;
    int pe_turn = __shfl(pe, lane);
    if (pe_turn == pe) {
      DISPATCH_PROVIDER_TYPE_COMPILE_TIME(ShmemPutMemNbiThreadKernelImpl, dest, destOffset, source,
                                          sourceOffset, bytes, pe, qpId);
      need_turn = false;
    }
    turns = __ballot(need_turn);
  }
}

ShmemPutMemNbiThreadKernelImpl

ShmemPutMemNbiThreadKernelImpl 是 MORI 库中 设备端非阻塞 Put 操作 的核心模板函数，负责将一个对称内存区域的数据通过 RDMA 发送到远端 PE。该函数具备以下关键特性：

• 支持静态堆与 VMM 分块堆；
• 支持多种网卡（MLX5/BNXT/PSD）；
• Warp 级协同预留 WQE 槽位，减少原子竞争；
• 自动流控，当发送队列满时主动回收完成事件。

初始化与循环控制

bool needsChunking = globalGpuStates->useVMMHeap;
size_t currentOffset = 0;
size_t remaining = bytes;
while (true) {
    // 检查本线程是否还有剩余数据
    bool has_remaining = (remaining > 0);
    uint64_t activemask = __ballot(has_remaining);
    if (activemask == 0) break;
    ...
}

_ballot 同步 Warp 内所有线程，获得仍有数据待发送的线程掩码。

若掩码为 0，则所有线程已完成，退出循环。

只有掩码中的线程参与本轮发送，其他线程跳过。

计算传输大小与密钥

根据 needsChunking 分两种情况：

静态堆模式

• lkey = source->lkey（全局统一
• srcAddr = source->localPtr + sourceOffset + currentOffset
• raddr = dest->peerPtrspe + destOffset + currentOffset
• rkey = dest->peerRkeyspe
• transfer_size = remaining（一次性发送剩余全部）

VMM 堆模式

• 查询源端 Chunk：VmmQueryLocalKey(srcAddr, remaining, lkey, src_chunk_size)
• 查询目标端 Chunk：VmmQueryRemoteAddr(dstAddr, pe, remaining, raddr, rkey, dst_chunk_size)
• transfer_size = min(src_chunk_size, dst_chunk_size, remaining)

注意：VmmQuery* 会返回当前地址所在 Chunk 的剩余大小及对应的 Key。

关键点：VMM 模式下，每次循环只传输一个 Chunk 内的数据，确保 lkey 和 rkey 有效。

对于 BNXT 网卡，需要额外计算每个线程的 PSN（包序列号）数量：

if constexpr (PrvdType == BNXT) {
    psnCnt = (transfer_size + wq->mtuSize - 1) / wq->mtuSize;
    my_psn_excl = WarpActivePsnPrefix(psnCnt, activemask, &warp_total_psn);
}

WarpActivePsnPrefix 计算每个活跃线程之前的 PSN 总数，用于后续分配 PSN 偏移。

Warp 级 WQE 槽位预留

步骤说明：

1. 确定本次需要预留的 WQE 数量 num_wqes：

• MLX5/PSD：每个活跃线程 1 个 WQE → num_wqes = num_active_lanes
• BNXT：每个线程需要的 WQE 数为 psnCnt（因为每个 WQE 可能携带多个 MTU 数据）→ num_wqes = warp_total_psn

2. 由 Leader 线程（最后一个活跃线程）执行原子预留：

if (is_leader) {
    if constexpr (MLX5) warp_sq_counter = atomicAdd(&wq->postIdx, num_active_lanes);
    else if constexpr (BNXT) {
        // 同时增加 msn 和 psn
        atomic_add_packed_msn_and_psn(&wq->msnPack, num_active_lanes, warp_total_psn, ...);
        warp_sq_counter = warp_msntbl_counter; // MSN 即 SQ 索引
        atomicMax(&wq->postIdx, warp_sq_counter + num_active_lanes);
    }
}

3. 广播起始槽位：warp_sq_counter = __shfl(warp_sq_counter, leader_phys_lane_id)

4. 计算每个线程的私有 WQE 索引：

   MLX5：my_sq_counter = warp_sq_counter + my_logical_lane_id

   BNXT：my_sq_counter = warp_sq_counter + my_logical_lane_id（但 my_msntbl_counter 和 my_psn_counter 另有偏移）

设计思想：将多个线程的预留合并为一次原子操作，降低争用。

发送队列流控

while (true) {
    uint64_t db_touched = wq->dbTouchIdx;  // 已提交硬件（但可能未完成）
    uint64_t db_done = wq->doneIdx;        // 已确认完成
    uint64_t active = db_touched - db_done;
    uint64_t free = wq->sqWqeNum - active;
    uint64_t need_until_end = warp_sq_counter + num_wqes - db_touched;
    if (free > need_until_end) break;
    // 队列不够，主动等待完成
    ShmemQuietThreadKernelImpl<PrvdType>(pe, qpId);
}

dbTouchIdx：网卡已读取的 WQE 最大索引（但可能未完成）。

doneIdx：已完成（CQE 已回收）的最大索引。

若空闲槽位不足以容纳本次预留，则调用 ShmemQuietThreadKernelImpl 尝试从 CQ 中回收完成事件，更新 doneIdx。

循环直至满足条件。

填充 WQE 与更新门铃

调用 core::PostWrite。该函数根据网卡类型在 SQ 内存中写入 WQE，返回一个门铃值（包含需要通知硬件的状态）。

uint64_t dbr_val;
if constexpr (MLX5) {
    dbr_val = PostWrite(wq, my_sq_counter, my_sq_counter, my_sq_counter, is_leader, qpn,
                        srcAddr, lkey, raddr, rkey, transfer_size);
} else if constexpr (BNXT) {
    dbr_val = PostWrite(wq, my_sq_counter, my_msntbl_counter, my_psn_counter, is_leader, qpn,
                        srcAddr, lkey, raddr, rkey, transfer_size);
} else if constexpr (PSD) {
    dbr_val = PostWrite(wq, my_sq_counter, my_sq_counter, my_sq_counter, is_leader, qpn,
                        srcAddr, lkey, raddr, rkey, transfer_size);
}

不同网卡需要不同的索引（MSN/PSN）。

is_leader 决定是否在 WQE 中设置"完成通知"标志（通常只有最后一个 WQE 需要产生 CQE，减少中断开销）。

内存屏障与门铃提交

__threadfence_system();  // 确保所有 WQE 写入全局可见
if (is_leader) {
    // 等待 dbTouchIdx 追上预留起始位置（确保前面的 WQE 已被硬件读取）
    while (wq->dbTouchIdx != warp_sq_counter) {}
    // 更新网卡侧的 dbr 记录
    UpdateSendDbrRecord(wq->dbrRecAddr, warp_sq_counter + num_wqes);
    __threadfence_system();
    // 敲响门铃
    RingDoorbell(wq->dbrAddr, dbr_val);
    __threadfence_system();
    // 更新 needConsIdx 并设置 dbTouchIdx
    atomicAdd(&cq->needConsIdx, 1);
    atomicStore(&wq->dbTouchIdx, warp_sq_counter + num_wqes);
}

UpdateSendDbrRecord 更新网卡侧的内存指针，告诉硬件新的 WQE 范围。

RingDoorbell 写入 PCIe 门铃寄存器，触发硬件 DMA 读取 WQE 并执行。

dbTouchIdx 的更新让其他线程知道这些槽位已被提交。

特殊优化与细节

1. 多线程协同预留
   只有 Leader 线程执行原子操作，其他线程通过 __shfl 获取索引，避免所有线程争抢同一个原子变量。这种做法充分利用了 Warp 内通信的低延迟优势。
2. 流控与完成回收的原子锁
   ShmemQuietThreadKernelImpl 内部会尝试获取 CQ 锁（pollCqLock），只有获取到的线程执行 PollCq，其他线程自旋等待或返回（非最终静默场景）。这保证了 CQ 处理的串行化。

Reference

1 rocm mori

2 vLLM distributed inference with MoRI

3 ROCmPD add moriio kv connector

4 nccl分析(二)------RDMA带外建链过程

5 RDMA带外建联过程