07: DeepEP Intra/Inter-node 协调 — NVLink Scale-up + RDMA Scale-out 三阶段流水线的代码实证

enum class WarpRole {
    kRDMASender,              // IB Sending — 写入 RDMA send buffer
    kRDMASenderCoordinator,   // 协调 RDMA 发送 (博客未提及)
    kRDMAAndNVLForwarder,     // IB-to-NVLink Forwarding — 从 RDMA buffer 转发到 NVLink
    kForwarderCoordinator,    // 协调转发进度 (博客未提及)
    kNVLReceivers             // NVLink Receiving — 从 NVLink 接收并写入 recv_x
};

const auto role_meta = [=]() -> std::pair<WarpRole, int> {
    if (is_forwarder) {
        // SM 0,2,4,... (偶数) = 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};       // 奇数 SM 的前几个 warp
    } else if (warp_id == kNumDispatchRDMASenderWarps) {
        return {WarpRole::kRDMASenderCoordinator, -1>;  // 紧接着的一个 warp
    } else {
        return {WarpRole::kNVLReceivers, (warp_id + channel_id - kNumDispatchRDMASenderWarps) % LEGACY_NUM_MAX_NVL_PEERS};
    }
}();

flowchart LR
    subgraph Stage1["Stage 1: IB Sending"]
        A[Source GPU SM 奇数] -->|ld_nc_global + st_na_global| B[RDMA Send Buffer]
        B -->|nvshmemi_ibgda_put_nbi_warp| C[NIC HCA]
    end

    subgraph Stage2["Stage 2: IB-to-NVLink Forwarding"]
        C -->|RDMA write| D[RDMA Recv Buffer]
        D -->|tma_load_1d| E[Forwarder SM 偶数]
        E -->|src_meta.is_token_in_nvl_rank| F[过滤目标 NVL rank]
        F -->|tma_store_1d| G[NVL Channel Buffer]
    end

    subgraph Stage3["Stage 3: NVLink Receiving"]
        G -->|ld_acquire_sys_global tail| H[NVL Receiver SM 奇数]
        H -->|tma_load + tma_store| I[recv_x Expert Buffer]
    end

    C -->|Stage 1| D
    E -->|Stage 2| G
    H -->|Stage 3| I

// ibgda_device.cuh:128-141 — GPU SM 直接 post RDMA send request
__device__ static __forceinline__ void ibgda_post_send(nvshmemi_ibgda_device_qp_t* qp, uint64_t new_prod_idx) {
    nvshmemi_ibgda_device_qp_management_t* mvars = &qp->mvars;
    uint64_t old_prod_idx;

    ibgda_lock_acquire(&mvars->post_send_lock);           // 1. 获取 per-QP 锁
    old_prod_idx = atomicMax(reinterpret_cast<unsigned long long int*>(&mvars->tx_wq.prod_idx), new_prod_idx);
    if (new_prod_idx > old_prod_idx) {
        ibgda_update_dbr(qp, new_prod_idx);                // 2. 更新 Doorbell Record
        ibgda_ring_db(qp, new_prod_idx);                   // 3. Ring Doorbell (通知 NIC)
    }
    ibgda_lock_release(&mvars->post_send_lock);
}

// ibgda_device.cuh:336-381 — Warp 协作发送
template <bool kAlwaysDoPostSend = false>
__device__ static __forceinline__ void nvshmemi_ibgda_put_nbi_warp(
    uint64_t req_rptr, uint64_t req_lptr, size_t bytes, int dst_pe, int qp_id, int lane_id, int message_idx) {
    auto qp = ibgda_get_rc(dst_pe, qp_id);

    // 1. 计算 WQE 数量 (每个 lane 处理一个 chunk, 最多 3 WQE)
    auto remaining_bytes = bytes;
    while (remaining_bytes > 0) {
        if (lane_id == num_wqes) {
            my_chunk_size = min(remaining_bytes,
                ibgda_get_lkey_and_rkey(my_laddr = req_lptr, &my_lkey, req_rptr, dst_pe, &my_raddr, &my_rkey, qp->dev_idx));
        }
        auto chunk_size = __shfl_sync(0xffffffff, my_chunk_size, static_cast<int>(num_wqes));
        remaining_bytes -= chunk_size;
        ++num_wqes;
    }

    // 2. Lane 0 预留 WQE 槽位, 广播给所有 lane
    if (lane_id == 0)
        base_wqe_idx = ibgda_reserve_wqe_slots(qp, num_wqes);
    base_wqe_idx = __shfl_sync(0xffffffff, base_wqe_idx, 0);

    // 3. 每个 lane 写入自己的 WQE (ctrl + raddr + data segment)
    if (lane_id < num_wqes) {
        auto wqe_ptr = ibgda_get_wqe_ptr(qp, base_wqe_idx + lane_id);
        ibgda_write_rdma_write_wqe(qp, my_laddr, my_lkey, my_raddr, my_rkey, my_chunk_size, wqe_idx, &wqe_ptr);
    }

    // 4. Lane 0 ring doorbell
    if (lane_id == 0)
        ibgda_submit_requests<kAlwaysDoPostSend>(qp, base_wqe_idx, num_wqes, message_idx);
}

// 逐 token 处理
for (token_idx = token_start_idx; token_idx < token_end_idx; ++token_idx) {
    // 1. 读取 token 属于哪些 RDMA rank
    is_token_in_rank_uint64 = __ldg(...);

    // 2. 等待远端 buffer 有空位 (flow control)
    while (rdma_tail_idx - cached_rdma_channel_head >= num_max_rdma_chunked_recv_tokens) {
        cached_rdma_channel_head = ld_volatile_global(rdma_channel_head.buffer(lane_id));
    }

    // 3. 写入 symmetric send buffer (x + scales + src_meta + topk)
    UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4, ld_nc_global, st_broadcast);
    // ... scales, src_meta, topk_idx, topk_weights

    // 4. 更新发送窗口 (32-bit 位图, 追踪 in-flight 事务)
    acquire_lock(rdma_send_channel_lock + lane_id);
    window = rdma_send_channel_window[lane_id] | (1u << offset);
    if (offset == 0) {
        num_empty_slots = (~window) == 0 ? 32 : __ffs(~window) - 1;
        st_release_cta(rdma_send_channel_tail + lane_id, latest_tail + num_empty_slots);
    }
    release_lock(rdma_send_channel_lock + lane_id);
}

