DeepEP 整体架构分析

flowchart LR
    subgraph 目标["DeepEP 设计目标"]
        A["高吞吐 all-to-all<br/>逼近硬件带宽极限"]
        B["低延迟<br/>RDMA 直连"]
        C["极低 SM 占用<br/>4-64 SMs for comm"]
        D["FP8/BF16 混合精度"]
    end

    subgraph 场景["应用场景"]
        E["MoE Training"]
        F["MoE Inference Prefill"]
        G["MoE Inference Decode"]
        H["实验: PP/CP/Engram"]
    end

    A --> E
    B --> G
    C --> E
    D --> F

flowchart TB
    subgraph L1["Layer 1: Python API"]
        P1["ElasticBuffer / Buffer"]
        P2["EPHandle / EventOverlap"]
        P3["Utils: envs, event, math, comm"]
    end

    subgraph L2["Layer 2: C++ Binding (pybind11)"]
        B1["python_api.cpp"]
        B2["jit::register_apis"]
        B3["elastic::register_apis"]
        B4["legacy::register_apis"]
    end

    subgraph L3["Layer 3: JIT Compilation"]
        J1["Compiler / NVCCCompiler"]
        J2["KernelRuntimeCache"]
        J3["KernelRuntime (cubin load)"]
        J4["IncludeParser / DeviceRuntime"]
    end

    subgraph L4["Layer 4: CUDA Kernels + Backend"]
        K1["elastic/dispatch.hpp · combine.hpp"]
        K2["elastic/engram.hpp · pp_send_recv.hpp"]
        K3["legacy/internode.cu · intranode.cu"]
        K4["NCCL Gin Backend / NVSHMEM Backend"]
    end

    L1 --> L2 --> L3 --> L4

# __init__.py 核心初始化流程
check_nccl_so()          # 1. 确保 NCCL 版本一致
init_jit()                # 2. 初始化 JIT 编译器路径
from .buffers.legacy import Buffer          # V1
from .buffers.elastic import ElasticBuffer, EPHandle  # V2
from .utils.event import EventOverlap

flowchart TB
    subgraph V1["V1: Legacy (NVSHMEM-based)"]
        V1P["Buffer (Python)"]
        V1C["legacy::Buffer (C++)"]
        V1K["legacy/internode.cu<br/>legacy/intranode.cu"]
        V1B["NVSHMEM Backend<br/>+ IBGDA"]
        V1P --> V1C --> V1K --> V1B
    end

    subgraph V2["V2: Elastic (NCCL Gin-based)"]
        V2P["ElasticBuffer (Python)"]
        V2C["elastic::ElasticBuffer (C++)"]
        V2K["elastic/dispatch.hpp<br/>elastic/combine.hpp"]
        V2B["NCCL Gin Backend<br/>+ Symmetric Memory"]
        V2P --> V2C --> V2K --> V2B
    end

// V1 的 Config 查表是硬编码的
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),
    // ... 一直到 160
}

classDiagram
    class ElasticBuffer {
        +group: ProcessGroup
        +rank_idx, num_ranks: int
        +num_bytes, num_cpu_buffer_bytes: int64
        +allow_hybrid_mode: bool
        +allow_multiple_reduction: bool
        +prefer_overlap_with_compute: bool
        +num_allocated_qps: int
        +runtime: C++ ElasticBuffer
        +nccl_comm_handle: NCCLCommHandle
        +num_scaleout/num_scaleup_ranks: int
        +num_rdma/num_nvl_ranks: int

        +__init__(group, num_max_tokens_per_rank, hidden, ...)
        +destroy()
        +dispatch(x, topk_idx, ...) → (recv_x, handle, event)
        +combine(x, handle, ...) → (combined_x, event)
        +barrier()
        +engram_write(storage) [实验]
        +engram_fetch(indices) → hook [实验]
        +pp_send/pp_recv [实验]
        +all_gather [实验]
        +get_theoretical_num_sms() → int
        +get_buffer_size_hint() → int
    }

    class NCCLSymmetricMemoryContext {
        +rank_idx, num_ranks: int
        +num_scaleout/num_scaleup_ranks: int
        +num_rdma/num_nvl_ranks: int
        +is_scaleup_nvlink: bool
        +comm: ncclComm_t
        +dev_comm: ncclDevComm_t
        +window: ncclWindow_t
        +mapped_window_ptr: void*
        +num_allocated_qps: int
    }

    ElasticBuffer --> NCCLSymmetricMemoryContext : nccl_context
    ElasticBuffer --> EPHandle : dispatch returns
    ElasticBuffer --> EventOverlap : returns

