From 6be80da7119ed1a978a46934429dc798129cbd47 Mon Sep 17 00:00:00 2001 From: yinding Date: Wed, 17 Jun 2026 10:06:23 +0800 Subject: [PATCH 1/5] Enable SM90 FP8 MegaMoE swapAB Add swapAB code paths for small-batch SM90 FP8 MegaMoE and remove ptxas C7510 sources from hot device code. (cherry picked from commit 00749387fcb3379332b9f8f883da5f6269c9aff8) --- csrc/jit_kernels/heuristics/sm90_mega_moe.hpp | 39 +- csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp | 9 +- .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 564 ++++++++++++++---- 3 files changed, 484 insertions(+), 128 deletions(-) diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp index 9b9c8d9f3..537280427 100644 --- a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -84,11 +84,22 @@ static int get_num_experts_per_wave_for_mega_moe_sm90( num_ring_tokens, num_max_tokens_per_rank, num_ranks); } +static bool should_use_swap_ab_for_mega_moe_sm90( + const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, + const int& block_m, const int& num_epilogue_threads) { + const float expected_tokens_per_expert = + static_cast(num_tokens) * num_topk / num_experts_per_rank; + const bool decode_split_n_path = + block_m == 64 and num_epilogue_threads == 256; + return decode_split_n_path and num_tokens <= 128 and expected_tokens_per_expert > 0.0f; +} + static std::pair get_pipeline_config_for_mega_moe_sm90( const int& smem_capacity, const int& num_experts, const int& hidden, const int& block_m, const int& block_n, const int& block_k, - const int& num_dispatch_warps, const int& num_epilogue_warps) { + const int& num_dispatch_warps, const int& num_epilogue_warps, + const bool& use_swap_ab = false) { constexpr int kSmemAlignment = 1024; const int smem_expert_count_size = align( @@ -100,7 +111,13 @@ static std::pair get_pipeline_config_for_mega_moe_sm90( const int smem_cd_l1 = block_m * (block_n / 2); const int smem_cd_l2 = block_m * block_n * static_cast(sizeof(nv_bfloat16)); - const int smem_cd = align(std::max(smem_cd_l1, smem_cd_l2), kSmemAlignment); + const int smem_cd_swap_l1 = use_swap_ab + ? block_m * (block_n / 2) * + (static_cast(sizeof(float)) + static_cast(sizeof(uint8_t))) + : 0; + const int smem_cd = align( + std::max(std::max(smem_cd_l1, smem_cd_l2), smem_cd_swap_l1), + kSmemAlignment); const int smem_sfa_per_stage = align(2 * block_m * static_cast(sizeof(float)), 128); const int smem_sfb_per_stage = 0; @@ -134,9 +151,13 @@ static MegaMoESM90Config get_mega_moe_config_sm90( const bool decode_use_block_n_256 = decode_split_n_path and intermediate_hidden >= 3072 and expected_tokens_per_expert >= 0.25f and - (2 * intermediate_hidden) % 256 == 0; - const int block_n = auto_split_mn ? 256 - : (decode_use_block_n_256 ? 256 : 128); + (2 * intermediate_hidden) % 256 == 0 and hidden % 256 == 0; + const bool use_swap_ab = should_use_swap_ab_for_mega_moe_sm90( + num_experts_per_rank, num_tokens, num_topk, + block_m, num_epilogue_threads); + int block_n = use_swap_ab ? 128 + : (auto_split_mn ? 256 : + (decode_use_block_n_256 ? 256 : 128)); const int block_k = 128; const int cluster_size = 1; const int num_max_pool_tokens = layout::get_num_max_pool_tokens( @@ -166,7 +187,8 @@ static MegaMoESM90Config get_mega_moe_config_sm90( SM90ArchSpec::smem_capacity, num_experts, hidden, block_m, block_n, block_k, - num_dispatch_threads / 32, num_epilogue_threads / 32); + num_dispatch_threads / 32, num_epilogue_threads / 32, + use_swap_ab); const auto config = MegaMoESM90Config { block_m, block_n, block_k, @@ -180,8 +202,9 @@ static MegaMoESM90Config get_mega_moe_config_sm90( if (get_env("DG_JIT_DEBUG") or get_env("DG_PRINT_CONFIGS")) { const auto key = fmt::format( - "MegaMoESM90Config(num_ranks={}, num_experts={}, hidden={}, intermediate_hidden={}, num_max_tokens_per_rank={}, num_tokens={}, num_topk={})", - num_ranks, num_experts, hidden, intermediate_hidden, num_max_tokens_per_rank, num_tokens, num_topk); + "MegaMoESM90Config(num_ranks={}, num_experts={}, hidden={}, intermediate_hidden={}, num_max_tokens_per_rank={}, num_tokens={}, num_topk={}, swap_ab={})", + num_ranks, num_experts, hidden, intermediate_hidden, num_max_tokens_per_rank, num_tokens, num_topk, + use_swap_ab); static std::unordered_set printed; if (printed.count(key) == 0) { std::cout << key << ": " << config << std::endl; diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index a86ebe6f9..06e33d027 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -44,6 +44,7 @@ class SM90FP8MegaMoERuntime final : public LaunchRuntime bool l2_arrival_counter; bool l2_epilogue_requires_full_sync; bool split_phase_hot_path; + bool use_swap_ab; MegaMoESM90Config config; // Runtime arguments @@ -93,6 +94,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, + {}, {} >); }}; @@ -113,7 +115,8 @@ static void __instantiate_kernel() {{ args.reuse_accum_as_final ? "true" : "false", args.l2_arrival_counter ? "true" : "false", args.l2_epilogue_requires_full_sync ? "true" : "false", - args.split_phase_hot_path ? "true" : "false"); + args.split_phase_hot_path ? "true" : "false", + args.use_swap_ab ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -187,6 +190,9 @@ static void sm90_fp8_mega_moe( default_split_mn_barrier_opt or decode_l2_counter; const bool l2_epilogue_requires_full_sync = not l2_arrival_counter; + const bool use_swap_ab = should_use_swap_ab_for_mega_moe_sm90( + num_experts_per_rank, num_tokens, num_topk, + config.block_m, config.num_epilogue_threads); // Tensormap construction // Acts/weights: standard 2D TMA descriptors (FP8 K-major). @@ -280,6 +286,7 @@ static void sm90_fp8_mega_moe( .l2_arrival_counter = l2_arrival_counter, .l2_epilogue_requires_full_sync = l2_epilogue_requires_full_sync, .split_phase_hot_path = split_phase_hot_path, + .use_swap_ab = use_swap_ab, .config = config, .y = y.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, 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 a025815de..74bac96a3 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 @@ -113,13 +113,13 @@ CUTLASS_DEVICE void sm90_fp8_mega_moe_for_each_block_split( // * Math warpgroups (totalling kNumEpilogueThreads) consume each // stage with WGMMA, accumulate into registers, then run the epilogue: // - L1 (Linear1): SwiGLU with gate/up granularity-8 interleaved layout, -// per-row amax over each output-SF group, FP8 e4m3 quantize, STSM into -// SMEM, TMA store to local L1 output buffer. +// per-row amax over each output-SF group, FP8 e4m3 quantize, stage +// through SMEM, then TMA store to local L1 output buffer. // The per-row SF is written as a *float* into the L2-acts SF buffer at // per-64 K granularity (one SF per L1 N block), so each block is fully // self-contained and no cross-CTA amax synchronisation is needed. -// - L2 (Linear2): BF16 cast of the GEMM output, STSM into SMEM, then -// NVLink scatter to remote combine buffers. +// - L2 (Linear2): BF16 cast of the GEMM output, stage through SMEM, +// then NVLink scatter to remote combine buffers. // * After all GEMM blocks, the math warps run the COMBINE step (top-k // reduction in BF16) — ported verbatim from the SM100 kernel. // ============================================================================ @@ -143,6 +143,7 @@ template < bool kL2ArrivalCounter, bool kL2EpilogueRequiresFullSync, bool kSplitPhaseHotPath, + bool kFP8SwapAB = false, uint32_t L1_SHAPE_N = kIntermediateHidden * 2, uint32_t L1_SHAPE_K = kHidden, uint32_t L2_SHAPE_N = kHidden, @@ -289,6 +290,13 @@ sm90_fp8_mega_moe_impl(void* y, // (k_sf_idx == n_block_idx) instead of one per warpgroup. The amax that // feeds that shared SF must be reduced across both warpgroups. constexpr bool kSplitNSharesSF = kSplitNWarpgroups and (WG_L1_OUT_BLOCK_N < 64); + constexpr bool kSwapABEligible = + kFP8SwapAB and kSplitNWarpgroups and (BLOCK_M == 64) and (BLOCK_N == 128) and + (kWarpgroupSplitN == 2); + constexpr bool kSwapABActive = kSwapABEligible; + constexpr uint32_t kSwapABTokenChunks = BLOCK_M / 8; + DG_STATIC_ASSERT(not kSwapABEligible or (BLOCK_M % 8 == 0), + "swapAB epilogue token chunks assume BLOCK_M is a multiple of 8"); constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t); // 128 constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t); // 128 constexpr uint32_t kSwizzleCDMode = 128; @@ -314,14 +322,23 @@ sm90_fp8_mega_moe_impl(void* y, // warpgroup, so no SMEM is needed. constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = 0; - // CD output: max of L1 FP8 (BLOCK_M * (BLOCK_N/2) * 1 byte) and - // L2 BF16 (BLOCK_M * BLOCK_N * 2 bytes). Split-M warpgroups own disjoint - // row slices; shared-SF split-N warpgroups stage disjoint column slices - // into one CTA tile. + // CD output: max of L1 FP8 (BLOCK_M * (BLOCK_N/2) * 1 byte), L2 BF16 + // (BLOCK_M * BLOCK_N * 2 bytes), and the swapAB L1 FP32+FP8 staging + // buffers. Split-M warpgroups own disjoint row slices; shared-SF split-N + // warpgroups stage disjoint column slices into one CTA tile. constexpr uint32_t SMEM_CD_L1_SIZE = BLOCK_M * L1_OUT_BLOCK_N * sizeof(cutlass::float_e4m3_t); constexpr uint32_t SMEM_CD_L2_SIZE = BLOCK_M * BLOCK_N * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_CD_SWAP_L1_FP32_SIZE = + kSwapABActive ? BLOCK_M * L1_OUT_BLOCK_N * sizeof(float) : 0; + constexpr uint32_t SMEM_CD_SWAP_L1_FP8_SIZE = + kSwapABActive ? BLOCK_M * L1_OUT_BLOCK_N * sizeof(cutlass::float_e4m3_t) : 0; + constexpr uint32_t SMEM_CD_SWAP_L1_SIZE = + kSwapABActive ? (SMEM_CD_SWAP_L1_FP32_SIZE + SMEM_CD_SWAP_L1_FP8_SIZE) : 0; + constexpr uint32_t SMEM_CD_BASE_SIZE = + SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE; constexpr uint32_t SMEM_CD_SIZE = math::constexpr_align( - SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE, kSharedMemoryAlignment); + SMEM_CD_BASE_SIZE > SMEM_CD_SWAP_L1_SIZE ? SMEM_CD_BASE_SIZE : SMEM_CD_SWAP_L1_SIZE, + kSharedMemoryAlignment); constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE = SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_CD_SIZE + @@ -339,6 +356,9 @@ sm90_fp8_mega_moe_impl(void* y, // CD output is shared by L1 (FP8) and L2 (BF16); reinterpret-cast as needed. auto smem_cd_l1 = reinterpret_cast(smem_gemm_base); auto smem_cd_l2 = reinterpret_cast(smem_gemm_base); + auto smem_cd_swap_l1_fp32 = reinterpret_cast(smem_gemm_base); + auto smem_cd_swap_l1_fp8 = reinterpret_cast( + math::advance_ptr(smem_gemm_base, SMEM_CD_SWAP_L1_FP32_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); @@ -1338,93 +1358,252 @@ sm90_fp8_mega_moe_impl(void* y, } if (block_phase == sched::BlockPhase::Linear1) { - // Single per-128 K-block WGMMA group - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread; ++ 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] + smem_a_wg_offset + k * WGMMA::K, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); - WGMMA::wgmma(desc_a, desc_b, accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_wait<0>(); + if constexpr (kSwapABActive) { + auto run_swap_ab_l1 = [&]() { + using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; + float swap_accum[kSwapAccum]; + + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++ i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / SwapWGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * SwapWGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++ i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++ i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + const float scale_0 = token_0 < valid_m ? + ptx::ld_shared(smem_sfa[stage_idx] + token_0) : 0.0f; + const float scale_1 = token_1 < valid_m ? + ptx::ld_shared(smem_sfa[stage_idx] + token_1) : 0.0f; + final_accum[i * 4 + 0] += scale_0 * gate_sf * swap_accum[i * 4 + 0]; + final_accum[i * 4 + 2] += scale_0 * up_sf * swap_accum[i * 4 + 2]; + final_accum[i * 4 + 1] += scale_1 * gate_sf * swap_accum[i * 4 + 1]; + final_accum[i * 4 + 3] += scale_1 * up_sf * swap_accum[i * 4 + 3]; + } + }; + + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if constexpr (kIntermediateHidden <= 2048) { + if (n_swap <= 8) { + run_swap_ab_l1.template operator()<8>(); + } else if (n_swap <= 16) { + run_swap_ab_l1.template operator()<16>(); + } else if (n_swap <= 32) { + run_swap_ab_l1.template operator()<32>(); + } else { + run_swap_ab_l1.template operator()<64>(); + } + } else { + switch (n_swap) { + case 8: run_swap_ab_l1.template operator()<8>(); break; + case 16: run_swap_ab_l1.template operator()<16>(); break; + case 24: run_swap_ab_l1.template operator()<24>(); break; + case 32: run_swap_ab_l1.template operator()<32>(); break; + case 40: run_swap_ab_l1.template operator()<40>(); break; + case 48: run_swap_ab_l1.template operator()<48>(); break; + case 56: run_swap_ab_l1.template operator()<56>(); break; + default: run_swap_ab_l1.template operator()<64>(); break; + } + } + } else { + // Single per-128 K-block WGMMA group + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ 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] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); - if (lane_idx == 0) - empty_barriers[stage_idx]->arrive(); + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); - // L1: gate/up alternate at gran=8 along N; each `i` block of 8 - // cols belongs entirely to one of {gate, up}, so .x and .y - // share the same scalar. - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { - const float sb = (i & 1u) ? up_sf : gate_sf; - final_accum[i*4+0] += scale_a_0_lo * sb * accum[i*4+0]; - final_accum[i*4+1] += scale_a_0_lo * sb * accum[i*4+1]; - final_accum[i*4+2] += scale_a_1_lo * sb * accum[i*4+2]; - final_accum[i*4+3] += scale_a_1_lo * sb * accum[i*4+3]; + // L1: gate/up alternate at gran=8 along N; each `i` block of 8 + // cols belongs entirely to one of {gate, up}, so .x and .y + // share the same scalar. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float sb = (i & 1u) ? up_sf : gate_sf; + final_accum[i*4+0] += scale_a_0_lo * sb * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_lo * sb * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_lo * sb * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_lo * sb * accum[i*4+3]; + } } } else { - // L2: split BLOCK_K=128 into two halves (per-64 SFA), each 2 WGMMAs. - // First half: K=0..63, SFA = scale_a_*_lo - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { - auto desc_a = mma::sm90::make_smem_desc( - smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); - WGMMA::wgmma(desc_a, desc_b, accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_wait<0>(); + if constexpr (kSwapABActive) { + DG_STATIC_ASSERT(kL2ActsSFGranK == 64, + "L2 swapAB assumes per-64 activation scales"); + auto run_swap_ab_l2 = [&]() { + using SwapWGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kSwapAccum = SwapWGMMA::kNumAccum; + float swap_accum[kSwapAccum]; + + auto promote_swap_accum = [&](const uint32_t& sf_group) { + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum / 4; ++ i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + const float scale_0 = token_0 < valid_m ? + ptx::ld_shared(smem_sfa[stage_idx] + sf_group * BLOCK_M + token_0) : 0.0f; + const float scale_1 = token_1 < valid_m ? + ptx::ld_shared(smem_sfa[stage_idx] + sf_group * BLOCK_M + token_1) : 0.0f; + final_accum[i * 4 + 0] += scale_0 * l2_sf * swap_accum[i * 4 + 0]; + final_accum[i * 4 + 2] += scale_0 * l2_sf * swap_accum[i * 4 + 2]; + final_accum[i * 4 + 1] += scale_1 * l2_sf * swap_accum[i * 4 + 1]; + final_accum[i * 4 + 3] += scale_1 * l2_sf * swap_accum[i * 4 + 3]; + } + }; + + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++ i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * SwapWGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k * SwapWGMMA::K, 1); + SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++ i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_wait<0>(); + promote_swap_accum(0); + + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++ i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / SwapWGMMA::K; ++ k) { + const uint32_t k_off = (BLOCK_K / 2) + k * SwapWGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_a[stage_idx] + k_off, 1); + SwapWGMMA::wgmma(desc_a, desc_b, swap_accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kSwapAccum; ++ i) + ptx::warpgroup_fence_operand(swap_accum[i]); + ptx::warpgroup_wait<0>(); + promote_swap_accum(1); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + }; + + const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; + if constexpr (kIntermediateHidden <= 2048) { + if (n_swap <= 8) { + run_swap_ab_l2.template operator()<8>(); + } else if (n_swap <= 16) { + run_swap_ab_l2.template operator()<16>(); + } else if (n_swap <= 32) { + run_swap_ab_l2.template operator()<32>(); + } else { + run_swap_ab_l2.template operator()<64>(); + } + } else { + switch (n_swap) { + case 8: run_swap_ab_l2.template operator()<8>(); break; + case 16: run_swap_ab_l2.template operator()<16>(); break; + case 24: run_swap_ab_l2.template operator()<24>(); break; + case 32: run_swap_ab_l2.template operator()<32>(); break; + case 40: run_swap_ab_l2.template operator()<40>(); break; + case 48: run_swap_ab_l2.template operator()<48>(); break; + case 56: run_swap_ab_l2.template operator()<56>(); break; + default: run_swap_ab_l2.template operator()<64>(); break; + } + } + } else { + // L2: split BLOCK_K=128 into two halves (per-64 SFA), each 2 WGMMAs. + // First half: K=0..63, SFA = scale_a_*_lo + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); - // L2 first half: single scalar `l2_sf` broadcast across N. - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { - final_accum[i*4+0] += scale_a_0_lo * l2_sf * accum[i*4+0]; - final_accum[i*4+1] += scale_a_0_lo * l2_sf * accum[i*4+1]; - final_accum[i*4+2] += scale_a_1_lo * l2_sf * accum[i*4+2]; - final_accum[i*4+3] += scale_a_1_lo * l2_sf * accum[i*4+3]; - } + // L2 first half: single scalar `l2_sf` broadcast across N. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] += scale_a_0_lo * l2_sf * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_lo * l2_sf * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_lo * l2_sf * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_lo * l2_sf * accum[i*4+3]; + } - // Second half: K=64..127, SFA = scale_a_*_hi - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { - const uint32_t k_off = (BLOCK_K / 2) + k * WGMMA::K; - auto desc_a = mma::sm90::make_smem_desc( - smem_a[stage_idx] + smem_a_wg_offset + k_off, 1); - auto desc_b = mma::sm90::make_smem_desc( - smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); - WGMMA::wgmma(desc_a, desc_b, accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_wait<0>(); + // Second half: K=64..127, SFA = scale_a_*_hi + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + const uint32_t k_off = (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); - if (lane_idx == 0) - empty_barriers[stage_idx]->arrive(); + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); - // L2 second half: same broadcast scalar `l2_sf`. - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { - final_accum[i*4+0] += scale_a_0_hi * l2_sf * accum[i*4+0]; - final_accum[i*4+1] += scale_a_0_hi * l2_sf * accum[i*4+1]; - final_accum[i*4+2] += scale_a_1_hi * l2_sf * accum[i*4+2]; - final_accum[i*4+3] += scale_a_1_hi * l2_sf * accum[i*4+3]; + // L2 second half: same broadcast scalar `l2_sf`. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] += scale_a_0_hi * l2_sf * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_hi * l2_sf * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_hi * l2_sf * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_hi * l2_sf * accum[i*4+3]; + } } } } @@ -1443,6 +1622,122 @@ sm90_fp8_mega_moe_impl(void* y, } if (block_phase == sched::BlockPhase::Linear1) { + if constexpr (kSwapABActive) { + auto silu = [](float x) -> float { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; + }; + auto clamp_gate = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + }; + auto clamp_up = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + }; + + const uint32_t out_col_base = + wg_l1_out_n_offset + warp_idx_in_wg * 8 + row_idx; + auto store_l1_swap_chunk = [&](const uint32_t& i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + if (token_0 < valid_m) { + float g0 = final_accum[i * 4 + 0]; + float u0 = final_accum[i * 4 + 2]; + clamp_gate(g0); + clamp_up(u0); + const float weight_0 = *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_0) + .get_base_ptr(); + smem_cd_swap_l1_fp32[token_0 * L1_OUT_BLOCK_N + out_col_base] = + silu(g0) * u0 * weight_0; + } + if (token_1 < valid_m) { + float g1 = final_accum[i * 4 + 1]; + float u1 = final_accum[i * 4 + 3]; + clamp_gate(g1); + clamp_up(u1); + const float weight_1 = *l1_topk_weights_buffer + .get_data_buffer(m_idx + token_1) + .get_base_ptr(); + smem_cd_swap_l1_fp32[token_1 * L1_OUT_BLOCK_N + out_col_base] = + silu(g1) * u1 * weight_1; + } + }; + + const uint32_t num_swap_token_chunks = (valid_m + 7u) / 8u; + store_l1_swap_chunk(0); + if (valid_m > 8) { + #pragma unroll + for (uint32_t i = 1; i < kSwapABTokenChunks; ++ i) { + if (i < num_swap_token_chunks) + store_l1_swap_chunk(i); + } + } + + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + + for (uint32_t token = epilogue_thread_idx; token < valid_m; token += kNumEpilogueThreads) { + float amax = 0.0f; + #pragma unroll + for (uint32_t col = 0; col < L1_OUT_BLOCK_N; ++ col) { + const float v = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col]; + amax = cute::max(amax, cute::abs(v)); + } + float2 amax_pair = {amax, amax}; + float2 sf_pair, sf_inv_pair; + sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv(amax_pair, sf_pair, sf_inv_pair); + const float sf = sf_pair.x; + const float sf_inv = sf_inv_pair.x; + + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + const uint32_t token_idx = pool_block_idx * BLOCK_M + token; + sf_base_ptr[n_block_idx * kNumPaddedSFPoolTokens + token_idx] = sf; + + #pragma unroll + for (uint32_t col = 0; col < L1_OUT_BLOCK_N; col += 2) { + const float v0 = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col + 0] * sf_inv; + const float v1 = smem_cd_swap_l1_fp32[token * L1_OUT_BLOCK_N + col + 1] * sf_inv; + const __nv_fp8x2_e4m3 pair(make_float2(v0, v1)); + auto* ptr = reinterpret_cast( + smem_cd_swap_l1_fp8 + token * L1_OUT_BLOCK_N + col); + *ptr = pair.__x; + } + } + + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy( + &tensor_map_l1_output, + smem_cd_swap_l1_fp8, + n_block_idx * L1_OUT_BLOCK_N, + m_idx); + cute::tma_store_arrive(); + } + __syncwarp(); + ptx::tma_store_wait<0>(); + + if constexpr (kL2ArrivalCounter) { + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { + ptx::red_add_rel( + reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), + kWarpgroupSplitN); + } + } else { + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_or_rel_gpu( + workspace.get_l2_arrival_mask_ptr(pool_block_idx), + 1ull << n_block_idx); + } + } + __syncwarp(); + if constexpr (kL2ArrivalCounter) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } else { // ---------------- L1 EPILOGUE: activation + FP8 quantize + TMA store ---------------- // Layout in `final_accum`: @@ -1700,6 +1995,7 @@ sm90_fp8_mega_moe_impl(void* y, // that store has drained. if constexpr (kSplitNSharesSF) ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } } else { // ---------------- L2 EPILOGUE: BF16 cast + NVLink scatter ---------------- constexpr uint32_t kNumRowsPerWarp = WG_BLOCK_M / 8; @@ -1708,38 +2004,68 @@ sm90_fp8_mega_moe_impl(void* y, const uint32_t lane_in_row = lane_idx % 16; const uint32_t cols_per_lane = WG_BLOCK_N / 16; - // STSM into smem_cd_l2 (BF16). Reuse SM100 column-swizzle layout. - #pragma unroll - for (uint32_t i = 0; i < kAccumPerThread / 8; ++ i) { - // Each i consumes 8 floats (one 16x256b chunk in SM100 terms). - // For SM90 WGMMA layout, 8 floats per i correspond to 2 chunks of 4 floats: - // final_accum[i*8 + (0..3)] = chunk 2i: (r0c0, r0c1, r1c0, r1c1) - // final_accum[i*8 + (4..7)] = chunk 2i+1: same shape - const uint32_t chunk_lo = 2 * i, chunk_hi = 2 * i + 1; - - auto write_pair = [&](uint32_t row, uint32_t col, uint32_t packed) { - auto smem_ptr = smem_cd_l2 - + smem_cd_l2_wg_offset - + row * WG_BLOCK_N - + col; - // BF16 STS: 2 bf16 elements - *reinterpret_cast(smem_ptr) = packed; + if constexpr (kSwapABActive) { + auto store_bf16 = [&](const uint32_t& token, const uint32_t& col, float value) { + smem_cd_l2[smem_cd_l2_wg_offset + token * WG_BLOCK_N + col] = + __float2bfloat16_rn(value); }; - if (valid_r0) { - const uint32_t r0_lo = math::cast_into_bf16_and_pack( - final_accum[chunk_lo*4 + 0], final_accum[chunk_lo*4 + 1]); - const uint32_t r0_hi = math::cast_into_bf16_and_pack( - final_accum[chunk_hi*4 + 0], final_accum[chunk_hi*4 + 1]); - write_pair(r_0, chunk_lo * 8 + col_idx * 2, r0_lo); - write_pair(r_0, chunk_hi * 8 + col_idx * 2, r0_hi); + + auto store_l2_swap_chunk = [&](const uint32_t& i) { + const uint32_t token_0 = i * 8 + col_idx * 2; + const uint32_t token_1 = token_0 + 1; + if (token_0 < valid_m) { + store_bf16(token_0, r_0, final_accum[i * 4 + 0]); + store_bf16(token_0, r_1, final_accum[i * 4 + 2]); + } + if (token_1 < valid_m) { + store_bf16(token_1, r_0, final_accum[i * 4 + 1]); + store_bf16(token_1, r_1, final_accum[i * 4 + 3]); + } + }; + + const uint32_t num_swap_token_chunks = (valid_m + 7u) / 8u; + store_l2_swap_chunk(0); + if (valid_m > 8) { + #pragma unroll + for (uint32_t i = 1; i < kSwapABTokenChunks; ++ i) { + if (i < num_swap_token_chunks) + store_l2_swap_chunk(i); + } } - if (valid_r1) { - const uint32_t r1_lo = math::cast_into_bf16_and_pack( - final_accum[chunk_lo*4 + 2], final_accum[chunk_lo*4 + 3]); - const uint32_t r1_hi = math::cast_into_bf16_and_pack( - final_accum[chunk_hi*4 + 2], final_accum[chunk_hi*4 + 3]); - write_pair(r_1, chunk_lo * 8 + col_idx * 2, r1_lo); - write_pair(r_1, chunk_hi * 8 + col_idx * 2, r1_hi); + } else { + // STSM into smem_cd_l2 (BF16). Reuse SM100 column-swizzle layout. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 8; ++ i) { + // Each i consumes 8 floats (one 16x256b chunk in SM100 terms). + // For SM90 WGMMA layout, 8 floats per i correspond to 2 chunks of 4 floats: + // final_accum[i*8 + (0..3)] = chunk 2i: (r0c0, r0c1, r1c0, r1c1) + // final_accum[i*8 + (4..7)] = chunk 2i+1: same shape + const uint32_t chunk_lo = 2 * i, chunk_hi = 2 * i + 1; + + auto write_pair = [&](uint32_t row, uint32_t col, uint32_t packed) { + auto smem_ptr = smem_cd_l2 + + smem_cd_l2_wg_offset + + row * WG_BLOCK_N + + col; + // BF16 STS: 2 bf16 elements + *reinterpret_cast(smem_ptr) = packed; + }; + if (valid_r0) { + const uint32_t r0_lo = math::cast_into_bf16_and_pack( + final_accum[chunk_lo*4 + 0], final_accum[chunk_lo*4 + 1]); + const uint32_t r0_hi = math::cast_into_bf16_and_pack( + final_accum[chunk_hi*4 + 0], final_accum[chunk_hi*4 + 1]); + write_pair(r_0, chunk_lo * 8 + col_idx * 2, r0_lo); + write_pair(r_0, chunk_hi * 8 + col_idx * 2, r0_hi); + } + if (valid_r1) { + const uint32_t r1_lo = math::cast_into_bf16_and_pack( + final_accum[chunk_lo*4 + 2], final_accum[chunk_lo*4 + 3]); + const uint32_t r1_hi = math::cast_into_bf16_and_pack( + final_accum[chunk_hi*4 + 2], final_accum[chunk_hi*4 + 3]); + write_pair(r_1, chunk_lo * 8 + col_idx * 2, r1_lo); + write_pair(r_1, chunk_hi * 8 + col_idx * 2, r1_hi); + } } } @@ -1943,7 +2269,7 @@ sm90_fp8_mega_moe_impl(void* y, } #else if (blockIdx.x == 0 and threadIdx.x == 0) - DG_DEVICE_ASSERT(false and "This kernel only supports sm_90"); + DG_TRAP_ONLY_DEVICE_ASSERT(false and "This kernel only supports sm_90"); #endif } From c36c6662fd9a68801e63dda67e689724ebc70809 Mon Sep 17 00:00:00 2001 From: yinding Date: Wed, 17 Jun 2026 13:26:24 +0800 Subject: [PATCH 2/5] Fix SM90 FP8 MegaMoE swapAB synchronization (cherry picked from commit b23008580430d83813eef396c3cab419a205ffd0) --- csrc/jit_kernels/heuristics/sm90_mega_moe.hpp | 8 ++++---- .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 17 ++++++++++------- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp index 537280427..780cbaf90 100644 --- a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -44,8 +44,7 @@ struct MegaMoESM90Config { static std::tuple get_block_config_for_mega_moe_sm90( const int& num_ranks, const int& num_experts, - const int& num_max_tokens_per_rank, const int& num_topk, - const int& num_tokens) { + const int& num_topk, const int& num_tokens) { const float expected_tokens_per_expert = static_cast(num_tokens) * num_ranks * num_topk / num_experts; const bool auto_split_mn = expected_tokens_per_expert >= 64.0f; @@ -142,10 +141,11 @@ static MegaMoESM90Config get_mega_moe_config_sm90( const int& hidden, const int& intermediate_hidden, const int& num_padded_sf_pool_tokens) { const auto [block_m, num_epilogue_threads] = get_block_config_for_mega_moe_sm90( - num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens); + num_ranks, num_experts, num_topk, num_tokens); const float expected_tokens_per_expert = static_cast(num_tokens) * num_ranks * num_topk / num_experts; - const bool auto_split_mn = expected_tokens_per_expert >= 64.0f; + const bool auto_split_mn = + block_m == 128 and num_epilogue_threads == 512; const bool decode_split_n_path = block_m == 64 and num_epilogue_threads == 256; const bool decode_use_block_n_256 = 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 74bac96a3..b10edcbc7 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 @@ -1382,9 +1382,6 @@ sm90_fp8_mega_moe_impl(void* y, ptx::warpgroup_fence_operand(swap_accum[i]); ptx::warpgroup_wait<0>(); - if (lane_idx == 0) - empty_barriers[stage_idx]->arrive(); - #pragma unroll for (uint32_t i = 0; i < kSwapAccum / 4; ++ i) { const uint32_t token_0 = i * 8 + col_idx * 2; @@ -1398,6 +1395,9 @@ sm90_fp8_mega_moe_impl(void* y, final_accum[i * 4 + 1] += scale_1 * gate_sf * swap_accum[i * 4 + 1]; final_accum[i * 4 + 3] += scale_1 * up_sf * swap_accum[i * 4 + 3]; } + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); }; const uint32_t n_swap = ((valid_m + 7u) / 8u) * 8u; @@ -2069,10 +2069,13 @@ sm90_fp8_mega_moe_impl(void* y, } } - // Each warp writes and then scatters only its own 16-row - // slice, so a warp-level fence is enough before reading - // back from shared memory. - __syncwarp(); + // In the normal layout each warp writes and then scatters its + // own 16-row slice. swapAB writes by output-column ownership + // instead, so a token row is produced by the whole warpgroup. + if constexpr (kSwapABActive) + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + else + __syncwarp(); // Scatter to remote ranks via NVLink (one row per warp-pair) // Each warpgroup-warp covers 8 unique rows x 2 (r_0 + r_1 doubled by warps) From 8d5d398a050331d005f6aa251f98f61ad102fbc5 Mon Sep 17 00:00:00 2001 From: yinding Date: Thu, 25 Jun 2026 17:14:17 +0800 Subject: [PATCH 3/5] Add SM90 FP8 MegaMoE pre-dispatch kernel (cherry picked from commit 34fe473860227b95f2ced13cd06c4f8e5fd0f917) --- csrc/apis/sm90_mega.hpp | 19 +++ .../impls/sm90_mega_moe_pre_dispatch.hpp | 142 ++++++++++++++++++ csrc/tvm_ffi_api.cpp | 20 +++ .../impls/sm90_mega_moe_pre_dispatch.cuh | 104 +++++++++++++ deep_gemm/mega/__init__.py | 17 +++ sgl_deep_gemm/__init__.py | 18 +++ 6 files changed, 320 insertions(+) create mode 100644 csrc/jit_kernels/impls/sm90_mega_moe_pre_dispatch.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index 403747861..2e43626bb 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -4,6 +4,7 @@ #include "mega.hpp" #include "../jit_kernels/impls/sm90_fp8_mega_moe.hpp" +#include "../jit_kernels/impls/sm90_mega_moe_pre_dispatch.hpp" namespace deep_gemm::mega { @@ -11,6 +12,24 @@ static int get_token_alignment_for_sm90_mega_moe() { return layout::kLCMCandidateBlockM; } +static void mega_moe_pre_dispatch_sm90( + const torch::Tensor& x, + const torch::Tensor& topk_idx, + const torch::Tensor& topk_weights, + const torch::Tensor& buf_x, + const torch::Tensor& buf_x_sf, + const torch::Tensor& buf_topk_idx, + const torch::Tensor& buf_topk_weights, + const int& num_tokens, + const int& group_size, + const float& routed_scaling_factor) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + sm90_mega_moe_pre_dispatch( + x, topk_idx, topk_weights, + buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, + num_tokens, group_size, routed_scaling_factor); +} + static std::tuple(const torch::Tensor&)>> get_symm_buffer_size_for_sm90_mega_moe( const int& num_ranks, const int& num_experts, diff --git a/csrc/jit_kernels/impls/sm90_mega_moe_pre_dispatch.hpp b/csrc/jit_kernels/impls/sm90_mega_moe_pre_dispatch.hpp new file mode 100644 index 000000000..1e8dc91b4 --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_mega_moe_pre_dispatch.hpp @@ -0,0 +1,142 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" + +namespace deep_gemm { + +class SM90MegaMoEPreDispatchRuntime final : public LaunchRuntime { +public: + struct Args { + int group_size; + bool use_pdl; + + const void* x; + const void* topk_idx; + const void* topk_weights; + void* buf_x; + void* buf_x_sf; + void* buf_topk_idx; + void* buf_topk_weights; + uint32_t num_tokens; + uint32_t padded_max; + uint32_t hidden; + uint32_t num_groups; + uint32_t top_k; + float routed_scaling_factor; + + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_mega_moe_pre_dispatch_kernel< + {}, {} + >); +}} +)", args.group_size, args.use_pdl ? "true" : "false"); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.x, args.topk_idx, args.topk_weights, + args.buf_x, args.buf_x_sf, args.buf_topk_idx, args.buf_topk_weights, + args.num_tokens, args.padded_max, args.hidden, args.num_groups, + args.top_k, args.routed_scaling_factor)); + } +}; + +static void sm90_mega_moe_pre_dispatch( + const torch::Tensor& x, + const torch::Tensor& topk_idx, + const torch::Tensor& topk_weights, + const torch::Tensor& buf_x, + const torch::Tensor& buf_x_sf, + const torch::Tensor& buf_topk_idx, + const torch::Tensor& buf_topk_weights, + const int& num_tokens, + const int& group_size, + const float& routed_scaling_factor) { + DG_HOST_ASSERT(group_size == 128); + DG_HOST_ASSERT(x.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(x.is_contiguous()); + DG_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt32); + DG_HOST_ASSERT(topk_weights.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(topk_idx.is_contiguous() && topk_weights.is_contiguous()); + DG_HOST_ASSERT(x.dim() == 2 && topk_idx.dim() == 2 && topk_weights.dim() == 2); + DG_HOST_ASSERT(buf_x.dim() == 2 && buf_x_sf.dim() == 2); + DG_HOST_ASSERT(buf_topk_idx.dim() == 2 && buf_topk_weights.dim() == 2); + DG_HOST_ASSERT(buf_x.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(buf_x_sf.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(buf_topk_idx.scalar_type() == torch::kInt64); + DG_HOST_ASSERT(buf_topk_weights.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(buf_x.is_contiguous() && buf_x_sf.is_contiguous()); + + const auto m = static_cast(x.size(0)); + const auto hidden = static_cast(x.size(1)); + const auto top_k = static_cast(topk_idx.size(1)); + const auto padded_max = static_cast(buf_x.size(0)); + + DG_HOST_ASSERT(num_tokens == m); + DG_HOST_ASSERT(num_tokens <= padded_max); + DG_HOST_ASSERT(static_cast(topk_idx.size(0)) == m); + DG_HOST_ASSERT(static_cast(topk_weights.size(0)) == m); + DG_HOST_ASSERT(static_cast(topk_weights.size(1)) == top_k); + DG_HOST_ASSERT(static_cast(buf_x.size(1)) == hidden); + DG_HOST_ASSERT(static_cast(buf_topk_idx.size(0)) == padded_max); + DG_HOST_ASSERT(static_cast(buf_topk_idx.size(1)) == top_k); + DG_HOST_ASSERT(static_cast(buf_topk_weights.size(0)) == padded_max); + DG_HOST_ASSERT(static_cast(buf_topk_weights.size(1)) == top_k); + + DG_HOST_ASSERT(hidden % group_size == 0); + const auto num_groups = hidden / group_size; + DG_HOST_ASSERT(static_cast(buf_x_sf.size(0)) == padded_max); + DG_HOST_ASSERT(static_cast(buf_x_sf.size(1)) == num_groups); + DG_HOST_ASSERT(hidden % 8 == 0); + const auto num_threads = hidden / 8; + DG_HOST_ASSERT(num_threads <= 1024); + DG_HOST_ASSERT(num_threads >= top_k); + + const auto pad_slots = (padded_max - num_tokens) * top_k; + const auto num_pad_blocks = pad_slots == 0 ? 0 + : (pad_slots + num_threads - 1) / num_threads; + const auto num_total_blocks = num_tokens + num_pad_blocks; + if (num_total_blocks == 0) return; + + const bool use_pdl = device_runtime->get_pdl(); + SM90MegaMoEPreDispatchRuntime::Args args = { + .group_size = group_size, + .use_pdl = use_pdl, + .x = x.const_data_ptr(), + .topk_idx = topk_idx.const_data_ptr(), + .topk_weights = topk_weights.const_data_ptr(), + .buf_x = buf_x.data_ptr(), + .buf_x_sf = buf_x_sf.data_ptr(), + .buf_topk_idx = buf_topk_idx.data_ptr(), + .buf_topk_weights = buf_topk_weights.data_ptr(), + .num_tokens = static_cast(num_tokens), + .padded_max = static_cast(padded_max), + .hidden = static_cast(hidden), + .num_groups = static_cast(num_groups), + .top_k = static_cast(top_k), + .routed_scaling_factor = routed_scaling_factor, + .launch_args = LaunchArgs(num_total_blocks, num_threads, /*smem_size=*/0, + /*cluster_dim=*/1, /*enable_pdl=*/use_pdl) + }; + + const auto code = SM90MegaMoEPreDispatchRuntime::generate(args); + const auto runtime = compiler->build("sm90_mega_moe_pre_dispatch", code); + SM90MegaMoEPreDispatchRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index 82313cf9c..bee2310a7 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -813,6 +813,25 @@ void dg_mega_moe_pre_dispatch( ); } +void dg_mega_moe_pre_dispatch_sm90( + TensorView x, TensorView topk_idx, TensorView topk_weights, + TensorView buf_x, TensorView buf_x_sf, + TensorView buf_topk_idx, TensorView buf_topk_weights, + int64_t num_tokens, int64_t group_size, double routed_scaling_factor) { + mega::mega_moe_pre_dispatch_sm90( + convert_to_torch_tensor(x), + convert_to_torch_tensor(topk_idx), + convert_to_torch_tensor(topk_weights), + convert_to_torch_tensor(buf_x), + convert_to_torch_tensor(buf_x_sf), + convert_to_torch_tensor(buf_topk_idx), + convert_to_torch_tensor(buf_topk_weights), + static_cast(num_tokens), + static_cast(group_size), + static_cast(routed_scaling_factor) + ); +} + TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_token_alignment_for_mega_moe, dg_get_token_alignment_for_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_ring_limit_for_mega_moe, dg_get_ring_limit_for_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_symm_buffer_size_for_mega_moe, dg_get_symm_buffer_size_for_mega_moe); @@ -821,6 +840,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_mega_moe, dg_fp8_fp4_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_mega_moe, dg_bf16_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_mega_moe, dg_fp8_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(mega_moe_pre_dispatch, dg_mega_moe_pre_dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(mega_moe_pre_dispatch_sm90, dg_mega_moe_pre_dispatch_sm90); #endif // DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh new file mode 100644 index 000000000..8fa4a90b3 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh @@ -0,0 +1,104 @@ +#pragma once + +#include +#include +#include + +#include + +namespace deep_gemm { + +template +__launch_bounds__(1024, 2) +__global__ void sm90_mega_moe_pre_dispatch_kernel( + const __nv_bfloat16* __restrict__ x, + const int32_t* __restrict__ topk_idx, + const float* __restrict__ topk_weights, + __nv_fp8_e4m3* __restrict__ buf_x, + float* __restrict__ buf_x_sf, + int64_t* __restrict__ buf_topk_idx, + float* __restrict__ buf_topk_weights, + const uint32_t num_tokens, + const uint32_t padded_max, + const uint32_t hidden, + const uint32_t num_groups, + const uint32_t top_k, + const float routed_scaling_factor) { + static_assert(kGroupSize == 128, "SM90 mega-moe pre-dispatch requires per-128 SF"); + constexpr uint32_t kVecElems = 8; // 16-byte BF16 load per thread + static_assert(kGroupSize % kVecElems == 0, "kGroupSize must be a multiple of 8"); + constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems; + + const uint32_t bid = blockIdx.x; + const uint32_t tid = threadIdx.x; + + if constexpr (kUsePDL) { + cudaGridDependencySynchronize(); + } + + if (bid < num_tokens) { + const uint32_t token_id = bid; + const auto* token_in = x + static_cast(token_id) * hidden; + + uint4 in_bits = reinterpret_cast(token_in)[tid]; + const auto* bf16_pairs = reinterpret_cast(&in_bits); + + float vals[kVecElems]; + float local_max = 0.0f; + #pragma unroll + for (uint32_t i = 0; i < kVecElems / 2; ++i) { + float2 fp = __bfloat1622float2(bf16_pairs[i]); + vals[2 * i + 0] = fp.x; + vals[2 * i + 1] = fp.y; + local_max = fmaxf(local_max, fmaxf(fabsf(fp.x), fabsf(fp.y))); + } + + local_max = warp_reduce( + local_max, ReduceMax{}); + + const float absmax = fmaxf(local_max, 1e-10f); + const float raw_scale = absmax / 448.0f; + const float inv_scale = 1.0f / raw_scale; + + uint64_t packed = 0; + #pragma unroll + for (uint32_t i = 0; i < kVecElems / 2; ++i) { + const __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2( + make_float2(vals[2 * i + 0] * inv_scale, + vals[2 * i + 1] * inv_scale), + __NV_SATFINITE, __NV_E4M3); + packed |= static_cast(fp8x2) << (16u * i); + } + auto* row_out = reinterpret_cast(buf_x) + + static_cast(token_id) * (hidden / 8u); + row_out[tid] = packed; + + const uint32_t group_id = tid / kThreadsPerGroup; + const uint32_t within_group_id = tid % kThreadsPerGroup; + if (within_group_id == 0u && group_id < num_groups) { + const uint32_t off = token_id * num_groups + group_id; + buf_x_sf[off] = raw_scale; + } + + if (tid < top_k) { + const uint32_t off = token_id * top_k + tid; + buf_topk_idx[off] = static_cast(topk_idx[off]); + buf_topk_weights[off] = topk_weights[off] * routed_scaling_factor; + } + } else { + const uint32_t copy_bid = bid - num_tokens; + const uint32_t pad_base = num_tokens * top_k; + const uint32_t slot = pad_base + copy_bid * blockDim.x + tid; + const uint32_t total = padded_max * top_k; + if (slot < total) { + buf_topk_idx[slot] = static_cast(-1); + buf_topk_weights[slot] = 0.0f; + } + } + + if constexpr (kUsePDL) { + cudaTriggerProgrammaticLaunchCompletion(); + } +} + +} // namespace deep_gemm diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index ef21460cf..165ff769c 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -227,3 +227,20 @@ def mega_moe_pre_dispatch(x: torch.Tensor, buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, num_tokens, group_size, use_fp4_acts, ) + + +def mega_moe_pre_dispatch_sm90(x: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + buf_x: torch.Tensor, + buf_x_sf: torch.Tensor, + buf_topk_idx: torch.Tensor, + buf_topk_weights: torch.Tensor, + num_tokens: int, + group_size: int = 128, + routed_scaling_factor: float = 1.0) -> None: + _C.mega_moe_pre_dispatch_sm90( + x, topk_idx, topk_weights, + buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, + num_tokens, group_size, float(routed_scaling_factor), + ) diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index 4d94f6bdd..c3e73931e 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -271,6 +271,7 @@ def k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, grouped_layout, c=None, compi fp8_fp4_mega_moe, bf16_mega_moe, mega_moe_pre_dispatch, + mega_moe_pre_dispatch_sm90, ) @@ -414,6 +415,23 @@ def fp8_mega_moe(y: torch.Tensor, fast_math ) + +def mega_moe_pre_dispatch_sm90(x: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + buf_x: torch.Tensor, + buf_x_sf: torch.Tensor, + buf_topk_idx: torch.Tensor, + buf_topk_weights: torch.Tensor, + num_tokens: int, + group_size: int = 128, + routed_scaling_factor: float = 1.0) -> None: + _C.mega_moe_pre_dispatch_sm90( + x, topk_idx, topk_weights, + buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, + num_tokens, group_size, float(routed_scaling_factor), + ) + # Some utils from . import testing from . import utils From 6d2e9d584730ef1f9ba248d5fbf490b8785ba9ee Mon Sep 17 00:00:00 2001 From: yinding Date: Fri, 26 Jun 2026 12:01:24 +0800 Subject: [PATCH 4/5] Add SM90 mega-MoE pre-dispatch test (cherry picked from commit 15a6f42764702f824ab4851cdbc5b6a671122b86) --- sgl_deep_gemm/run_tests.sh | 1 + .../tests/test_mega_moe_pre_dispatch_sm90.py | 180 ++++++++++++++++++ 2 files changed, 181 insertions(+) create mode 100644 sgl_deep_gemm/tests/test_mega_moe_pre_dispatch_sm90.py diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh index 3b1913bcc..7c643cbda 100755 --- a/sgl_deep_gemm/run_tests.sh +++ b/sgl_deep_gemm/run_tests.sh @@ -135,6 +135,7 @@ MEGA_MOE_BLACKWELL=( ) MEGA_MOE_HOPPER=( test_mega_moe_hopper.py + test_mega_moe_pre_dispatch_sm90.py ) MEGA_MOE_ALL=( "${MEGA_MOE_BLACKWELL[@]}" diff --git a/sgl_deep_gemm/tests/test_mega_moe_pre_dispatch_sm90.py b/sgl_deep_gemm/tests/test_mega_moe_pre_dispatch_sm90.py new file mode 100644 index 000000000..2459fadc1 --- /dev/null +++ b/sgl_deep_gemm/tests/test_mega_moe_pre_dispatch_sm90.py @@ -0,0 +1,180 @@ +# Correctness probe for `deep_gemm.mega_moe_pre_dispatch_sm90` (SM90 / Hopper). +# +# SM90 variant of the fused mega-MoE pre-dispatch. Unlike the Blackwell +# `mega_moe_pre_dispatch` (covered by test_mega_moe_pre_dispatch.py), the SM90 +# kernel emits plain FP8 (no packed-FP4 branch) with per-(token, 128-channel) +# FP32 scales, and folds a `routed_scaling_factor` into the topk_weights write. +# This test verifies, against a pure-torch reference: +# 1) per-token group-128 FP8 quantization of x -> buf_x[:M], buf_x_sf[:M] +# (dequantised output within one e4m3 mantissa step), +# 2) buf.topk_idx[:M] == topk_idx ; buf.topk_weights[:M] == weights * alpha, +# 3) pad rows: buf.topk_idx[M:] == -1 ; buf.topk_weights[M:] == 0. +# +# Single-GPU; SM90-only (skips cleanly on non-Hopper or when the wheel lacks +# the kernel). Run directly: `python test_mega_moe_pre_dispatch_sm90.py`. + +import argparse +import sys + +import torch + +import deep_gemm + +# fp8-e4m3 representable max; per-group scale = amax / FP8_E4M3_MAX. +FP8_E4M3_MAX = 448.0 + + +def _has_hopper() -> bool: + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_capability()[0] == 9 + + +def _ref_quant_fp8_group128(x: torch.Tensor): + """Per-token, per-128-channel-group absmax FP8-e4m3 quant with fp32 scales. + + Returns (q_fp8 [M, H], scale_fp32 [M, H/128]) such that + ``q_fp8.float().reshape(M, G, 128) * scale[..., None] ≈ x``. + """ + M, H = x.shape + G = H // 128 + xf = x.float().reshape(M, G, 128) + amax = xf.abs().amax(dim=-1).clamp(min=1e-12) # (M, G) + scale = amax / FP8_E4M3_MAX # (M, G) fp32 + q = (xf / scale.unsqueeze(-1)).to(torch.float8_e4m3fn) + return q.reshape(M, H), scale + + +def _alloc(padded_max: int, hidden: int, top_k: int, device): + num_groups = hidden // 128 + buf_x = torch.empty((padded_max, hidden), dtype=torch.float8_e4m3fn, device=device) + buf_x_sf = torch.empty((padded_max, num_groups), dtype=torch.float32, device=device) + buf_idx = torch.empty((padded_max, top_k), dtype=torch.int64, device=device) + buf_w = torch.empty((padded_max, top_k), dtype=torch.float32, device=device) + return buf_x, buf_x_sf, buf_idx, buf_w + + +def _reference(x, topk_idx, topk_weights, padded_max, routed_scaling_factor): + M, _ = x.shape + buf_x, buf_x_sf, buf_idx, buf_w = _alloc(padded_max, x.shape[1], topk_idx.shape[1], x.device) + if M > 0: + qx, sf = _ref_quant_fp8_group128(x) + buf_x[:M].copy_(qx) + buf_x_sf[:M].copy_(sf) + buf_idx[:M].copy_(topk_idx) + buf_w[:M].copy_(topk_weights * routed_scaling_factor) + if M < padded_max: + buf_idx[M:].fill_(-1) + buf_w[M:].zero_() + return buf_x, buf_x_sf, buf_idx, buf_w + + +def _run_kernel(x, topk_idx, topk_weights, padded_max, routed_scaling_factor): + M, H = x.shape + K = topk_idx.shape[1] + buf_x, buf_x_sf, buf_idx, buf_w = _alloc(padded_max, H, K, x.device) + # Poison the padded region so a missing pad-write shows up as a hard failure. + if M < padded_max: + buf_idx[M:].fill_(0x4242) + buf_w[M:].fill_(float("nan")) + + if M > 0: + x_in, idx_in, w_in = x, topk_idx, topk_weights + else: + x_in = x.new_empty((0, H), dtype=x.dtype) + idx_in = x.new_empty((0, K), dtype=torch.int32) + w_in = x.new_empty((0, K), dtype=torch.float32) + + deep_gemm.mega_moe_pre_dispatch_sm90( + x_in, + idx_in, + w_in, + buf_x, + buf_x_sf, + buf_idx, + buf_w, + num_tokens=M, + group_size=128, + routed_scaling_factor=routed_scaling_factor, + ) + torch.cuda.synchronize() + return buf_x, buf_x_sf, buf_idx, buf_w + + +def _check_case(M, H, K, P, scale, seed): + assert H % 128 == 0 and H % 8 == 0, f"hidden {H} must be a multiple of 128" + assert M <= P, f"num_tokens {M} must be <= padded_max {P}" + device = torch.device("cuda") + torch.manual_seed(seed) + + # Mix of magnitudes so absmax / clamp logic is exercised. + x = torch.randn(M, H, dtype=torch.bfloat16, device=device) * 4.0 + topk_idx = torch.randint(0, 256, (M, K), dtype=torch.int32, device=device) + topk_weights = torch.randn(M, K, dtype=torch.float32, device=device) + + ref_x, ref_sf, ref_idx, ref_w = _reference(x, topk_idx, topk_weights, P, scale) + out_x, out_sf, out_idx, out_w = _run_kernel(x, topk_idx, topk_weights, P, scale) + + if M > 0: + G = H // 128 + # Scales: same amax/448 formula; allow small fp32 drift. + torch.testing.assert_close(out_sf[:M], ref_sf[:M], rtol=1e-3, atol=0) + + def _dequant(bx, bsf): + return bx[:M].float().reshape(M, G, 128) * bsf[:M].unsqueeze(-1) + + out_deq = _dequant(out_x, out_sf) + ref_deq = _dequant(ref_x, ref_sf) + # Kernel vs reference: at most one e4m3 mantissa step (1/8) apart. + torch.testing.assert_close(out_deq, ref_deq, rtol=1.0 / 8, atol=0) + # Sanity: dequantised values approximate the original input. + torch.testing.assert_close( + out_deq, x.float().reshape(M, G, 128), rtol=0.25, atol=1e-4 + ) + + # topk_idx pass-through + pad-fill (exact); weights scaled by alpha. + torch.testing.assert_close(out_idx, ref_idx, rtol=0, atol=0) + torch.testing.assert_close(out_w, ref_w, rtol=1e-6, atol=0) + + pad = "exact-fill" if M == P else f"{P - M} pad rows" + print(f" PASS M={M} H={H} K={K} P={P} alpha={scale} ({pad})") + + +# (M, H, K, padded_max, routed_scaling_factor) +CASES = [ + (0, 2048, 8, 32, 1.0), # zero tokens -> all pad + (1, 1024, 4, 8, 2.5), + (7, 2048, 8, 16, 1.0), + (7, 2048, 8, 16, 2.5), + (32, 4096, 8, 32, 2.5), + (128, 7168, 8, 256, 2.5), + (128, 7168, 8, 128, 1.0), # exact-fill, no pad rows +] + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--seed", type=int, default=0) + # Accept (and ignore) --num-processes so the shared test runner can pass it. + parser.add_argument("--num-processes", type=int, default=1) + args = parser.parse_args() + + if not torch.cuda.is_available(): + print("SKIP: no CUDA device") + return 0 + if not _has_hopper(): + major = torch.cuda.get_device_capability()[0] + print(f"SKIP: SM90/Hopper required (got arch major {major})") + return 0 + if not hasattr(deep_gemm, "mega_moe_pre_dispatch_sm90"): + print("SKIP: installed deep_gemm has no mega_moe_pre_dispatch_sm90") + return 0 + + for i, (M, H, K, P, s) in enumerate(CASES): + _check_case(M, H, K, P, s, seed=(args.seed ^ (0xC0FFEE + i))) + print("OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From dbbb0f4745af11c4c899ea7648888af687394ede Mon Sep 17 00:00:00 2001 From: yinding Date: Thu, 2 Jul 2026 13:17:28 +0800 Subject: [PATCH 5/5] =?UTF-8?q?SM90=20FP8=20MegaMoE=20swapAB=20=E2=80=94?= =?UTF-8?q?=20accuracy=20issue=20&=20mitigation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- csrc/jit_kernels/heuristics/sm90_mega_moe.hpp | 5 +++++ deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh | 8 +++++++- .../deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh | 7 ++++--- 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp index 780cbaf90..c449fc5b2 100644 --- a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -86,6 +86,11 @@ static int get_num_experts_per_wave_for_mega_moe_sm90( static bool should_use_swap_ab_for_mega_moe_sm90( const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, const int& block_m, const int& num_epilogue_threads) { + // swapAB is ENABLED by default (the L1 SF-pool stride bug that corrupted + // pool blocks >= 1 was fixed: BLOCK_M -> SF_BLOCK_M in the swapAB L1 epilogue). + // Kill-switch retained: set DG_SM90_FP8_SWAP_AB=0 to force the non-swap path. + if (get_env("DG_SM90_FP8_SWAP_AB", 1) == 0) + return false; const float expected_tokens_per_expert = static_cast(num_tokens) * num_topk / num_experts_per_rank; const bool decode_split_n_path = 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 b10edcbc7..a514803bf 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 @@ -1692,7 +1692,13 @@ sm90_fp8_mega_moe_impl(void* y, const float sf_inv = sf_inv_pair.x; auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); - const uint32_t token_idx = pool_block_idx * BLOCK_M + token; + // ROOT-CAUSE FIX: the L2-activation SF pool is strided by SF_BLOCK_M + // (=align(BLOCK_M,128)=128), which is how the L2 producer reads it + // (sfa_m_idx = pool_block_idx * SF_BLOCK_M) and how the non-swap L1 + // writes it. This swapAB path used BLOCK_M (64), so for pool_block_idx>=1 + // the SF landed in the wrong rows -> L2 read stale SF -> every pool block + // after the first was corrupted (block 0 was correct because 0*64==0*128). + const uint32_t token_idx = pool_block_idx * SF_BLOCK_M + token; sf_base_ptr[n_block_idx * kNumPaddedSFPoolTokens + token_idx] = sf; #pragma unroll diff --git a/deep_gemm/include/deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh b/deep_gemm/include/deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh index 8fa4a90b3..f26b4d117 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh @@ -4,7 +4,8 @@ #include #include -#include +#include +#include namespace deep_gemm { @@ -53,8 +54,8 @@ __global__ void sm90_mega_moe_pre_dispatch_kernel( local_max = fmaxf(local_max, fmaxf(fabsf(fp.x), fabsf(fp.y))); } - local_max = warp_reduce( - local_max, ReduceMax{}); + local_max = math::warp_reduce( + local_max, math::ReduceMax{}); const float absmax = fmaxf(local_max, 1e-10f); const float raw_scale = absmax / 448.0f;