diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index 0bd32f5..6e28683 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -16,7 +16,7 @@ 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>; static int get_token_alignment_for_sm90_mega_moe() { return layout::kLCMCandidateBlockM; @@ -164,11 +164,15 @@ get_symm_buffer_size_for_sm90_mega_moe( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_accum_debug_buffer.base)), {128, 128}, torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto combine_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(combine_token_buffer.base)), + {num_topk, num_max_tokens_per_rank, hidden}, + torch::TensorOptions().dtype(torch::kBFloat16).device(buffer.device())); return std::make_tuple(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, l2_arrival_mask, token_src_metadata, - l1_accum_debug); + l1_accum_debug, combine_acts); }; return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; } @@ -245,8 +249,9 @@ static void fp8_mega_moe( l1_acts, l1_acts_sf, l1_topk_weights, l2_acts, l2_acts_sf, expert_recv_count_sum, l1_arrival_count, l2_arrival_mask, token_src_metadata, - l1_accum_debug] = slice(sym_buffer); + l1_accum_debug, combine_acts] = slice(sym_buffer); (void)l2_arrival_mask; + (void)combine_acts; // 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 e112d7c..d487168 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 @@ -49,14 +49,14 @@ sm90_fp8_mega_moe_impl(void* y, const __grid_constant__ layout::SymBuffer sym_buffer, const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, - const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, - const void* l1_weights_sf, - const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, - const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, - const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, - const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, - const void* l2_weights_sf, - void* l1_accum_debug) { + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, + const void* l1_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, + const void* l2_weights_sf, + void* l1_accum_debug) { DG_STATIC_ASSERT(kNumThreads == 384, "SM90 MegaMoE expects 384 threads"); DG_STATIC_ASSERT(BLOCK_N == 128, "SM90 MegaMoE expects BLOCK_N=128"); DG_STATIC_ASSERT(BLOCK_K == 128, "SM90 MegaMoE expects BLOCK_K=128"); @@ -89,6 +89,8 @@ sm90_fp8_mega_moe_impl(void* y, constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t); constexpr uint32_t kNumL1WeightSFGroupsN = L1_SHAPE_N / 128; constexpr uint32_t kNumL1WeightSFGroupsK = L1_SHAPE_K / 128; + constexpr uint32_t kNumL2WeightSFGroupsN = L2_SHAPE_N / 128; + constexpr uint32_t kNumL2WeightSFGroupsK = L2_SHAPE_K / 128; constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; constexpr uint32_t kL2ActsGranK = 64; constexpr uint32_t kMathBarrierIdx = 2; @@ -106,6 +108,9 @@ sm90_fp8_mega_moe_impl(void* y, cute::prefetch_tma_descriptor(&tensor_map_l1_acts); cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf); cute::prefetch_tma_descriptor(&tensor_map_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights); } __syncwarp(); @@ -121,6 +126,7 @@ sm90_fp8_mega_moe_impl(void* y, sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk); constexpr auto fp8_token_layout = layout::Data(kHidden); + constexpr auto bf16_token_layout = layout::Data(kHidden * static_cast(sizeof(nv_bfloat16))); 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 / kL2ActsGranK * static_cast(sizeof(float)), false); @@ -156,7 +162,13 @@ sm90_fp8_mega_moe_impl(void* y, const auto l2_sf_buffer = layout::Buffer( fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens, l2_token_buffer.get_end_ptr()); - (void)l2_sf_buffer; + constexpr auto l1_accum_debug_layout = layout::Data(128 * static_cast(sizeof(float)), false); + const auto l1_accum_debug_buffer = layout::Buffer( + l1_accum_debug_layout, 1, 128, + l2_sf_buffer.get_end_ptr()); + const auto combine_token_buffer = layout::Buffer( + bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, + l1_accum_debug_buffer.get_end_ptr()); constexpr uint32_t kSharedMemoryAlignment = 1024; extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[]; @@ -466,13 +478,16 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t& num_k_blocks, const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - if (block_phase != sched::SM90BlockPhase::Linear1) - return; - const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; - const uint32_t valid_m = scheduler.get_valid_m(); - const auto arrival_ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); - while (ptx::ld_acq(arrival_ptr) != valid_m); + if (block_phase == sched::SM90BlockPhase::Linear1) { + const uint32_t valid_m = scheduler.get_valid_m(); + const auto arrival_ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); + while (ptx::ld_acq(arrival_ptr) != valid_m); + } else { + const auto arrival_ptr = workspace.get_l2_arrival_mask_ptr(pool_block_idx); + const uint64_t expected = ((1ull << num_k_blocks) << num_k_blocks) - 1; + while (ptx::ld_acq_gpu(arrival_ptr) != expected); + } #pragma unroll 1 for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { @@ -482,16 +497,26 @@ sm90_fp8_mega_moe_impl(void* y, auto& full_barrier = *full_barriers[stage_idx]; const uint32_t k_idx = k_block_idx * BLOCK_K; const uint32_t m_idx = pool_block_idx * BLOCK_M; - const uint32_t n_idx = local_expert_idx * L1_SHAPE_N + n_block_idx * BLOCK_N; - const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M; - tma::copy( - &tensor_map_l1_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx); - tma::copy( - &tensor_map_l1_acts_sf, &full_barrier, smem_sfa[stage_idx], sfa_m_idx, k_block_idx); - tma::copy( - &tensor_map_l1_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx); - full_barrier.arrive_and_expect_tx( - SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE); + if (block_phase == sched::SM90BlockPhase::Linear1) { + const uint32_t n_idx = local_expert_idx * L1_SHAPE_N + n_block_idx * BLOCK_N; + const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M; + tma::copy( + &tensor_map_l1_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx); + tma::copy( + &tensor_map_l1_acts_sf, &full_barrier, smem_sfa[stage_idx], sfa_m_idx, k_block_idx); + tma::copy( + &tensor_map_l1_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx); + full_barrier.arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE); + } else { + const uint32_t n_idx = local_expert_idx * L2_SHAPE_N + n_block_idx * BLOCK_N; + tma::copy( + &tensor_map_l2_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx); + tma::copy( + &tensor_map_l2_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx); + full_barrier.arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + } } __syncwarp(); } @@ -508,6 +533,7 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx; const uint32_t r_1 = r_0 + 8; const auto l1_weights_sf_ptr = reinterpret_cast(l1_weights_sf); + const auto l2_weights_sf_ptr = reinterpret_cast(l2_weights_sf); auto l1_accum_debug_ptr = reinterpret_cast(l1_accum_debug); auto l2_acts_ptr = l2_token_buffer.get_base_ptr<__nv_fp8_e4m3>(); auto l2_acts_sf_ptr = l2_sf_buffer.get_base_ptr(); @@ -543,11 +569,12 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t& num_k_blocks, const uint32_t& m_block_idx, const uint32_t& n_block_idx) { - if (block_phase != sched::SM90BlockPhase::Linear1) - return; - const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; const uint32_t valid_m = scheduler.get_valid_m(); + const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0; + const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1; + const uint32_t pool_token_idx_0 = pool_block_idx * BLOCK_M + row_0; + const uint32_t pool_token_idx_1 = pool_block_idx * BLOCK_M + row_1; float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0}; const auto empty_barrier_arrive = [&](const uint32_t& s) { @@ -559,52 +586,121 @@ sm90_fp8_mega_moe_impl(void* y, for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { full_barriers[stage_idx]->wait(phase); - const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_0); - const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_1); + if (block_phase == sched::SM90BlockPhase::Linear1) { + const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_0); + const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_1); - #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) - ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { - auto desc_a = mma::sm90::make_smem_desc( - smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_b[stage_idx] + k * WGMMA::K, 1); - WGMMA::wgmma(desc_a, desc_b, accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) - ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_wait<0>(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); - empty_barrier_arrive(stage_idx); + empty_barrier_arrive(stage_idx); - #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { - const uint32_t col_0 = i * 8 + col_idx * 2; - const uint32_t col_1 = col_0 + 1; - const uint32_t n_0 = n_block_idx * BLOCK_N + col_0; - const uint32_t n_1 = n_block_idx * BLOCK_N + col_1; - const auto sf_base = l1_weights_sf_ptr + - local_expert_idx * kNumL1WeightSFGroupsN * kNumL1WeightSFGroupsK; - const float scale_b_0 = __ldg( - sf_base + get_l1_weight_sf_group(n_0) * kNumL1WeightSFGroupsK + k_block_idx); - const float scale_b_1 = __ldg( - sf_base + get_l1_weight_sf_group(n_1) * kNumL1WeightSFGroupsK + k_block_idx); - final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; - final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; - final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; - final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t col_0 = i * 8 + col_idx * 2; + const uint32_t col_1 = col_0 + 1; + const uint32_t n_0 = n_block_idx * BLOCK_N + col_0; + const uint32_t n_1 = n_block_idx * BLOCK_N + col_1; + const auto sf_base = l1_weights_sf_ptr + + local_expert_idx * kNumL1WeightSFGroupsN * kNumL1WeightSFGroupsK; + const float scale_b_0 = __ldg( + sf_base + get_l1_weight_sf_group(n_0) * kNumL1WeightSFGroupsK + k_block_idx); + const float scale_b_1 = __ldg( + sf_base + get_l1_weight_sf_group(n_1) * kNumL1WeightSFGroupsK + k_block_idx); + final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; + } + } else { + #pragma unroll + for (uint32_t half = 0; half < 2; ++ half) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < 2; ++ k) { + const uint32_t smem_k = (half * 2 + k) * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + smem_k, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_k, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + const uint32_t sf_group = k_block_idx * 2 + half; + const uint32_t sf_pool_token_idx_0 = pool_block_idx * SF_BLOCK_M + row_0; + const uint32_t sf_pool_token_idx_1 = pool_block_idx * SF_BLOCK_M + row_1; + const float scale_a_0 = __ldg( + l2_acts_sf_ptr + sf_group * kNumPaddedSFPoolTokens + sf_pool_token_idx_0); + const float scale_a_1 = __ldg( + l2_acts_sf_ptr + sf_group * kNumPaddedSFPoolTokens + sf_pool_token_idx_1); + const auto sf_base = l2_weights_sf_ptr + + local_expert_idx * kNumL2WeightSFGroupsN * kNumL2WeightSFGroupsK; + + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t col_0 = i * 8 + col_idx * 2; + const uint32_t col_1 = col_0 + 1; + const uint32_t n_group_0 = (n_block_idx * BLOCK_N + col_0) / 128; + const uint32_t n_group_1 = (n_block_idx * BLOCK_N + col_1) / 128; + const float scale_b_0 = __ldg(sf_base + n_group_0 * kNumL2WeightSFGroupsK + k_block_idx); + const float scale_b_1 = __ldg(sf_base + n_group_1 * kNumL2WeightSFGroupsK + k_block_idx); + final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; + } + } + + empty_barrier_arrive(stage_idx); } } - const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0; - const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1; - const uint32_t pool_token_idx_0 = pool_block_idx * BLOCK_M + row_0; - const uint32_t pool_token_idx_1 = pool_block_idx * BLOCK_M + row_1; + if (block_phase == sched::SM90BlockPhase::Linear2) { + const auto scatter_row = [&](const uint32_t& row, const uint32_t& pool_token_idx, const uint32_t& accum_offset) { + if (row >= valid_m) + return; + const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_token_idx); + const auto dst_token = combine_token_buffer.get_rank_buffer(src_metadata.topk_idx) + .get_data_buffer(src_metadata.token_idx); + auto dst_ptr = dst_token.get_base_ptr(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t out_col = n_block_idx * BLOCK_N + i * 8 + col_idx * 2; + *sym_buffer.map(dst_ptr + out_col + 0, src_metadata.rank_idx) = + __float2bfloat16(final_accum[i * 4 + accum_offset + 0]); + *sym_buffer.map(dst_ptr + out_col + 1, src_metadata.rank_idx) = + __float2bfloat16(final_accum[i * 4 + accum_offset + 1]); + } + }; + + scatter_row(row_0, pool_token_idx_0, 0); + scatter_row(row_1, pool_token_idx_1, 2); + return; + } + const float topk_weight_0 = row_0 < valid_m ? *l1_topk_weights_buffer.get_data_buffer(pool_token_idx_0).get_base_ptr() : 0.0f; const float topk_weight_1 = row_1 < valid_m ? diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index c048205..b82f25a 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -70,7 +70,8 @@ class SymmBuffer: self.l1_arrival_count, self.l2_arrival_mask, self.token_src_metadata, - self.l1_accum_debug) = buffer_views + self.l1_accum_debug, + self.combine_acts) = buffer_views else: (self.x, self.x_sf, self.topk_idx, self.topk_weights, @@ -82,6 +83,7 @@ class SymmBuffer: self.l2_arrival_mask = None self.token_src_metadata = None self.l1_accum_debug = None + self.combine_acts = None def destroy(self): self.handle = None @@ -101,6 +103,7 @@ class SymmBuffer: self.l2_arrival_mask = None self.token_src_metadata = None self.l1_accum_debug = None + self.combine_acts = None def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, diff --git a/megamoe_dev_test_scripts/phase1/interface_smoke.py b/megamoe_dev_test_scripts/phase1/interface_smoke.py index 3d71db1..e3e8c08 100644 --- a/megamoe_dev_test_scripts/phase1/interface_smoke.py +++ b/megamoe_dev_test_scripts/phase1/interface_smoke.py @@ -85,6 +85,8 @@ def main() -> None: assert buffer.l2_acts.shape[1] == args.intermediate_hidden assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 64 assert buffer.l2_acts_sf.dtype == torch.float32 + assert buffer.combine_acts.shape == (args.num_topk, buffer.num_max_tokens_per_rank, args.hidden) + assert buffer.combine_acts.dtype == torch.bfloat16 num_tokens = args.num_tokens buffer.x[:num_tokens].copy_(torch.randn((num_tokens, args.hidden), device='cuda').to(torch.float8_e4m3fn)) diff --git a/megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py b/megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py new file mode 100644 index 0000000..de3dcff --- /dev/null +++ b/megamoe_dev_test_scripts/phase5/l2_gemm_scatter_correctness.py @@ -0,0 +1,201 @@ +import argparse +import inspect +import os +import pathlib +import sys +from typing import 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 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', '8364')) + + 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 make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int): + torch.manual_seed(2345 + dist.get_rank()) + l1_weights = (torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn) + l2_weights = (torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn) + + l1_weights_sf = torch.empty( + (num_experts_per_rank, ceil_div(intermediate_hidden * 2, 128), hidden // 128), + dtype=torch.float32, device='cuda') + l2_weights_sf = torch.empty( + (num_experts_per_rank, ceil_div(hidden, 128), intermediate_hidden // 128), + dtype=torch.float32, device='cuda') + + for expert in range(num_experts_per_rank): + for n_group in range(l1_weights_sf.shape[1]): + for k_group in range(l1_weights_sf.shape[2]): + l1_weights_sf[expert, n_group, k_group] = 0.5 + 0.125 * expert + 0.0625 * n_group + 0.03125 * k_group + for n_group in range(l2_weights_sf.shape[1]): + for k_group in range(l2_weights_sf.shape[2]): + l2_weights_sf[expert, n_group, k_group] = 0.75 + 0.125 * expert + 0.0625 * n_group + 0.03125 * k_group + + return deep_gemm.transform_weights_for_mega_moe( + (l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf)) + + +def dequant_l2_acts(buffer: deep_gemm.SymmBuffer, + start: int, + count: int, + intermediate_hidden: int) -> torch.Tensor: + x = buffer.l2_acts[start:start + count, :intermediate_hidden].to(torch.float32) + out = torch.empty_like(x) + for sf_group in range(intermediate_hidden // 64): + col_start = sf_group * 64 + col_end = col_start + 64 + sf = buffer.l2_acts_sf[start:start + count, sf_group].to(torch.float32) + out[:, col_start:col_end] = x[:, col_start:col_end] * sf[:, None] + return out + + +def verify_l2_scatter(buffer: deep_gemm.SymmBuffer, + l2_weights: torch.Tensor, + l2_weights_sf: torch.Tensor, + hidden: int, + intermediate_hidden: int, + num_tokens: int, + atol: float, + rtol: float) -> None: + block_m = 128 + ref = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + counts = [int(v) & 0xffffffff for v in buffer.expert_recv_count_sum.to(torch.int64).cpu().tolist()] + pool_block_offset = 0 + expected_mask = (1 << (intermediate_hidden // 64)) - 1 + + for expert, count in enumerate(counts): + if count == 0: + continue + + pool_start = pool_block_offset * block_m + x = dequant_l2_acts(buffer, pool_start, count, intermediate_hidden) + expected = torch.zeros((count, hidden), dtype=torch.float32, device='cuda') + + for k_group in range(intermediate_hidden // 128): + k_start = k_group * 128 + k_end = k_start + 128 + w = l2_weights[expert, :, k_start:k_end].to(torch.float32) + partial = x[:, k_start:k_end] @ w.t() + for n_group in range(hidden // 128): + n_start = n_group * 128 + n_end = n_start + 128 + sfb = l2_weights_sf[expert, n_group, k_group].to(torch.float32) + expected[:, n_start:n_end] += partial[:, n_start:n_end] * sfb + + metadata = buffer.token_src_metadata[pool_start:pool_start + count].to(torch.int64) + for row in range(count): + rank_idx, token_idx, topk_idx = [int(v) for v in metadata[row].tolist()] + assert rank_idx == dist.get_rank() + assert topk_idx == 0 + ref[token_idx] = expected[row].to(torch.bfloat16) + + for block in range(ceil_div(count, block_m)): + mask = int(buffer.l2_arrival_mask[pool_block_offset + block].item()) + assert mask == expected_mask, (expert, block, mask, expected_mask) + + pool_block_offset += ceil_div(count, block_m) + + actual = buffer.combine_acts[0, :num_tokens, :hidden] + torch.testing.assert_close(actual.cpu(), ref.cpu(), rtol=rtol, atol=atol) + + +def run_case(args: argparse.Namespace, group: dist.ProcessGroup, rank_idx: int, num_ranks: int) -> None: + assert num_ranks == 1, 'Phase 5 milestone verifies multi-expert single-rank scatter first' + hidden = args.hidden + intermediate_hidden = args.intermediate_hidden + num_tokens = args.num_tokens + num_topk = 1 + num_experts = args.num_experts + num_experts_per_rank = num_experts // num_ranks + assert num_experts % num_ranks == 0 + + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, args.num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden) + weights = make_weights(num_experts_per_rank, hidden, intermediate_hidden) + + torch.manual_seed(6789 + rank_idx) + x = (torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda') * 0.25).to(torch.float8_e4m3fn) + x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') * 0.25 + 0.875 + topk_idx = (torch.arange(num_tokens, dtype=torch.long, device='cuda') // 128).reshape(num_tokens, 1) + topk_idx = torch.clamp(topk_idx, max=num_experts - 1) + topk_weights = torch.linspace(0.75, 1.25, num_tokens, dtype=torch.float32, device='cuda').reshape(num_tokens, 1) + + buffer.x[:num_tokens].copy_(x) + 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) + + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer, + activation_clamp=args.activation_clamp, + fast_math=False) + torch.cuda.synchronize() + + verify_l2_scatter(buffer, weights[1][0], weights[1][1], hidden, intermediate_hidden, + num_tokens, args.atol, args.rtol) + + dist.barrier(group=group) + if rank_idx == 0: + print('[PASSED] Phase 5 L2 GEMM scatter correctness', flush=True) + buffer.destroy() + + +def main() -> None: + parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 5 L2 GEMM scatter correctness') + parser.add_argument('--num-tokens', type=int, default=256) + parser.add_argument('--num-max-tokens-per-rank', type=int, default=384) + parser.add_argument('--hidden', type=int, default=256) + parser.add_argument('--intermediate-hidden', type=int, default=128) + parser.add_argument('--num-experts', type=int, default=2) + parser.add_argument('--activation-clamp', type=float, default=None) + parser.add_argument('--local-rank', type=int, default=None) + parser.add_argument('--atol', type=float, default=3e-2) + parser.add_argument('--rtol', type=float, default=5e-2) + args = parser.parse_args() + + rank_idx, num_ranks, group = init_test_dist(args.local_rank) + assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9 + assert args.num_tokens == 256 + assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0 + + run_case(args, group, rank_idx, num_ranks) + dist.destroy_process_group() + + +if __name__ == '__main__': + main()