09: Symmetric Memory 演进深度对比 — DeepEP V1 / V2 vs DeepGEMM Mega MoE

timeline
    title Symmetric Memory 在 MoE 通信中的演进
    2024 : DeepEP V1 (NVSHMEM)
         : nvshmem_align + IBGDA put
         : 显式 Push 模型 + 三阶段流水线
         : SM90 NVLink + RDMA
    2025 : DeepEP V2 (NCCL Gin)
         : ncclMemAlloc + Window Register
         : 仍 Push 但用 QP 抽象
         : SM90/SM100 NVLink + RDMA
    2025 : DeepGEMM Mega MoE
         : torch.distributed._symmetric_memory
         : Pull 模型 + TMA load/store
         : SM100 NVLink Symmetric Memory only

// csrc/kernels/backend/nvshmem.cu (line 23-25)
void* alloc(const size_t& size, const size_t& alignment) {
    return nvshmem_align(alignment, size);  // NVSHMEM 对称内存分配
}

# deep_ep/buffers/legacy.py (line 33-36)
def __init__(self, group, num_nvl_bytes=0, num_rdma_bytes=0, ...):
    # 两个 buffer 分别用于 NVLink 和 RDMA

// csrc/kernels/legacy/internode.cu (line 617-623)
nvshmemi_ibgda_put_nbi_warp<true>(
    reinterpret_cast<uint64_t>(rdma_channel_meta.recv_buffer(rdma_rank)),  // 远端目标
    reinterpret_cast<uint64_t>(rdma_channel_meta.send_buffer(dst_rdma_rank)), // 本地源
    sizeof(int) * (LEGACY_NUM_MAX_NVL_PEERS * 2 + 2),
    translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank),
    channel_id, lane_id, 0);

// csrc/kernels/legacy/internode.cu (line 93-95)
template <bool kLowLatencyMode>
__forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx(const nvshmem_team_t& rdma_team) {
    kLowLatencyMode ? nvshmem_sync(rdma_team) : nvshmem_sync_all();
}

// csrc/kernels/backend/symmetric.hpp (line 123-140)
// GPUSymmetricMemory: 纯 GPU，当前默认
class GPUSymmetricMemory final : public SymmetricMemory {
    explicit GPUSymmetricMemory(const int64_t& num_bytes) {
        NCCL_CHECK(ncclMemAlloc(&ptr, num_bytes));  // NCCL 分配器
        this->num_gpu_bytes = num_bytes;
    }
};

// csrc/kernels/backend/nccl.cu (line 150-153)
void* NCCLSymmetricMemoryContext::get_sym_ptr(void* ptr, const int& dst_rank_idx) const {
    const auto offset = static_cast<uint8_t*>(ptr) - static_cast<uint8_t*>(mapped_window_ptr);
    return static_cast<uint8_t*>(nvl_window_ptrs[dst_rank_idx]) + offset;
}

# deep_gemm/mega/__init__.py (line 42-49)
allocator = torch if group.size() == 1 else symm_mem
self.buffer = allocator.empty(num_bytes, dtype=torch.int8, device='cuda')
self.handle = (
    types.SimpleNamespace(buffer_ptrs=[self.buffer.data_ptr()])
    if group.size() == 1
    else symm_mem.rendezvous(self.buffer, group=group)
)

// deep_gemm/layout/sym_buffer.cuh (line 34-40)
template <typename ptr_t>
CUTLASS_DEVICE ptr_t map(const ptr_t& ptr, const uint32_t& dst_rank_idx) const {
    if constexpr (kNumRanks == 1) return ptr;
    int64_t mapped_ptr = offsets[dst_rank_idx] + reinterpret_cast<int64_t>(ptr);
    return *reinterpret_cast<ptr_t*>(&mapped_ptr);
}

flowchart TB
    subgraph V1["DeepEP V1: Buffer(group, num_nvl_bytes, num_rdma_bytes)"]
        V1A["nvshmem_align() 分配 RDMA buffer"] --> V1B["CUDA IPC handle 交换 NVLink buffer"]
        V1B --> V1C["nvshmemx_init_attr() 初始化 NVSHMEM"]
    end

    subgraph V2["DeepEP V2: ElasticBuffer(group, num_bytes, num_cpu_bytes)"]
        V2A["symmetric::alloc() 分配 symmetric memory"] --> V2B["ncclCommWindowRegister() 注册窗口"]
        V2B --> V2C["ncclDevCommCreate() 创建 Gin QP"]
    end

    subgraph MM["DeepGEMM: SymmBuffer(group, ...)"]
        MMA["symm_mem.empty() 分配 buffer"] --> MMB["symm_mem.rendezvous() 注册到进程组"]
        MMB --> MMC["handle.buffer_ptrs 获取所有 rank 地址"]
    end