while (__any_sync(0xffffffff, num_tokens_to_send > 0)) {
    for (int i = 0; i < kNumRDMARanks; ++i) {
        int dst_rdma_rank = (i + channel_id + rdma_rank) % kNumRDMARanks; // incast 避免

        // 检查该 rank 是否有待发数据
        processed_tail = ld_acquire_cta(rdma_send_channel_tail + dst_rdma_rank);
        num_tokens_processed = processed_tail - synced_last_issued_tail;
        if (num_tokens_processed != synced_num_tokens_to_send and num_tokens_processed < num_max_rdma_chunked_send_tokens)
            continue;

        // 发起 RDMA put
        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);
        } else {
            memory_fence();  // 本地 rank 只需 memory fence
        }

        // 通过 RDMA atomic 通知远端 tail 更新
        nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_tail.buffer(rdma_rank),
            num_tokens_to_issue, ..., dst_rdma_rank == rdma_rank);
    }
}

// 逐 token 写入 receiver 的 channel buffer
for (int64_t token_idx = token_start_idx; token_idx < token_end_idx;) {
    // 1. 等待 receiver 释放 buffer 槽位
    while (true) {
        num_used_slots = cached_channel_tail_idx - ld_volatile_global(channel_head_idx.buffer());
        if (num_recv_buffer_tokens - num_used_slots >= num_max_send_tokens) break;
    }

    // 2. 写入 receiver 的 channel_x_buffers (通过 NVLink 可见)
    UNROLLED_WARP_COPY(5, lane_id, hidden_int4, shifted_channel_x_buffers, shifted_x, __ldg, st_na_global);
    channel_src_idx_buffers[dst_slot_idx] = token_idx;
    channel_topk_idx_buffers[dst_slot_idx * num_topk + lane_id] = idx_value;
    // ... topk_weights, x_scales

    // 3. 更新 tail (receiver 通过 ld_acquire_sys_global 看到)
    st_release_sys_global(channel_tail_idx.buffer(), cached_channel_tail_idx);
}

while (num_tokens_to_recv > 0) {
    // 1. 等待 tail 更新 (sender 通过 NVLink 写入)
    while (true) {
        cached_channel_tail_idx = __shfl_sync(0xffffffff, ld_acquire_sys_global(channel_tail_idx.buffer()), 0);
        if (cached_channel_head_idx != cached_channel_tail_idx) break;
    }

    // 2. 从 channel buffer 搬到 recv_x (TMA load + store on SM90)
    for (int chunk_idx = 0; chunk_idx < num_recv_tokens; ++chunk_idx) {
        if (elect_one_sync()) {
            tma_load_1d(tma_buffer, shifted_buffer_x_int4 + i * half_hidden_int4, tma_mbarrier, half_hidden_bytes);
            mbarrier_arrive_and_expect_tx(tma_mbarrier, half_hidden_bytes);
            mbarrier_wait(tma_mbarrier, tma_phase);
            tma_store_1d(tma_buffer, shifted_recv_x_int4 + i * half_hidden_int4, half_hidden_bytes, false);
        }
    }

    // 3. 更新 head (通知 sender 可以复用 buffer)
    st_relaxed_sys_global(channel_head_idx.buffer(), cached_channel_head_idx);
}

