从 DeepEP 到 Mega MoE：MoE Runtime 演进分析

DeepEP: Communication + Data Movement Runtime
   ↓
Mega MoE: Communication + Compute Fusion Runtime
   ↓
DeepEP + DeepGEMM + Fusion Kernel → Unified MoE Dataflow Runtime

timeline
    title MoE Runtime 演进路径
    2024 : DeepEP v1
         : All-to-All 通信库
         : Dispatch/Combine 分离
         : 通信与计算解耦
    2025 : DeepEP v2 (Low-Latency Kernel)
         : 三阶段流水线
         : Warp Specialization
         : NVLink + RDMA 统一
    2025 : DeepGEMM Mega MoE
         : 通信+计算融合
         : 单 Kernel 完成全部 MoE 计算
         : Symmetric Memory (NVLink)
    2026 : Unified MoE Dataflow Runtime (愿景)
         : DeepEP + DeepGEMM + Fusion Kernel
         : 统一数据流执行

graph TB
    subgraph "Mega MoE Kernel (Single CUDA Kernel)"
        W0["Warp 0..N-1<br/>Dispatch Warps<br/>Token 路由 + NVLink Pull"]
        WN["Warp N<br/>TMA Load (Tokens + SFA)"]
        WN1["Warp N+1<br/>TMA Load (Weights + SFB)"]
        WN2["Warp N+2<br/>MMA Issue<br/>UMMA 指令发射"]
        WE["Warp N+3..End<br/>Epilogue Warps<br/>SwiGLU + FP8 Cast<br/>+ Combine Write-back"]
    end

    W0 -->|"Token 就绪通知"| WN
    WN -->|"Smem A 就绪"| WN2
    WN1 -->|"Smem B 就绪"| WN2
    WN2 -->|"TMEM Accumulator"| WE

// 不同 warp 角色划分
if (warp_idx < kNumDispatchWarps) {
    // Dispatch warps: Token 路由 + 跨 Rank 拉取
    ...
} else if (warp_idx == kNumDispatchWarps) {
    // GEMM TMA load warp for tokens with SFA
    ...
} else if (warp_idx == kNumDispatchWarps + 1) {
    // GEMM TMA load warp for weights with SF
    ...
} else if (warp_idx == kNumDispatchWarps + 2) {
    // GEMM MMA issue warp (UMMA 指令发射)
    ...
} else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) {
    // Epilogue warps: SwiGLU + Combine
    ...
}

# deep_gemm/mega/__init__.py
def transform_weights_for_mega_moe(l1_weights, l2_weights):
    # L1: interleave gate/up, then transpose SF for UTCCP
    l1_interleaved = _interleave_l1_weights(l1_weights)
    l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1]))
    # L2: only transpose SF for UTCCP
    l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1]))
    return l1_weights, l2_weights

[Router] → [DeepEP Dispatch] → [Expert Buffer] → [GEMM] → [DeepEP Combine] → [Output]
                ↑                    ↑                           ↑
            NVLink/RDMA          HBM 写入                    NVLink/RDMA

[Router] → [Mega MoE Kernel] → [Output]
                │
                ├─ Dispatch Warps: NVLink Pull → Smem → L1 Buffer (HBM local)
                ├─ MMA Warps: L1 Buffer → UMMA → TMEM
                ├─ Epilogue: TMEM → SwiGLU → FP8 Cast → L2 Buffer (HBM local)
                ├─ MMA Warps: L2 Buffer → UMMA → TMEM
                └─ Epilogue: TMEM → NVLink Write-back → Combine Reduce → Output

// sm100_fp8_fp4_mega_moe.cuh - Dispatch Warps
// 1. 统计每个 Expert 的 Token 数量
read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) {
    atomicAdd_block(smem_expert_count + expert_idx, 1);
});

// 2. 计算全局偏移（跨 Rank 累加）
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
    const uint64_t send_value = (1ull << 32) | static_cast<uint64_t>(smem_expert_count[i]);
    smem_expert_count[i] = static_cast<uint32_t>(
        ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value));
}

// 3. 写入源 Token 索引到 Symmetric Buffer（跨 Rank 可见）
*sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;

// 4. 从远程 Rank 拉取 Token 数据（TMA Load via NVLink）
ptx::tma_load_1d(
    pull_buffer.get_base_ptr(),
    sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(),
                   current_rank_in_expert_idx),
    pull_mbarrier, kHidden);

// sm100_fp8_fp4_mega_moe.cuh - Combine (Epilogue Warps)
// 1. 从 combine_token_buffer 读取各 Top-K 专家结果
const auto src_ptr = math::advance_ptr<uint8_t>(
    combine_token_buffer.get_rank_buffer(slot_idx)
                        .get_data_buffer(token_idx).get_base_ptr(),
    chunk_byte_offset);
ptx::tma_load_1d(combine_load_buffer[i], src_ptr, combine_load_barriers[i], kNumChunkBytes);

// 2. 累加所有 Top-K 贡献
float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {};
while (do_reduce) {
    do_reduce = move_mask_and_load(load_stage_idx ^ 1);
    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]);
    }
}

