06_06: NVLink Scale-up + RDMA Scale-out — DeepEP 三阶段流水线 vs DeepGEMM 对称内存

// csrc/kernels/legacy/internode.cu, line 487
enum class WarpRole {
    kRDMASender,              // IB Sending 阶段
    kRDMASenderCoordinator,   // RDMA 发送协调
    kRDMAAndNVLForwarder,     // IB-to-NVLink Forwarding 阶段
    kForwarderCoordinator,    // Forwarder 协调
    kNVLReceivers             // NVLink Receiving 阶段
};

// csrc/kernels/legacy/internode.cu, line 499-514
const auto role_meta = [=]() -> std::pair<WarpRole, int> {
    if (is_forwarder) {
        if (warp_id < LEGACY_NUM_MAX_NVL_PEERS) {
            return {WarpRole::kRDMAAndNVLForwarder, (warp_id + channel_id) % LEGACY_NUM_MAX_NVL_PEERS};
        } else {
            return {WarpRole::kForwarderCoordinator, warp_id - LEGACY_NUM_MAX_NVL_PEERS};
        }
    } else if (warp_id < kNumDispatchRDMASenderWarps) {
        return {WarpRole::kRDMASender, -1};
    } else if (warp_id == kNumDispatchRDMASenderWarps) {
        return {WarpRole::kRDMASenderCoordinator, -1};
    } else {
        return {WarpRole::kNVLReceivers, (warp_id + channel_id - kNumDispatchRDMASenderWarps) % LEGACY_NUM_MAX_NVL_PEERS};
    }
}();

// csrc/kernels/legacy/internode.cu, line 587-757
if (warp_role == WarpRole::kRDMASender) {
    // 获取 token 任务范围
    get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx);

    // 遍历 token 并拷贝到 symmetric send buffer
    for (token_idx = token_start_idx; token_idx < token_end_idx; ++token_idx) {
        // 读取 RDMA rank 存在性
        // 拷贝 x, x_scales, source meta, topk_idx, topk_weights 到 send buffer
        UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4, ld_nc_global, st_broadcast);

        // 释放事务槽 (流控)
        acquire_lock(rdma_send_channel_lock + lane_id);
        // ... 窗口管理 ...
        release_lock(rdma_send_channel_lock + lane_id);
    }
}

// csrc/kernels/legacy/internode.cu, line 758-848
} else if (warp_role == WarpRole::kRDMASenderCoordinator) {
    // 迭代所有 RDMA ranks
    while (__any_sync(0xffffffff, num_tokens_to_send > 0)) {
        for (int i = 0; i < kNumRDMARanks; ++i) {
            // 实际发起 RDMA send
            if (dst_rdma_rank != rdma_rank) {
                nvshmemi_ibgda_put_nbi_warp<true>(
                    dst_ptr, src_ptr, num_bytes_per_msg,
                    translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank),
                    channel_id, lane_id, 0);
            }
            // 更新 tail
            nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_tail.buffer(rdma_rank), ...);
        }
    }
}

// csrc/kernels/legacy/internode.cu, line 849-1013
} else if (warp_role == WarpRole::kRDMAAndNVLForwarder) {
    const auto dst_nvl_rank = target_rank;

    // 1. 等待 RDMA meta 数据到达
    while (true) {
        auto meta_0 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + dst_nvl_rank);
        // ... 等待 4 个 meta 字段就绪 ...
        if (meta_0 < 0 and meta_1 < 0 and meta_2 < 0 and meta_3 < 0) {
            // 通知 NVL ranks
            st_relaxed_sys_global(nvl_channel_prefix_start.buffer() + lane_id, -start_sum - 1);
            st_relaxed_sys_global(nvl_channel_prefix_end.buffer() + lane_id, -end_sum - 1);
            break;
        }
    }

    // 2. 从 RDMA buffer 转发 token 到 NVL
    while (__any_sync(0xffffffff, num_tokens_to_recv_from_rdma > 0)) {
        // 轮询源 RDMA rank
        src_rdma_rank = (src_rdma_rank + 1) % kNumRDMARanks;

        // 遍历 RDMA buffer 中的 token
        for (int i = src_rdma_head; i < src_rdma_tail; ++i) {
            auto shifted = rdma_channel_data.recv_buffer(src_rdma_rank) + rdma_slot_idx * num_bytes_per_token;
            auto src_meta = ld_nc_global(reinterpret_cast<SourceMeta*>(shifted + hidden_bytes + scale_bytes));

            // 检查是否属于目标 NVL rank
            bool is_in_dst_nvl_rank = src_meta.is_token_in_nvl_rank(dst_nvl_rank);
            if (not is_in_dst_nvl_rank) continue;

            // TMA load from RDMA buffer, TMA store to NVL buffer
            if (elect_one_sync()) {
                tma_load_1d(tma_buffer, shifted, tma_mbarrier, num_bytes_per_token, false);
                mbarrier_arrive_and_expect_tx(tma_mbarrier, num_bytes_per_token);
            }
            mbarrier_wait(tma_mbarrier, tma_phase);
            if (elect_one_sync())
                tma_store_1d(tma_buffer, dst_shifted, num_bytes_per_token);
        }
    }
}