# deep_ep/buffers/legacy.py
buffer = Buffer(group, num_nvl_bytes=1024**3, num_rdma_bytes=1024**3)
# 内部: nvshmem_align() + IPC handle exchange

# deep_ep/buffers/elastic.py (line 346-354)
runtime = _C.ElasticBuffer(group.rank(), group.size(),
    nccl_comm_handle.get(), cpu_comm,
    num_bytes, num_cpu_bytes, ...)
# 内部: ncclMemAlloc/cuMemCreate + ncclCommWindowRegister

# deep_gemm/mega/__init__.py (line 42-49)
self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')
self.handle = symm_mem.rendezvous(self.buffer, group=group)
# handle.buffer_ptrs = [addr_rank0, addr_rank1, ...]

NVLink buffer (num_nvl_bytes): [send_buf | recv_buf] per channel per rank
RDMA buffer (num_rdma_bytes):  [send_buf | recv_buf] per channel per rank

Symmetric Memory: [[Workspace] GPU buffer [CPU buffer]]
                  ↑ mapped_window_ptr

SymmBuffer: [Workspace | Input Token | L1 Pool | L2 Pool | Combine]
            ↑ 所有 rank 相同布局，map() 通过偏移访问

// csrc/kernels/legacy/internode.cu (line 688-693)
// RDMA Sender Warp: 将本地 token 写入远端 RDMA buffer
auto st_broadcast = [=](const int key, const int4& value) {
    for (int j = 0; j < num_topk_ranks; ++j)
        st_na_global(reinterpret_cast<int4*>(dst_send_buffers[j]) + key, value);
    // ↑ 写入远端 send_buffer (通过 NVLink/RDMA)
};
UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4, ld_nc_global, st_broadcast);

// csrc/elastic/buffer.hpp (line 469-489)
// all_gather: 通过 cudaMemcpyBatchAsync 推送到远端
for (int i = 0; i < nccl_context->num_ranks; ++i) {
    for (int j = 0; j < num_tensors; ++j) {
        void* dst_ptr = nccl_context->get_sym_ptr(
            math::advance_ptr(buffer, offset[j] + x.nbytes() * nccl_context->rank_idx),
            dst_rank_idx);  // 远端目标地址
        dst_ptrs[count] = dst_ptr;
        src_ptrs[count] = x.data_ptr();  // 本地源
    }
}
cudaMemcpyBatchAsync(dst_ptrs.data(), src_ptrs.data(), sizes.data(), num_copies, attrs, comm_stream);

// sm100_fp8_fp4_mega_moe.cuh (line 533-555)
// Dispatch Warp: 从远端 rank 拉取 token 到本地 smem
const auto src_base_ptr = sym_buffer.map(
    buffer.input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(),
    current_rank_in_expert_idx);  // ← 远端源地址
const auto dst_base_ptr = buffer.l1_token_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).get_base_ptr();

// TMA load: 远端 → 本地 smem → 本地 L1 buffer
for (uint32_t i = 0; i < kNumChunks; ++i) {
    ptx::tma_load_1d(
        pull_buffer.get_base_ptr(),                    // 本地 smem 目标
        math::advance_ptr(src_base_ptr, i * kNumBytesPerPull),  // 远程源
        pull_mbarrier, kNumBytesPerPull);
    ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kNumBytesPerPull);
    issue_and_wait_pull_store(i);  // smem → L1 buffer (TMA store)
}

// V1: Forwarder warp 将 NVLink recv_buffer 数据写入远端 combine buffer
// V2: 类似，通过 get_sym_ptr 定位远端 combine buffer

// sm100_fp8_fp4_mega_moe.cuh (line 1293-1299)
// 从本地 smem 读取，写入远端 combine buffer
const auto packed = ptx::ld_shared(reinterpret_cast<float4*>(smem_ptr));
const auto dst_token = buffer.combine_token_buffer.get_rank_buffer(dst_topk_idx)
                       .get_data_buffer(dst_token_idx);
