feat: implement sm90 megamoe phase2 dispatch-only

csrc/apis/sm90_mega.hpp

@@ -13,7 +13,9 @@ namespace deep_gemm::mega {
 
 using SM90MegaMoEBufferViews = std::tuple<
     torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
-    torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>;
+    torch::Tensor, torch::Tensor, torch::Tensor,
+    torch::Tensor, torch::Tensor,
+    torch::Tensor, torch::Tensor, torch::Tensor>;
 
 static int get_token_alignment_for_sm90_mega_moe() {
     return layout::kLCMCandidateBlockM;
@@ -92,8 +94,11 @@ get_symm_buffer_size_for_sm90_mega_moe(
     DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
     DG_HOST_ASSERT(num_max_padded_sf_pool_tokens % 4 == 0);
 
-    // Slice function: creates input and L1/L2 pool tensor views.
+    // Slice function: creates input, L1/L2 pool, and Phase 2 dispatch-result tensor views.
     auto slice_input_buffers = [=](const torch::Tensor& buffer) {
+        static_assert(sizeof(layout::TokenSrcMetadata) == 3 * sizeof(uint32_t));
+        const auto runtime_workspace = layout::Workspace(
+            buffer.data_ptr(), num_ranks, num_experts, num_max_tokens_per_rank, num_topk);
         auto x = torch::from_blob(
             math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_token_buffer.base)),
             {num_max_tokens_per_rank, hidden},
@@ -119,6 +124,10 @@ get_symm_buffer_size_for_sm90_mega_moe(
             {num_max_padded_sf_pool_tokens, hidden / 128},
             {1, num_max_padded_sf_pool_tokens},
             torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
+        auto l1_topk_weights = torch::from_blob(
+            math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_topk_weights_buffer.base)),
+            {num_max_pool_tokens},
+            torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
         auto l2_acts = torch::from_blob(
             math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_token_buffer.base)),
             {num_max_pool_tokens, intermediate_hidden},
@@ -128,8 +137,22 @@ get_symm_buffer_size_for_sm90_mega_moe(
             {num_max_padded_sf_pool_tokens, intermediate_hidden / 128},
             {1, num_max_padded_sf_pool_tokens},
             torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
+        auto expert_recv_count_sum = torch::from_blob(
+            runtime_workspace.get_expert_recv_count_sum_ptr(),
+            {num_experts / num_ranks},
+            torch::TensorOptions().dtype(torch::kInt64).device(buffer.device()));
+        auto l1_arrival_count = torch::from_blob(
+            runtime_workspace.get_l1_arrival_count_ptr(),
+            {static_cast<int>(runtime_workspace.num_max_pool_blocks)},
+            torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
+        auto token_src_metadata = torch::from_blob(
+            reinterpret_cast<int32_t*>(runtime_workspace.get_token_src_metadata_ptr()),
+            {num_max_pool_tokens, 3},
+            torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
         return std::make_tuple(x, x_sf, topk_idx, topk_weights,
-                               l1_acts, l1_acts_sf, l2_acts, l2_acts_sf);
+                               l1_acts, l1_acts_sf, l1_topk_weights,
+                               l2_acts, l2_acts_sf,
+                               expert_recv_count_sum, l1_arrival_count, token_src_metadata);
     };
     return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
 }
@@ -202,7 +225,10 @@ static void fp8_mega_moe(
     DG_HOST_ASSERT(num_experts == num_experts_);
 
     // Already registered tensors
-    const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer);
+    const auto [x, x_sf, topk_idx, topk_weights,
+                l1_acts, l1_acts_sf, l1_topk_weights,
+                l2_acts, l2_acts_sf,
+                expert_recv_count_sum, l1_arrival_count, token_src_metadata] = slice(sym_buffer);
 
     // Dispatch into SM90 path
     DG_HOST_ASSERT(arch_major == 9);

deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh

@@ -2,11 +2,15 @@
 
 #include <cutlass/cutlass.h>
 
+#include <deep_gemm/common/math.cuh>
 #include <deep_gemm/common/exception.cuh>
 #include <deep_gemm/common/tma_utils.cuh>
 #include <deep_gemm/common/types.cuh>
+#include <deep_gemm/comm/barrier.cuh>
 #include <deep_gemm/layout/mega_moe.cuh>
 #include <deep_gemm/layout/sym_buffer.cuh>
+#include <deep_gemm/ptx/ld_st.cuh>
+#include <deep_gemm/ptx/utils.cuh>
 #include <deep_gemm/scheduler/sm90_mega_moe.cuh>
 
 namespace deep_gemm {
@@ -51,7 +55,310 @@ sm90_fp8_mega_moe_impl(void* y,
     DG_STATIC_ASSERT(BLOCK_K == 128, "SM90 MegaMoE expects BLOCK_K=128");
     DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks");
 
-    // Phase 1 only validates the host/JIT/API path and launches an empty kernel.
+#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
+    DG_STATIC_ASSERT(kNumDispatchThreads == 64, "SM90 dispatch-only path expects 64 dispatch threads");
+    DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of top-k experts");
+    DG_STATIC_ASSERT(kHidden % 128 == 0, "SM90 activation SF expects per-128K groups");
+    DG_STATIC_ASSERT(kHidden % sizeof(uint4) == 0, "Token copy expects 16-byte alignment");
+
+    constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32;
+    constexpr uint32_t kNumTokensPerWarp = 32 / kNumTopk;
+    constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk;
+    constexpr uint32_t SF_BLOCK_M = math::constexpr_align(BLOCK_M, 128u);
+    constexpr uint32_t kNumMaxPoolBlocks = kNumMaxPoolTokens / layout::kMinCandidateBlockM;
+    constexpr uint32_t kNumTokenUint4 = kHidden / sizeof(uint4);
+    constexpr uint32_t kNumSFValues = kHidden / 128;
+    DG_STATIC_ASSERT(kNumTokensPerWarp > 0, "Invalid number of top-k experts");
+    DG_STATIC_ASSERT(kNumPaddedSFPoolTokens % SF_BLOCK_M == 0, "Invalid padded SF pool size");
+
+    // Only the first two warps participate in Phase 2 dispatch. TMA/math warps stay idle.
+    const uint32_t thread_idx = threadIdx.x;
+    if (thread_idx >= kNumDispatchThreads)
+        return;
+
+    const uint32_t sm_idx = blockIdx.x;
+    const uint32_t warp_idx = thread_idx / 32;
+    const uint32_t lane_idx = ptx::get_lane_idx();
+
+    constexpr uint32_t kDispatchBarrierIdx = 0;
+    constexpr uint32_t kDispatchGridSyncIndex = 0;
+    constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 1;
+    constexpr uint32_t kBeforeDispatchPullBarrierTag = 2;
+    const auto dispatch_sync = []() {
+        ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
+    };
+
+    const auto workspace = layout::Workspace(
+        sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk);
+
+    constexpr auto fp8_token_layout = layout::Data(kHidden);
+    constexpr auto fp8_sf_layout = layout::Data(kHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
+    constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden);
+    constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
+    constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(int64_t)), false);
+    constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(float)), false);
+    constexpr auto l1_topk_weights_layout = layout::Data(static_cast<uint32_t>(sizeof(float)), false);
+
+    const auto input_token_buffer = layout::Buffer(
+        fp8_token_layout, 1, kNumMaxTokensPerRank,
+        workspace.get_end_ptr());
+    const auto input_sf_buffer = layout::Buffer(
+        fp8_sf_layout, 1, kNumMaxTokensPerRank,
+        input_token_buffer.get_end_ptr());
+    const auto input_topk_idx_buffer = layout::Buffer(
+        input_topk_idx_layout, 1, kNumMaxTokensPerRank,
+        input_sf_buffer.get_end_ptr());
+    const auto input_topk_weights_buffer = layout::Buffer(
+        input_topk_weights_layout, 1, kNumMaxTokensPerRank,
+        input_topk_idx_buffer.get_end_ptr());
+
+    const auto l1_token_buffer = layout::Buffer(
+        fp8_token_layout, 1, kNumMaxPoolTokens,
+        input_topk_weights_buffer.get_end_ptr());
+    const auto l1_sf_buffer = layout::Buffer(
+        fp8_sf_layout, 1, kNumPaddedSFPoolTokens,
+        l1_token_buffer.get_end_ptr());
+    const auto l1_topk_weights_buffer = layout::Buffer(
+        l1_topk_weights_layout, 1, kNumMaxPoolTokens,
+        l1_sf_buffer.get_end_ptr());
+    const auto l2_token_buffer = layout::Buffer(
+        fp8_intermediate_token_layout, 1, kNumMaxPoolTokens,
+        l1_topk_weights_buffer.get_end_ptr());
+    const auto l2_sf_buffer = layout::Buffer(
+        fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens,
+        l2_token_buffer.get_end_ptr());
+    (void)l2_sf_buffer;
+
+    constexpr uint32_t kSharedMemoryAlignment = 1024;
+    extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[];
+    const auto smem_expert_count = reinterpret_cast<uint32_t*>(smem_buffer);
+
+    // Clean local dispatch workspace from any previous call. Barrier state is intentionally
+    // persistent and must not be reset because grid/NVLink barriers use phase counters.
+    for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
+        *workspace.get_expert_send_count_ptr(i) = 0;
+    for (uint32_t i = thread_idx; i < kNumRanks * kNumExpertsPerRank; i += kNumDispatchThreads)
+        *workspace.get_expert_recv_count_ptr(i / kNumExpertsPerRank, i % kNumExpertsPerRank) = 0;
+    for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads)
+        *workspace.get_expert_recv_count_sum_ptr(i) = 0;
+    for (uint32_t i = thread_idx; i < kNumMaxPoolBlocks; i += kNumDispatchThreads) {
+        *workspace.get_l1_arrival_count_ptr(i) = 0;
+        *workspace.get_l2_arrival_mask_ptr(i) = 0;
+    }
+    comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
+                         kDispatchGridSyncIndex, kAfterWorkspaceCleanBarrierTag>(
+        workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync);
+
+    for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
+        smem_expert_count[i] = 0;
+    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
+
+    const auto read_topk_idx = [&](const auto& process) {
+        #pragma unroll 1
+        for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp;
+             i < num_tokens;
+             i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) {
+            const uint32_t token_offset = lane_idx / kNumTopk;
+            const uint32_t token_idx = i + token_offset;
+            int expert_idx = -1;
+            if (token_idx < num_tokens and lane_idx < kNumActivateLanes) {
+                expert_idx = static_cast<int>(
+                    __ldg(input_topk_idx_buffer.get_base_ptr<int64_t>() + i * kNumTopk + lane_idx));
+                if (expert_idx >= 0 and expert_idx < static_cast<int>(kNumExperts))
+                    process(i * kNumTopk + lane_idx, static_cast<uint32_t>(expert_idx));
+            }
+            __syncwarp();
+        }
+    };
+
+    // Count local outgoing token-topk entries per global expert.
+    read_topk_idx([&](const uint32_t&, const uint32_t& expert_idx) {
+        atomicAdd_block(smem_expert_count + expert_idx, 1);
+    });
+    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
+
+    // Convert per-SM counts into global per-rank offsets. High 32 bits count arrived SMs.
+    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));
+    }
+    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
+
+    // Write source token-topk indices into destination ranks' local-expert tables.
+    read_topk_idx([&](const uint32_t& token_topk_idx, const uint32_t& expert_idx) {
+        const uint32_t dst_rank_idx = expert_idx / kNumExpertsPerRank;
+        const uint32_t dst_local_expert_idx = expert_idx - dst_rank_idx * kNumExpertsPerRank;
+        const uint32_t dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1);
+        const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(
+            dst_local_expert_idx, sym_buffer.rank_idx, dst_slot_idx);
+        *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;
+    });
+
+    // Wait until all local SMs have finished filling local send-count/source-index data.
+    comm::grid_sync<kNumSMs, kDispatchGridSyncIndex>(
+        workspace, sm_idx, thread_idx, dispatch_sync);
+
+    // Publish per-rank expert counts and summed ready/count words to destination ranks.
+    if (sm_idx == 0) {
+        for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
+            const uint32_t dst_rank_idx = i / kNumExpertsPerRank;
+            const uint32_t dst_local_expert_idx = i - dst_rank_idx * kNumExpertsPerRank;
+            const uint64_t expert_status = *workspace.get_expert_send_count_ptr(i);
+            *sym_buffer.map(
+                workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx),
+                dst_rank_idx) = expert_status & 0xffffffffu;
+            ptx::atomic_add_sys(
+                sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx),
+                expert_status);
+        }
+    }
+    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
+
+    // All ranks must observe the complete count/source-index tables before pulling.
+    comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
+                         kDispatchGridSyncIndex, kBeforeDispatchPullBarrierTag>(
+        workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync,
+        /* sync_prologue */ false, /* sync_epilogue */ true);
+
+    auto scheduler = sched::SM90MegaMoEScheduler<
+        BLOCK_M, BLOCK_N, BLOCK_K,
+        L1_SHAPE_N, L1_SHAPE_K,
+        L2_SHAPE_N, L2_SHAPE_K,
+        kNumExpertsPerRank,
+        kNumExpertsPerWave,
+        kNumSMs, kNumRanks,
+        kUseNMajorL2>(workspace);
+    scheduler.fetch_expert_recv_count();
+
+    if (sm_idx == 0 and cumulative_local_expert_recv_stats != nullptr) {
+        for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads) {
+            const auto num_recv_tokens = static_cast<uint32_t>(*workspace.get_expert_recv_count_sum_ptr(i));
+            ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast<int>(num_recv_tokens));
+        }
+    }
+    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
+
+    constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u);
+    int current_expert_idx = -1;
+    uint32_t stored_rank_count[kNumRanksPerLane] = {};
+    uint32_t expert_start_idx = 0, expert_end_idx = 0;
+    uint32_t expert_pool_block_offset = 0;
+
+    constexpr uint32_t kNumGlobalDispatchWarps = kNumSMs * kNumDispatchWarps;
+    for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalDispatchWarps) {
+        int old_expert_idx = current_expert_idx;
+        while (token_idx >= expert_end_idx) {
+            if (++ current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
+                break;
+
+            expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M);
+            expert_start_idx = expert_end_idx;
+            expert_end_idx += scheduler.get_num_tokens(static_cast<uint32_t>(current_expert_idx));
+        }
+
+        if (current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
+            break;
+
+        if (old_expert_idx != current_expert_idx) {
+            old_expert_idx = current_expert_idx;
+            #pragma unroll
+            for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
+                const uint32_t j = i * 32 + lane_idx;
+                stored_rank_count[i] = j < kNumRanks ?
+                    static_cast<uint32_t>(*workspace.get_expert_recv_count_ptr(j, static_cast<uint32_t>(current_expert_idx))) : 0;
+            }
+        }
+
+        uint32_t current_rank_in_expert_idx = 0;
+        uint32_t remaining[kNumRanksPerLane];
+        #pragma unroll
+        for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
+            remaining[i] = stored_rank_count[i];
+
+        uint32_t offset = 0;
+        const uint32_t token_idx_in_expert = token_idx - expert_start_idx;
+        uint32_t slot_idx = token_idx_in_expert;
+        uint32_t token_idx_in_rank = 0;
+        while (true) {
+            uint32_t num_actives_in_lane = 0;
+            uint32_t min_in_lane = 0xffffffffu;
+            #pragma unroll
+            for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
+                num_actives_in_lane += remaining[i] > 0;
+                if (remaining[i] > 0)
+                    min_in_lane = cute::min(min_in_lane, remaining[i]);
+            }
+            const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane);
+            const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane);
+            const uint32_t num_round_tokens = length * num_active_ranks;
+
+            if (slot_idx < num_round_tokens) {
+                const uint32_t slot_idx_in_round = slot_idx % num_active_ranks;
+                uint32_t num_seen_ranks = 0;
+                #pragma unroll
+                for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
+                    const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0);
+                    const uint32_t num_active_lanes = __popc(mask);
+                    if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes)
+                        current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1);
+                    num_seen_ranks += num_active_lanes;
+                }
+                token_idx_in_rank = offset + slot_idx / num_active_ranks;
+                break;
+            }
+
+            slot_idx -= num_round_tokens;
+            offset += length;
+            #pragma unroll
+            for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
+                remaining[i] -= cute::min(remaining[i], length);
+        }
+
+        const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr(
+            static_cast<uint32_t>(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 - src_token_idx * kNumTopk;
+        const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert;
+        const uint32_t sf_pool_token_idx =
+            (expert_pool_block_offset + token_idx_in_expert / BLOCK_M) * SF_BLOCK_M +
+            (token_idx_in_expert % BLOCK_M);
+
+        const auto remote_token_ptr = sym_buffer.map(
+            input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr<uint4>(),
+            current_rank_in_expert_idx);
+        const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr<uint4>();
+        #pragma unroll
+        for (uint32_t i = lane_idx; i < kNumTokenUint4; i += 32)
+            local_token_ptr[i] = remote_token_ptr[i];
+
+        const auto remote_sf_ptr = sym_buffer.map(
+            input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr<float>(),
+            current_rank_in_expert_idx);
+        const auto local_sf_ptr = l1_sf_buffer.get_base_ptr<float>();
+        #pragma unroll
+        for (uint32_t i = lane_idx; i < kNumSFValues; i += 32)
+            local_sf_ptr[i * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[i];
+
+        __syncwarp();
+        if (lane_idx == 0) {
+            const float weight = *sym_buffer.map(
+                input_topk_weights_buffer.get_base_ptr<float>() + src_token_topk_idx,
+                current_rank_in_expert_idx);
+            *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr<float>() = weight;
+            *workspace.get_token_src_metadata_ptr(pool_token_idx) = {
+                current_rank_in_expert_idx, src_token_idx, src_topk_idx};
+        }
+        __syncwarp();
+        __threadfence();
+        __syncwarp();
+        if (lane_idx == 0) {
+            ptx::red_add_rel(
+                workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1);
+        }
+        __syncwarp();
+    }
+#endif
     return;
 }
 

deep_gemm/mega/__init__.py

@@ -60,14 +60,24 @@ class SymmBuffer:
         torch.cuda.synchronize()
 
         # Create input buffer views
-        (self.x, self.x_sf,
+        buffer_views = slice_input_buffers(self.buffer)
-         self.topk_idx, self.topk_weights,
+        if _is_sm90():
-         self.l1_acts, self.l1_acts_sf,
+            (self.x, self.x_sf,
-         self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer)
+             self.topk_idx, self.topk_weights,
-        self.l1_topk_weights = None
+             self.l1_acts, self.l1_acts_sf, self.l1_topk_weights,
-        self.expert_recv_count_sum = None
+             self.l2_acts, self.l2_acts_sf,
-        self.l1_arrival_count = None
+             self.expert_recv_count_sum,
-        self.token_src_metadata = None
+             self.l1_arrival_count,
+             self.token_src_metadata) = buffer_views
+        else:
+            (self.x, self.x_sf,
+             self.topk_idx, self.topk_weights,
+             self.l1_acts, self.l1_acts_sf,
+             self.l2_acts, self.l2_acts_sf) = buffer_views
+            self.l1_topk_weights = None
+            self.expert_recv_count_sum = None
+            self.l1_arrival_count = None
+            self.token_src_metadata = None
 
     def destroy(self):
         self.handle = None

megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py

@@ -0,0 +1,266 @@
+import argparse
+import inspect
+import os
+import pathlib
+import random
+import sys
+from typing import List, Tuple
+
+import torch
+import torch.distributed as dist
+
+
+REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+import deep_gemm
+from deep_gemm.utils.math import ceil_div
+
+
+def parse_tokens_list(value: str) -> List[int]:
+    return [int(item) for item in value.split(',') if item]
+
+
+def gather_same_shape(tensor: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
+    gathered = [torch.empty_like(tensor) for _ in range(dist.get_world_size(group))]
+    dist.all_gather(gathered, tensor.contiguous(), group=group)
+    return torch.stack(gathered, dim=0)
+
+
+def init_test_dist(local_rank_arg: int = None) -> Tuple[int, int, dist.ProcessGroup]:
+    local_rank = local_rank_arg if local_rank_arg is not None else int(os.environ.get('LOCAL_RANK', '0'))
+    rank = int(os.environ.get('RANK', '0'))
+    world_size = int(os.environ.get('WORLD_SIZE', '1'))
+    master_addr = os.environ.get('MASTER_ADDR', '127.0.0.1')
+    master_port = int(os.environ.get('MASTER_PORT', '8361'))
+
+    torch.cuda.set_device(local_rank)
+    sig = inspect.signature(dist.init_process_group)
+    params = {
+        'backend': 'nccl',
+        'init_method': f'tcp://{master_addr}:{master_port}',
+        'world_size': world_size,
+        'rank': rank,
+    }
+    if 'device_id' in sig.parameters:
+        params['device_id'] = torch.device(f'cuda:{local_rank}')
+    dist.init_process_group(**params)
+    torch.set_default_device('cuda')
+    return rank, world_size, dist.new_group(list(range(world_size)))
+
+
+def get_block_m(num_tokens: int, num_ranks: int, num_topk: int, num_experts: int) -> int:
+    expected = num_tokens * num_ranks * num_topk / num_experts
+    if expected <= 16.5:
+        return 32
+    if expected <= 64.5:
+        return 64
+    return 128
+
+
+def make_topk(num_tokens: int, num_experts: int, num_topk: int, rank_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    token_idx = torch.arange(num_tokens, device='cuda', dtype=torch.long).unsqueeze(1)
+    topk_slot = torch.arange(num_topk, device='cuda', dtype=torch.long).unsqueeze(0)
+    topk_idx = (token_idx * num_topk + topk_slot + rank_idx * 3) % num_experts
+    topk_weights = (token_idx.float() * 0.125 + topk_slot.float() * 0.25 + rank_idx + 1.0).contiguous()
+    return topk_idx.contiguous(), topk_weights.contiguous()
+
+
+def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int):
+    l1_weights = torch.randn(
+        (num_experts_per_rank, intermediate_hidden * 2, hidden),
+        dtype=torch.float32, device='cuda').to(torch.float8_e4m3fn)
+    l2_weights = torch.randn(
+        (num_experts_per_rank, hidden, intermediate_hidden),
+        dtype=torch.float32, device='cuda').to(torch.float8_e4m3fn)
+    l1_weights_sf = torch.ones(
+        (num_experts_per_rank, ceil_div(intermediate_hidden * 2, 128), hidden // 128),
+        dtype=torch.float32, device='cuda')
+    l2_weights_sf = torch.ones(
+        (num_experts_per_rank, ceil_div(hidden, 128), intermediate_hidden // 128),
+        dtype=torch.float32, device='cuda')
+    return deep_gemm.transform_weights_for_mega_moe(
+        (l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf))
+
+
+def build_expected_entries(all_topk_idx: torch.Tensor,
+                           rank_idx: int,
+                           num_experts_per_rank: int) -> Tuple[List[List[Tuple[int, int, int]]], torch.Tensor]:
+    num_ranks, num_tokens, num_topk = all_topk_idx.shape
+    local_start = rank_idx * num_experts_per_rank
+    local_end = local_start + num_experts_per_rank
+    entries: List[List[Tuple[int, int, int]]] = [[] for _ in range(num_experts_per_rank)]
+    counts = [0 for _ in range(num_experts_per_rank)]
+
+    all_topk_idx_cpu = all_topk_idx.cpu()
+    for src_rank in range(num_ranks):
+        for token_idx in range(num_tokens):
+            for topk_idx in range(num_topk):
+                expert_idx = int(all_topk_idx_cpu[src_rank, token_idx, topk_idx])
+                if local_start <= expert_idx < local_end:
+                    local_expert = expert_idx - local_start
+                    entries[local_expert].append((src_rank, token_idx, topk_idx))
+                    counts[local_expert] += 1
+    return entries, torch.tensor(counts, dtype=torch.int64, device='cuda')
+
+
+def verify_case(buffer: deep_gemm.SymmBuffer,
+                all_x_f32: torch.Tensor,
+                all_x_sf: torch.Tensor,
+                all_topk_idx: torch.Tensor,
+                all_topk_weights: torch.Tensor,
+                cumulative_stats: torch.Tensor,
+                rank_idx: int,
+                num_ranks: int,
+                num_experts: int,
+                num_topk: int,
+                block_m: int) -> None:
+    del num_topk
+    num_experts_per_rank = num_experts // num_ranks
+    sf_block_m = ((block_m + 127) // 128) * 128
+    num_sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
+    entries_by_expert, counts = build_expected_entries(all_topk_idx, rank_idx, num_experts_per_rank)
+
+    torch.testing.assert_close(cumulative_stats.cpu(), counts.to(torch.int32).cpu(), rtol=0, atol=0)
+
+    expert_status = buffer.expert_recv_count_sum[:num_experts_per_rank].detach().cpu().tolist()
+    arrival = buffer.l1_arrival_count.detach().cpu()
+    metadata = buffer.token_src_metadata.detach().cpu()
+
+    pool_block_offset = 0
+    local_start = rank_idx * num_experts_per_rank
+    for local_expert, expected_entries in enumerate(entries_by_expert):
+        count = len(expected_entries)
+        status = int(expert_status[local_expert])
+        assert (status & 0xffffffff) == count, (local_expert, status, count)
+        assert (status >> 32) == num_sms * num_ranks, (local_expert, status, num_sms, num_ranks)
+
+        expected_set = set(expected_entries)
+        actual_set = set()
+        base_pool_token = pool_block_offset * block_m
+        expected_global_expert = local_start + local_expert
+
+        for token_in_expert in range(count):
+            pool_token_idx = base_pool_token + token_in_expert
+            src_rank, src_token_idx, src_topk_idx = [int(v) for v in metadata[pool_token_idx].tolist()]
+            actual_entry = (src_rank, src_token_idx, src_topk_idx)
+            assert actual_entry in expected_set, (local_expert, token_in_expert, actual_entry, expected_set)
+            assert actual_entry not in actual_set, (local_expert, token_in_expert, actual_entry)
+            actual_set.add(actual_entry)
+
+            expert_idx = int(all_topk_idx[src_rank, src_token_idx, src_topk_idx].item())
+            assert expert_idx == expected_global_expert, (pool_token_idx, expert_idx, expected_global_expert)
+
+            torch.testing.assert_close(
+                buffer.l1_acts[pool_token_idx].to(torch.float32).cpu(),
+                all_x_f32[src_rank, src_token_idx].cpu(),
+                rtol=0, atol=0)
+            torch.testing.assert_close(
+                buffer.l1_topk_weights[pool_token_idx].cpu(),
+                all_topk_weights[src_rank, src_token_idx, src_topk_idx].cpu(),
+                rtol=0, atol=0)
+
+            sf_pool_token_idx = (pool_block_offset + token_in_expert // block_m) * sf_block_m + token_in_expert % block_m
+            torch.testing.assert_close(
+                buffer.l1_acts_sf[sf_pool_token_idx].cpu(),
+                all_x_sf[src_rank, src_token_idx].cpu(),
+                rtol=0, atol=0)
+
+        assert actual_set == expected_set, (local_expert, actual_set, expected_set)
+
+        for block_idx in range(ceil_div(count, block_m)):
+            expected_arrivals = min(block_m, count - block_idx * block_m)
+            actual_arrivals = int(arrival[pool_block_offset + block_idx].item())
+            assert actual_arrivals == expected_arrivals, (local_expert, block_idx, actual_arrivals, expected_arrivals)
+
+        pool_block_offset += ceil_div(count, block_m)
+
+
+def run_case(args: argparse.Namespace,
+             group: dist.ProcessGroup,
+             rank_idx: int,
+             num_ranks: int,
+             buffer: deep_gemm.SymmBuffer,
+             weights,
+             num_tokens: int) -> None:
+    hidden = args.hidden
+    num_topk = args.num_topk
+    num_experts = args.num_experts
+
+    x_f32 = torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda')
+    x_fp8 = x_f32.to(torch.float8_e4m3fn)
+    x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') + 0.5
+    topk_idx, topk_weights = make_topk(num_tokens, num_experts, num_topk, rank_idx)
+
+    all_x_f32 = gather_same_shape(x_fp8.to(torch.float32), group)
+    all_x_sf = gather_same_shape(x_sf, group)
+    all_topk_idx = gather_same_shape(topk_idx, group)
+    all_topk_weights = gather_same_shape(topk_weights, group)
+
+    buffer.x[:num_tokens].copy_(x_fp8)
+    buffer.x_sf[:num_tokens].copy_(x_sf)
+    buffer.topk_idx[:num_tokens].copy_(topk_idx)
+    buffer.topk_weights[:num_tokens].copy_(topk_weights)
+    torch.cuda.synchronize()
+    dist.barrier(group=group)
+
+    cumulative_stats = torch.zeros((num_experts // num_ranks,), dtype=torch.int32, device='cuda')
+    y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
+    deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer,
+                           cumulative_local_expert_recv_stats=cumulative_stats)
+    torch.cuda.synchronize()
+
+    block_m = get_block_m(num_tokens, num_ranks, num_topk, num_experts)
+    verify_case(buffer, all_x_f32, all_x_sf, all_topk_idx, all_topk_weights,
+                cumulative_stats, rank_idx, num_ranks, num_experts, num_topk, block_m)
+    dist.barrier(group=group)
+    if rank_idx == 0:
+        print(f'[PASSED] tokens={num_tokens}, block_m={block_m}', flush=True)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 2 dispatch-only correctness')
+    parser.add_argument('--tokens-list', type=str, default='0,8,48,192')
+    parser.add_argument('--num-max-tokens-per-rank', type=int, default=384)
+    parser.add_argument('--hidden', type=int, default=512)
+    parser.add_argument('--intermediate-hidden', type=int, default=256)
+    parser.add_argument('--num-experts', type=int, default=16)
+    parser.add_argument('--num-topk', type=int, default=6)
+    parser.add_argument('--local-rank', type=int, default=None)
+    args = parser.parse_args()
+
+    local_rank = args.local_rank if args.local_rank is not None else int(os.environ.get('LOCAL_RANK', '0'))
+    rank_idx, num_ranks, group = init_test_dist(local_rank)
+    assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
+    assert args.num_experts % num_ranks == 0
+    assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
+    assert args.num_topk <= args.num_experts
+
+    torch.manual_seed(1234 + rank_idx)
+    random.seed(1234 + rank_idx)
+
+    tokens_list = parse_tokens_list(args.tokens_list)
+    max_tokens = max([args.num_max_tokens_per_rank] + tokens_list)
+    buffer = deep_gemm.get_symm_buffer_for_mega_moe(
+        group, args.num_experts, max_tokens, args.num_topk,
+        args.hidden, args.intermediate_hidden)
+    weights = make_weights(args.num_experts // num_ranks, args.hidden, args.intermediate_hidden)
+
+    if rank_idx == 0:
+        print(f'[Phase 2] ranks={num_ranks}, tokens_list={tokens_list}, '
+              f'hidden={args.hidden}, intermediate={args.intermediate_hidden}, '
+              f'experts={args.num_experts}, topk={args.num_topk}', flush=True)
+
+    for num_tokens in tokens_list:
+        run_case(args, group, rank_idx, num_ranks, buffer, weights, num_tokens)
+
+    dist.barrier(group=group)
+    buffer.destroy()
+    dist.destroy_process_group()
+    if rank_idx == 0:
+        print('[PASSED] SM90 MegaMoE Phase 2 dispatch-only correctness', flush=True)
+
+
+if __name__ == '__main__':
+    main()