DeepEP → Mega MoE：通信与计算的边界演进分析

# 核心通信 API
from .buffers.legacy import Buffer                    # 传统 Buffer
from .buffers.elastic import ElasticBuffer, EPHandle  # 弹性 Buffer（v2）

# 工具 API
from .utils.event import EventOverlap, EventHandle     # 事件/异步
from .utils.envs import get_physical_domain_size, get_logical_domain_size

# C++ 扩展
from deep_ep._C import Config, topk_idx_t

// dispatch.cuh - Warp 角色划分
if (warp_idx < kNumNotifyWarps) {
    // ===== Notify Warps：元数据计算（非 GEMM） =====
    // 1. 统计每个 Expert 接收多少 Token
    const auto dst_expert_idx = __ldg(topk_idx + i * kNumTopk + lane_idx);
    atomicAdd_block(expert_count + dst_expert_idx, 1);
    atomicAdd_block(rank_count + dst_rank_idx, 1);

    // 2. 跨 Rank 交换 count 信息（NVLink/RDMA put）
    gin.put_value<team_t>(dst_rank_counter, rank_count[i], i);

    // 3. Prefix Sum 计算偏移
    do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, kNumRanks, 0);
} else {
    // ===== Dispatch Warps：纯数据搬运 =====
    // 1. TMA Load Token 数据到 Smem
    ptx::tma_load_1d(tma_buffer.get_hidden_ptr(), x + token_i64_idx * kNumHiddenBytes, ...);

    // 2. TMA Store 到 Send Buffer
    ptx::tma_store_1d(send_buffer_ptr, tma_buffer.get_base_ptr(), ...);

    // 3. NVLink Store 到远程 Rank
    ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), ...);

    // 4. RDMA Put 到远程 Rank
    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);
}

// combine.cuh - 三种模式
if (no_local_reduce) {
    // 模式 1：无本地归约，直接发送
    ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, ...);
    ptx::tma_store_1d(master_token_buffer.get_base_ptr(), tma_buffer.get_base_ptr(), ...);
} else if constexpr (kAllowMultipleReduction) {
    // 模式 2：本地 Top-K 归约（累加）
    combine_reduce<kHiddenVec, kUnrollFactor>(..., tma_buffer.get_base_ptr(),
        /* Get source */ [=](const int& slot_idx) { return x + slot_idx * kNumHiddenBytes; },
        /* Wait */ [=]() { ptx::tma_store_wait(); });
} else {
    // 模式 3：扩展发送（每个 Top-K 独立发送）
    for (int k = 0; k < kNumTopk; ++ k) {
        ptx::tma_load_1d(tma_buffer.get_base_ptr(), src_token_ptr, ...);
        ptx::tma_store_1d(..., tma_buffer.get_base_ptr(), ...);
        gin.put<team_t>(token_buffer.get_base_ptr(), send_buffer_ptr, ...);
    }
}

// combine_utils.cuh - Top-K 归约（bf16 累加）
float2 reduced[kUnrollFactor * kNumElemsPerVec / 2] = {};
for (int k = 0; k < num_topk; ++ k) {
    const auto bf162_view = reinterpret_cast<nv_bfloat162*>(values);
    for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++ j)
        ptx::accumulate(reduced[j], bf162_view[j]);  // bf16 → float 累加
}
// 转回 bf16 输出
bf162_view[l] = __float22bfloat162_rn(reduced[j]);

┌─────────────────────────────────────────────────────────┐
│                   DeepEP 的能力边界                      │
├─────────────────────────────────────────────────────────┤
│ ✅ 通信：NVLink (intra-node) + RDMA (inter-node)        │
│ ✅ 数据搬运：TMA Load/Store, cp.async                   │
│ ✅ 元数据计算：Expert 计数, Prefix Sum, Slot 分配        │
│ ✅ Top-K 归约：bf16 累加（Combine 语义恢复）             │
│ ❌ GEMM：无任何矩阵乘法                                  │
│ ❌ Activation：无任何激活函数（SwiGLU, ReLU, etc.）      │
│ ❌ 量化：无 FP8/FP4 重量化                              │
└─────────────────────────────────────────────────────────┘

