DeepEP Warp Specialization 深度分析: 博客描述 vs 源码实现

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

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

__launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + LEGACY_NUM_MAX_NVL_PEERS) * 32), 1)
// 即: (7 + 1 + 8) * 32 = 512 threads per CTA

┌─────────────────────────────────────────────────────────────────────────────────────────┐
│                    DeepEP Legacy Dispatch CTA (1 SM per channel)                         │
├────────────────┬─────────────────┬─────────────────────┬───────────────┬────────────────┤
│ RDMA Senders   │ Sender Coord    │ RDMA-NVL Forwarders │ Forward Coord │ NVL Receivers  │
│ (7 warps)      │ (1 warp)        │ (8 warps)           │ (1 warp)      │ (8 warps)      │
├────────────────┼─────────────────┼─────────────────────┼───────────────┼────────────────┤
│ • Token 遍历    │ • 流控窗口维护   │ • RDMA buf 消费      │ • Head 协调   │ • NVL buf 消费  │
│ • 数据打包      │ • Issue put_nbi │ • NVL buf 生产       │ • RDMA release│ • 数据解包      │
│ • RDMA buf 生产 │ • Tail 同步     │ • TMA load/store    │ • 流控        │ • TopK 写入     │
│ • 窗口流控      │                 │ • Round-robin       │               │ • TMA store     │
└────────────────┴─────────────────┴─────────────────────┴───────────────┴────────────────┘
         │                 │               │                   │              │
         └────── Producer ─┘               └───── Producer ────┘              │
                (RDMA buffer)                  (NVL buffer)                   │
                                     Consumer ─────────────────────────────────┘

enum class WarpRole {
    kNVLSender,          // NVL 发送
    kNVLAndRDMAForwarder, // NVL-RDMA 转发
    kRDMAReceiver,       // RDMA 接收
    kCoordinator         // 协调
};

Dispatch: RDMA Send → Forward → NVL Receive
Combine:  NVL Send → Forward → RDMA Receive

} else if (warp_role == WarpRole::kRDMAAndNVLForwarder) {
    // RDMA consumers and NVL producers
    const auto dst_nvl_rank = target_rank;

    // 1. 等待 RDMA 元数据到达
    int num_tokens_to_recv_from_rdma = 0, src_rdma_channel_prefix = 0;
    if (lane_id < kNumRDMARanks) {
        while (true) {
            auto meta_0 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + dst_nvl_rank);
            auto meta_1 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + LEGACY_NUM_MAX_NVL_PEERS + dst_nvl_rank);
            // ... 等待 4 个 meta 都为负值 (编码后的计数)
            if (meta_0 < 0 and meta_1 < 0 and meta_2 < 0 and meta_3 < 0) {
                int start_sum = -meta_0 - 1, end_sum = -meta_1 - 1;
                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. Round-robin 从各 RDMA rank 取数据转发
    int src_rdma_rank = sm_id % kNumRDMARanks;
    while (__any_sync(0xffffffff, num_tokens_to_recv_from_rdma > 0)) {
        // 轮询找下一个有数据的 RDMA rank
        src_rdma_rank = (src_rdma_rank + 1) % kNumRDMARanks;
        if (__shfl_sync(0xffffffff, num_tokens_to_recv_from_rdma, src_rdma_rank) > 0) { ... }

        // 3. 从 RDMA buffer 读取并写入 NVL buffer
        for (int i = src_rdma_head; i < src_rdma_tail; ++i) {
            auto rdma_slot_idx = i % num_max_rdma_chunked_recv_tokens;
            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
            if (not src_meta.is_token_in_nvl_rank(dst_nvl_rank))
                continue;

            // 4. TMA load from RDMA buffer, TMA store to NVL buffer
            int dst_slot_idx = (cached_nvl_channel_tail++) % num_max_nvl_chunked_recv_tokens;
            auto dst_shifted = nvl_channel_x.buffer() + dst_slot_idx * num_bytes_per_token;

            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);
            }
            __syncwarp();
            mbarrier_wait(tma_mbarrier, tma_phase);
            if (elect_one_sync())
                tma_store_1d(tma_buffer, dst_shifted, num_bytes_per_token);
        }
    }
}