// hybrid_dispatch.cuh:28-32 — 三种 warp 角色
int kNumSMs,
int kNumNotifyWarps, int kNumScaleoutWarps, int kNumForwardWarps,
int kNumScaleoutRanks, int kNumScaleupRanks,

static constexpr int kDeviceBarrierTag = 0;
static constexpr int kKernelBarrierTag = 1;
static constexpr int kDispatchTag0 = 2;
static constexpr int kDispatchTag1 = 3;
static constexpr int kCombineTag0 = 4;
static constexpr int kCombineTag1 = 5;
static constexpr int kHybridDispatchTag0 = 6;   // Hybrid 专用 tag
static constexpr int kHybridDispatchTag1 = 7;
static constexpr int kHybridCombineTag0 = 8;
static constexpr int kHybridCombineTag1 = 9;

team_world(ncclTeamWorld(nccl_dev_comm)),    // 全局 team (scale-out + scale-up)
team_lsa(ncclTeamLsa(nccl_dev_comm)),        // scale-up team (NVLink 域)
team_rail(ncclTeamRail(nccl_dev_comm));      // scale-out team (RDMA 域)

template <int kNumSMs, int kNumQPs, int kNumChannelsPerSM, bool kWithNotifyWarps = false>
__device__ __forceinline__ std::pair<int, ncclGinResourceSharingMode> get_qp_mode(
    const int& sm_idx, const int& channel_in_sm_idx, const bool& is_notify_warp = false) {
    constexpr auto kSharingCTA = NCCL_GIN_RESOURCE_SHARING_CTA;
    constexpr auto kSharingGrid = kNumSMs == 1 ? NCCL_GIN_RESOURCE_SHARING_CTA : NCCL_GIN_RESOURCE_SHARING_GPU;

    // 1. 只有一个 QP: 所有 SM 共享
    if constexpr (kNumQPs == 1)
        return {0, kSharingGrid};

    // 2. Notify warp: 固定用 QP 0, CTA 级共享
    if (is_notify_warp)
        return {0, kSharingCTA};

    // 3. 数据 channel 的 QP 分配
    constexpr int kQPStartIdx = static_cast<int>(kWithNotifyWarps);
    if constexpr (kNumSMs <= kNumAvailableQPs) {
        // SM 数 ≤ QP 数: 一个 SM 独占多个 QP
        // 例: 3 SMs, 10 QPs → SM0: 0,3,6,9; SM1: 1,4,7; SM2: 2,5,8
        const int num_qps_in_sm = (kNumAvailableQPs / kNumSMs) + (sm_idx < (kNumAvailableQPs % kNumSMs));
        return {kQPStartIdx + sm_idx + (channel_in_sm_idx % num_qps_in_sm) * kNumSMs, kSharingCTA};
    } else {
        // SM 数 > QP 数: 所有 SM 共享所有 QP
        const auto global_channel_idx = sm_idx * kNumChannelsPerSM + channel_in_sm_idx;
        return {kQPStartIdx + (global_channel_idx % kNumAvailableQPs), kSharingGrid};
    }
}