[Workspace (2MB-aligned)] [GPU Buffer] [CPU Buffer (可选)]
         ↑ mapped_window_ptr

class EPHandle:
    do_expand: bool                    # 是否使用 expand 布局
    num_experts: int
    expert_alignment: int              # expert token 对齐
    num_max_tokens_per_rank: int
    num_sms: int
    topk_idx: Tensor                   # [num_tokens, num_topk]，克隆的 expert 索引
    psum_num_recv_tokens_per_scaleup_rank: Tensor  # scaleup rank 前缀和
    psum_num_recv_tokens_per_expert: Tensor        # expert 前缀和
    num_unaligned_recv_tokens_perpert: Tensor      # expand 模式下未对齐计数
    recv_src_metadata: Tensor          # [num_recv_tokens, topk+2]，源 token 索引 + slot
    dst_buffer_slot_idx: Tensor        # 目标 buffer slot
    token_metadata_at_forward          # hybrid 模式元数据
    channel_linked_list                # hybrid 模式链表
    num_recv_tokens: int
    num_expanded_tokens: int

# 典型用法：先做不依赖 recv_x 的计算，再等待通信完成
recv_x, ..., handle, event = buffer.dispatch(...)

with event:                    # __exit__ 时 current_stream_wait()
    do_other_compute()         # 与通信重叠执行

# 退出 with 后，recv_x 安全可用
recv_x = forward_experts(recv_x, ...)

csrc/kernels/
├── elastic/                    # V2 kernels
│   ├── api.hpp                 # kernel 总入口 (launch 函数)
│   ├── dispatch.hpp            # dispatch + dispatch_copy_epilogue
│   ├── combine.hpp             # combine + combine_reduce_epilogue
│   ├── engram.hpp              # Engram fetch (实验)
│   ├── pp_send_recv.hpp        # PP send/recv (实验)
│   ├── barrier.hpp             # GPU barrier
│   └── utils.hpp               # kernel utilities
├── legacy/                     # V1 kernels
│   ├── api.cuh                 # V1 kernel 入口
│   ├── internode.cu            # 节点间 (RDMA + NVLink)
│   ├── intranode.cu            # 节点内 (NVLink)
│   ├── buffer.cuh              # V1 buffer 管理
│   ├── compiled.cuh            # V1 编译配置
│   ├── ibgda_device.cuh        # IBGDA low-latency
│   ├── launch.cuh              # V1 launch 逻辑
│   ├── layout.cu               # V1 layout 计算
│   └── utils.cuh               # V1 utilities
└── backend/                    # 后端抽象
    ├── api.cuh                 # nvshmem / nccl / cuda_driver API
    ├── symmetric.hpp           # SymmetricMemory 分配器
    ├── cuda_driver.cu          # CUDA Driver API 封装
    ├── nccl.cu                 # NCCL 初始化/销毁
    └── nvshmem.cu              # NVSHMEM 初始化/销毁

flowchart LR
    subgraph Dispatch["dispatch (2-phase)"]
        D1["Phase 1: dispatch_impl<br/>notify warps + dispatch warps<br/>写对称内存 + 计数"]
        D2["Phase 2: dispatch_copy_epilogue<br/>从对称内存拷贝到用户 tensor<br/>full SMs"]
        D1 -->|"named barrier"| D2
    end

    subgraph Combine["combine (2-phase)"]
        C1["Phase 1: combine_impl<br/>从本地 push 数据到远端"]
        C2["Phase 2: combine_reduce_epilogue<br/>reduce + bias + 写回"]
        C1 -->|"named barrier"| C2
    end