RDMA Buffer (per rank) → TMA Load → Shared Memory → TMA Store → NVL Buffer (per rank)
                                ↓
                       检查 SourceMeta.is_token_in_nvl_rank(dst_nvl_rank)
                       只转发属于本 NVL rank 的 tokens

} else if (warp_role == WarpRole::kForwarderCoordinator) {
    // Forward warp coordinator
    if (target_rank > 0) return;  // 只有第一个 forwarder SM 的 coordinator 工作

    // 清理共享内存
    for (int i = lane_id; i < kNumRDMARanks * LEGACY_NUM_MAX_NVL_PEERS; i += 32)
        forward_channel_head[i % LEGACY_NUM_MAX_NVL_PEERS][i / LEGACY_NUM_MAX_NVL_PEERS] = 0;
    sync_forwarder_smem();

    while (true) {
        // 找所有 forwarder 的最小 head
        int min_head = std::numeric_limits<int>::max();
        for (int i = 0; i < LEGACY_NUM_MAX_NVL_PEERS; ++i)
            if (not forward_channel_retired[i])
                min_head = min(min_head, forward_channel_head[i][target_rdma]);

        // 更新远端 RDMA head (释放 RDMA buffer 空间)
        if (min_head >= last_head + num_max_rdma_chunked_send_tokens and lane_id < kNumRDMARanks) {
            nvshmemi_ibgda_amo_nonfetch_add(rdma_channel_head.buffer(rdma_rank),
                                            min_head - last_head,
                                            translate_dst_rdma_rank<kLowLatencyMode>(lane_id, nvl_rank),
                                            channel_id + num_channels,
                                            lane_id == rdma_rank);
            last_head = min_head;
        }
        __nanosleep(LEGACY_NUM_WAIT_NANOSECONDS);  // 让出 SM 给其他 warps
    }
}

// 21-29 行: 模板参数定义
template <bool kIsScaleupNVLink,
          bool kDoCPUSync,
          bool kReuseSlotIndices,
          int kNumSMs,
          int kNumNotifyWarps, int kNumDispatchWarps,  // ← 两种 warp 角色
          int kNumRanks, ...>
__global__ void __launch_bounds__(kNumThreads, 1)
dispatch_impl(...) {
    // 79 行: warp 分派
    if (warp_idx < kNumNotifyWarps) {
        // ====== Notify Warps ======
        // 1. atomicAdd 统计 expert_count / rank_count
        // 2. full-grid reduction
        // 3. 写入 peer 的 rank_expert_count
        // 4. 计算 prefix sum
        // 5. 写入 host workspace
    } else {
        // ====== Dispatch Warps ======
        // 1. TMA load token from global memory
        // 2. TMA store to send buffer
        // 3. NVLink store (tma_store_1d to sym_ptr)
        // 4. RDMA put (gin.put, 若非 NVLink 可达)
    }
}

┌───────────────────────────────────────────────────────────────┐
│              DeepEP Scale-up Only CTA                          │
├────────────────────────────┬──────────────────────────────────┤
│ Notify Warps               │ Dispatch Warps                   │
│ (warp_idx < kNumNotifyWarps)│ (warp_idx >= kNumNotifyWarps)    │
├────────────────────────────┼────────────────────────────────┤
│ • Expert/Rank 计数          │ • TMA Load tokens              │
│ • Prefix sum 计算           │ • TMA Store to send buffer     │
│ • Peer 元数据同步           │ • NVLink Store (对称内存)       │
│ • Host workspace 写入       │ • RDMA Put (若非 NVLink 可达)   │
└────────────────────────────┴────────────────────────────────┘

// 16-32 行: 模板参数定义**三种**数据路径 warp 角色
template <bool kDoCPUSync,
          bool kReuseSlotIndices,
          int kNumSMs,
          int kNumNotifyWarps, int kNumScaleoutWarps, int kNumForwardWarps,  // ← 三种!
          int kNumScaleoutRanks, int kNumScaleupRanks,
          ...>
