06.10 Counting Sort / Bucketization 三方对比分析：博客 ↔ DeepEP ↔ DeepGEMM Mega MoE

┌─────────────────────────────────────────────────────────────────────────────┐
│                    dispatch_impl (主 Dispatch Kernel)                        │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │ Notify Warp (warp_idx < kNumNotifyWarps)                             │   │
│  │                                                                      │   │
│  │  Step 1: COUNT          Step 2: PREFIX SUM         Step 3: 输出      │   │
│  │  ┌──────────────┐       ┌───────────────────┐      ┌──────────────┐  │   │
│  │  │ atomicAdd to │  ──▶  │ SM 0 warp-level   │──▶  │ psum_num_    │  │   │
│  │  │ smem expert_ │       │ inclusive/exclusive│     │ recv_tokens_ │  │   │
│  │  │ count[expert]│       │ scan              │      │ per_expert   │  │   │
│  │  └──────────────┘       └───────────────────┘      └──────────────┘  │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │ Dispatch Warp (warp_idx >= kNumNotifyWarps)                          │   │
│  │                                                                      │   │
│  │  Step 4: SCATTER                                                     │   │
│  │  ┌──────────────────────────────────────────────────────────────┐    │   │
│  │  │ atomicAdd(scaleup_atomic_sender_counter + dst_rank_idx, 1)   │    │   │
│  │  │ → 获得 dst_slot_idx                                           │    │   │
│  │  │ → TMA store 到目标 rank 的对称内存                              │    │   │
│  │  └──────────────────────────────────────────────────────────────┘    │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────────────────────┘
                                      │
                                      ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│              dispatch_copy_epilogue_impl (Copy Epilogue Kernel)              │
│                                                                             │
│  Step 5: EXPAND MODE SCATTER (可选)                                         │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │ atomicAdd(psum_num_recv_tokens_per_expert + dst_expert_idx, 1)       │   │
│  │ → 获得 expert 内的 slot 位置                                          │   │
│  │ → TMA store 到 recv_x（Expert-major layout）                          │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────────────────────┘

// Atomic add on shared memory
EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes");
const auto global_warp_idx = warp_idx * kNumSMs + sm_idx;
for (int i = global_warp_idx; i < num_tokens; i += kNumNotifyWarps * kNumSMs) {
    // Expert choice can not be redundant
    const auto dst_expert_idx = lane_idx < kNumTopk ?
        static_cast<int>(__ldg(topk_idx + i * kNumTopk + lane_idx)) : -1;
    if (dst_expert_idx >= 0)
        atomicAdd_block(expert_count + dst_expert_idx, 1);

    // Rank choice should do deduplication here
    const auto dst_rank_idx = dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1;
    if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0)
        atomicAdd_block(rank_count + dst_rank_idx, 1);
}
ptx::named_barrier<kNumNotifyThreads>(kNotifyBarrierIndex);

// Do full-grid reduction
#pragma unroll
for (int i = thread_idx; i < kNumRanks + kNumExperts; i += kNumNotifyThreads) {
    const int64_t counter = (1ll << 32ll) | rank_expert_count[i];
    ptx::red_add(workspace_layout.get_notify_reduction_workspace_ptr() + i, counter);
}

// Do prefix sum by the warps
const auto do_psum = [=](const int* count, int* out, const int n, const int is_exclusive) {
    int psum = 0;
    #pragma unroll
    for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++ i) {
        const auto idx = i * 32 + lane_idx;
        const auto mem_idx = idx - is_exclusive;
        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);

        // Store into global memory
        if (idx < n + is_exclusive)
            out[idx] = sum;

        // Update `psum` by using the last lane's value
        psum = ptx::exchange(sum, 31);
    }
};
if (warp_idx == 0) {
    // Inclusive prefix sum
    do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, kNumRanks, 0);
} else if (warp_idx == 1) {
    // Exclusive prefix sum for later expanding
    do_psum(expert_count, psum_num_recv_tokens_per_expert, kNumExpertsPerRank, 1);
}

// Deduplicate ranks and assign slots
int stored_dst_slot_idx = -1;
if constexpr (kReuseSlotIndices) {
    if (lane_idx < kNumTopk)
        stored_dst_slot_idx = __ldg(dst_buffer_slot_idx + token_idx * kNumTopk + lane_idx);
    stored_dst_slot_idx = stored_dst_slot_idx >= 0 ?
        (stored_dst_slot_idx - rank_idx * kNumMaxTokensPerRank) : -1;
} else {
    if (ptx::deduplicate(stored_dst_rank_idx, lane_idx) and stored_dst_rank_idx >= 0)
        stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_rank_idx, 1);
    if (lane_idx < kNumTopk) {
        const auto value = stored_dst_slot_idx >= 0 ?
            rank_idx * kNumMaxTokensPerRank + stored_dst_slot_idx : -1;
        dst_buffer_slot_idx[token_idx * kNumTopk + lane_idx] = value;
    }
}

