Warp Specialization: DeepEP vs Mega MoE 深度对比分析

Warp Group A: IB Sending        (GPU → NIC)
Warp Group B: IB-NVLink Forward (NIC → NVLink 中继)
Warp Group C: NVLink Receive    (NVLink → GPU Memory)

// 148-176 行
const int num_dispatch_threads = 128;      // 4 warps
const int num_non_epilogue_threads = 128;  // 4 warps (MMA issue + TMA load)
const int num_epilogue_threads = 256;      // 8 warps
// 总计: 512 threads = 16 warps per CTA

// 356-877 行: 主 kernel 中的 warp 分派
if (warp_idx < kNumDispatchWarps) {
    // Warp 0-3: Dispatch warps (48 regs)
    cutlass::arch::warpgroup_reg_dealloc<kNumDispatchRegisters>();
    // ... token routing, expert counting, NVLink pull ...
} else if (warp_idx == kNumDispatchWarps) {
    // Warp 4: TMA load for tokens + SFA (40 regs)
    cutlass::arch::warpgroup_reg_dealloc<kNumNonEpilogueRegisters>();
} else if (warp_idx == kNumDispatchWarps + 1) {
    // Warp 5: TMA load for weights + SFB (40 regs)
    cutlass::arch::warpgroup_reg_dealloc<kNumNonEpilogueRegisters>();
} else if (warp_idx == kNumDispatchWarps + 2) {
    // Warp 6: MMA issue warp (40 regs)
    cutlass::arch::warpgroup_reg_dealloc<kNumNonEpilogueRegisters>();
} else if (warp_idx == kNumDispatchWarps + 3) {
    // Warp 7: 预留 (40 regs)
    cutlass::arch::warpgroup_reg_dealloc<kNumNonEpilogueRegisters>();
} else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) {
    // Warp 8-15: Epilogue warps (208 regs)
    cutlass::arch::warpgroup_reg_alloc<kNumEpilogueRegisters>();
    // ... SwiGLU, scaling, combine ...
}

┌─────────────────────────────────────────────────────────────────┐
│                    Mega MoE CTA (512 threads)                    │
├──────────┬──────────┬──────────┬─────────────────────────────────┤
│ Warp 0-3 │ Warp 4-5 │ Warp 6-7 │ Warp 8-15                      │
│ Dispatch │ TMA Load │ MMA Issue│ Epilogue (8 warps = 2 WGs)     │
│ 4 warps  │ 2 warps  │ 2 warps  │ 8 warps                        │
│ 48 regs  │ 40 regs  │ 40 regs  │ 208 regs                       │
└──────────┴──────────┴──────────┴─────────────────────────────────┘

graph TB
    subgraph "Mega MoE Warp Pipeline"
        A[Warp 0-3: Dispatch<br/>Token Routing + NVLink Pull<br/>48 regs] -->|l1_arrival_count| B
        B[Warp 4: TMA Load<br/>Tokens + SFA → smem<br/>40 regs] -->|full_barriers| D
        C[Warp 5: TMA Load<br/>Weights + SFB → smem<br/>40 regs] -->|full_barriers| D
        D[Warp 6: MMA Issue<br/>UMMA → Tensor Core<br/>40 regs] -->|tmem_full_barriers| E
        E[Warp 8-15: Epilogue<br/>SwiGLU + Scale + Combine<br/>208 regs] -->|tmem_empty_barriers| D
    end

    style A fill:#e1f5fe
    style B fill:#fff3e0
    style C fill:#fff3e0
    style D fill:#f3e5f5
    style E fill:#e8f5e9