__global__ void __launch_bounds__(kNumThreads, 1)
hybrid_dispatch_impl(...) {

int kNumNotifyThreads = kNumNotifyWarps * 32,
int kNumScaleoutSendThreads = kNumScaleoutWarps * 32,
int kNumForwardThreads = kNumForwardWarps * 32,
int kNumThreads = kNumNotifyThreads + kNumScaleoutSendThreads + kNumForwardThreads

if (warp_idx < kNumNotifyWarps) {
    // ====== Notify Warps ======
    // 功能: 元数据收集与同步
    // 1. atomicAdd 统计 expert_count / rank_count
    // 2. full-grid reduction
    // 3. 写入 scale-out peers 的 rank_expert_count (RDMA)
    // 4. 写入 scale-up peers 的 rank_expert_count (NVLink)
    // 5. 计算 prefix sum

} else if (warp_idx < kNumNotifyWarps + kNumScaleoutWarps) {
    // ====== Scaleout Warps (IBGDA/RDMA Sending) ======
    // 功能: 跨节点数据发送
    for (int token_idx = channel_idx; token_idx < num_tokens; token_idx += kNumChannels) {
        // TMA load token
        // TMA store to scaleout_send_buffer
        gin.put<ncclTeamTagRail>(..., stored_dst_scaleout_rank_idx);  // IBGDA!
        // Issue scale-out tail update
    }

} else {
    // ====== Forward Warps (NVLink Forwarding) ======
    // 功能: 从 scale-out 接收 buffer 转发到 scale-up 目标
    const int forward_warp_idx = warp_idx - (kNumNotifyWarps + kNumScaleoutWarps);
    // ...
    // Forward tokens from scale-out ranks
    for (int slot_idx = start_slot_idx; slot_idx < end_slot_idx; ++slot_idx) {
        const auto token_buffer = recv_buffer.get_token_buffer(slot_idx);
        // TMA load into shared memory
        // TMA store to scaleup_buffer (NVLink sym_ptr)
    }
}

┌─────────────────────────────────────────────────────────────────────────────────┐
│                    DeepEP Hybrid Dispatch CTA                                    │
├──────────────────┬──────────────────────────┬───────────────────────────────────┤
│ Notify Warps     │ Scaleout Warps           │ Forward Warps                     │
│ (元数据同步)      │ (IBGDA/RDMA 跨节点发送)   │ (NVLink 节点内转发)                │
├──────────────────┼──────────────────────────┼─────────────────────────────────┤
│ • Expert 计数     │ • Token TMA Load         │ • 从 scaleout_recv 读取           │
│ • Rank 计数       │ • Send buffer 写入        │ • TMA Load → smem                 │
│ • Prefix sum     │ • IBGDA put (RDMA)       │ • TMA Store → scaleup_buffer     │
│ • Peer 元数据同步  │ • Tail update            │ • 按 scale-up rank 分发            │
│ • Grid sync      │ • Channel 流控            │ • Linked list 索引维护             │
└──────────────────┴──────────────────────────┴─────────────────────────────────┘

// 25 行: 两种 warp 角色
int kNumWarps = kNumScaleupWarps + kNumForwardWarps;

// 106 行: warp 分派
if (warp_idx < kNumScaleupWarps) {
    // ====== Scaleup Warps (节点内发送) ======
    // 功能: 从本地 buffer 发送到 scale-up peers
    // 241 行: TMA load from global memory
    // 328 行: TMA store to remote scale-up buffer
    // 127 行: update_tails (通知 forward warps)

} else {
    // ====== Forward Warps (节点间转发) ======
    // 功能: 从 scale-up buffer 转发到 scale-out peers
    const auto forward_warp_idx = warp_idx - kNumScaleupWarps;
    // 467 行: TMA load from scaleup_buffer
    // 477 行: TMA store to scaleout_send_buffer
    // 484 行: gin.put<ncclTeamTagRail> (RDMA)
}

const bool kAdjustRegisters = (kNumChannelsPerSM == 4 or kNumChannelsPerSM == 8) and not kUseExpandedLayout;
constexpr int kNumRegistersForScaleupWarps = 40;      // 少
constexpr int kNumRegistersForForwardWarps = 256 - kNumRegistersForScaleupWarps;  // 多 (216)

┌──────────────────────────────────────────────────────────────────────────────────────────┐
│                          Legacy Dispatch Pipeline                                         │
│                                                                                          │
│  ┌──────────────┐    RDMA Buffer    ┌──────────────────┐    NVL Buffer    ┌────────────┐│
│  │ RDMA Senders │ ───────────────→  │ RDMA-NVL Forwarder│ ──────────────→  │ NVL Receiver││
│  │ (Producer)   │   (per rank)      │ (Consumer/Prod)   │  (per rank)      │ (Consumer)  ││
│  └──────────────┘                   └──────────────────┘                  └────────────┘│
│         ↑                               ↑                                     ↑        │
│         │                               │                                     │        │
│  ┌──────────────┐                ┌──────────────┐                      ┌────────────┐  │
│  │ Sender Coord │                │ Forward Coord│                      │ (implicit) │  │
│  │ (流控/Issue) │                │ (Head 协调)  │                      │            │  │
│  └──────────────┘                └──────────────┘                      └────────────┘  │
└──────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────┐
│                       Elastic Hybrid Dispatch Pipeline                                    │
│                                                                                          │
│  ┌────────────┐                                                    ┌─────────────────┐  │
│  │ Notify     │ ──────────────────────────────────────────────────→ │ (Barrier/Sync)  │  │
│  │ Warps      │  (元数据: rank_count, expert_count, prefix sum)     │                 │  │
│  └────────────┘                                                    └─────────────────┘  │
│                                                                                          │
│  ┌────────────────┐   scaleout_send_buffer   ┌─────────────────┐   scaleup_buffer  ┌────┐│
│  │ Scaleout Warps │ ──────────────────────→  │  Forward Warps   │ ───────────────→ │Done││
│  │ (RDMA Producer)│   (per channel)          │ (NVL Producer)   │  (per rank)       │    ││
│  └────────────────┘                          └─────────────────┘                   └────┘│
│         ↑                                          ↑                                    │
│         │                                          │                                    │
│  ┌────────────────┐                        ┌─────────────────┐                          │
│  │ scaleout_tail  │                        │  scaleup_tail   │                          │
│  │ (signaled)     │                        │  (linked list)  │                          │
│  └────────────────┘                        └─────────────────┘                          │
└──────────────────────────────────────────────────────────────────────────────────────────┘

if (lane_id < kNumRDMARanks) {
    while (true) {
        auto meta_0 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + dst_nvl_rank);
        auto meta_1 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + LEGACY_NUM_MAX_NVL_PEERS + dst_nvl_rank);
        auto meta_2 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + LEGACY_NUM_MAX_NVL_PEERS * 2);
        auto meta_3 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + LEGACY_NUM_MAX_NVL_PEERS * 2 + 1);
        if (meta_0 < 0 and meta_1 < 0 and meta_2 < 0 and meta_3 < 0) {
            // 解码: -value - 1 编码
            int start_sum = -meta_0 - 1, end_sum = -meta_1 - 1;
            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;
        }
    }
}