// Issue TMA NVLink stores
const auto dst_ptr = stored_dst_slot_idx >= 0 ?
    gin.get_sym_ptr<team_t>(recv_buffer.get_token_buffer(stored_dst_slot_idx).get_base_ptr(), stored_dst_rank_idx) :
    nullptr;
if (dst_ptr != nullptr)
    ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), tma_buffer.get_num_bytes<false>());
ptx::tma_store_commit();

// Issue RDMA put
if constexpr (not kIsScaleupNVLink) {
    ptx::tma_store_wait<1>();
    if (stored_dst_slot_idx >= 0 and dst_ptr == nullptr) {
        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);
    }
}

} else if (kDoExpand and not kCachedMode and dst_expert_idx >= 0) {
    dst_tensor_idx = atomicAdd(psum_num_recv_tokens_per_expert + dst_expert_idx, 1);
}

┌─────────────────────────────────────────────────────────────────────────────┐
│                    sm100_fp8_fp4_mega_moe_impl                               │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │ Dispatch Warp (warp_idx < kNumDispatchWarps)                         │   │
│  │                                                                      │   │
│  │  Step 1: COUNT          Step 2: PREFIX SUM         Step 3: SCATTER   │   │
│  │  ┌──────────────┐       ┌───────────────────┐      ┌──────────────┐  │   │
│  │  │ atomicAdd to │  ──▶  │ atomicAdd to      │──▶  │ atomicAdd to │  │   │
│  │  │ smem_expert  │       │ global expert_    │      │ get dst slot │  │   │
│  │  │ _count[expert]│      │ send_count →      │      │ write src_   │  │   │
│  │  │              │       │ get global offset │      │ token_topk_idx│  │   │
│  │  └──────────────┘       └───────────────────┘      └──────────────┘  │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │ Pull Warp (同一 dispatch warp，角色切换)                               │   │
│  │                                                                      │   │
│  │  Step 4: PULL + 连续化                                                │   │
│  │  ┌──────────────────────────────────────────────────────────────┐    │   │
│  │  │ 读取 scatter 表 → TMA pull token → 写入 pool buffer          │    │   │
│  │  │ (expert_pool_block_offset 决定 pool 中的位置)                 │    │   │
│  │  └──────────────────────────────────────────────────────────────┘    │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │ Scheduler (GEMM 调度)                                                │   │
│  │                                                                      │   │
│  │  Step 5: POOL BLOCK OFFSET (Prefix Sum)                              │   │
│  │  ┌──────────────────────────────────────────────────────────────┐    │   │
│  │  │ get_pool_block_offset: 对 stored_num_tokens_per_expert       │    │   │
│  │  │ 做 prefix sum → expert 在 pool 中的起始 block 位置            │    │   │
│  │  └──────────────────────────────────────────────────────────────┘    │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────────────────────┘

// Count experts' tokens
read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) {
   atomicAdd_block(shared_storage.expert_token_count + expert_idx, 1);
});
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);

// Get SM offset (~6.5 us)
#pragma unroll
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
    const uint64_t send_value = (1ull << 32) | static_cast<uint64_t>(shared_storage.expert_token_count[i]);
    shared_storage.expert_token_count[i] = static_cast<uint32_t>(
        ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value));
}
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);

// Write source indices (~2 us with 512 tokens)
read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) {
    const auto dst_rank_idx = expert_idx / kNumExpertsPerRank;
    const auto dst_slot_idx = atomicAdd_block(shared_storage.expert_token_count + expert_idx, 1);
    const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(
        expert_idx % kNumExpertsPerRank, sym_buffer.rank_idx, dst_slot_idx);
    *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;
});

// Read source token-topk index (written by remote dispatch via NVLink)
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);
const uint32_t src_token_idx = src_token_topk_idx / kNumTopk;
const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk;

// 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),
    pull_mbarrier, kHidden);

// Store token to local L1 buffer via TMA
const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert;
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());

CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) const {
    uint32_t num_blocks = 0;
    #pragma unroll
    for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
        if (i * 32 + ptx::get_lane_idx() < expert_idx)
            num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], BLOCK_M);
    }
    return __reduce_add_sync(0xffffffff, num_blocks);
}

