DeepEP Buffer 系统深度分析：代码实证 vs 博客"第一性原理"模型

class ElasticBuffer {
    // Buffer bytes = GPU buffer + CPU buffer (excludes workspace)
    // Memory layout: [[[Workspace] GPU buffer] CPU buffer]
    int64_t num_buffer_bytes;
    int64_t num_gpu_buffer_bytes;
    int64_t num_cpu_buffer_bytes;
    void* buffer;
    void *workspace;
    void *host_workspace, *mapped_host_workspace;
    // ...
};

ElasticBuffer(...) {
    // Workspace is aligned to 2 MB so that it sits cleanly at the front of the GPU segment
    const auto num_workspace_bytes = math::align<int64_t>(
        layout::WorkspaceLayout::get_num_bytes(), symmetric::kNumAlignmentBytes);

    // Create NCCL symmetric memory context
    // Symmetric memory layout: [[[Workspace] GPU buffer] CPU buffer]
    const auto num_sym_bytes = num_workspace_bytes + num_buffer_bytes;
    this->nccl_context = std::make_shared<nccl::NCCLSymmetricMemoryContext>(
        nccl_comm, cpu_comm, num_ranks, rank_idx,
        num_sym_bytes, num_cpu_buffer_bytes, ...);

    // Assign workspaces and buffers
    workspace = this->nccl_context->mapped_window_ptr;
    workspace_layout_wo_expert = std::make_shared<layout::WorkspaceLayout>(
        workspace, nccl_context->num_scaleout_ranks, nccl_context->num_scaleup_ranks, 0);
    buffer = static_cast<uint8_t*>(workspace) + num_workspace_bytes;
    CUDA_RUNTIME_CHECK(cudaMemset(workspace, 0, num_workspace_bytes));

    // Allocate host workspaces
    CUDA_RUNTIME_CHECK(cudaMallocHost(&host_workspace, layout::WorkspaceLayout::get_num_bytes(), cudaHostAllocMapped));
    CUDA_RUNTIME_CHECK(cudaHostGetDevicePointer(&mapped_host_workspace, host_workspace, 0));
    std::memset(host_workspace, 0, layout::WorkspaceLayout::get_num_bytes());
}

template <bool kWithMBarrier>
struct BufferLayout {
    TokenLayout token_layout;
    int num_ranks;
    int num_max_tokens_per_rank;
    void* base;

    __forceinline__ __device__ __host__
    BufferLayout(const TokenLayout& token_layout,
                 const int& num_ranks,
                 const int& max_num_tokens_per_rank,
                 void* base = nullptr) :
        token_layout(token_layout),
        num_ranks(num_ranks), num_max_tokens_per_rank(max_num_tokens_per_rank),
        base(base) {}

    int64_t get_num_bytes_per_rank() const {
        return num_max_tokens_per_rank * get_num_bytes_per_token();
    }

    int64_t get_num_bytes() const {
        return get_num_bytes_per_rank() * num_ranks;
    }

    BufferLayout get_rank_buffer(const int& rank_idx) const {
        return BufferLayout(token_layout,
                            1, num_max_tokens_per_rank,
                            static_cast<int8_t*>(base) + get_num_bytes_per_rank() * rank_idx);
    }

    TokenLayout get_token_buffer(const int& token_idx, const bool& global = false) const {
        EP_UNIFIED_ASSERT(num_ranks == 1 or global);
        return TokenLayout(token_layout.num_hidden_bytes, token_layout.num_sf_bytes,
                           token_layout.num_topk, token_layout.with_metadata,
                           static_cast<int8_t*>(base) + token_layout.get_num_bytes<kWithMBarrier, int64_t>() * token_idx);
    }
};

BufferLayout = [num_ranks] × [num_max_tokens_per_rank] × [TokenLayout]

struct TokenLayout {
    int num_hidden_bytes, num_sf_bytes;
    bool with_metadata;
    int num_topk, num_metadata_bytes;
    void* base;

    __forceinline__ __device__ __host__
    TokenLayout(const int& num_hidden_bytes, const int& num_sf_bytes,
                const int& num_topk, const bool& with_metadata, void* base = nullptr) :
        num_hidden_bytes(num_hidden_bytes),
        num_sf_bytes(num_sf_bytes),
        with_metadata(with_metadata),
        num_topk(num_topk),
        num_metadata_bytes(num_topk * (sizeof(int) + sizeof(float)) +
                           (with_metadata ? (1 + num_topk) * sizeof(int) : 0)),
        base(base) {}