// sm100_fp8_fp4_mega_moe.cuh - Warp 角色划分
if (warp_idx < kNumDispatchWarps) {
    // ===== Dispatch Warps：Token 路由 + NVLink Pull =====
    atomicAdd_block(smem_expert_count + expert_idx, 1);  // 计数
    *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;  // 元数据
    ptx::tma_load_1d(pull_buffer, sym_buffer.map(input_token, src_rank), ...);  // 拉取
} else if (warp_idx == kNumDispatchWarps) {
    // ===== TMA Load Warp：Token + Scale Factor =====
    tma::copy<BLOCK_K, LOAD_BLOCK_M>(tensor_map_a_ptr, ..., smem_a[stage_idx], ...);
} else if (warp_idx == kNumDispatchWarps + 1) {
    // ===== TMA Load Warp：Weight + Scale Factor =====
    tma::copy<BLOCK_K, LOAD_BLOCK_N>(tensor_map_b_ptr, ..., smem_b[stage_idx], ...);
} else if (warp_idx == kNumDispatchWarps + 2) {
    // ===== MMA Issue Warp：UMMA 指令发射 =====
    ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma(b_desc, a_desc, ..., runtime_instr_desc, ...);
} else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) {
    // ===== Epilogue Warps：SwiGLU + FP8 Cast + Combine =====
    // SwiGLU
    auto gate = __bfloat1622float2(bf16_gate);
    gate = __fmul2_rn(gate, {math::fast_rcp(denom.x), math::fast_rcp(denom.y)});
    const auto up = __bfloat1622float2(bf16_up);
    activation_values[i][k] = __fmul2_rn(__fmul2_rn(gate, up), weights);
    // FP8 重量化
    const auto fp8x4_values = __nv_fp8x4_e4m3(make_float4(upper.x, upper.y, lower.x, lower.y));
    // Combine 写回
    *sym_buffer.map(dst_ptr, dst_rank_idx) = packed;
}

// sm100_fp8_fp4_mega_moe.cuh - L1 Epilogue (SwiGLU 后)
// 输出直接留在 Smem，作为 L2 的输入
alignas(kSharedMemoryAlignment) cutlass::float_e4m3_t l1[...];  // L1 输出
alignas(kSharedMemoryAlignment) nv_bfloat16 l2[...];            // L2 输出

DeepEP 模式（分离式）：
┌──────────┐    ┌──────────┐    ┌──────────┐    ┌──────────┐
│ Dispatch │───▶│  Expert  │───▶│   GEMM   │───▶│ Combine  │
│ (NVLink/ │    │  Buffer  │    │(外部库)  │    │(NVLink/  │
│  RDMA)   │    │  (HBM)   │    │          │    │  RDMA)   │
└──────────┘    └──────────┘    └──────────┘    └──────────┘
   Kernel 1       HBM Write      Kernel 2        Kernel 3

Mega MoE 模式（融合式）：
┌──────────────────────────────────────────────────────────┐
│                   Mega MoE Kernel                         │
│  ┌────────┐   ┌────────┐   ┌────────┐   ┌────────┐     │
│  │Dispatch│──▶│Linear1 │──▶│SwiGLU  │──▶│Linear2 │──▶  │
│  │(NVLink)│   │(UMMA)  │   │(Reg)   │   │(UMMA)  │     │
│  └────────┘   └────────┘   └────────┘   └────────┘     │
│       Smem/TMEM 流水线（不经过 HBM）                      │
└──────────────────────────────────────────────────────────┘

Unified MoE Dataflow Runtime 需要：
┌─────────────────────────────────────────────────────────────┐
│  1. 跨节点支持：RDMA + NVLink 统一调度                       │
│     - Mega MoE 目前仅支持节点内 NVLink                       │
│     - DeepEP 的 RDMA 能力未集成                              │
│                                                             │
│  2. 自适应通信路径选择                                       │
│     - 节点内：Symmetric Memory (低延迟)                      │
│     - 节点间：RDMA (高带宽)                                  │
│     - 当前两者是独立系统                                     │
│                                                             │
│  3. 全局任务调度器                                           │
│     - 通信任务 vs 计算任务的动态平衡                         │
│     - 内存带宽 vs SM 计算的 Roofline 感知                    │
│                                                             │
│  4. Decode 优化                                              │
│     - DeepEP Low-Latency Kernel 的 Decode 优化               │
│     - Mega MoE 目前仅针对 Training/Prefill                   │
└─────────────────────────────────────────────────────────────┘

