[SM90] Optimize FP8 MegaMoE small-batch decode path with swapAB

csrc/apis/sm90_mega.hpp

@@ -4,13 +4,32 @@
 
 #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 {
 
 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<int64_t, std::function<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>(const torch::Tensor&)>>
 get_symm_buffer_size_for_sm90_mega_moe(
     const int& num_ranks, const int& num_experts,

csrc/jit_kernels/heuristics/sm90_mega_moe.hpp

@@ -44,8 +44,7 @@ struct MegaMoESM90Config {
 
 static std::tuple<int, int> 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<float>(num_tokens) * num_ranks * num_topk / num_experts;
     const bool auto_split_mn = expected_tokens_per_expert >= 64.0f;
@@ -84,11 +83,27 @@ 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) {
+    // 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<int>("DG_SM90_FP8_SWAP_AB", 1) == 0)
+        return false;
+    const float expected_tokens_per_expert =
+        static_cast<float>(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<int, int> 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 +115,13 @@ static std::pair<int, int> 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<int>(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<int>(sizeof(float)) + static_cast<int>(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<int>(sizeof(float)), 128);
     const int smem_sfb_per_stage = 0;
@@ -125,18 +146,23 @@ 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<float>(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 =
         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 +192,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 +207,9 @@ static MegaMoESM90Config get_mega_moe_config_sm90(
 
     if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("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<std::string> printed;
         if (printed.count(key) == 0) {
             std::cout << key << ": " << config << std::endl;

csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp

@@ -44,6 +44,7 @@ class SM90FP8MegaMoERuntime final : public LaunchRuntime<SM90FP8MegaMoERuntime>
         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,

csrc/jit_kernels/impls/sm90_mega_moe_pre_dispatch.hpp

@@ -0,0 +1,142 @@
+#pragma once
+
+#include <torch/python.h>
+
+#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<SM90MegaMoEPreDispatchRuntime> {
+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 <deep_gemm/impls/sm90_mega_moe_pre_dispatch.cuh>
+
+using namespace deep_gemm;
+
+static void __instantiate_kernel() {{
+    auto ptr = reinterpret_cast<void*>(&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<int>(x.size(0));
+    const auto hidden = static_cast<int>(x.size(1));
+    const auto top_k = static_cast<int>(topk_idx.size(1));
+    const auto padded_max = static_cast<int>(buf_x.size(0));
+
+    DG_HOST_ASSERT(num_tokens == m);
+    DG_HOST_ASSERT(num_tokens <= padded_max);
+    DG_HOST_ASSERT(static_cast<int>(topk_idx.size(0)) == m);
+    DG_HOST_ASSERT(static_cast<int>(topk_weights.size(0)) == m);
+    DG_HOST_ASSERT(static_cast<int>(topk_weights.size(1)) == top_k);
+    DG_HOST_ASSERT(static_cast<int>(buf_x.size(1)) == hidden);
+    DG_HOST_ASSERT(static_cast<int>(buf_topk_idx.size(0)) == padded_max);
+    DG_HOST_ASSERT(static_cast<int>(buf_topk_idx.size(1)) == top_k);
+    DG_HOST_ASSERT(static_cast<int>(buf_topk_weights.size(0)) == padded_max);
+    DG_HOST_ASSERT(static_cast<int>(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<int>(buf_x_sf.size(0)) == padded_max);
+    DG_HOST_ASSERT(static_cast<int>(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<uint32_t>(num_tokens),
+        .padded_max = static_cast<uint32_t>(padded_max),
+        .hidden = static_cast<uint32_t>(hidden),
+        .num_groups = static_cast<uint32_t>(num_groups),
+        .top_k = static_cast<uint32_t>(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

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<int>(num_tokens),
+        static_cast<int>(group_size),
+        static_cast<float>(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