    template <bool kWithMBarrier, typename dtype_t = int>
    __forceinline__ __device__ __host__ dtype_t get_num_bytes() const {
        const auto num_bytes = math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes) +
                               math::align(num_sf_bytes, ptx::kNumTMAAlignBytes) +
                               math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes) +
                               math::align<int>(kWithMBarrier ? sizeof(ptx::mbarrier) : 0, ptx::kNumTMAAlignBytes);
        return static_cast<dtype_t>(num_bytes);
    }

    void* get_hidden_ptr() const { return get_base_ptr(); }
    sf_pack_t* get_sf_ptr() const {
        return math::advance_ptr<sf_pack_t>(base, math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes));
    }
    int* get_topk_idx_ptr() const { return get_metadata_ptr(); }
    float* get_topk_weights_ptr() const {
        return math::advance_ptr<float>(get_metadata_ptr(), num_topk * sizeof(int));
    }
    int* get_src_token_global_idx_ptr() const {
        return math::advance_ptr<int>(get_topk_weights_ptr(), num_topk * sizeof(float));
    }
    int* get_linked_list_idx_ptr() const {
        return get_src_token_global_idx_ptr() + 1;
    }
    ptx::mbarrier* get_mbarrier_ptr() const {
        return math::advance_ptr<ptx::mbarrier>(get_metadata_ptr(), math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes));
    }
};

[hidden_bytes | SF_bytes | metadata_bytes | mbarrier?]
     ↓              ↓            ↓
   hidden     scale factors   topk_idx[num_topk]
                              topk_weights[num_topk]
                              src_token_global_idx
                              linked_list_idx[num_topk]

static int64_t get_dispatch_buffer_size(const int& num_max_tokens_per_rank,
                                        const int& hidden, const int& num_sf_packs, const int& num_topk,
                                        const int& elem_size,
                                        const int& num_scaleout_ranks, const int& num_scaleup_ranks,
                                        const bool& is_scaleup_nvlink) {
    const auto num_ranks = num_scaleup_ranks * num_scaleout_ranks;
    const auto token_layout = get_dispatch_token_layout(hidden, elem_size, num_sf_packs, num_topk);

    if (num_scaleout_ranks == 1) {
        // Direct dispatch
        const auto send_buffer_layout = layout::BufferLayout<false>(
            token_layout, is_scaleup_nvlink ? 0 : 1, num_max_tokens_per_rank);
        const auto recv_buffer_layout = layout::BufferLayout<false>(
            token_layout, num_ranks, num_max_tokens_per_rank);
        return send_buffer_layout.get_num_bytes() + recv_buffer_layout.get_num_bytes();
    }
    // ...
}

GPU buffer = [recv_buffer (num_ranks × max_tokens) | send_buffer (1 × max_tokens)]

// Buffer layouts
const auto token_layout = layout::TokenLayout(kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true);
const auto tma_buffer = layout::BufferLayout<true>(token_layout, kNumDispatchWarps, 1,
    math::advance_ptr<int>(smem, kNumSmemBytesForNotify)).get_rank_buffer(dispatch_warp_idx).get_token_buffer(0);
auto recv_buffer = layout::BufferLayout<false>(token_layout, kNumRanks, kNumMaxTokensPerRank, buffer);
auto send_buffer = layout::BufferLayout<false>(token_layout, 1, kNumMaxTokensPerRank, recv_buffer.get_buffer_end_ptr());
recv_buffer = recv_buffer.get_rank_buffer(rank_idx);

输入 x (用户 Tensor)
    ↓ TMA load
smem tma_buffer (warp 本地 staging)
    ↓ TMA store
recv_buffer[dst_rank][dst_slot]  (NVLink 直达目标 rank 的对称内存)
    ↓ gin.put (RDMA)
recv_buffer[...]  (对于 NVLink 不可达的 rank，先写入 send_buffer 再 RDMA)