// comm::get_qp_mode: 决定每个 SM 的每个 channel 使用哪个 QP
if (kNumSMs <= kNumAvailableQPs) {
    // SM 少：每个 SM 独占 QP，channel 间 round-robin
    // e.g., 3 SMs, 10 QPs: SM0→{0,3,6,9}, SM1→{1,4,7}, SM2→{2,5,8}
} else {
    // SM 多：所有 SM 共享所有 QP
    qp_idx = kQPStartIdx + (global_channel_idx % kNumAvailableQPs)
}

flowchart TB
    subgraph HybridDispatch["Hybrid Dispatch 数据流"]
        T0["Token on Rank 0"]
        T1["Notify Warps: 计算目标 slot"]
        T2["Scaleout Warps: RDMA 写到<br/>远端 scaleout rank 的 buffer"]
        T3["Forward Warps: NVLink 转发<br/>到目标 scaleup rank"]
        T4["最终到达目标 expert 所在 rank"]
    end
    T0 --> T1 --> T2 --> T3 --> T4

flowchart TB
    subgraph Init["初始化阶段"]
        I1["__init__.py: init_jit()"]
        I2["jit::init(library_root, cuda_home, nccl_root)"]
        I3["Compiler::prepare_init()"]
        I4["IncludeParser::prepare_init()"]
        I5["DeviceRuntime::prepare_init()"]
        I1 --> I2 --> I3 & I4 & I5
    end

    subgraph Build["内核构建阶段"]
        B1["generate_impl(args) → 生成 kernel.cu 源码"]
        B2["计算 kernel_signature = name$$sig$$flags$$code"]
        B3["计算 cache 路径: ~/.deep_ep/cache/kernel.{name}.{hash}/"]
        B4{"KernelRuntimeCache 命中?"}
        B5["是: 返回缓存的 KernelRuntime"]
        B6["否: NVCCCompiler::compile() → kernel.cubin"]
        B7["cuobjdump 反汇编 → kernel.sass (可选)"]
        B8["原子 rename tmp_dir → cache_dir"]
        B9["KernelRuntime 加载 cubin"]
        B1 --> B2 --> B3 --> B4
        B4 -->|Yes| B5
        B4 -->|No| B6 --> B7 --> B8 --> B9
    end

    subgraph Launch["内核启动阶段"]
        L1["construct_launch_config()"]
        L2["launch_kernel(kernel, config, args...)"]
        L1 --> L2
    end

    Init --> Build --> Launch

// dispatch.hpp 生成的 kernel.cu 示例
#include <deep_ep/impls/dispatch.cuh>
using namespace deep_ep::elastic;
static void __instantiate_kernel() {
    auto ptr = reinterpret_cast<void*>(&dispatch_impl<
        true,    // is_scaleup_nvlink
        false,   // do_cpu_sync
        false,   // reuse_slot_indices
        12,      // num_sms
        4,       // num_notify_warps
        28,      // num_dispatch_warps
        16,      // num_scaleup_ranks
        14336,   // num_hidden_bytes (7168*2)
        0,       // num_sf_packs
        2048,    // num_max_tokens_per_rank
        256,     // num_experts
        8,       // num_topk
        128,     // expert_alignment
        65,      // num_qps
        cycles   // num_timeout_cycles
    >);
}

kernel_signature = "{name}$${signature}$${flags}$${code}"

~/.deep_ep/
├── cache/
│   └── kernel.dispatch.{hex_digest}/
│       ├── kernel.cu        # 生成的源码
│       ├── kernel.cubin     # 编译产物
│       └── kernel.sass      # 可选的反汇编
└── tmp/                     # 临时目录（原子 rename 用）

flowchart LR
    subgraph Backends["DeepEP 后端"]
        direction TB
        N1["NCCL Gin Backend<br/>(V2 主力)"]
        N2["NVSHMEM Backend<br/>(V1 专用)"]
        N3["CUDA Driver API<br/>(辅助)"]
    end

    N1 --- N1D["ncclMemAlloc / ncclCommWindowRegister<br/>ncclGin (QP 抽象)<br/>Header-only, 复用 NCCL communicator"]
    N2 --- N2D["nvshmem_init / ibgda<br/>NVSHMEM team / QP<br/>独立初始化"]
    N3 --- N3D["cuMemCreate / cuMemMap<br/>cuMemcpyBatchAsync<br/>batched_write_and_wait"]

