07 - DeepEP FIFO/Streaming Pipeline 机制深度分析: Blog 论述 vs 源码实现

enum class WarpRole {
    kRDMASender,           // RDMA 发送
    kRDMASenderCoordinator,// RDMA 发送协调
    kRDMAAndNVLForwarder,  // RDMA-NVL 转发
    kForwarderCoordinator, // 转发协调
    kNVLReceivers          // NVL 接收
};

template <int kNumRanks, bool kSyncOnly = false>
__forceinline__ __device__ void barrier_block(int** barrier_signal_ptrs, int rank) {
    // 1. 内存 fence + block 同步
    if constexpr (not kSyncOnly) {
        memory_fence();          // fence.acq_rel.sys
        __syncthreads();         // block 级同步
    }

    // 2. 原子操作: 通知其他 rank, 减去自己的信号
    if (thread_id < kNumRanks) {
        atomicAdd_system(barrier_signal_ptrs[rank] + thread_id, LEGACY_FINISHED_SUM_TAG);
        atomicSub_system(barrier_signal_ptrs[thread_id] + rank, LEGACY_FINISHED_SUM_TAG);
    }

    // 3. 轮询等待所有 rank 完成
    while (true) {
        auto value = ld_volatile_global(barrier_signal_ptrs[rank] + thread_id);
        if (__all_sync(0xffffffff, value <= 0))
            break;  // 所有 rank 完成
        // timeout 检查...
    }
    __syncthreads();
}

// 1. 等待所有之前的 RDMA 操作完成
for (int i = thread_id; i < qps_per_rdma_rank * (kNumRDMARanks - 1); ++i)
    nvshmemi_ibgda_quiet(dst_rdma_rank, qp_id);
__syncthreads();

// 2. 跨节点 barrier
if (thread_id == 32)
    nvshmem_sync_with_same_gpu_idx(rdma_team);

// 3. 节点内 barrier
barrier_block<LEGACY_NUM_MAX_NVL_PEERS, true>(barrier_signal_ptrs, nvl_rank);

// ... 执行发送 ...

// 4. 再次 barrier
if (thread_id == 0)
    nvshmem_sync_with_same_gpu_idx(rdma_team);
__syncthreads();

// ... NVL 操作 ...

// 5. 最终 barrier
if (thread_id == 32)
    nvshmem_sync_with_same_gpu_idx(rdma_team);
barrier_block<LEGACY_NUM_MAX_NVL_PEERS>(barrier_signal_ptrs, nvl_rank);

// RDMA sender warp: 使用 head/tail 管理发送窗口
while (is_token_in_rank_uint64 != 0 and 
       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));
    // timeout 检查...
}

// 发送完成后, 使用锁保护更新 tail
acquire_lock(rdma_send_channel_lock + lane_id);
auto window = rdma_send_channel_window[lane_id] | (1u << offset);
if (offset == 0) {
    auto num_empty_slots = (~window) == 0 ? 32 : __ffs(~window) - 1;
    st_release_cta(rdma_send_channel_tail + lane_id, latest_tail + num_empty_slots);
    window >>= num_empty_slots;
}
rdma_send_channel_window[lane_id] = window;
release_lock(rdma_send_channel_lock + lane_id);

// Forwarder 轮询等待 RDMA 数据
while (true) {
    src_rdma_rank = (src_rdma_rank + 1) % kNumRDMARanks;
    if (__shfl_sync(0xffffffff, num_tokens_to_recv_from_rdma, src_rdma_rank) > 0) {
        if (lane_id == src_rdma_rank and cached_rdma_channel_head == cached_rdma_channel_tail)
            cached_rdma_channel_tail = ld_acquire_sys_global(rdma_channel_tail.buffer(src_rdma_rank));
        if (__shfl_sync(0xffffffff, cached_rdma_channel_tail > cached_rdma_channel_head, src_rdma_rank))
            break;
    }
    // timeout 检查...
}

template <bool kIsScaleupNVLink, int kNumScaleoutRanks, int kNumScaleupRanks,
          int kNumSMs, int kNumThreads, int kNumQPs,
          int64_t kNumTimeoutCycles, int kTag = kDeviceBarrierTag,
          bool kFlushStores = true, bool kSyncAtStart = true, bool kSyncAtEnd = true>