else {
    // Hybrid dispatch
    const auto scaleup_recv_buffer = layout::BufferLayout<false>(
        token_layout, num_scaleup_ranks, num_scaleout_ranks * num_max_tokens_per_rank);
    const auto scaleout_send_buffer = layout::BufferLayout<false>(
        token_layout, 1, num_max_tokens_per_rank);
    const auto scaleout_recv_buffer = layout::BufferLayout<false>(
        token_layout, num_scaleout_ranks,
        /* kNumChannels * kNumMaxTokensPerChannel */ num_max_tokens_per_rank + kNumMaxChannels);
    return scaleup_recv_buffer.get_num_bytes() +
           scaleout_send_buffer.get_num_bytes() +
           scaleout_recv_buffer.get_num_bytes();
}

GPU buffer = [scaleup_recv_buffer | scaleout_send_buffer | scaleout_recv_buffer]

// All the buffers
auto scaleup_buffer = layout::BufferLayout<false>(
    token_layout, kNumScaleupRanks, kNumScaleoutRanks * kNumMaxTokensPerRank, buffer);
auto scaleout_send_buffer = layout::BufferLayout<false>(
    token_layout, 1, kNumMaxTokensPerRank, scaleup_buffer.get_buffer_end_ptr());
auto scaleout_recv_buffer = layout::BufferLayout<false>(
    token_layout, kNumScaleoutRanks, kNumChannels * kNumMaxTokensPerChannel, scaleout_send_buffer.get_buffer_end_ptr());

输入 x
    ↓ (scaleout warps)
scaleout_send_buffer → gin.put(RDMA) → 远端 scaleout_recv_buffer
    ↓ (forward warps 轮询各 scaleout peer)
scaleout_recv_buffer → TMA store → scaleup_buffer[dst_scaleup_rank][slot]
    ↓ (copy epilogue)
scaleup_buffer → recv_x (用户 Tensor)

struct Buffer {
    // Low-latency mode buffer
    int low_latency_buffer_idx = 0;
    bool low_latency_mode = false;

    // NVLink Buffer
    int64_t num_nvl_bytes;
    void* buffer_ptrs[LEGACY_NUM_MAX_NVL_PEERS] = {nullptr};
    void** buffer_ptrs_gpu = nullptr;

    // NVSHMEM Buffer
    int64_t num_rdma_bytes;
    void* rdma_buffer_ptr = nullptr;

    // Workspace
    void* workspace = nullptr;

    // Host-side MoE info
    volatile int* moe_recv_counter = nullptr;
    int* moe_recv_counter_mapped = nullptr;
    volatile int* moe_recv_expert_counter = nullptr;
    volatile int* moe_recv_rdma_counter = nullptr;
    // ...
};

size_t get_nvl_buffer_size_hint(size_t hidden_bytes, int num_ranks) const {
    const int num_channels = num_sms / 2;
    size_t num_bytes = 0;
    num_bytes += num_channels * num_nvl_ranks * (2 * num_rdma_ranks + 3) * sizeof(int);  // metadata
    num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * hidden_bytes;  // data
    num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * internode::get_source_meta_bytes();
    num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(topk_idx_t);
    num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(float);
    num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxScales * sizeof(float);
    return num_bytes;
}

num_nvl_bytes = Σ_channels Σ_ranks [
    metadata (start/end/head/tail)
    + data[num_max_nvl_chunked_recv_tokens]
    + source_meta[...]
    + topk_idx[...]
    + topk_weights[...]
    + scales[...]
]

// 获取某个 rank 的子 BufferLayout（维度退化为 1 × num_max_tokens_per_rank）
BufferLayout get_rank_buffer(const int& rank_idx) const {
    return BufferLayout(token_layout,
                        1, num_max_tokens_per_rank,
                        static_cast<int8_t*>(base) + get_num_bytes_per_rank() * rank_idx);
}

// 获取某个 token 的 TokenLayout（维度退化为单个 token）
TokenLayout get_token_buffer(const int& token_idx, const bool& global = false) const {
    EP_UNIFIED_ASSERT(num_ranks == 1 or global);
    return TokenLayout(token_layout.num_hidden_bytes, token_layout.num_sf_bytes,
                       token_layout.num_topk, token_layout.with_metadata,
                       static_cast<int8_t*>(base) + token_layout.get_num_bytes<kWithMBarrier, int64_t>() * token_idx);
}

BufferLayout(num_ranks, max_tokens)
    ├── .get_rank_buffer(i) → BufferLayout(1, max_tokens)
    │       └── .get_token_buffer(j) → TokenLayout(single token)
    └── .get_channel_buffer<kTokensPerChannel>(c) → BufferLayout(num_ranks, max_tokens)
            // 注意：channel 不是降维，而是 stride 访问

