DeepEP Chunk Streaming 机制深度分析：博客第一性原理 vs 源码实现

struct Config {
    int num_sms;
    int num_max_nvl_chunked_send_tokens;      // NVL 发送 chunk 大小
    int num_max_nvl_chunked_recv_tokens;      // NVL 接收 chunk 大小
    int num_max_rdma_chunked_send_tokens;     // RDMA 发送 chunk 大小
    int num_max_rdma_chunked_recv_tokens;     // RDMA 接收 chunk 大小
    // ...
};

config_map = {
    2:   Config(Buffer.num_sms, 24, 256, 6, 128),
    4:   Config(Buffer.num_sms, 6, 256, 6, 128),
    8:   Config(Buffer.num_sms, 6, 256, 6, 128),
    16:  Config(Buffer.num_sms, 36, 288, 20, 128),
    24:  Config(Buffer.num_sms, 32, 288, 8, 128),
    32:  Config(Buffer.num_sms, 32, 288, 8, 128),
    48:  Config(Buffer.num_sms, 32, 288, 8, 128),
    64:  Config(Buffer.num_sms, 32, 288, 8, 128),
    96:  Config(Buffer.num_sms, 20, 480, 12, 128),
    128: Config(Buffer.num_sms, 20, 560, 12, 128),
    # Config 参数顺序：num_sms, nvl_chunked_send, nvl_chunked_recv, rdma_chunked_send, rdma_chunked_recv
}

// Issue RDMA send
auto num_tokens_to_issue = min(num_tokens_processed, num_max_rdma_chunked_send_tokens);
EP_DEVICE_ASSERT(num_tokens_to_issue >= 0 and num_tokens_to_issue <= synced_num_tokens_to_send);
if (dst_rdma_rank != rdma_rank) {
    auto dst_slot_idx = synced_last_issued_tail % num_max_rdma_chunked_recv_tokens;
    EP_DEVICE_ASSERT(dst_slot_idx + num_tokens_to_issue <= num_max_rdma_chunked_recv_tokens);
    const size_t num_bytes_per_msg = num_bytes_per_token * num_tokens_to_issue;
    nvshmemi_ibgda_put_nbi_warp<true>(dst_ptr, src_ptr, num_bytes_per_msg, ...);
}

// Iterate over all tokens and send by chunks
int current_rdma_idx = channel_id % kNumRDMARanks;
while (true) {
    // ...
    int num_tokens_in_chunk =
        __shfl_sync(0xffffffff, min(num_max_nvl_chunked_send_tokens, token_end_idx - token_start_idx), current_rdma_idx);

    // Send by chunk
    for (int chunk_idx = 0; chunk_idx < num_tokens_in_chunk; ++chunk_idx, ++token_idx) {
        int dst_slot_idx = (cached_channel_tail_idx++) % num_max_nvl_chunked_recv_tokens_per_rdma;
        // ... TMA store
    }
}

flowchart LR
    TokenStream["Token Stream<br/>调度粒度"] -->|"token 逐个处理"| Accumulate["Accumulate<br/>计数达到 chunk 大小"]
    Accumulate -->|"num_tokens_processed<br/>≥ chunk_size"| Chunk["Chunk<br/>通信粒度"]
    Chunk -->|"nvshmemi_ibgda_put_nbi<br/>一次性发送"| RDMA["RDMA Network"]
    Chunk -->|"tma_store_1d<br/>一次性发送"| NVL["NVLink"]

// Iterate all tokens
const auto token_start = dispatch_warp_idx * kNumSMs + sm_idx;
const auto token_stride = kNumDispatchWarps * kNumSMs;
for (int token_idx = token_start; token_idx < num_tokens; token_idx += token_stride) {
    // ... TMA load hidden + SF + topk_idx + topk_weights
    // ... TMA store to send buffer
    // ... TMA NVLink store
    // Issue RDMA put
    if constexpr (not kIsScaleupNVLink) {
        ptx::tma_store_wait<1>();
        __syncwarp();
        if (stored_dst_slot_idx >= 0 and dst_ptr == nullptr) {
            gin.put<team_t>(recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(),
                            send_buffer_ptr, tma_buffer.get_num_bytes<false>(), stored_dst_rank_idx);
        }
        __syncwarp();
    }
}

// Issue scaleup rank count writes to peers
for (int i = thread_idx; i < kNumRanks; i += kNumNotifyThreads) {
    const auto dst_rank_counter =
        workspace_layout.get_scaleup_rank_count_ptr<false>() + rank_idx;
    gin.put_value<team_t>(dst_rank_counter, static_cast<int64_t>(rank_count[i]), i,
                          ncclGinOptFlagsAggregateRequests);
}

int kScaleoutUpdateInterval = 6,
int kNumSlotsPerForwardChunk = kScaleoutUpdateInterval,

// Process one chunk from the current rank
const auto start_slot_idx = ptx::exchange(stored_scaleout_old_tail_idx, recv_scaleout_rank_idx);
const auto end_slot_idx = std::min(
    ptx::exchange(stored_scaleout_tail_idx, recv_scaleout_rank_idx),
    start_slot_idx + kNumSlotsPerForwardChunk
);
if (lane_idx == recv_scaleout_rank_idx)
    stored_scaleout_old_tail_idx = end_slot_idx;