const auto dst_ptr = math::advance_ptr<float4>(dst_token.get_base_ptr(), ...);
*sym_buffer.map(dst_ptr, dst_rank_idx) = packed;  // ← NVLink store 写回

sequenceDiagram
    participant S as Source GPU
    participant T as Target GPU

    Note over S,T: Push Model (DeepEP V1/V2)
    S->>S: 读取本地 token
    S->>T: RDMA/NVLink 写入 Target buffer
    T->>T: 轮询等待数据到达
    T->>T: 从 buffer 读取处理

    Note over S,T: Pull Model (DeepGEMM)
    T->>T: 计算远端 source 地址
    T->>S: TMA Load 从 Source HBM 读取
    S->>S: (被动响应)
    T->>T: 本地 L1 buffer 写入

# deep_gemm/mega/__init__.py
import torch.distributed._symmetric_memory as symm_mem
# PyTorch SymMem 内部: 同节点 NVLink, 跨节点 RDMA (NCCL 自动选择)

// csrc/kernels/backend/nccl.cu (line 143-147)
// 仅获取 NVLink domain 内 peer 的指针
nvl_window_ptrs.resize(num_nvl_ranks);
for (int i = 0; i < num_nvl_ranks; ++i)
    NCCL_CHECK(ncclGetLsaDevicePointer(window, 0, i, &nvl_window_ptrs[i]));

// csrc/kernels/legacy/internode.cu (line 147-149)
// 两阶段同步: NVLink intra-node + RDMA inter-node
if (thread_id == 32)
    nvshmem_sync_with_same_gpu_idx<kLowLatencyMode>(rdma_team);
barrier_block<LEGACY_NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);

// csrc/elastic/buffer.hpp (line 191-199)
void barrier(const bool& use_comm_stream, const bool& with_cpu_sync, const bool& sequential) {
    // GPU barrier via NCCL
    launch_barrier(nccl_context->dev_comm, nccl_context->window,
                   workspace,
                   nccl_context->scaleout_rank_idx, nccl_context->scaleup_rank_idx,
                   nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks,
                   num_gpu_timeout_cycles, ...);
}

// deep_gemm/comm/barrier.cuh (line 46-89)
// NVLink barrier: 跨 rank 全局同步
void nvlink_barrier(...) {
    grid_sync<kNumSMs, kGridSyncIndex>(...);  // 1. 节点内 SM 同步
    if (sm_idx == 0) {
        // SM 0 发送信号到远端 rank
        ptx::red_add_rel_sys(sym_buffer.map(signal_ptr, thread_idx), signal_sign ? -1 : 1);
        // 等待所有 rank 到达
        while (ptx::ld_acq_sys(signal_ptr) != target) { /* spin */ }
    }
    grid_sync<kNumSMs, kGridSyncIndex>(...);  // 3. 再次节点内同步
}

// sm100_fp8_fp4_mega_moe.cuh (line 309-311)
constexpr uint32_t kBeforeDispatchPullBarrierTag = 1;      // Dispatch 拉取前
constexpr uint32_t kBeforeCombineReduceBarrierTag = 2;     // Combine 归约前
constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3;     // Workspace 清理后

// sm100_fp8_fp4_mega_moe.cuh (line 329-416)
if (warp_idx < kNumDispatchWarps) {
    // Dispatch Warps: 远程拉取 token
    cutlass::arch::warpgroup_reg_dealloc<kNumDispatchRegisters>();  // 48 regs
} else if (warp_idx < kNumDispatchWarps + kNumMMANonEpilogueWarps) {
    // MMA Warps: GEMM 计算
    cutlass::arch::warpgroup_reg_dealloc<kNumNonEpilogueRegisters>();  // 40 regs
} else {
    // Epilogue Warps: SwiGLU + Combine 写回
    // 208 regs (高寄存器用于 TMEM load/STSM)
}