sequenceDiagram
    participant D as Dispatch Warps (0-3)
    participant L1 as L1 Token Buffer
    participant T4 as TMA Load Warp 4 (tokens)
    participant T5 as TMA Load Warp 5 (weights)
    participant S as Shared Memory (smem_a/b)
    participant M as MMA Issue Warp 6
    participant T as Tensor Memory (TMEM)
    participant E as Epilogue Warps (8-15)
    participant C as Combine Buffer

    D->>L1: Pull tokens via NVLink (TMA load)
    D->>L1: red_add_rel(l1_arrival_count)
    L1->>T4: Wait arrival → TMA copy to smem_a
    T4->>S: full_barriers arrive
    T5->>S: TMA copy weights to smem_b
    T5->>S: full_barriers arrive
    S->>M: Wait full_barriers → Issue UMMA
    M->>T: Accumulate in TMEM
    M->>E: tmem_full_barriers arrive
    E->>T: TMEM load → SwiGLU → Scale
    E->>M: tmem_empty_barriers arrive (release TMEM)
    E->>C: Write combine buffer (NVLink)

GPU0 要发送到 NIC1，但 NIC1 绑定在 GPU4 上
→ 需要 GPU0 → NVLink → GPU4 (Forward) → PCIe → NIC1
→ 需要 Warp Group B 做 Forwarding

// sym_buffer.cuh - 对称内存映射
CUTLASS_DEVICE ptr_t map(const ptr_t& ptr, const uint32_t& dst_rank_idx) const {
    int64_t mapped_ptr = offsets[dst_rank_idx] + reinterpret_cast<int64_t>(ptr);
    return *reinterpret_cast<ptr_t*>(&mapped_ptr);
}

// Mega MoE Dispatch Warp 的核心操作 (441-599 行)
// 1. 读取远端 token 的 topk_idx
const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr(
    current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank);

// 2. TMA load token from remote rank into shared memory
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),  // 远端 rank 的对称地址
    pull_mbarrier, kHidden);

// 3. 写入本地 L1 buffer
ptx::tma_store_1d(
    l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(),
    pull_buffer.get_base_ptr(), pull_buffer.get_num_bytes());

// Warp 6: MMA Issue (763-872 行)
if (warp_idx == kNumDispatchWarps + 2) {
    // 只有 leader CTA 执行
    if (is_leader_cta) {
        // 创建 UMMA 指令描述符
        auto instr_desc = cute::UMMA::make_instr_desc_block_scaled<...>();

        // 持久化调度所有 block
        scheduler.for_each_block([&](...) {
            // 等待 TMEM empty barrier
            tmem_empty_barriers[accum_stage_idx]->wait(accum_phase ^ 1);

            // 等待 TMA load 完成
            full_barriers[stage_idx]->wait(phase);

            // UTCCP copy SFA/SFB to TMEM
            cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4);

            // Issue UMMA
            ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma(
                b_desc, a_desc, accum_stage_idx * UMMA_N, ...);

            // Commit to empty barrier
            empty_barrier_arrive(k_block_idx == num_k_blocks - 1);
        });
    }
}

// Warp 8-15: Epilogue (877-1357 行)
// L1 Epilogue: SwiGLU + FP8 量化 + TMA Store
if (block_phase == sched::BlockPhase::Linear1) {
    // 1. TMEM Load
    cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, ...);

    // 2. SwiGLU: silu(gate) * up
    gate = __fmul2_rn(gate, {math::fast_rcp(denom.x), ...});
    swiglu_values[i * 2 + k] = __fmul2_rn(__fmul2_rn(gate, up), weights);

    // 3. Amax reduction + SF 计算
    math::get_e4m3_sf_and_sf_inv(amax_values[i], sf, sf_inv);

    // 4. Cast to FP8 + STSM store
    ptx::SM100_U8x4_STSM_T<__nv_fp8x4_e4m3>::copy(fp8x4_values, smem_ptr);

    // 5. TMA Store
    cute::SM90_TMA_STORE_2D::copy(&tensor_map_l1_output, ...);
}

// L2 Epilogue: BF16 写回远端 Combine Buffer
else {
    // 1. TMEM Load
    // 2. Store to shared memory
    // 3. Write to remote combine buffer via NVLink
    *sym_buffer.map(dst_ptr, dst_rank_idx) = packed;  // 对称内存写回
}