// 3. 写回最终输出
ptx::tma_store_1d(
    math::advance_ptr(y, static_cast<uint64_t>(token_idx) * kNumHiddenBytes + chunk_byte_offset),
    combine_store_buffer, kNumChunkBytes);

// MMA Issue Warp - Linear1 Phase
if (block_phase == sched::BlockPhase::Linear1) {
    // 等待 Token 数据到达
    const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx);
    while (ptx::ld_acq(ptr) != expected);

    // 发射 UMMA 指令
    ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma(
        b_desc, a_desc, accum_stage_idx * UMMA_N,
        k_block_idx > 0 or k > 0, runtime_instr_desc,
        kTmemStartColOfSFB, kTmemStartColOfSFA);
}

// L1 Epilogue: SwiGLU in-place with weight application
if (block_phase == sched::BlockPhase::Linear1) {
    // 从 TMEM 加载 MMA 结果
    cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, values[0], ..., values[7]);

    // 应用 SwiGLU: silu(gate) * up * weight
    auto gate = __bfloat1622float2(bf16_gate);
    auto neg_gate_exp = make_float2(expf(-gate.x), 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);
    swiglu_values[i * 2 + k] = __fmul2_rn(__fmul2_rn(gate, up), weights);

    // 重量化为 FP8 并写回 Shared Memory
    const auto fp8x4_values = __nv_fp8x4_e4m3(make_float4(upper.x, upper.y, lower.x, lower.y));
    ptx::SM100_U8x4_STSM_T<__nv_fp8x4_e4m3>::copy(fp8x4_values, smem_ptr);
}

// L2 Epilogue: BF16 output write-back via NVLink
} else {  // block_phase == BlockPhase::Linear2
    // 从 TMEM 加载 L2 MMA 结果
    cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, values[0], ..., values[7]);

    // 写入 Shared Memory (BF16)
    ptx::SM90_U32x4_STSM_T<uint32_t>::copy(
        math::cast_into_bf16_and_pack(values[0], values[1]), ...);

    // 通过 NVLink 写入远程 Rank 的 Combine Buffer
    const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx)
                           .get_data_buffer(dst_token_idx);
    *sym_buffer.map(dst_ptr, dst_rank_idx) = packed;
}

graph LR
    subgraph "DeepEP (分离)"
        D1["Dispatch"] -->|"HBM Write"| B1["Expert Buffer"]
        B1 -->|"HBM Read"| G1["GEMM"]
        G1 -->|"HBM Write"| B2["Output Buffer"]
        B2 -->|"HBM Read"| C1["Combine"]
    end

    subgraph "Mega MoE (融合)"
        D2["Dispatch"] -->|"Smem/TMEM"| G2["Linear1"]
        G2 -->|"Register"| S2["SwiGLU"]
        S2 -->|"Smem/TMEM"| G3["Linear2"]
        G3 -->|"Smem/NVLink"| C2["Combine"]
    end

# tests/test_mega_moe.py
# TFLOPS: 3 matmuls (L1 left, L1 right, L2), each 2 * M * N * K
tflops = safe_div(2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_fused)

# HBM bytes: weights (FP4 packed = 0.5 bytes) + activations (FP8 = 1 byte) + output (BF16 = 2 bytes)
num_hbm_bytes = (
    num_experts_per_rank * intermediate_hidden * 2 * hidden // 2 +  # L1 weights (FP4)
    num_experts_per_rank * hidden * intermediate_hidden // 2 +      # L2 weights (FP4)
    num_recv_tokens * hidden +                                      # L1 acts read (FP8)
    num_recv_tokens * intermediate_hidden +                         # L1 output write (FP8)
    num_recv_tokens * intermediate_hidden +                         # L2 acts read (FP8)
    num_recv_tokens * hidden * 2                                    # L2 output write (BF16)
)

# NVLink bytes: dispatch pull + combine write-back
num_nvlink_bytes = num_recv_tokens * hidden * 3

# tests/test_mega_moe.py
# Non-overlapped baseline: EP dispatch + GEMM + EP combine
def run_baseline():
    recv_x, _, recv_topk_weights, handle, _ = ep_buffer.dispatch(...)
    deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(recv_x, l1_weights, l1_y, ...)
    l1_y = tilelang_ops.swiglu_apply_weight_to_fp8(...)
    deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(l1_y, l2_weights, l2_y, ...)
    return ep_buffer.combine(l2_y, handle=handle)[0]

# Run fused mega MoE
def run_fused():
    deep_gemm.fp8_fp4_mega_moe(y, transformed_l1_weights, transformed_l2_weights, buffer, ...)
    return y

Phase 1 (当前): Mega MoE
├── 节点内 NVLink 通信 + 计算融合
├── 单 Kernel 完成 MoE 全流水线
└── 依赖 Symmetric Memory (PyTorch)

Phase 2 (愿景): Unified MoE Dataflow Runtime
├── 集成 DeepEP 的跨节点 RDMA 能力
├── 统一调度器（通信 + 计算 + 内存）
├── 自适应选择通信路径（NVLink vs RDMA）
└── 全局最优的 Dataflow 执行