diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index ef21327..203e39e 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -15,7 +15,8 @@ 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; @@ -85,10 +86,16 @@ get_symm_buffer_size_for_sm90_mega_moe( fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens, l2_token_buffer.get_end_ptr()); + // Phase 3 verification output: one scaled FP32 L1 accumulator tile. + const 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()); + // Combine input buffer const auto combine_token_buffer = layout::Buffer( bf16_token_layout, num_topk, num_max_tokens_per_rank, - l2_sf_buffer.get_end_ptr()); + l1_accum_debug_buffer.get_end_ptr()); // Check SF buffer requirements DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); @@ -149,10 +156,15 @@ get_symm_buffer_size_for_sm90_mega_moe( reinterpret_cast(runtime_workspace.get_token_src_metadata_ptr()), {num_max_pool_tokens, 3}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())); + auto l1_accum_debug = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_accum_debug_buffer.base)), + {128, 128}, + torch::TensorOptions().dtype(torch::kFloat32).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, token_src_metadata); + expert_recv_count_sum, l1_arrival_count, token_src_metadata, + l1_accum_debug); }; return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; } @@ -228,7 +240,8 @@ static void fp8_mega_moe( 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); + expert_recv_count_sum, l1_arrival_count, token_src_metadata, + l1_accum_debug] = slice(sym_buffer); // Dispatch into SM90 path DG_HOST_ASSERT(arch_major == 9); @@ -237,6 +250,7 @@ static void fp8_mega_moe( l2_acts, l2_acts_sf, l1_weights, l2_weights, l1_weights_sf, l2_weights_sf, + l1_accum_debug, cumulative_local_expert_recv_stats, sym_buffer_ptrs, rank_idx, num_max_tokens_per_rank, diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index c826118..6043cca 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -29,6 +29,7 @@ public: // Runtime arguments void* y; + void* l1_accum_debug; int* cumulative_local_expert_recv_stats; int num_tokens; layout::SymBuffer<> sym_buffer_ptrs; @@ -101,7 +102,8 @@ static void __instantiate_kernel() {{ args.tensor_map_l2_acts, args.tensor_map_l2_acts_sf, args.tensor_map_l2_weights, - args.l2_weights_sf + args.l2_weights_sf, + args.l1_accum_debug )); } }; @@ -112,6 +114,7 @@ static void sm90_fp8_mega_moe( const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf, const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf, + const torch::Tensor& l1_accum_debug, const std::optional cumulative_local_expert_recv_stats, const std::vector& sym_buffer_ptrs, const int& rank_idx, const int& num_max_tokens_per_rank, @@ -184,6 +187,7 @@ static void sm90_fp8_mega_moe( .fast_math = fast_math, .config = config, .y = y.data_ptr(), + .l1_accum_debug = l1_accum_debug.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, .num_tokens = num_tokens, .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), 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 cd8a27c..6bd9d21 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 @@ -1,16 +1,22 @@ #pragma once +#include #include #include #include +#include #include #include +#include #include #include #include +#include #include +#include #include +#include #include namespace deep_gemm { @@ -43,44 +49,63 @@ 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) { + 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"); DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); #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(kNumDispatchThreads == 64, "SM90 MegaMoE expects 64 dispatch threads"); + DG_STATIC_ASSERT(kNumTMAThreads == 64, "SM90 MegaMoE expects 64 TMA threads"); + DG_STATIC_ASSERT(kNumMathThreads == 256, "SM90 MegaMoE expects 256 math 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"); + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using a_dtype_t = __nv_fp8_e4m3; + using b_dtype_t = __nv_fp8_e4m3; + using WGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32; + constexpr uint32_t kNumTMAWarps = kNumTMAThreads / 32; + constexpr uint32_t kNumMathWarps = kNumMathThreads / 32; + constexpr uint32_t kMathWarpStart = kNumDispatchWarps + kNumTMAWarps; 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; + constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t); + 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; DG_STATIC_ASSERT(kNumTokensPerWarp > 0, "Invalid number of top-k experts"); DG_STATIC_ASSERT(kNumPaddedSFPoolTokens % SF_BLOCK_M == 0, "Invalid padded SF pool size"); + DG_STATIC_ASSERT(BLOCK_N == WGMMA::N and BLOCK_K % WGMMA::K == 0, "Invalid WGMMA tile shape"); - // 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; + if (warp_idx == kNumDispatchWarps and cute::elect_one_sync()) { + 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); + } + __syncwarp(); + + constexpr uint32_t kDispatchBarrierIdx = 1; constexpr uint32_t kDispatchGridSyncIndex = 0; constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 1; constexpr uint32_t kBeforeDispatchPullBarrierTag = 2; @@ -131,23 +156,78 @@ sm90_fp8_mega_moe_impl(void* y, 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; + constexpr uint32_t SMEM_EXPERT_COUNT_SIZE = + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment); + constexpr uint32_t SMEM_SEND_BUFFER_SIZE = + math::constexpr_align(fp8_token_layout.get_num_bytes() * kNumDispatchWarps, kSharedMemoryAlignment); + constexpr uint32_t SMEM_CD_L1_SIZE = + (kCooperativeMode ? 2u : 1u) * kStoreBlockM * (BLOCK_N / 2) * 2; + constexpr uint32_t SMEM_CD_L2_SIZE = BLOCK_M * BLOCK_N * static_cast(sizeof(nv_bfloat16)); + constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE; + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(a_dtype_t); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(b_dtype_t); + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(float); + DG_STATIC_ASSERT(SMEM_A_SIZE_PER_STAGE % kSharedMemoryAlignment == 0, "Invalid A smem alignment"); + DG_STATIC_ASSERT(SMEM_B_SIZE_PER_STAGE % kSharedMemoryAlignment == 0, "Invalid B smem alignment"); + + const auto smem_expert_count = reinterpret_cast(smem_buffer); + auto smem_gemm_base = math::advance_ptr( + smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE); + auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); + }); + auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE) + i * SMEM_SFA_SIZE_PER_STAGE); + }); + auto barrier_start_ptr = reinterpret_cast(math::advance_ptr( + smem_gemm_base, + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE))); + auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { + return barrier_start_ptr + i; + }); + auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { + return barrier_start_ptr + kNumStages + i; + }); + + if constexpr (BLOCK_M == 128) { + if (sm_idx == 0) { + auto debug_ptr = reinterpret_cast(l1_accum_debug); + for (uint32_t i = thread_idx; i < 128 * 128; i += kNumThreads) + debug_ptr[i] = 0.0f; + } + + if (warp_idx == kNumDispatchWarps + 1 and cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(kNumMathWarps); + } + cutlass::arch::fence_barrier_init(); + } + __syncthreads(); } - comm::nvlink_barrier( - workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync); + + if (thread_idx < kNumDispatchThreads) { + // 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; @@ -358,6 +438,181 @@ sm90_fp8_mega_moe_impl(void* y, } __syncwarp(); } + } else if (thread_idx < kNumDispatchThreads + kNumTMAThreads) { + if constexpr (BLOCK_M == 128) { + if (warp_idx == kNumDispatchWarps) { + 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); + + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + scheduler.for_each_block([&](const sched::SM90BlockPhase& block_phase, + const uint32_t& local_expert_idx, + 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); + + #pragma unroll 1 + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + if (cute::elect_one_sync()) { + 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); + } + __syncwarp(); + } + }); + } + } + } else { + if constexpr (BLOCK_M == 128) { + const uint32_t math_warp_idx = warp_idx - kMathWarpStart; + const uint32_t math_wg_idx = math_warp_idx / 4; + const uint32_t warp_idx_in_wg = math_warp_idx % 4; + const uint32_t row_idx = lane_idx / 4; + const uint32_t col_idx = lane_idx % 4; + 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); + auto l1_accum_debug_ptr = reinterpret_cast(l1_accum_debug); + + const auto get_l1_weight_sf_group = [](const uint32_t& interleaved_n) { + constexpr uint32_t kInterleaveGran = 8; + const uint32_t pair_group = interleaved_n / (2 * kInterleaveGran); + const uint32_t offset = interleaved_n - pair_group * (2 * kInterleaveGran); + const uint32_t natural_n = offset < kInterleaveGran ? + pair_group * kInterleaveGran + offset : + (L1_SHAPE_N / 2) + pair_group * kInterleaveGran + offset - kInterleaveGran; + return natural_n / 128; + }; + + 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); + + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + scheduler.for_each_block([&](const sched::SM90BlockPhase& block_phase, + const uint32_t& local_expert_idx, + 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 valid_m = scheduler.get_valid_m(); + float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0}; + + const auto empty_barrier_arrive = [&](const uint32_t& s) { + if (lane_idx == 0) + empty_barriers[s]->arrive(); + }; + + #pragma unroll 1 + 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); + + #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); + + #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]; + } + } + + if (local_expert_idx == 0 and m_block_idx == 0 and n_block_idx == 0) { + const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0; + const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t col = i * 8 + col_idx * 2; + if (row_0 < valid_m) { + l1_accum_debug_ptr[row_0 * BLOCK_N + col + 0] = final_accum[i * 4 + 0]; + l1_accum_debug_ptr[row_0 * BLOCK_N + col + 1] = final_accum[i * 4 + 1]; + } + if (row_1 < valid_m) { + l1_accum_debug_ptr[row_1 * BLOCK_N + col + 0] = final_accum[i * 4 + 2]; + l1_accum_debug_ptr[row_1 * BLOCK_N + col + 1] = final_accum[i * 4 + 3]; + } + } + } + }); + } + } #endif return; } diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index b25c629..428bb7c 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -68,7 +68,8 @@ class SymmBuffer: self.l2_acts, self.l2_acts_sf, self.expert_recv_count_sum, self.l1_arrival_count, - self.token_src_metadata) = buffer_views + self.token_src_metadata, + self.l1_accum_debug) = buffer_views else: (self.x, self.x_sf, self.topk_idx, self.topk_weights, @@ -78,6 +79,7 @@ class SymmBuffer: self.expert_recv_count_sum = None self.l1_arrival_count = None self.token_src_metadata = None + self.l1_accum_debug = None def destroy(self): self.handle = None @@ -95,6 +97,7 @@ class SymmBuffer: self.expert_recv_count_sum = None self.l1_arrival_count = None self.token_src_metadata = None + self.l1_accum_debug = None def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, diff --git a/megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py b/megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py new file mode 100644 index 0000000..823fa8a --- /dev/null +++ b/megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py @@ -0,0 +1,169 @@ +import argparse +import inspect +import os +import pathlib +import random +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', '8362')) + + 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 interleaved_to_natural_n(n: torch.Tensor, half_n: int, gran: int = 8) -> torch.Tensor: + pair_group = n // (2 * gran) + offset = n - pair_group * (2 * gran) + gate_n = pair_group * gran + offset + up_n = half_n + pair_group * gran + offset - gran + return torch.where(offset < gran, gate_n, up_n) + + +def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int): + torch.manual_seed(2027 + dist.get_rank()) + l1_weights = (torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.float32, device='cuda') * 0.5).to(torch.float8_e4m3fn) + l2_weights = (torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.float32, device='cuda') * 0.5).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') + 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.75 + 0.125 * n_group + 0.0625 * k_group + + 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 reference_l1_accum(buffer: deep_gemm.SymmBuffer, + l1_weights: torch.Tensor, + l1_weights_sf: torch.Tensor, + hidden: int, + intermediate_hidden: int) -> torch.Tensor: + block_m = 128 + block_n = 128 + x_fp8 = buffer.l1_acts[:block_m] + x_sf = buffer.l1_acts_sf[:block_m, :hidden // 128] + w_fp8 = l1_weights[0, :block_n] + + n = torch.arange(block_n, device='cuda') + natural_n = interleaved_to_natural_n(n, intermediate_hidden) + n_groups = natural_n // 128 + + expected = torch.zeros((block_m, block_n), dtype=torch.float32, device='cuda') + for k_group in range(hidden // 128): + k_start = k_group * 128 + k_end = k_start + 128 + x = x_fp8[:, k_start:k_end].to(torch.float32) + w = w_fp8[:, k_start:k_end].to(torch.float32) + sfa = x_sf[:, k_group].to(torch.float32) + sfb = l1_weights_sf[0, n_groups, k_group].to(torch.float32) + expected += (x @ w.t()) * sfa[:, None] * sfb[None, :] + return expected + + +def run_case(args: argparse.Namespace, group: dist.ProcessGroup, rank_idx: int, num_ranks: int) -> None: + hidden = args.hidden + intermediate_hidden = args.intermediate_hidden + num_tokens = args.num_tokens + num_topk = 1 + num_experts = args.num_experts if args.num_experts is not None else num_ranks + num_experts_per_rank = num_experts // num_ranks + assert num_experts % num_ranks == 0 + assert num_experts_per_rank >= 1 + + 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(1234 + rank_idx) + x = (torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda') * 0.5).to(torch.float8_e4m3fn) + x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') + 0.75 + local_global_expert = rank_idx * num_experts_per_rank + topk_idx = torch.full((num_tokens, num_topk), local_global_expert, dtype=torch.long, device='cuda') + topk_weights = torch.ones((num_tokens, num_topk), dtype=torch.float32, device='cuda') + + 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) + torch.cuda.synchronize() + + expected = reference_l1_accum(buffer, weights[0][0], weights[0][1], hidden, intermediate_hidden) + actual = buffer.l1_accum_debug + torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=args.rtol, atol=args.atol) + + dist.barrier(group=group) + if rank_idx == 0: + max_diff = (actual - expected).abs().max().item() + print(f'[PASSED] Phase 3 L1 WGMMA single tile, max_diff={max_diff:.6f}', flush=True) + buffer.destroy() + + +def main() -> None: + parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 3 L1 WGMMA single-tile correctness') + parser.add_argument('--num-tokens', type=int, default=128) + 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=None) + parser.add_argument('--local-rank', type=int, default=None) + parser.add_argument('--atol', type=float, default=1e-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 == 128 + assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0 + + random.seed(5678 + rank_idx) + run_case(args, group, rank_idx, num_ranks) + dist.destroy_process_group() + + +if __name__ == '__main__': + main()