// Record metadata at forward
if constexpr (not kReuseSlotIndices) {
    const auto metadata_ptr = token_metadata_at_forward +
        num_tokens_processed * kNumForwardMetadataDims;

    // Source token index and last token index flag
    if (ptx::elect_one_sync()) {
        metadata_ptr[0] = tma_buffer.get_src_token_global_idx_ptr()[0];
        metadata_ptr[1] = slot_idx == (end_slot_idx - 1);  // ← is_token_last_in_chunk
    }
    // ...
}

const auto flush_last_tma_and_issue_rdma = [&]() {
    if (last_src_scaleout_rank_idx >= 0 and ptx::elect_one_sync()) {
        ptx::tma_store_wait();

        // Issue only if not local rank
        if (last_src_scaleout_rank_idx != scaleout_rank_idx) {
            gin.put<ncclTeamTagRail>(
                last_recv_token_buffer_ptr,
                last_send_token_buffer_ptr,
                token_layout.get_num_bytes<false>(),
                last_src_scaleout_rank_idx,
                last_is_token_last_in_chunk ? 0 : ncclGinOptFlagsAggregateRequests
                //  ↑↑↑ 关键逻辑 ↑↑↑
            );
        }
    }
    __syncwarp();
};

if is_token_last_in_chunk:
    flags = 0              → 立即 flush，不聚合
else:
    flags = ncclGinOptFlagsAggregateRequests  → 聚合等待

gin.put<ncclTeamTagRail>(
    recv_buffer_ptr,
    send_buffer_ptr,
    token_layout.get_num_bytes<false>(),
    src_scaleout_rank_idx,
    topk_valid_mask == 0 and is_token_last_in_chunk ? 0 : ncclGinOptFlagsAggregateRequests
);

int kNumChannelsPerSM = kNumScaleoutWarps,
int kNumChannels = kNumScaleoutWarps * kNumSMs,
int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, kNumChannels),

// Write the per-scaleup channel index for this token
int linked_list_idx = -1;
#pragma unroll
for (int j = 0; j < kNumScaleupRanksPerLane; ++ j) {
    const auto src_lane_idx = stored_dst_scaleup_rank_idx - j * 32;
    const bool valid = 0 <= src_lane_idx and src_lane_idx < 32;
    const auto exchanged = ptx::exchange(
        stored_scaleup_send_counters[j], valid ? src_lane_idx : 0);
    linked_list_idx = valid ? exchanged : linked_list_idx;
}
if (not kReuseSlotIndices and lane_idx < kNumTopk) {
    tma_buffer.get_linked_list_idx_ptr()[lane_idx] = transform_linked_list_idx(linked_list_idx);
    ptx::tma_store_fence();
}

flowchart TB
    subgraph Scaleout["Scale-Out Dispatch"]
        SOT["Scaleout Warp<br/>Channel c"] -->|"token_idx = c + k"| LOAD["TMA Load Token"]
        LOAD --> DEDUP["Deduplicate<br/>scaleout ranks"]
        LOAD --> META["Record Metadata<br/>src_token_global_idx<br/>is_token_last_in_chunk"]
        DEDUP --> SEND["Send Buffer<br/>scaleout_send_buffer"]
        SEND -->|"ncclGinOptFlagsAggregateRequests"| IBGDA["IBGDA Put<br/>RDMA Scale-Out"]
    end

    subgraph Forward["Forward Warp"]
        FWD["Forward Warp<br/>Channel c"] -->|"wait signaled_tail"| FR["Forward from<br/>scaleout_recv_buffer"]
        FR --> CHUNK["Process Chunk<br/>≤ kNumSlotsPerForwardChunk=6"]
        CHUNK --> NVL["TMA Store →<br/>scaleup_buffer NVLink"]
        CHUNK --> LNK["Linked List<br/>per scaleup rank"]
        CHUNK --> META2["metadata_at_forward<br/>[src_global_idx, is_last_in_chunk, dst_rank, dst_slot]"]
    end

    subgraph Combine["Combine"]
        CU["Scaleup Combine Warp"] -->|"replay linked_list"| RDX["Reduce from<br/>scaleup_buffer"]
        RDX -->|"is_token_last_in_chunk?<br/>flush : aggregate"| RDMA2["RDMA Put →<br/>scaleout_recv_buffer"]
    end

    IBGDA --> FWD
    LNK --> CU
    META2 --> CU

V1: 应用层 Chunk → 固定阈值 flush → RDMA
V2: Token 流式 → NCCL AggregateRequests → NCCL 运行时决定 flush 时机

flowchart LR
    subgraph V1["Legacy: 应用层 Chunk"]
        T1["Token 1"] -->|等待| C1["Chunk 1<br/>6 tokens"]
        T2["Token 2"] -->|等待| C1
        T3["Token 3"] -->|等待| C1
        T4["Token 4"] -->|等待| C1
        T5["Token 5"] -->|等待| C1
        T6["Token 6"] -->|等待| C1
        C1 -->|一次性发送| R1["RDMA PUT"]
    end

    subgraph V2["Elastic: NCCL 聚合"]
        TT1["Token 1"] -->|"put(aggregate)"| Q["NCCL QP<br/>Pending Queue"]
        TT2["Token 2"] -->|"put(aggregate)"| Q
        TT3["Token 3"] -->|"put(aggregate)"| Q
        TT4["Token 4"] -->|"put(aggregate)"| Q
        TT5["Token 5"] -->|"put(aggregate)"| Q
        TT6["Token 6"] -->|"put(flush)"| Q
        Q -->|doorbell batch| R2["RDMA PUT"]
    end