__forceinline__ __device__ void gpu_barrier(...) {
    // 1. TMA store 等待 (确保数据写入完成)
    if constexpr (kFlushStores) {
        ptx::tma_store_commit();
        ptx::tma_store_wait();
        __syncwarp();
    }

    // 2. 网格级同步 (所有 SM 等待)
    if constexpr (kSyncAtStart) {
        cooperative_groups::this_grid().sync();  // 网格级 barrier!
    }

    // 3. 执行 scaleout/scaleup barrier
    if (do_scaleup and do_scaleout) {
        // 并行执行: SM 0 做 scaleup, 其他 SM 做 scaleout
        if (sm_idx == 0) {
            scaleup_barrier_wo_local_sync(...);
            if constexpr (kFlushStores) 
                cooperative_groups::this_grid().sync();
        } else {
            scaleout_barrier_wo_local_sync(...);
        }
    }

    // 4. 最终网格级同步
    if constexpr (kSyncAtEnd)
        cooperative_groups::this_grid().sync();
}

template <int kNumRanks, int kNumSMs, int kNumThreads, int64_t kNumTimeoutCycles, int kTag>
__forceinline__ __device__ void nvlink_barrier_wo_local_sync(...) {
    // 只使用 SM 0
    if (kNumSMs > 1 and sm_idx > 0) return;

    // 读取当前 barrier phase
    const int status = static_cast<int>((*workspace.get_nvl_barrier_counter_ptr()) & 3);
    const int phase = status & 1, sign = status >> 1;

    // 每个 thread 原子操作信号
    if (thread_idx < kNumRanks) {
        const auto dst_ptr = gin.get_sym_ptr<ncclTeamTagLsa>(
            workspace.get_nvl_barrier_signal_ptr(phase), thread_idx);
        ptx::red_add_rel_sys(dst_ptr, sign ? -1 : 1);  // 原子加减
    }
    __syncthreads();

    // 增加 phase counter
    if (thread_idx == 0)
        atomicAdd(workspace.get_nvl_barrier_counter_ptr(), 1);

    // 轮询等待信号到达目标值
    const auto target = sign ? 0 : kNumRanks;
    timeout_while<kNumTimeoutCycles>(thread_idx == 0, [=](const bool& is_last_check) {
        const auto signal = ptx::ld_acquire_sys<int>(workspace.get_nvl_barrier_signal_ptr(phase));
        if (signal == target) return true;  // barrier 完成
        // timeout 处理...
    });
}

template <int kNumRanks, int kNumSMs, int kNumThreads, int kNumQPs, 
          int64_t kNumTimeoutCycles, typename team_t, int kTag, bool kFlushStores>
__forceinline__ __device__ void gin_barrier_wo_local_sync(...) {
    // 1. 刷新所有 QP (确保 RDMA 操作完成)
    if constexpr (kFlushStores) {
        for (int i = global_warp_idx; i < num_qps; i += kNumSMs * kNumWarps) {
            ncclGin(nccl_dev_comm, i, NCCL_GIN_RESOURCE_SHARING_CTA).flush(ncclCoopWarp());
        }
        (gridDim.x > 1) ? cooperative_groups::this_grid().sync() : __syncthreads();
    }

    // 2. 使用 QP 0 做 barrier
    if (sm_idx == 0) {
        const auto team = (std::is_same_v<team_t, ncclTeamTagWorld>) ?
            ncclTeamWorld(nccl_dev_comm) : ncclTeamRail(nccl_dev_comm);
        const ncclGin gin(nccl_dev_comm, 0, NCCL_GIN_RESOURCE_SHARING_CTA);

        // 发送 signal 给所有 rank
        for (int i = thread_idx; i < kNumRanks; i += kNumThreads)
            gin.signal(team, i, ncclGin_SignalInc{static_cast<ncclGinSignal_t>(rank_idx)});

        // 等待所有 rank 的 signal
        for (int i = thread_idx; i < kNumRanks; i += kNumThreads) {
            const auto signal_idx = static_cast<ncclGinSignal_t>(i);
            const auto shadow_ptr = gin.getSignalShadowPtr(signal_idx);
            const auto target = ++(*shadow_ptr);
            const auto signal_ptr = reinterpret_cast<uint64_t*>(...);

            timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
                const auto signal = ptx::ld_acquire_sys<uint64_t>(signal_ptr);
                if (signal >= target) return true;  // signal 到达
                // timeout 处理...
            });
        }
    }
}

// dispatch_impl 开始: 全局 barrier (kDispatchTag0)
comm::gpu_barrier<kIsScaleupNVLink, 1, kNumRanks, kNumSMs, kNumThreads, 
                  kNumQPs, kNumTimeoutCycles, comm::kDispatchTag0, 
                  false, false, true>(...);  // kSyncAtStart=true

// ... 执行 dispatch 操作 ...

// dispatch_impl 结束: 全局 barrier (kDispatchTag1)
comm::gpu_barrier<kIsScaleupNVLink, 1, kNumRanks, kNumSMs, kNumThreads, 
                  kNumQPs, kNumTimeoutCycles, comm::kDispatchTag1, 
                  true, true, false>(...);  // kFlushStores=true, kSyncAtStart=true