int dst_slot_idx = (cached_nvl_channel_tail++) % num_max_nvl_chunked_recv_tokens;
auto dst_shifted = nvl_channel_x.buffer() + dst_slot_idx * num_bytes_per_token;

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);
}
__syncwarp();
mbarrier_wait(tma_mbarrier, tma_phase);
if (elect_one_sync())
    tma_store_1d(tma_buffer, dst_shifted, num_bytes_per_token);

// Issue IBGDA requests
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);
}

// TMA load into shared memory
if (ptx::elect_one_sync()) {
    ptx::tma_load_1d(tma_buffer.get_base_ptr(), token_buffer.get_base_ptr(),
                     mbarrier_ptr, token_layout.get_num_bytes<false>());
    ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, token_layout.get_num_bytes<false>());
    ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase);
}

// Issue TMAs to scale-up buffer (NVLink)
if (stored_dst_slot_idx >= 0) {
    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>());
    ptx::tma_store_commit();
}

┌─────────────────────────────────────────────────────────────────────────────┐
│ Legacy Dispatch (5 roles)                                                   │
├─────────────┬──────────────┬────────────────┬──────────────┬────────────────┤
│ RDMA Sender │ Sender Coord │ RDMA-NVL Fwd   │ Forward Coord│ NVL Receiver   │
│ 7 warps     │ 1 warp       │ 8 warps        │ 1 warp       │ 8 warps        │
├─────────────┴──────────────┴────────────────┴──────────────┴────────────────┤
│ 同步: rdma_channel_tail/head, nvl_channel_tail/head, rdma_channel_meta     │
│ 特点: 显式 Coordinator, 基于 nvshmem put_nbi, 共享内存 barrier.sync         │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│ Elastic Scale-up Dispatch (2 roles)                                         │
├────────────────────────────┬────────────────────────────────────────────────┤
│ Notify Warps               │ Dispatch Warps                                 │
│ (元数据 + 同步)             │ (数据搬运)                                      │
├────────────────────────────┴────────────────────────────────────────────────┤
│ 同步: gpu_barrier, NCCL Gin, tma_store_wait                                │
│ 特点: 无 Forward (全 NVLink 直达), Notify 独立化, TMA load/store            │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│ Elastic Hybrid Dispatch (3 roles)                                           │
├────────────────┬─────────────────────────┬──────────────────────────────────┤
│ Notify Warps   │ Scaleout Warps          │ Forward Warps                    │
│ (元数据同步)    │ (RDMA 跨节点发送)         │ (NVLink 节点内转发)               │
├────────────────┴─────────────────────────┴──────────────────────────────────┤
│ 同步: scaleout_channel_signaled_tail, channel_scaleup_tail, linked list     │
│ 特点: 3 角色, 显式 Notify, Round-robin 调度, 寄存器重分配 (Combine)          │
└─────────────────────────────────────────────────────────────────────────────┘