// 写入 scale-out send buffer (TMA store)
if (scaleout_rank_mask ^ (1 << scaleout_rank_idx)) {  // 有非本地 rank
    ptx::tma_store_1d(scaleout_send_buffer.get_token_buffer(token_idx).get_base_ptr(),
                      tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
}

// 本地 rank 直接写入 recv buffer (bypass)
if (stored_dst_slot_idx >= 0 and stored_dst_scaleout_rank_idx == scaleout_rank_idx) {
    ptx::tma_store_1d(scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(),
                      tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
}

// 发起 RDMA put (通过 NCCL Gin)
if (stored_dst_slot_idx >= 0 and stored_dst_scaleout_rank_idx != scaleout_rank_idx) {
    gin.put<ncclTeamTagRail>(
        scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(),
        scaleout_send_buffer.get_token_buffer(token_idx).get_base_ptr(),
        tma_buffer.get_num_bytes<false>(),
        stored_dst_scaleout_rank_idx,
        ncclGinOptFlagsAggregateRequests);
}

// Forward warp: 从 scale-out recv buffer 转发到 scale-up buffer
while ((wip_mask = ptx::gather(stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or stored_finish_flag == 0))) {
    // 1. Round-robin 选择下一个 scale-out rank
    recv_scaleout_rank_idx = hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask);

    // 2. 等待该 rank 有数据到达 (通过 signaled_tail)
    comm::timeout_while<kNumTimeoutCycles>([&](const bool& is_last_check) {
        arrived_or_finished = stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or stored_finish_flag > 0;
        ...
    });

    // 3. TMA load 从 scale-out recv buffer
    ptx::tma_load_1d(tma_buffer.get_base_ptr(), token_buffer.get_base_ptr(),
                     mbarrier_ptr, token_layout.get_num_bytes<false>());

    // 4. 解析 top-k, 确定目标 scale-up rank
    stored_dst_scaleup_rank_idx = dst_expert_idx / kNumExpertsPerRank;

    // 5. TMA store 到 scale-up buffer (通过 NVLink)
    const auto dst_ptr = gin.get_sym_ptr<ncclTeamTagLsa>(
        scaleup_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(),
        stored_dst_scaleup_rank_idx);
    ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
}

template <typename team_t, typename dtype_t>
__device__ __forceinline__
void red_add_rel(dtype_t* sym_ptr, const dtype_t& value, const int& dst_rank_idx, const int& extra_options = 0) const {
    const auto dst_ptr = get_sym_ptr<team_t>(sym_ptr, dst_rank_idx);
    // 优先使用对称指针 (NVLink), 否则走 RDMA
    if (dst_ptr != nullptr) {
        if (std::is_same_v<team_t, ncclTeamTagRail> or dst_ptr == sym_ptr) {
            ptx::red_add_rel_gpu(dst_ptr, value);   // NVLink: GPU scope atomic
        } else {
            ptx::red_add_rel_sys(dst_ptr, value);   // 跨端: system scope
        }
    } else {
        // 走 NCCL Gin signal (RDMA path)
        gin.signal(TEAM_WORLD_RAIL(), dst_rank_idx,
            ncclGin_VASignalAdd(nccl_window, reinterpret_cast<int64_t>(sym_ptr) - lsa_base_ptr, static_cast<uint64_t>(value)),
            ...);
    }
}

template <typename team_t>
__device__ __forceinline__ bool is_nvlink_accessible(const int& dst_rank_idx) const {
    IS_TEAM_LSA({
        return true;  // LSA team 内总是 NVLink 可达
    })
    IS_TEAM_WORLD({
        // 检查 dst_rank 是否在当前 rail (NVLink 域) 内
        return team_rail.rank * team_lsa.nRanks <= dst_rank_idx and
               dst_rank_idx < (team_rail.rank + 1) * team_lsa.nRanks;
    })
    IS_TEAM_RAIL({
        return team_rail.rank == dst_rank_idx;  // Rail 内只有本 rank
    })
}

@functools.lru_cache()
def get_nvlink_gbs(factor: float = 0.9) -> float:
    """通过 nvidia-smi nvlink -s 获取 NVLink 总带宽"""
    result = subprocess.run(['nvidia-smi', 'nvlink', '-s'], ...)
    link_pattern = r'Link \d+:\s*([\d\.]+) GB/s'
    link_matches = re.findall(link_pattern, gpu_block)
    return sum(float(bw) for bw in link_matches) * factor

@functools.lru_cache()
def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float:
    """通过 ibstat 获取 RDMA 带宽"""
    result = subprocess.run(['ibstat'], ...)
    pattern = rf"CA '{nic_name}'.*?Port \d+:\s*.*?Rate:\s*(\d+)"
    rate = int(match.group(1))
    return rate / 8  # Gbps → GB/s

def approximate_num_sms(self, ...):
    # 1. 获取带宽
    rdma_gbs = get_rdma_gbs() if rdma_gbs == 0 else rdma_gbs
    nvlink_gbs = get_nvlink_gbs() if nvlink_gbs == 0 else nvlink_gbs

    # 2. 估算 traffic (基于 top-k 概率模型)
    if self.num_scaleout_ranks > 1:
        rdma_traffic += (1 / num_expected_topk) * (num_expected_scaleout_topk * (1 - 1 / self.num_scaleout_ranks))
        nvlink_traffic += 1 - (1 / self.num_scaleup_ranks)  # Forward 阶段
    else:
        nvlink_traffic += self.num_nvlink_ranks / self.num_ranks * (1 - 1 / self.num_nvlink_ranks)
        rdma_traffic += (self.num_ranks - self.num_nvlink_ranks) / self.num_ranks

    # 3. 找到瓶颈
    if (rdma_traffic / rdma_gbs) > (nvlink_traffic / nvlink_gbs):
        bounded_traffic, bounded_gbs = rdma_traffic, rdma_gbs  # RDMA 瓶颈
    else:
        bounded_traffic, bounded_gbs = nvlink_traffic, nvlink_gbs  # NVLink 瓶颈

    # 4. 计算所需 SM 数 (基于 SM read/write 带宽)
    num_sms = max(
        bounded_gbs / bounded_traffic * sm_read / sm_read_gbs,
        bounded_gbs / bounded_traffic * sm_write / sm_write_gbs,
    )
    return min(align(max(4, ceil(num_sms * 1.25)), 2), num_device_sms)

template <bool kLowLatencyMode>
__forceinline__ __device__ int translate_dst_rdma_rank(const int dst_rdma_rank, const int nvl_rank) {
    return kLowLatencyMode ? (dst_rdma_rank * LEGACY_NUM_MAX_NVL_PEERS + nvl_rank) : dst_rdma_rank;
}

template <bool kLowLatencyMode>
__forceinline__ __device__ void nvshmem_sync_with_same_gpu_idx(const nvshmem_team_t& rdma_team) {
    kLowLatencyMode ? void(nvshmem_sync(rdma_team)) : nvshmem_sync_all();
}

template <bool kLowLatencyMode, int kNumRDMARanks, bool kCachedMode, ...>
__global__ void dispatch(...) {
    enum class WarpRole { kRDMASender, kRDMASenderCoordinator, kRDMAAndNVLForwarder, kForwarderCoordinator, kNVLReceivers };
    // ... 完全相同的三阶段流水线
}

// internode.cu:823-829 (kRDMASenderCoordinator 发起 RDMA put)
nvshmemi_ibgda_put_nbi_warp<true>(
    dst_ptr,                                          // 远端 RDMA recv buffer
    src_ptr,                                          // 本地 RDMA send buffer
    num_bytes_per_msg,                                // 批量发送的字节数
    translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank),  // 目标 QP
    channel_id,                                       // channel ID (用于 QP 选择)
    lane_id,
    0);                                               // message_idx