// csrc/kernels/legacy/internode.cu, line 1061-1196
} else {
    // NVL consumers
    int src_nvl_rank = target_rank;

    // 1. 等待 NVL channel prefix 数据
    while (lane_id < kNumRDMARanks) {
        start_offset = ld_volatile_global(nvl_channel_prefix_start.buffer() + lane_id);
        end_offset = ld_volatile_global(nvl_channel_prefix_end.buffer() + lane_id);
        if (start_offset < 0 and end_offset < 0) {
            start_offset = -start_offset - 1, end_offset = -end_offset - 1;
            break;
        }
    }

    // 2. 从 NVL buffer 读取 token
    while (num_tokens_to_recv > 0) {
        // 等待数据到达
        if (cached_channel_head_idx != cached_channel_tail_idx)
            break;
        cached_channel_tail_idx = ld_acquire_sys_global(nvl_channel_tail.buffer());

        // 拷贝数据
        for (int chunk_idx = 0; chunk_idx < num_recv_tokens; ++chunk_idx, --num_tokens_to_recv) {
            auto shifted = nvl_channel_x.buffer() + token_idx_in_buffer * num_bytes_per_token;

            // TMA load from NVL buffer, TMA store to recv_x
            if (elect_one_sync()) {
                tma_load_1d(tma_buffer, shifted, tma_mbarrier, tma_load_bytes);
                mbarrier_arrive_and_expect_tx(tma_mbarrier, tma_load_bytes);
            }
            mbarrier_wait(tma_mbarrier, tma_phase);
            if (elect_one_sync()) {
                tma_store_1d(tma_buffer, recv_x + recv_token_idx * hidden_int4, hidden_bytes, false);
            }
        }
    }
}

// csrc/kernels/legacy/internode.cu, line 1746
enum class WarpRole { kNVLSender, kNVLAndRDMAForwarder, kRDMAReceiver, kCoordinator };

// deep_gemm/include/deep_gemm/layout/sym_buffer.cuh
template <uint32_t kNumRanks = kNumMaxRanks>
struct SymBuffer {
    int64_t base;                    // 本地 buffer 基地址
    int64_t offsets[kNumMaxRanks];   // 各 rank 相对于本 rank 的偏移
    uint32_t rank_idx;               // 当前 rank 编号

    // 关键: 将本地指针映射为可远程访问的指针
    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);
    }
};

// deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh, line 533-556

// 1. 读取源 token-topk 索引 (由远端 dispatch 通过 NVLink 写入)
const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr(
    current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank);

// 2. 通过 sym_buffer.map 计算远程地址
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);  // ← 远程 rank 地址

// 3. TMA load token 从远端 rank 到本地 smem
if (cute::elect_one_sync()) {
    for (uint32_t i = 0; i < kNumChunks; ++ i) {
        ptx::tma_load_1d(
            pull_buffer.get_base_ptr(),
            math::advance_ptr(src_base_ptr, i * kNumBytesPerPull),
            pull_mbarrier, kNumBytesPerPull
        );
        ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kNumBytesPerPull);
    }
}

// deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh, line 1293-1299

// 从 shared memory 读取
const auto packed = ptx::ld_shared(reinterpret_cast<float4*>(smem_ptr));

// 写入远端 combine buffer
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(),
    n_idx * sizeof(nv_bfloat16) + (lane_idx % 16) * sizeof(float4));

*sym_buffer.map(dst_ptr, dst_rank_idx) = packed;  // ← 远程 NVLink/RDMA 写