// Combine: Top-K reduce (1226-1356 行)
for (uint32_t token_idx = sm_idx * kNumEpilogueWarps + epilogue_warp_idx;
     token_idx < num_tokens;
     token_idx += kNumSMs * kNumEpilogueWarps) {
    // 加载所有 top-k 贡献
    // 累加: reduced += bf16_values[l]
    // Cast + TMA Store 到输出
}

// 343-349 行: 寄存器预算
constexpr uint32_t kNumDispatchRegisters = 48;      // Dispatch: 少
constexpr uint32_t kNumNonEpilogueRegisters = 40;   // MMA/TMA: 最少
constexpr uint32_t kNumEpilogueRegisters = 208;     // Epilogue: 最多
DG_STATIC_ASSERT(kNumDispatchRegisters * kNumDispatchThreads +
                 kNumNonEpilogueRegisters * kNumNonEpilogueThreads +
                 kNumEpilogueRegisters * kNumEpilogueThreads <= 64512,
                 "Too many registers");

DeepEP:
|---- IB Send ----|---- Forward ----|---- Receive ----|
     Warp A            Warp B            Warp C

Mega MoE:
|-- Dispatch --|-- TMA Load --|-- MMA Compute --|-- Epilogue --|
   Warp 0-3       Warp 4-5        Warp 6          Warp 8-15
      ↓              ↓               ↓                ↓
   NVLink        smem fill      Tensor Core      SwiGLU +
   Pull          (tokens+wt)    (UMMA)           Combine

// 331-335 行: Intra-SM Barrier 索引
constexpr uint32_t kDispatchBarrierIdx = 0;
constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1;
constexpr uint32_t kEpilogueFullBarrierIdx = 2;
constexpr uint32_t kEpilogueWGBarrierStartIdx = 3;

// 263-268 行: Barrier 分配
auto dispatch_barriers      = ...;  // kNumDispatchWarps 个
auto full_barriers          = ...;  // kNumStages 个 (TMA → MMA)
auto empty_barriers         = ...;  // kNumStages 个 (MMA → TMA)
auto tmem_full_barriers     = ...;  // kNumEpilogueStages 个 (MMA → Epilogue)
auto tmem_empty_barriers    = ...;  // kNumEpilogueStages 个 (Epilogue → MMA)
auto combine_barriers       = ...;  // kNumEpilogueWarps * 2 个

graph LR
    subgraph "Mega MoE = 整个 MoE 层的 Warp Specialization"
        A[Token Routing] --> B[NVLink Pull]
        B --> C[TMA Load Tokens]
        B --> D[TMA Load Weights]
        C --> E[MMA Issue L1]
        D --> E
        E --> F[SwiGLU + Quantize]
        F --> G[MMA Issue L2]
        G --> H[Write Combine Buffer]
        H --> I[Top-K Reduce]
    end

    style A fill:#e1f5fe
    style B fill:#e1f5fe
    style C fill:#fff3e0
    style D fill:#fff3e0
    style E fill:#f3e5f5
    style F fill:#e8f5e9
    style G fill:#f3e5f5
    style H fill:#e8f5e9
    style I fill:#e8f5e9

gantt
    title Mega MoE Warp Pipeline (单 CTA 视角)
    dateFormat X
    axisFormat %L

    section Dispatch (Warp 0-3)
    Token Routing + Expert Count :active, d1, 0, 10
    NVLink Pull (TMA Load)       :active, d2, after d1, 20
    Workspace Clean              :d3, after d2, 5

    section TMA Load (Warp 4-5)
    Wait L1 Arrival              :crit, t0, 5, 5
    TMA Copy Tokens+Weights      :active, t1, after d2, 15

    section MMA Issue (Warp 6)
    Wait full_barriers           :crit, m0, 15, 3
    Issue UMMA (L1)              :active, m1, after t1, 20
    Issue UMMA (L2)              :active, m2, after m1, 20

    section Epilogue (Warp 8-15)
    Wait tmem_full_barriers      :crit, e0, 20, 3
    SwiGLU + Quantize (L1)       :active, e1, after m1, 15
    Write Combine Buffer (L2)    :active, e2, after m2, 15
    Top-K Reduce + TMA Store     :active, e3, after e2, 10