flowchart TB
    subgraph NotifyWarp ["Notify Warp (专用)"]
        A[遍历 topk_idx] --> B[atomicAdd_block to smem expert_count]
        B --> C[sync barrier]
        C --> D[red_add to global notify_reduction_workspace]
        D --> E[SM 0 等待所有 SM 完成]
        E --> F[Warp-level inclusive/exclusive scan]
        F --> G[输出 psum_num_recv_tokens_per_scaleup_rank]
        F --> H[输出 psum_num_recv_tokens_per_expert]
    end

    subgraph DispatchWarp ["Dispatch Warp"]
        I[遍历 token] --> J[deduplicate rank]
        J --> K[atomicAdd scaleup_atomic_sender_counter]
        K --> L[获得 dst_slot_idx]
        L --> M[TMA store to recv_buffer]
    end

    subgraph CopyEpilogue ["Copy Epilogue Kernel (可选 Expand)"]
        N[遍历 recv_buffer] --> O[atomicAdd psum_num_recv_tokens_per_expert]
        O --> P[获得 expert 内 slot]
        P --> Q[TMA store to recv_x Expert-major]
    end

    NotifyWarp --> DispatchWarp
    DispatchWarp --> CopyEpilogue

flowchart TB
    subgraph DispatchWarp ["Dispatch Warp (Count + Scatter)"]
        A[read_topk_idx] --> B[atomicAdd_block to smem expert_token_count]
        B --> C[sync barrier]
        C --> D[atomic_add to global expert_send_count]
        D --> E[返回值写回 smem 作为全局偏移]
        E --> F[sync barrier]
        F --> G[read_topk_idx again]
        G --> H[atomicAdd_block expert_token_count 获取 slot]
        H --> I[src_token_topk_idx = token_topk_idx]
        I --> J[NVLink map to dst_rank]
        J --> K[写入 dst_rank 的 src_token_topk_idx 表]
    end

    subgraph PullPhase ["Pull Phase (同一 warp 角色切换)"]
        L[scheduler.fetch_expert_recv_count] --> M[等待所有 rank 完成]
        M --> N[读取 src_token_topk_idx]
        N --> O[解码 src_token_idx, src_topk_idx]
        O --> P[TMA load from remote rank]
        P --> Q[pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert]
        Q --> R[TMA store to l1_token_buffer]
    end

    subgraph Scheduler ["Scheduler (GEMM 调度)"]
        S[fetch_expert_recv_count] --> T[get_pool_block_offset: prefix sum]
        T --> U[输出 expert 在 pool 中的起始 block]
    end

    DispatchWarp --> PullPhase
    PullPhase --> Scheduler

DeepEP (Push 范式):
  Source Rank → Count → Prefix Sum → Scatter → TMA Store → Dest Rank's Buffer
                                                      ↓
                                             Copy Epilogue (Expand)
                                                      ↓
                                              Expert Buffer (GEMM Input)

Mega MoE (Pull 范式):
  Source Rank → Count → Prefix Sum → Scatter → Write Scatter Table (Dest Rank)
                                                      ↓
                                              Pull Phase (Dest Rank)
                                                      ↓
                                              Pool Buffer (GEMM Input)

// smem 中的计数数组（Notify Warp 使用）
int* rank_count = rank_expert_count;        // [kNumRanks]
int* expert_count = rank_expert_count + kNumRanks;  // [kNumExperts]

// global workspace 中的 reduction 数组
int64_t* notify_reduction_workspace;        // [kNumRanks + kNumExperts]

// 输出 prefix sum
int* psum_num_recv_tokens_per_scaleup_rank; // [kNumRanks] (inclusive)
int* psum_num_recv_tokens_per_expert;       // [kNumExpertsPerRank] (exclusive)

// scatter 使用的全局原子计数器
int* scaleup_atomic_sender_counter;         // [kNumRanks]

// smem 中的计数数组（Dispatch Warp 使用）
uint32_t expert_token_count[kNumExperts];   // 计数 → 全局偏移（复用）

// global workspace 中的 expert count
uint64_t* expert_send_count;                // [kNumExperts] (64-bit 编码)
uint64_t* expert_recv_count;                // [kNumRanks][kNumExpertsPerRank]
uint64_t* expert_recv_count_sum;            // [kNumExpertsPerRank]

// scatter 索引表
uint32_t* src_token_topk_idx;               // [kNumExpertsPerRank][kNumRanks][max_recv_tokens]

// combine 使用的源 metadata
TokenSrcMetadata* token_src_metadata;       // [num_max_pool_tokens]