graph LR
    subgraph "Legacy Dispatch CTA"
        A[RDMA Senders<br/>7 warps] -->|rdma_channel_tail| B[RDMA Buffer]
        B -->|rdma_channel_head| C[RDMA-NVL Forwarder<br/>8 warps]
        C -->|nvl_channel_tail| D[NVL Buffer]
        D -->|nvl_channel_head| E[NVL Receiver<br/>8 warps]

        F[Sender Coord<br/>1 warp] -->|issue put_nbi| A
        G[Forward Coord<br/>1 warp] -->|release head| C
    end

    subgraph "同步信号"
        H[rdma_channel_meta] -->|通知到达| C
        I[nvl_channel_prefix] -->|偏移通知| E
    end

graph TB
    subgraph "Elastic Hybrid Dispatch CTA"
        A[Notify Warps<br/>元数据同步] -->|rank_count, expert_count, prefix_sum| B[Global Sync]

        C[Scaleout Warps<br/>RDMA 发送] -->|gin.put + tail update| D[scaleout_recv_buffer]
        D -->|signaled_tail| E[Forward Warps<br/>NVLink 转发]
        E -->|TMA store + linked_list| F[scaleup_buffer]
    end

    subgraph "数据流"
        G[Token Buffer] -->|TMA load| C
        F -->|arrival| H[Scale-up Peers]
    end

graph TB
    subgraph "Legacy V1"
        L1[RDMA Sender] --> L2[RDMA-NVL Forwarder] --> L3[NVL Receiver]
        L1 -.-> LC1[Sender Coord]
        L2 -.-> LC2[Forward Coord]
    end

    subgraph "Elastic V2 Scale-up"
        E1[Notify Warps] --> E2[Dispatch Warps]
    end

    subgraph "Elastic V2 Hybrid"
        H1[Notify Warps] --> H2[Scaleout Warps] --> H3[Forward Warps]
    end