template <bool kWithMBarrier, typename dtype_t = int>
__forceinline__ __device__ __host__ dtype_t get_num_bytes() const {
    const auto num_bytes = math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes) +
                           math::align(num_sf_bytes, ptx::kNumTMAAlignBytes) +
                           math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes) +
                           math::align<int>(kWithMBarrier ? sizeof(ptx::mbarrier) : 0, ptx::kNumTMAAlignBytes);
    return static_cast<dtype_t>(num_bytes);
}

template <int kNumTokensPerChannel>
__forceinline__ __device__ __host__
BufferLayout get_channel_buffer(const int& channel_idx) const {
    EP_UNIFIED_ASSERT(num_max_tokens_per_rank % kNumTokensPerChannel == 0);
    return BufferLayout(token_layout,
                        num_ranks, num_max_tokens_per_rank,
                        static_cast<int8_t*>(base) + get_num_bytes_per_token() * kNumTokensPerChannel * channel_idx);
}

struct WorkspaceLayout {
    void* workspace;
    int num_ranks;
    int num_scaleout_ranks, num_scaleup_ranks;
    int num_experts, num_experts_per_rank;

    static constexpr int kNumMaxRanks = 1024;
    static constexpr int kNumMaxExperts = 2048;
    static constexpr int kNumMaxExpertsPerRank = 256;
    static constexpr int kNumMaxInflightAGRS = 32;
    static constexpr int64_t kNumBarrierSignalBytes = 16;

    static int64_t get_num_bytes() {
        int64_t num_bytes = 0;
        num_bytes += kNumBarrierSignalBytes;                                    // NVLink barrier
        num_bytes += (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t);        // Notify reduction
        num_bytes += kNumMaxRanks * sizeof(int64_t) * 2;                       // Scaleup rank count
        num_bytes += kNumMaxExperts * sizeof(int64_t) * 2;                     // Scaleup expert count
        num_bytes += kNumMaxRanks * sizeof(int);                                // Scaleup atomic sender
        num_bytes += kNumMaxRanks * sizeof(int) * 2;                           // Scaleout rank count
        num_bytes += kNumMaxExperts * sizeof(int) * 2;                         // Scaleout expert count
        num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int64_t);         // Scaleout channel tails
        num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int);             // Channel scaleup tails
        num_bytes += 2 * 2 * sizeof(int64_t);                                   // PP counts
        num_bytes += (kNumMaxInflightAGRS + 1) * kNumMaxRanks * sizeof(int);    // AGRS signals
        return num_bytes;
    }
    // ...
};

// Allocate host workspaces
CUDA_RUNTIME_CHECK(cudaMallocHost(&host_workspace, layout::WorkspaceLayout::get_num_bytes(), cudaHostAllocMapped));
CUDA_RUNTIME_CHECK(cudaHostGetDevicePointer(&mapped_host_workspace, host_workspace, 0));

Python: buffer.dispatch(x, topk_idx, topk_weights, ...)
    ↓
C++: ElasticBuffer::dispatch (csrc/elastic/buffer.hpp:702)
    ├── 计算 num_sms, num_qps
    ├── 分配 psum_num_recv_tokens_per_scaleup_rank
    ├── 分配 psum_num_recv_tokens_per_expert
    ├── 分配 dst_buffer_slot_idx (non-hybrid) 或 [num_channels][num_scaleout][max_tokens_per_channel][topk] (hybrid)
    ├── 分配 token_metadata_at_forward (hybrid only)
    ├── 分配 channel_linked_list (hybrid only)
    ├── 清理 host_workspace 计数器
    ├── launch_dispatch(...) → dispatch_impl kernel
    │       ├── notify warps: 统计 expert/rank 接收数 → 写 workspace
    │       ├── dispatch warps: TMA load x → TMA store recv_buffer/send_buffer → gin.put (RDMA)
    │       └── cudaTriggerProgrammaticLaunchCompletion (触发 epilogue)
    ├── CPU 等待 host_workspace 计数器 (do_cpu_sync)
    ├── 分配 recv_x, recv_sf, recv_topk_idx, recv_topk_weights, recv_src_metadata
    └── launch_dispatch_copy_epilogue → dispatch_copy_epilogue_impl
            ├── 从 scaleup_buffer TMA load → smem
            ├── expand 模式: atomicAdd psum expert 计数器 → 分配 slot
            └── TMA store → recv_x

