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

# deep_gemm/mega/__init__.py
import torch.distributed._symmetric_memory as symm_mem

# 每个 rank 分配一块 buffer
self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')

# rendezvous 让 buffer 对所有 rank 可见
self.handle = symm_mem.rendezvous(self.buffer, group=group)

// 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 {
        int64_t mapped_ptr = offsets[dst_rank_idx] + reinterpret_cast<int64_t>(ptr);
        return *reinterpret_cast<ptr_t*>(&mapped_ptr);
    }
};

// sm100_fp8_fp4_mega_moe.cuh — Dispatch Warp 拉取远程 token
const auto pull_buffer = smem_send_buffers.get_rank_buffer(warp_idx).get_data_buffer(0);

// 关键: 通过 sym_buffer.map 将本地 input_token_buffer 指针映射到远程 rank
ptx::tma_load_1d(
    pull_buffer.get_base_ptr(),                                          // 本地 smem 目标
    sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx)     // 远程 rank 源地址
                   .get_base_ptr(), current_rank_in_expert_idx),
    pull_mbarrier, kHidden);

// 同一段代码，不同 dst_rank_idx → 不同底层传输
*sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;  // 跨节点写 → RDMA

Source GPU → IB Sending → RDMA Network → IB-to-NVLink Forwarding → NVLink Receiving → Target GPU

┌─────────────────────────────────────────────────────────┐
│                    Symmetric Memory                      │
│  ┌──────────┐    ┌──────────┐    ┌──────────┐          │
│  │ Rank 0   │◄──►│ Rank 1   │◄──►│ Rank N   │          │
│  │ Buffer   │ NVLink│ Buffer │ NVLink│ Buffer  │          │
│  └──────────┘    └──────────┘    └──────────┘          │
│       ▲              ▲              ▲                   │
│       │ TMA          │ TMA          │ TMA               │
│  ┌──────────┐    ┌──────────┐    ┌──────────┐          │
│  │ SM 0     │    │ SM 1     │    │ SM N     │          │
│  └──────────┘    └──────────┘    └──────────┘          │
└─────────────────────────────────────────────────────────┘

graph TB
    subgraph DeepEP["DeepEP: 三阶段推送 + 转发"]
        S[Source GPU] -->|读取 Dispatch Buffer| IB[IB Sending]
        IB -->|RDMA 网络| NET[RDMA Network]
        NET -->|接收| FWD[IB-to-NVLink Forwarding]
        FWD -->|NVLink| RECV[NVLink Receiving]
        RECV -->|写入 Receive Buffer| T[Target GPU]
    end

    subgraph MegaMoE["Mega MoE: 对称内存直接访问"]
        S2[Dispatch Warp] -->|TMA Load + sym_buffer.map| LB[Local L1 Buffer]
        S3[Epilogue Warp] -->|TMA Store + sym_buffer.map| RB[Remote Combine Buffer]
    end

// sm100_fp8_fp4_mega_moe.cuh (line 544-598)

// 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. TMA load token 从远端 rank 到本地 smem
if (cute::elect_one_sync()) {
    ptx::tma_load_1d(
        pull_buffer.get_base_ptr(),
        sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(),
                       current_rank_in_expert_idx),  // ← 远程 rank 地址
        pull_mbarrier, kHidden);
}

// 3. 加载并存储 SF (与 TMA token load 重叠)
const auto remote_sf_ptr = sym_buffer.map(
    input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr<uint32_t>(),
    current_rank_in_expert_idx);
// ... 直接读取 remote_sf_ptr

// 4. 存储到本地 L1 buffer
ptx::tma_store_1d(
    l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(),
    pull_buffer.get_base_ptr(), pull_buffer.get_num_bytes());

// sm100_fp8_fp4_mega_moe.cuh (line 1196-1202)

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

// 写入远端 combine buffer
const auto dst_token = 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/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);
}

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

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);

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

    S->>S: 读取 Dispatch Buffer
    S->>NIC: IB Sending (GPU→NIC)
    NIC->>NET: RDMA 发送
    NET->>FWD: 接收 (可能非目标 GPU)
    FWD->>FWD: IB-to-NVLink Forwarding
    FWD->>T: NVLink 直写
    T->>T: 写入 Receive Buffer

sequenceDiagram
    participant D as Dispatch Warp (SM)
    participant TMA as TMA Engine
    participant R as Remote Rank Buffer
    participant L as Local L1 Buffer
    participant E as Epilogue Warp (SM)
    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 归约