graph LR
    subgraph Metrics["性能关键指标"]
        L["延迟"]
        B["带宽"]
        S["SM 利用率"]
        O["计算通信重叠"]
    end

    subgraph V1["DeepEP V1"]
        L1["中 (RDMA put)"]
        B1["中 (FIFO 开销)"]
        S1["低 (专用 warp group)"]
        O1["无"]
    end

    subgraph V2["DeepEP V2"]
        L2["中 (Gin QP)"]
        B2["中-高 (QP 合并)"]
        S2["低 (分离 kernel)"]
        O2["部分 (stream)"]
    end

    subgraph MM["DeepGEMM"]
        L3["低 (TMA pull)"]
        B3["高 (TMA 硬件)"]
        S3["高 (warp 特化)"]
        O3["完全 (persistent)"]
    end

DeepEP V1: "在哪里分配?" — NVSHMEM 提供跨 rank 内存分配
    ↓
DeepEP V2: "怎么暴露?" — NCCL Window Register 暴露远程指针
    ↓
DeepGEMM: "怎么用?" — SymBuffer.map() + TMA 直接远程 load/store

// csrc/kernels/backend/symmetric.hpp (line 188-290)
// HybridElasticSymmetricMemory: GPU + 所有 rank 的 CPU segment
// Layout: [GPU VRAM (front)] [CPU rank0 | CPU rank1 | ... | CPU rank(N-1) (back)]
class HybridElasticSymmetricMemory final : public SymmetricMemory {
    // 每个 rank 创建 NUMA-local CPU segment
    // 通过 POSIX FD (pidfd_open + pidfd_getfd) 跨进程导入
    for (int i = 0; i < num_scaleup_ranks; ++ i) {
        auto [pid, fd] = cpu_comm[num_scaleup_ranks * scaleout_rank_idx + i];
        if (pid != local_pid) {
            int pidfd = syscall(SYS_pidfd_open, pid, 0);
            local_fd = syscall(SYS_pidfd_getfd, pidfd, fd, 0);
        }
        CUDA_DRIVER_CHECK(lazy_cuMemImportFromShareableHandle(&cpu_handles[i], ...));
        CUDA_DRIVER_CHECK(lazy_cuMemMap(addr + offset, num_cpu_bytes, 0, cpu_handles[i], 0));
    }
};

// sm100_fp8_fp4_mega_moe.cuh (line 1274-1299)
// Epilogue warp 将 L2 输出写回远程 combine buffer
uint32_t dst_rank_idx, dst_token_idx, dst_topk_idx;
if (task_info.is_shared()) {
    dst_rank_idx = sym_buffer.rank_idx;
    dst_token_idx = pool_m_idx + m_idx_in_block;
    dst_topk_idx = kNumTopk;
} else {
    const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_m_idx + m_idx_in_block);
    dst_rank_idx = src_metadata.rank_idx;
    dst_token_idx = src_metadata.token_idx;
    dst_topk_idx = src_metadata.topk_idx;
}
// 从本地 smem 读取，通过 sym_buffer.map 写入远端
const auto packed = ptx::ld_shared(reinterpret_cast<float4*>(smem_ptr));
const auto dst_token = buffer.combine_token_buffer.get_rank_buffer(dst_topk_idx)
                       .get_data_buffer(dst_token_idx);
const auto dst_ptr = math::advance_ptr<float4>(dst_token.get_base_ptr(), ...);
*sym_buffer.map(dst_ptr, dst_rank_idx) = packed;  // NVLink/RDMA store

// deep_gemm/comm/barrier.cuh (line 46-89)
// 三层同步: grid_sync → NVLink signal → grid_sync
template <uint32_t kNumRanks, uint32_t kNumSMs, ...>
void nvlink_barrier(...) {
    if (sync_prologue)
        grid_sync<kNumSMs, kGridSyncIndex>(...);  // 节点内同步

    if (sm_idx == 0) {  // 仅 SM 0 参与跨节点
        auto* signal_ptr = workspace.get_nvl_barrier_signal_ptr(signal_phase);
        if (thread_idx < kNumRanks)
            ptx::red_add_rel_sys(sym_buffer.map(signal_ptr, thread_idx), signal_sign ? -1 : 1);
        if (thread_idx == 0) {
            ptx::red_add(counter_ptr, 1);
            while (ptx::ld_acq_sys(signal_ptr) != target) { /* spin */ }
        }
    }

    if (sync_epilogue)
        grid_sync<kNumSMs, kGridSyncIndex>(...);  // 再次节点内同步
}