// 触发 copy epilogue kernel
cudaTriggerProgrammaticLaunchCompletion();

// 阶段 0: Notify warps - 统计 token 数量
if (warp_idx < kNumNotifyWarps) {
    // atomic add 统计 expert_count / rank_count
    // 全局 reduction
    // 写入 scaleout rank count
}

// 阶段 1: Scaleout warps - RDMA 发送
else if (warp_idx < kNumNotifyWarps + kNumScaleoutWarps) {
    // TMA load token
    // RDMA put 到远端
    // 更新 signaled_tail
}

// 阶段 2: Forward warps - NVLink 转发
else {
    // 轮询 signaled_tail 检查新数据
    // TMA load from scaleout_recv_buffer
    // TMA store to scaleup_buffer
}

// Forward warp 轮询等待 scaleout 数据
while ((wip_mask = ptx::gather(
    stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or stored_finish_flag == 0))) {

    // 等待该 rank 有数据到达 (或完成)
    comm::timeout_while<kNumTimeoutCycles>([&](const bool& is_last_check) {
        const uint32_t arrived_or_finished =
            stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or stored_finish_flag > 0;
        if (ptx::exchange(arrived_or_finished, recv_scaleout_rank_idx))
            return true;  // 数据到达

        // 读取新的 signaled tails
        if (lane_idx < kNumScaleoutRanks) {
            const auto signaled_tail = ptx::ld_acquire_sys<int64_t>(
                workspace_layout.get_scaleout_channel_signaled_tail_ptr(channel_idx, lane_idx));
            math::unpack2<int, int64_t>(signaled_tail, stored_finish_flag, stored_scaleout_tail_idx);
        }
        __syncwarp();
        return false;
    });
    // ... 处理数据 ...
}

// ptx.cuh 中的 mbarrier 原语
__forceinline__ __device__ void mbarrier_init_with_fence(mbarrier* ptr, const int& arrive_count) {
    asm volatile("mbarrier.init.shared::cta.b64 [%1], %0;" :: "r"(arrive_count), "r"(...));
    asm volatile("fence.mbarrier_init.release.cluster;");
}

__forceinline__ __device__ void mbarrier_arrive_and_set_tx(mbarrier* ptr, const int& num_bytes) {
    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0;" :: 
                 "r"(num_bytes), "r"(...));
}

__forceinline__ __device__ void mbarrier_wait_and_flip_phase(mbarrier* ptr, arrival_phase& phase) {
    asm volatile(
        "{\n\t"
        ".reg .pred P1; \n\t"
        "LAB_WAIT: \n\t"
        "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t"
        "@P1 bra DONE; \n\t"
        "bra LAB_WAIT; \n\t"
        "DONE: \n\t"
        "}" ::"r"(...), "r"(phase), "r"(0x989680));
    phase ^= 1;  // flip phase
}

// dispatch_impl 中的 TMA 流水线
ptx::arrival_phase phase = 0;
const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr();
if (ptx::elect_one_sync())
    ptx::mbarrier_init_with_fence(mbarrier_ptr, 1);
__syncwarp();

for (int token_idx = token_start; token_idx < num_tokens; token_idx += token_stride) {
    // 1. 等待 TMA store 完成
    ptx::tma_store_wait();
    __syncwarp();

    // 2. Issue TMA load (使用 mbarrier)
    if (ptx::elect_one_sync()) {
        ptx::tma_load_1d(tma_buffer.get_hidden_ptr(), 
                         math::advance_ptr(x, token_i64_idx * kNumHiddenBytes),
                         mbarrier_ptr, kNumHiddenBytes);  // cp.async.bulk with mbarrier
    }
    __syncwarp();

    // 3. 等待 TMA load 完成
    if (ptx::elect_one_sync()) {
        ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes);
        ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);  // 等待数据到达
    }
    __syncwarp();

    // 4. 使用数据: TMA store 到 send buffer
    ptx::tma_store_1d(send_buffer_ptr, tma_buffer.get_base_ptr(), ...);
    ptx::tma_store_commit();
    __syncwarp();
}

// combine_impl 中的 TMA 流水线
for (int i = token_start_idx; i < token_end_idx; ++i) {
    // 1. 等待 TMA store 完成
    ptx::tma_store_wait();
    __syncwarp();

    // 2. TMA load
    if (ptx::elect_one_sync()) {
        ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, mbarrier_ptr, kNumHiddenBytes);
        ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes);
        ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
    }
    __syncwarp();

    // 3. TMA store 到目标 buffer
    ptx::tma_store_1d(master_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), kNumHiddenBytes);
    ptx::tma_store_commit();
    __syncwarp();
}