// deep_gemm/include/deep_gemm/comm/barrier.cuh
template <uint32_t kNumRanks, uint32_t kNumSMs, uint32_t kNumThreads, ...>
CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace,
                                   const layout::SymBuffer<kNumRanks>& sym_buffer, ...) {
    // 1. Grid sync (节点内所有 SM 同步)
    if (sync_prologue)
        grid_sync<kNumSMs, kGridSyncIndex>(workspace, sm_idx, thread_idx, sync_scope);

    // 2. NVLink 跨 rank barrier (仅 SM 0 参与)
    if (sm_idx == 0) {
        // 发送信号到远端 rank
        if (thread_idx < kNumRanks)
            ptx::red_add_rel_sys(sym_buffer.map(signal_ptr, thread_idx), signal_sign ? -1 : 1);

        // 等待所有 rank 到达
        if (thread_idx == 0) {
            ptx::red_add(counter_ptr, 1);
            while (ptx::ld_acq_sys(signal_ptr) != target) { /* spin */ }
        }
    }

    // 3. Grid sync (确保所有 SM 看到 barrier 完成)
    if (sync_epilogue)
        grid_sync<kNumSMs, kGridSyncIndex>(workspace, sm_idx, thread_idx, sync_scope);
}

graph TB
    subgraph DeepEP["DeepEP: 三阶段推送 + 转发"]
        S[Source GPU<br/>kRDMASender] -->|写 Symmetric Buffer| SB[Symmetric RDMA Buffer]
        SB -->|ibgda_put_nbi| NIC[NIC / RDMA]
        NIC -->|RDMA 网络| NET[RDMA Network]
        NET -->|接收| FWD[Forwarding GPU<br/>kRDMAAndNVLForwarder]
        FWD -->|TMA load + TMA store| NVL[NVL Buffer]
        NVL -->|NVLink| RECV[Target GPU<br/>kNVLReceivers]
        RECV -->|TMA store| T[Receive Buffer]
    end

    subgraph MegaMoE["Mega MoE: 对称内存直接访问"]
        D[Dispatch Warp] -->|sym_buffer.map + TMA load| LB[Local L1 Buffer]
        LB -->|直接远程访问| R[Remote input_token_buffer]
        L[Epilogue Warp] -->|sym_buffer.map + TMA store| C[Remote combine_token_buffer]
    end

sequenceDiagram
    participant S as Source GPU (DeepEP)
    participant NIC as NIC (IB)
    participant NET as RDMA Network
    participant FWD as Forwarding GPU
    participant T as Target GPU

    S->>S: kRDMASender 写 Symmetric Buffer
    S->>NIC: Coordinator 发起 ibgda_put_nbi
    NIC->>NET: RDMA 发送
    NET->>FWD: 接收 (可能非目标 GPU)
    FWD->>FWD: kRDMAAndNVLForwarder TMA 中继
    FWD->>T: NVLink 直写到 NVL Buffer
    T->>T: kNVLReceivers TMA 读取

sequenceDiagram
    participant D as Dispatch Warp (Mega MoE)
    participant TMA as TMA Engine
    participant R as Remote Rank Buffer
    participant L as Local L1 Buffer
    participant E as Epilogue Warp
    participant C as Combine Buffer

    D->>D: 计算远程 rank + slot
    D->>TMA: 发起 TMA Load (sym_buffer.map)
    TMA->>R: NVLink/RDMA 直接读取
    R->>L: 写入本地 L1 buffer
    Note over L: mbarrier 同步
    L->>E: MMA + SwiGLU + Epilogue
    E->>C: TMA Store 写回远程 (sym_buffer.map)
    Note over C: Combine 归约

DeepEP (2025)                          Mega MoE (2026)
─────────────────────────────────────────────────────────────────
IBGDA (GPU Direct Async)         →     Symmetric Memory (PyTorch)
显式 RDMA put_nbi                →     TMA load/store (硬件自动选择 NVLink/RDMA)
GPU SM 发起 RDMA                 →     TMA 硬件引擎发起传输
对称 buffer (手动管理偏移)        →     SymBuffer.map (统一地址空间)
三阶段 (Send→Forward→Receive)    →    单阶段 (直接远程访问)

// DeepEP: 5 种 Warp 角色
enum class WarpRole {
    kRDMASender,              // Stage 1: IB Sending
    kRDMASenderCoordinator,   // Stage 1.5: RDMA 协调
    kRDMAAndNVLForwarder,     // Stage 2: Forwarding
    kForwarderCoordinator,    // Stage 2.5: Forwarder 协调
    kNVLReceivers             // Stage 3: NVLink Receiving
};