timeline
    title MoE Runtime 演进路径（愿景 vs 现实）
    2024 : DeepEP v1
         : All-to-All 通信库
         : Dispatch/Combine 分离
         : 通信与计算解耦
    2025 : DeepEP v2 (Elastic Buffer)
         : 三阶段流水线
         : Warp Specialization
         : NVLink + RDMA 统一
    2025 : DeepGEMM Mega MoE (当前)
         : 节点内通信+计算融合
         : 单 Kernel 完成 MoE 全流水线
         : Symmetric Memory (NVLink only)
    2026 : Unified MoE Dataflow Runtime (愿景)
         : DeepEP + DeepGEMM 统一
         : NVLink + RDMA 自适应
         : 全局最优 Dataflow 执行

// sm100_fp8_fp4_mega_moe.cuh - 编译期硬件检查
#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__)
    // Kernel 主体
#else
    if (blockIdx.x == 0 and threadIdx.x == 0)
        DG_DEVICE_ASSERT(false and "This kernel only support sm_100f");
#endif

# deep_gemm/mega/__init__.py - Symmetric Memory 依赖
import torch.distributed._symmetric_memory as symm_mem

# Mega MoE 必须使用 Symmetric Memory
allocator = torch if group.size() == 1 else symm_mem
self.buffer = allocator.empty(num_bytes, dtype=torch.int8, device='cuda')
self.handle = symm_mem.rendezvous(self.buffer, group=group)

// dispatch.cuh - DeepEP 使用 NCCL 抽象
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, sharing_mode);
gin.put<team_t>(dst_ptr, src_ptr, size, dst_rank_idx);  // 通用 API

# tests/test_mega_moe.py - DeepEP + DeepGEMM 作为 Baseline
def run_baseline():
    # 1. DeepEP Dispatch
    recv_x, _, recv_topk_weights, handle, _ = ep_buffer.dispatch(...)
    # 2. DeepGEMM GEMM
    deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(recv_x, l1_weights, l1_y, ...)
    # 3. 外部 SwiGLU
    l1_y = tilelang_ops.swiglu_apply_weight_to_fp8(...)
    # 4. DeepGEMM GEMM
    deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(l1_y, l2_weights, l2_y, ...)
    # 5. DeepEP Combine
    return ep_buffer.combine(l2_y, handle=handle)[0]

# Mega MoE 融合版本
def run_fused():
    deep_gemm.fp8_fp4_mega_moe(y, transformed_l1_weights, transformed_l2_weights, buffer, ...)
    return y

graph TB
    subgraph "DeepEP 能力域"
        D1["Intra-node NVLink"]
        D2["Inter-node RDMA"]
        D3["Low-Latency Decode"]
        D4["Pipeline Parallel"]
    end

    subgraph "Mega MoE 能力域"
        M1["Intra-node NVLink"]
        M2["FP8×FP4 GEMM"]
        M3["SwiGLU Activation"]
        M4["FP8 Rescale"]
    end

    subgraph "Unified Runtime 愿景"
        U1["NVLink + RDMA 自适应"]
        U2["通信 + 计算统一调度"]
        U3["Training + Decode 统一"]
    end

    D1 ---|"共享"| M1
    D2 -."->|"缺失"| M1
    M2 -."->|"缺失"| D1

    U1 --> D1
    U1 --> D2
    U1 --> M1
    U2 --> M2
    U2 --> M3
    U3 --> D3
    U3 --> M4

// dispatch.cuh - 完整 Kernel 中的"计算"操作
// 全部是整数运算（计数、索引、原子操作）

// 1. Expert 计数（整数 atomicAdd）
atomicAdd_block(expert_count + dst_expert_idx, 1);

// 2. Rank 计数（整数 atomicAdd + deduplicate）
if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0)
    atomicAdd_block(rank_count + dst_rank_idx, 1);

// 3. Prefix Sum（整数加法）
int psum = 0;
for (int i = 0; i < ceil_div(n + is_exclusive, 32); ++ i) {
    const auto value = (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0;
    const auto sum = psum + ptx::warp_inclusive_sum(value, lane_idx);
    out[idx] = sum;
}

// 4. Slot 分配（整数 atomicAdd）
stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_rank_idx, 1);