// ptx.cuh
template <int kNumThreads>
__forceinline__ __device__ void named_barrier(const int& idx) {
    asm volatile("bar.sync %0, %1;" ::"r"(idx), "r"(kNumThreads));
}

// dispatch_impl 中的使用
constexpr int kNotifyBarrierIndex = 1;

// Notify warps 内部同步
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);

// 统计 expert count
for (int i = global_warp_idx; i < num_tokens; i += kNumNotifyWarps * kNumSMs) {
    atomicAdd_block(expert_count + dst_expert_idx, 1);
}
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);

// 全局 reduction
for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) {
    ptx::red_add(workspace_layout.get_notify_reduction_workspace_ptr() + i, counter);
}

// handle.cuh 中的 signal 封装
template <typename team_t, typename remote_action_t>
__device__ __forceinline__
void signal(const int& dst_rank_idx, const remote_action_t& remote_action) const {
    gin.signal(TEAM_WORLD_RAIL(), dst_rank_idx, remote_action);
}

// gin_barrier_wo_local_sync 中的使用
gin.signal(team, i, ncclGin_SignalInc{static_cast<ncclGinSignal_t>(rank_idx)});

// gin_barrier_wo_local_sync 中的等待
for (int i = thread_idx; i < kNumRanks; i += kNumThreads) {
    const auto signal_idx = static_cast<ncclGinSignal_t>(i);
    const auto shadow_ptr = gin.getSignalShadowPtr(signal_idx);
    const auto target = ++(*shadow_ptr);  // 期望的 signal 值
    const auto signal_ptr = reinterpret_cast<uint64_t*>(...);

    timeout_while<kNumTimeoutCycles>([=](const bool& is_last_check) {
        const auto signal = ptx::ld_acquire_sys<uint64_t>(signal_ptr);
        if (signal >= target) return true;  // signal 到达
        // timeout 处理...
    });
}

sequenceDiagram
    participant S as Sender
    participant F as Forwarder
    participant R as Receiver

    S->>S: Send Batch
    S->>F: Barrier (等待 Send 完成)
    F->>F: Forward Batch
    F->>R: Barrier (等待 Forward 完成)
    R->>R: Receive Batch

sequenceDiagram
    participant S as Sender
    participant F as Forwarder
    participant R as Receiver

    loop Streaming
        S->>S: Send Chunk
        S->>F: FIFO Write
        F->>F: Forward Chunk
        F->>R: FIFO Write
        R->>R: Receive Chunk
    end

sequenceDiagram
    participant N as Notify Warp
    participant S as RDMA Sender
    participant F as Forwarder
    participant R as NVL Receiver

    N->>N: Count Tokens
    N->>N: Barrier (barrier_block)
    N->>S: Send Metadata
    S->>S: barrier_block
    loop Per Token
        S->>S: Send to RDMA Buffer
        S->>F: Update tail (FIFO)
        F->>F: Poll tail (FIFO)
        F->>F: Forward to NVL Buffer
        F->>R: Update tail (FIFO)
        R->>R: Poll tail (FIFO)
        R->>R: Receive
    end
    R->>R: barrier_block (final)

sequenceDiagram
    participant N as Notify Warp
    participant SO as Scaleout Warp
    participant FW as Forward Warp
    participant RE as Receiver

    Note over N,RE: gpu_barrier (kDispatchTag0)

    par Parallel Execution
        N->>N: Count & Reduce
    and
        loop Per Token
            SO->>SO: TMA Load
            SO->>SO: RDMA Put
            SO->>FW: signaled_tail++ (FIFO)
            FW->>FW: Poll signaled_tail
            FW->>FW: TMA Load/Store
            FW->>RE: NVLink Write
        end
    end

    Note over N,RE: gpu_barrier (kDispatchTag1)

┌─────────────────────────────────────────────────────────────┐
│                    网格级 Barrier                            │
│  gpu_barrier / cooperative_groups::this_grid().sync()       │
│  使用场景: dispatch/combine 首尾, 阶段间大同步               │
├─────────────────────────────────────────────────────────────┤
│                    节点级 Barrier                            │
│  nvlink_barrier / gin_barrier / barrier_block               │
│  使用场景: 跨 SM 同步, 跨节点同步                            │
├─────────────────────────────────────────────────────────────┤
│                    Warp 组内 Barrier                         │
│  named_barrier / __syncthreads / __syncwarp                 │
│  使用场景: Notify warps 内部, TMA 操作前后                   │
├─────────────────────────────────────────────────────────────┤
│                    生产者-消费者 FIFO                        │
│  head/tail 指针 / signaled_tail / mbarrier                  │
│  使用场景: 缓冲区管理, Token 级流水线                        │
└─────────────────────────────────────────────────────────────┘