From 540e5aeadc1b21dca8222a838ed6e5b6c4de93e8 Mon Sep 17 00:00:00 2001 From: Xinyi Liu <94362768+XinyiLiu577086410@users.noreply.github.com> Date: Thu, 18 Jun 2026 00:00:36 +0800 Subject: [PATCH] feat: implement sm90 megamoe phase2 dispatch-only --- csrc/apis/sm90_mega.hpp | 34 +- .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 309 +++++++++++++++++- deep_gemm/mega/__init__.py | 26 +- .../phase2/dispatch_only_correctness.py | 266 +++++++++++++++ 4 files changed, 622 insertions(+), 13 deletions(-) create mode 100644 megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index a68a554..ef21327 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/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(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(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(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(runtime_workspace.num_max_pool_blocks)}, + torch::TensorOptions().dtype(torch::kInt).device(buffer.device())); + auto token_src_metadata = torch::from_blob( + reinterpret_cast(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(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); diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh index 74cbd72..cd8a27c 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -2,11 +2,15 @@ #include +#include #include #include #include +#include #include #include +#include +#include #include 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(sizeof(float)), false); + constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden); + constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 128 * static_cast(sizeof(float)), false); + constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * static_cast(sizeof(int64_t)), false); + constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * static_cast(sizeof(float)), false); + constexpr auto l1_topk_weights_layout = layout::Data(static_cast(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(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( + 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( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0 and expert_idx < static_cast(kNumExperts)) + process(i * kNumTopk + lane_idx, static_cast(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(smem_expert_count[i]); + smem_expert_count[i] = static_cast( + 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( + 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( + 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(*workspace.get_expert_recv_count_sum_ptr(i)); + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(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(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(current_expert_idx)); + } + + if (current_expert_idx >= static_cast(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(*workspace.get_expert_recv_count_ptr(j, static_cast(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(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(), + current_rank_in_expert_idx); + const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(); + #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(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + #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() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = 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; } diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 4c7e19d..b25c629 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -60,14 +60,24 @@ class SymmBuffer: torch.cuda.synchronize() # Create input buffer views - (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) = slice_input_buffers(self.buffer) - self.l1_topk_weights = None - self.expert_recv_count_sum = None - self.l1_arrival_count = None - self.token_src_metadata = None + buffer_views = slice_input_buffers(self.buffer) + if _is_sm90(): + (self.x, self.x_sf, + self.topk_idx, self.topk_weights, + self.l1_acts, self.l1_acts_sf, self.l1_topk_weights, + self.l2_acts, self.l2_acts_sf, + self.expert_recv_count_sum, + 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 diff --git a/megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py b/megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py new file mode 100644 index 0000000..83dbb90 --- /dev/null +++ b/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()