// 5. 数据搬运（TMA，无计算）
ptx::tma_load_1d(tma_buffer.get_hidden_ptr(), x + token_i64_idx * kNumHiddenBytes, ...);
ptx::tma_store_1d(send_buffer_ptr, tma_buffer.get_base_ptr(), ...);
gin.put<team_t>(recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(), ...);

// combine_utils.cuh - 唯一的"计算"：bf16 累加
float2 reduced[kUnrollFactor * kNumElemsPerVec / 2] = {};
for (int k = 0; k < num_topk; ++ k) {
    const auto bf162_view = reinterpret_cast<nv_bfloat162*>(values);
    for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++ j)
        ptx::accumulate(reduced[j], bf162_view[j]);  // 唯一浮点运算
}
// 转回 bf16
bf162_view[l] = __float22bfloat162_rn(reduced[j]);

// sm100_fp8_fp4_mega_moe.cuh - 完整计算流水线

// ===== 1. GEMM (UMMA 指令) =====
ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma(
    b_desc, a_desc, accum_stage_idx * UMMA_N,
    k_block_idx > 0 or umma_k_block_idx > 0 or k > 0, runtime_instr_desc,
    kTmemStartColOfSFB, kTmemStartColOfSFA);

// ===== 2. SwiGLU Activation =====
auto gate = __bfloat1622float2(bf16_gate);
auto neg_gate_exp = make_float2(
    kFastMath ? __expf(-gate.x) : expf(-gate.x),
    kFastMath ? __expf(-gate.y) : expf(-gate.y));
const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_gate_exp);
gate = __fmul2_rn(gate, {math::fast_rcp(denom.x), math::fast_rcp(denom.y)});
const auto up = __bfloat1622float2(bf16_up);
activation_values[i][k] = __fmul2_rn(__fmul2_rn(gate, up), weights);

// ===== 3. Amax Reduction =====
thread_local_amax.x = cute::max(thread_local_amax.x, cute::abs(activation_values[i][k].x));
amax_values[i].x = math::warp_reduce<4, true>(thread_local_amax.x, math::ReduceMax<float>());

// ===== 4. FP8 重量化 =====
const auto fp8x4_values = __nv_fp8x4_e4m3(make_float4(upper.x, upper.y, lower.x, lower.y));

// ===== 5. Combine Top-K 归约 =====
float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {};
while (do_reduce) {
    combine_load_barriers[load_stage_idx]->wait(combine_phase);
    for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) {
        const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx];
        const auto bf16_values = reinterpret_cast<const nv_bfloat162*>(&uint4_values);
        for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l)
            ptx::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]);
    }
}

timeline
    title MoE Runtime 演进：从分离到融合
    section 2024
        DeepEP v1 : All-to-All 通信库
                  : Dispatch/Combine 分离
                  : 通信与计算解耦
    section 2025
        DeepEP v2 : Elastic Buffer
                  : Warp Specialization
                  : NVLink + RDMA 统一
                  : Low-Latency Decode Kernel
    section 2025
        DeepGEMM Mega MoE : 节点内通信+计算融合
                          : 单 Kernel 全流水线
                          : Symmetric Memory + UMMA + TMEM
    section 2026 (愿景)
        Unified Runtime : DeepEP + DeepGEMM 统一
                         : NVLink + RDMA 自适应
                         : 全局最优 Dataflow

DeepEP 时代：通信与计算分离
┌─────────┐   ┌─────────┐   ┌─────────┐
│ 通信 Kernel │ → │ 计算 Kernel │ → │ 通信 Kernel │
│ (HBM Write) │   │ (HBM Read)  │   │ (HBM Read)  │
└─────────┘   └─────────┘   └─────────┘

Mega MoE 时代：通信与计算融合
┌─────────────────────────────────────────┐
│          单 Kernel 内部流水线             │
│  Dispatch → GEMM → SwiGLU → GEMM → Combine │
│  (Smem/TMEM 直连，不经过 HBM)            │
└─────────────────────────────────────────┘

Unified Runtime 愿景：全局最优调度
┌─────────────────────────────────────────────────────────┐
│  自适应通信路径 (NVLink vs RDMA) + 全局任务调度器        │
│  Training/Decode 统一 + 跨节点/节点内统一               │
└─────────────────────────────────────────────────────────┘