Python: buffer.combine(x, handle, topk_weights, ...)
    ↓
C++: ElasticBuffer::combine (csrc/elastic/buffer.hpp:1180)
    ├── launch_combine(...) → combine_impl kernel
    │       ├── 读取 src_metadata 获取 (src_token_idx, src_rank_idx, src_topk_idx)
    │       ├── NVLink bypass: TMA store 到对端 scaleup_buffer
    │       ├── RDMA: TMA store 到 send_buffer → gin.put
    │       └── barrier
    ├── 分配 combined_x, combined_topk_weights
    └── launch_combine_reduce_epilogue → combine_reduce_epilogue_impl
            ├── 从 comm_buffer TMA load 多个 rank 的数据
            ├── combine_reduce: 本地归约 (allow_multiple_reduction)
            ├── 加 bias_0, bias_1
            └── TMA store → combined_x

graph TB
    subgraph "NCCL Symmetric Memory (num_sym_bytes)"
        W["Workspace<br/>(barrier + counters + tails)<br/>2MB aligned"]
        GPU["GPU Buffer<br/>(dispatch/combine 复用)"]
        CPU["CPU Buffer<br/>(Engram storage)"]
    end

    subgraph "GPU Buffer: Direct Mode"
        RBUF["recv_buffer<br/>[num_ranks × max_tokens]<br/>TokenLayout"]
        SBUF["send_buffer<br/>[1 × max_tokens]<br/>RDMA staging"]
    end

    subgraph "GPU Buffer: Hybrid Mode"
        SRBUF["scaleup_recv_buffer<br/>[num_scaleup_ranks × scaleout×max_tokens]"]
        SSOUT["scaleout_send_buffer<br/>[1 × max_tokens]<br/>RDMA staging"]
        SRECV["scaleout_recv_buffer<br/>[num_scaleout_ranks × channels×tokens_per_channel]"]
    end

    subgraph "TokenLayout (per token)"
        HID["hidden_bytes"]
        SF["sf_bytes<br/>(FP8 scale factors)"]
        META["metadata<br/>topk_idx + topk_weights<br/>+ src_token_global_idx<br/>+ linked_list_idx"]
        MBAR["mbarrier<br/>(仅 smem staging)"]
    end

    W --> GPU --> CPU
    GPU --> RBUF & SBUF
    GPU --> SRBUF --> SSOUT --> SRECV
    RBUF --> HID --> SF --> META --> MBAR

flowchart LR
    subgraph "User Space"
        X["x<br/>[num_tokens, hidden]"]
        RX["recv_x<br/>[num_recv, hidden]"]
    end

    subgraph "EP Buffer (Direct)"
        RB["recv_buffer<br/>[rank][slot]"]
        SB["send_buffer<br/>[0][token]"]
    end

    subgraph "Hardware"
        NVL["NVLink<br/>(对称内存直通)"]
        RDMA["RDMA<br/>gin.put"]
    end

    X -->|"TMA load → smem"| RB
    X -->|"TMA load → smem"| SB
    RB -->|"TMA store → gin.get_sym_ptr"| NVL
    SB -->|"gin.put"| RDMA
    NVL -->|"远端 RB"| RB
    RDMA -->|"远端 RB"| RB
    RB -->|"copy epilogue"| RX

flowchart TB
    subgraph "Stage 1: Scaleout Send (RDMA)"
        X["x<br/>输入"]
        SSOUT["scaleout_send_buffer<br/>RDMA staging"]
        X -->|"TMA store"| SSOUT
        SSOUT -->|"gin.put(ncclTeamTagRail)"| NET["RDMA Network"]
    end

    subgraph "Stage 2: Forward (NVLink)"
        SRECV["scaleup_recv_buffer<br/>[scaleout_rank][slot]"]
        SRBUF["scaleup_buffer<br/>[scaleup_rank][slot]"]
        NET -->|"远端写入"| SRECV
        SRECV -->|"TMA store<br/>轮询各 scaleout peer"| SRBUF
    end

    subgraph "Stage 3: Copy Epilogue"
        RX["recv_x"]
        SRBUF -->|"TMA load → expand/sort<br/>TMA store"| RX
    end

用户 x → smem staging → recv_buffer[rank][slot] (NVLink/RDMA) → smem staging → 用户 recv_x
         ↑ TMA load          ↑ TMA store / gin.put        ↑ TMA load     ↑ TMA store