// csrc/kernels/legacy/internode_ll.cu, line 253-278
// Issue IBGDA sends (直接 RDMA, 无 Forwarding)
if (dst_expert_idx >= 0) {
    const auto dst_rank = dst_expert_idx / num_local_experts;
    const auto dst_ptr = reinterpret_cast<uint64_t>(rdma_recv_x) + ...;
    const auto dst_p2p_ptr = nvshmemi_get_p2p_ptr(dst_ptr, rank, dst_rank);

    if (dst_p2p_ptr == 0) {
        // 跨节点: 直接 RDMA (无 Forwarding)
        nvshmemi_ibgda_put_nbi_warp(dst_ptr, src_ptr, num_bytes_per_msg, dst_rank, ...);
    } else {
        // 同节点: 直接 P2P 写 (无 Forwarding)
        UNROLLED_WARP_COPY(8, lane_id, num_int4_per_msg, dst_int4_ptr, src_int4_ptr, ld_nc_global, st_na_global);
    }
}

// csrc/kernels/legacy/internode.cu
// Inter-node: RDMA 网络
nvshmemi_ibgda_put_nbi_warp<true>(dst_ptr, src_ptr, ...);

// csrc/kernels/legacy/intranode.cu
// Intra-node: NVLink 直写
barrier_block<kNumRanks>(barrier_signal_ptrs, rank);

// Mega MoE: 同一段代码，不同 dst_rank_idx → 不同底层传输
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);  // 同节点 → NVLink, 跨节点 → RDMA

NVLink + RDMA + GPU SM → 共同构成数据路径

Symmetric Memory + TMA + GPU SM → 统一数据路径

// Mega MoE 的通信由以下三者协作:
// 1. Dispatch Warps (GPU SM): 计算远程地址, 发起 TMA
// 2. TMA (Tensor Memory Accelerator): 执行实际数据传输
// 3. SymBuffer.map(): 提供远程地址映射

// 关键代码: SM 直接发起远程 TMA load
ptx::tma_load_1d(
    local_smem_ptr,
    sym_buffer.map(remote_ptr, dst_rank_idx),  // SM 计算远程地址
    mbarrier, size);

// IBGDA (InfiniBand GPU Direct Async)
nvshmemi_ibgda_put_nbi_warp<true>(
    dst_ptr, src_ptr, num_bytes,
    dst_rdma_rank, channel_id, lane_id, slot_idx);

// NVLink barrier
barrier_block<kNumRanks>(barrier_signal_ptrs, rank);

// NVLink P2P 访问
st_release_sys_global(reinterpret_cast<int*>(dst_p2p_ptr), cnt);

// SymBuffer 地址映射
sym_buffer.map(ptr, dst_rank_idx);

// TMA load/store (硬件自动选择 NVLink/RDMA)
ptx::tma_load_1d(dst_ptr, src_ptr, mbarrier, size);
ptx::tma_store_1d(dst_ptr, src_ptr, size);

// NVLink barrier (跨 rank 同步)
ptx::red_add_rel_sys(sym_buffer.map(signal_ptr, thread_idx), value);

graph LR
    subgraph Stage1["Stage 1: IB Sending"]
        S1[Source GPU SM] -->|写 Dispatch Buffer| SB[Symmetric RDMA Buffer]
        SB -->|ibgda_put_nbi_warp| NIC1[NIC]
    end

    subgraph Stage2["Stage 2: RDMA Network"]
        NIC1 -->|RDMA 网络| NIC2[Remote NIC]
    end

    subgraph Stage3["Stage 3: IB-to-NVLink Forwarding"]
        NIC2 -->|接收| FWD[Forwarding GPU SM]
        FWD -->|TMA load from RDMA Buffer| TMA1[TMA Engine]
        TMA1 -->|TMA store to NVL Buffer| NVL[NVLink Buffer]
    end

    subgraph Stage4["Stage 4: NVLink Receiving"]
        NVL -->|NVLink| RECV[Target GPU SM]
        RECV -->|TMA load| TMA2[TMA Engine]
        TMA2 -->|TMA store| RB[Receive Buffer]
    end

    Stage1 --> Stage2 --> Stage3 --> Stage4

graph LR
    subgraph Pull["Dispatch: 直接远程拉取"]
        D[Dispatch Warp] -->|sym_buffer.map| ADDR[地址计算]
        ADDR -->|TMA load| TMA1[TMA Engine]
        TMA1 -->|NVLink/RDMA| R1[Remote input_token_buffer]
        R1 -->|数据| L1[l1_token_buffer]
    end

    subgraph Compute["计算"]
        L1 -->|MMA| GEMM[GEMM + SwiGLU]
    end

    subgraph Push["Combine: 直接远程写回"]
        GEMM -->|Epilogue| E[Epilogue Warp]
        E -->|sym_buffer.map| ADDR2[地址计算]
        ADDR2 -->|TMA store| TMA2[TMA Engine]
        TMA2 -->|NVLink/RDMA| C[Remote combine_token_buffer]
    end

    Pull --> Compute --> Push