// internode.cu:966-994 (kRDMAAndNVLForwarder)
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 的 token
    if (not src_meta.is_token_in_nvl_rank(dst_nvl_rank))
        continue;

    // TMA load 从 RDMA buffer, TMA store 到 NVL channel 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);  // NVLink 写入

    tma_store_wait<0>();
}

// internode.cu:1129-1152 (kNVLReceivers)
for (int chunk_idx = 0; chunk_idx < num_recv_tokens; ++chunk_idx) {
    auto shifted = nvl_channel_x.buffer() + token_idx_in_buffer * num_bytes_per_token;

    if (elect_one_sync()) {
        tma_load_1d(tma_buffer, shifted, tma_mbarrier, tma_load_bytes);      // NVLink load
        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);  // 写入 expert buffer
        tma_store_1d(tma_buffer + hidden_bytes, recv_x_scales + recv_token_idx * num_scales, scale_bytes, false);
    }
    tma_store_wait<0>();
}

flowchart TB
    subgraph Source["Source GPU (rank S)"]
        T[Token Buffer x] -->|ld_nc_global| SB[RDMA Send Buffer]
        SB -->| SenderWarp 写入 | SB2[Symmetric Send Buffer]
    end

    subgraph RDMA["RDMA Domain (scale-out)"]
        SB2 -->|CoordinatorWarp| IBGDA[nvshmemi_ibgda_put_nbi_warp]
        IBGDA -->|GPU SM 直接 ring doorbell| NIC[NIC HCA]
        NIC -->|RDMA Write| RB[Remote RDMA Recv Buffer]
    end

    subgraph Forward["Forwarding GPU (rank F, 偶数 SM)"]
        RB -->|tma_load_1d| TMA1[TMA Buffer]
        TMA1 -->|src_meta 过滤| FILT{属于本 NVL rank?}
        FILT -->|Yes| TMA2[TMA Store]
        TMA2 -->|NVLink| NB[NVL Channel Buffer]
    end

    subgraph Target["Target GPU (rank T, 奇数 SM)"]
        NB -->|ld_acquire_sys_global tail| TMA3[tma_load_1d]
        TMA3 -->|tma_store_1d| RX[recv_x Expert Buffer]
        RX --> GEMM[Expert GEMM]
    end

    Source -->|Stage 1: IB Sending| RDMA
    RDMA -->|Stage 2: Forwarding| Forward
    Forward -->|Stage 3: NVLink Receiving| Target