[SM90] Optimize FP8 MegaMoE across L2 epilogue, 4WG launch heuristic and decode swap-AB

csrc/jit_kernels/heuristics/sm90_mega_moe.hpp

@@ -44,10 +44,32 @@ 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_topk, const int& num_tokens) {
+    const int& num_topk, const int& num_tokens, const int &intermediate_hidden) {
     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;
+    // The relaxed 2-WG threshold enables the block_m=128 / 4-WG path only
+    // above a higher tokens/expert bar (instead of the original >= 64),
+    // trading two extra warpgroups for fewer register spills. On H20 the
+    // smaller SM count (78 vs 132 on H100/H200) makes the extra warpgroups
+    // costly, so the relaxation applies in two intermediate_hidden regimes:
+    //   * pro   (>= 3072): 4-WG only when expected_tokens_per_expert > 512
+    //   * flash (<= 2048): 4-WG only when expected_tokens_per_expert > 576,
+    //     because 2-WG + BLOCK_N=256 outperforms 4-WG in part of the flash
+    //     batch range -- 4-WG is reserved for the heaviest flash batches.
+    // On H200/H100 the larger SM count makes the extra warpgroups always win,
+    // so the original 4-WG-first (>= 64) threshold is kept for every shape,
+    // as well as for the H20 mid-range (2048 < intermediate_hidden < 3072).
+    const int num_sms = device_runtime->get_num_sms();
+    const bool is_h20 = num_sms <= 84;
+    const bool apply_h20_pro_relaxation   = is_h20 and intermediate_hidden >= 3072;
+    const bool apply_h20_flash_relaxation = is_h20 and intermediate_hidden <= 2048;
+    bool auto_split_mn;
+    if (apply_h20_pro_relaxation)
+        auto_split_mn = expected_tokens_per_expert > 512.0f;
+    else if (apply_h20_flash_relaxation)
+        auto_split_mn = expected_tokens_per_expert > 576.0f;
+    else
+        auto_split_mn = expected_tokens_per_expert >= 64.0f;
     if (auto_split_mn)
         return {128, 512};
 
@@ -146,21 +168,21 @@ 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_topk, num_tokens);
+        num_ranks, num_experts, num_topk, num_tokens, intermediate_hidden);
     const float expected_tokens_per_expert =
         static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;
     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
+        decode_split_n_path and
         expected_tokens_per_expert >= 0.25f and
         (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
+    int block_n = use_swap_ab ? 256
                               : (auto_split_mn ? 256 :
                                  (decode_use_block_n_256 ? 256 : 128));
     const int block_k = 128;

csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp

@@ -242,11 +242,12 @@ static void sm90_fp8_mega_moe(
     const int wg_l1_out_block_n = wg_block_n / 2;
     const bool split_n_shares_sf =
         split_n_warpgroups and wg_l1_out_block_n < kL2ActsSFGranK;
+    const bool l1_output_full_tile = split_n_shares_sf or use_swap_ab;
     const int l1_output_swizzle_mode = 0;
     const int l1_output_box_n =
-        split_n_shares_sf ? config.block_n / 2 : wg_l1_out_block_n;
+        l1_output_full_tile ? config.block_n / 2 : wg_l1_out_block_n;
     const int l1_output_box_m =
-        split_n_shares_sf ? config.block_m : wg_block_m;
+        l1_output_full_tile ? config.block_m : wg_block_m;
     const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts,
                                                        intermediate_hidden, config.num_max_pool_tokens,
                                                        l1_output_box_n, l1_output_box_m,