// 1. 获取 Gin handle（per-warp）
const auto [qp_idx, sharing_mode] = comm::get_qp_mode<...>(sm_idx, channel_idx, is_notify);
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode);

// 2. 发起 RDMA 请求
ncclGinRequest_t req;
gin.get(dst_ptr, src_ptr, size, dst_rank, req);  // RDMA get

// 3. 等待完成
gin.wait(req);

# 核心公式（简化）
# 假设 balanced gate，计算 expected top-k 数量
num_expected_topk = num_groups * (1 - C(num_experts - num_experts//num_groups, num_topk) / C(num_experts, num_topk))

# 计算各层流量
sm_read += 1 / num_expected_topk
sm_write += ...
rdma_traffic += ...
nvlink_traffic += ...

# SM 数 = ceil(max(sm_read/sm_read_bw, sm_write/sm_write_bw, rdma_traffic/rdma_bw, nvlink_traffic/nvlink_bw))

// dispatch.hpp: 根据共享内存消耗计算每个 SM 的 channel 数
num_channels_per_sm = min(
    (smem_bytes - notify_smem_bytes) / dispatch_token_layout_size,
    32 - kNumNotifyWarps
);
num_channels_per_sm = min(
    smem_bytes / combine_token_layout_size,
    num_channels_per_sm
);
num_channels_per_sm = min(num_channels_per_sm / 2, kNumMaxChannelsPerSM);

flowchart LR
    E1["engram_write(storage)<br/>写入 CPU buffer"]
    E2["engram_fetch(indices)<br/>发起 RDMA get"]
    E3["返回 hook 闭包<br/>延迟等待"]
    E4["hook() → 阻塞等待<br/>RDMA 完成"]
    E5["返回 (data, sf)"]
    E1 --> E2 --> E3 --> E4 --> E5

buffer.pp_set_config(num_max_tensor_bytes, num_max_inflight_tensors)
buffer.pp_send(t, dst_rank_idx, num_sms=0)  # 0 = 全部 SM
buffer.pp_recv(t, src_rank_idx, num_sms=0)

buffer.agrs_set_config(num_max_session_bytes, num_max_all_gathers)
buffer.create_agrs_session()
t = buffer.agrs_get_inplace_tensor(shape, dtype)  # zero-copy
gathered, handle = buffer.all_gather(t)
handle()  # wait for data arrival
buffer.destroy_agrs_session()

flowchart TB
    subgraph DispatchFlow["Dispatch 数据流"]
        direction TB
        D0["输入: x[num_tokens, hidden], topk_idx, topk_weights"]
        D1["notify warps: 解析 topk_idx<br/>原子累加 rank/expert 计数<br/>写入 host_workspace (mapped)"]
        D2["dispatch warps: 打包 token data + sf<br/>通过 Gin QP 发起 RDMA/NVLink get"]
        D3["数据到达远端对称内存 buffer"]
        D4["CPU sync 或 cached mode<br/>确定 num_recv_tokens"]
        D5["dispatch_copy_epilogue: 从 buffer 拷贝<br/>到 recv_x, recv_topk_idx, recv_topk_weights"]
        D6["输出: recv_x, EPHandle, EventOverlap"]
    end

    subgraph CombineFlow["Combine 数据流"]
        direction TB
        C0["输入: expert 计算结果 + EPHandle"]
        C1["combine warps: 读取 handle 中的路由元数据<br/>通过 Gin push 数据到远端 buffer"]
        C2["combine_reduce_epilogue: 从 buffer reduce<br/>按 topk_weights 加权 + bias"]
        C3["输出: combined_x, EventOverlap"]
    end

    D6 -->|"EPHandle 包含路由元数据"| C0
    DispatchFlow --> CombineFlow