Simplify SM90 FP4 MegaMoE heuristics to FP8 parity

The FP4 path had ~49 tuned conditions (17 shape x e-band API-default
bands, 9+12 wave-table rules, 11 stage-cap rules) fitted point-by-point
on benchmark batches; on real shapes most bands collapsed to constants
or slivers. Retire them:

- API defaults: constants + one decode/prefill split + a single swapAB
  threshold (measured crossover on H20: swapAB wins at e<=12, ties at
  e~16, loses beyond -> bound 16). wide_load_decode on for all decode.
- Wave scheduling: delegate to the FP8 rule
  (get_num_experts_per_wave_for_mega_moe_sm90).
- Stage cap: removed; always use as many pipeline stages as SMEM allows.
- CPU heuristic mirrors in the test synced to the same closed forms.

Verified on 8xH20 EP8: accuracy suites pass (diff <= 0.0006, tol 0.05);
GLM5.2 flat-or-better (b2 -12%, b96 -9%, prefill bit-identical);
DSV4-Flash within noise; DSV4-Pro avg ~+2%, worst b1 +8.6%.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: yinding

csrc/apis/sm90_mega.hpp

@@ -60,123 +60,35 @@ struct FP4SM90APIDefaults {
 static FP4SM90APIDefaults get_fp4_sm90_api_defaults(
     const int& num_experts_per_rank, const int& num_tokens, const int& num_topk,
     const int& intermediate_hidden) {
+    // Simplified, shape-agnostic defaults, mirroring the FP8 path's style
+    // (`should_use_swap_ab_for_mega_moe_sm90`): one decode/prefill split plus a
+    // single swapAB threshold. The historical per-(shape x e-band) table was
+    // tuned point-by-point on benchmark batches; on the shapes that matter it
+    // collapsed to constants plus a few sliver bands, so it is retired.
+    (void)intermediate_hidden;
     const float expected_tokens_per_expert =
         static_cast<float>(num_tokens) * num_topk / num_experts_per_rank;
-    // Shape bands exclude kernel tile/thread constraints; JIT heuristics add those as kernel bands.
-    const bool fp4_flash_shape = intermediate_hidden <= 2048;
-    const bool fp4_pro_shape = intermediate_hidden >= 3072;
-    const bool fp4_middle_shape = !fp4_flash_shape and !fp4_pro_shape;
-    const bool fp4_decode_lookahead_shape_band =
-        expected_tokens_per_expert >= 3.0f and expected_tokens_per_expert <= 6.0f;
-    const bool fp4_bigband_lookahead_shape_band =
-        expected_tokens_per_expert >= 12.0f and expected_tokens_per_expert <= 24.0f;
-    const bool fp4_b4_skip_decode_shape_band =
-        expected_tokens_per_expert >= 0.5f and expected_tokens_per_expert < 1.0f;
-    const bool fp4_pro_single_token_per_expert_shape_band =
-        fp4_pro_shape and
-        expected_tokens_per_expert >= 1.0f and expected_tokens_per_expert < 1.5f and
-        num_experts_per_rank % 8 == 0;
-    const bool fp4_pro_split_n_mbarrier_shape_band =
-        fp4_pro_shape and
-        expected_tokens_per_expert > 0.0f and expected_tokens_per_expert < 64.0f;
-    const bool fp4_pro_two_tokens_per_expert_shape_band =
-        fp4_pro_shape and
-        expected_tokens_per_expert >= 1.5f and expected_tokens_per_expert < 3.0f;
-    const bool fp4_pro_mid_decode_assist_shape_band =
-        fp4_pro_shape and
-        expected_tokens_per_expert >= 6.0f and expected_tokens_per_expert < 12.0f;
-    const bool fp4_pro_large_decode_assist_shape_band =
-        fp4_pro_shape and
-        expected_tokens_per_expert >= 24.0f and expected_tokens_per_expert < 64.0f;
-    const bool fp4_flash_two_tokens_per_expert_shape_band =
-        fp4_flash_shape and
-        expected_tokens_per_expert >= 1.5f and expected_tokens_per_expert < 2.0f;
-    const bool fp4_flash_half_token_per_expert_shape_band =
-        fp4_flash_shape and
-        expected_tokens_per_expert >= 0.375f and expected_tokens_per_expert < 0.5f;
-    const bool fp4_flash_decode_lookahead_shape_band =
-        fp4_flash_shape and
-        expected_tokens_per_expert >= 3.0f and expected_tokens_per_expert < 6.0f;
-    const bool fp4_flash_wide_load_decode_shape_band =
-        fp4_flash_shape and
-        expected_tokens_per_expert >= 6.0f and expected_tokens_per_expert < 64.0f;
-    const bool fp4_pro_wide_load_decode_shape_band =
-        fp4_pro_shape and
-        expected_tokens_per_expert > 0.0f and expected_tokens_per_expert < 64.0f;
-    const bool fp4_flash_split_n_mbarrier_shape_band =
-        fp4_flash_shape and
-        expected_tokens_per_expert >= 0.75f and expected_tokens_per_expert < 64.0f;
-    const bool fp4_flash_small_mbarrier_shape_band =
-        fp4_flash_shape and
-        expected_tokens_per_expert > 0.0f and expected_tokens_per_expert < 0.5f;
-    const bool fp4_2wg_decode_offload_shape_band =
-        expected_tokens_per_expert >= 64.0f;
-    const bool fp4_shared_decode_assist_shape_band =
-        ((expected_tokens_per_expert > 0.0f and expected_tokens_per_expert < 0.375f) or
-         fp4_flash_half_token_per_expert_shape_band or
-         fp4_b4_skip_decode_shape_band or fp4_decode_lookahead_shape_band or
-         fp4_flash_split_n_mbarrier_shape_band or
-         fp4_pro_mid_decode_assist_shape_band or fp4_pro_large_decode_assist_shape_band or
-         fp4_bigband_lookahead_shape_band or fp4_2wg_decode_offload_shape_band);
-    const bool default_math_wg_decode =
-        fp4_shared_decode_assist_shape_band or
-        (expected_tokens_per_expert >= 1.0f and expected_tokens_per_expert < 2.0f) or
-        fp4_pro_two_tokens_per_expert_shape_band;
-    const bool math_wg_participates_in_decode =
-        !default_math_wg_decode;
-    const bool default_skip_loader_decode_assist =
-        fp4_shared_decode_assist_shape_band or
-        fp4_pro_single_token_per_expert_shape_band or
-        (expected_tokens_per_expert >= 1.5f and expected_tokens_per_expert < 3.0f);
-    const bool default_wide_load_decode =
-        fp4_pro_wide_load_decode_shape_band or
-        fp4_flash_half_token_per_expert_shape_band or
-        fp4_flash_two_tokens_per_expert_shape_band or
-        fp4_flash_wide_load_decode_shape_band;
-    const bool default_ss_early_b_decode =
-        ((expected_tokens_per_expert >= 1.5f and expected_tokens_per_expert <= 3.0f and
-          !fp4_pro_two_tokens_per_expert_shape_band and
-          !fp4_flash_two_tokens_per_expert_shape_band and
-          !fp4_flash_decode_lookahead_shape_band) or
-         fp4_2wg_decode_offload_shape_band);
-    const bool fp4_middle_decode_lookahead_mbarrier_shape_band =
-        fp4_middle_shape and fp4_decode_lookahead_shape_band;
-    const bool fp4_middle_bigband_mbarrier_shape_band =
-        fp4_middle_shape and fp4_bigband_lookahead_shape_band;
-    const bool default_decode_done_mbarrier =
-        fp4_pro_split_n_mbarrier_shape_band or
-        fp4_flash_split_n_mbarrier_shape_band or
-        fp4_flash_small_mbarrier_shape_band or
-        fp4_middle_decode_lookahead_mbarrier_shape_band or
-        fp4_middle_bigband_mbarrier_shape_band or
-        fp4_2wg_decode_offload_shape_band;
-    const bool default_l2_arrival_counter =
-        ((fp4_flash_shape and
-          expected_tokens_per_expert >= 0.375f and expected_tokens_per_expert < 0.75f) or
-         (fp4_pro_shape and
-          expected_tokens_per_expert >= 0.25f and expected_tokens_per_expert < 0.375f));
+    const bool prefill_band = expected_tokens_per_expert >= 64.0f;
+    const bool decode_band =
+        expected_tokens_per_expert > 0.0f and !prefill_band;
     // swapAB on/off kill-switch (default ON). Set DG_SM90_FP4_SWAP_AB=0 to force
     // the non-swap path for A/B accuracy comparison.
     const bool swap_ab_env_enabled = get_env<int>("DG_SM90_FP4_SWAP_AB", 1) != 0;
-    const bool default_swap_ab =
-        swap_ab_env_enabled and
-        (fp4_flash_shape or fp4_pro_shape) and
-        expected_tokens_per_expert > 0.0f and expected_tokens_per_expert <= 24.0f;
-    const bool default_swap_ab_fast_amax =
-        swap_ab_env_enabled and
-        fp4_pro_shape and
-        expected_tokens_per_expert >= 12.0f and expected_tokens_per_expert <= 24.0f;
     return {
-        math_wg_participates_in_decode,
-        math_wg_participates_in_decode ? 4 : 0,
-        default_skip_loader_decode_assist ? 2 : 0,
-        default_wide_load_decode,
-        default_ss_early_b_decode,
-        default_decode_done_mbarrier,
-        default_l2_arrival_counter,
-        expected_tokens_per_expert >= 64.0f,
-        default_swap_ab,
-        default_swap_ab_fast_amax
+        /*math_wg_participates_in_decode=*/ false,
+        /*num_math_wg_decode_warps=*/ 0,
+        /*first_decode_assist_warp=*/ 2,
+        /*wide_load_decode=*/ decode_band,
+        /*early_b_decode=*/ prefill_band,
+        /*decode_done_mbarrier=*/ expected_tokens_per_expert > 0.0f,
+        /*l2_arrival_counter=*/ false,
+        /*ss_nsplit=*/ prefill_band,
+        // Measured crossover on H20: swapAB wins clearly at e<=12, ties at
+        // e~16 and loses beyond, so the bound is 16 (FP8 uses 30; the FP4
+        // kernel pays extra decode work on the swapped path).
+        /*swap_ab=*/ swap_ab_env_enabled and decode_band and
+                     expected_tokens_per_expert < 16.0f,
+        /*swap_ab_fast_amax=*/ false
     };
 }
 

csrc/jit_kernels/heuristics/sm90_mega_moe.hpp

@@ -144,169 +144,14 @@ static std::tuple<int, int> get_block_config_for_mega_moe_sm90_fp4(
     return {block_m, num_epilogue_warpgroups * 128};
 }
 
-struct FP4SM90WaveRule {
-    float min_tokens_per_expert;
-    float max_tokens_per_expert;
-    bool include_min;
-    int required_expert_divisor;
-    int num_experts_per_wave;
-};
-
-enum class FP4SM90StageShape {
-    Any,
-    Flash,
-    Pro,
-    NotPro,
-};
-
-struct FP4SM90StageCapRule {
-    float min_tokens_per_expert;
-    float max_tokens_per_expert;
-    bool include_min;
-    bool include_max;
-    FP4SM90StageShape shape;
-    int num_stages_cap;
-};
-
-static bool try_get_num_experts_per_wave_for_sm90_fp4(
-    const FP4SM90WaveRule* rules, const int& num_rules,
-    const float& expected_tokens_per_expert, const int& num_experts_per_rank,
-    int& num_experts_per_wave) {
-    for (int i = 0; i < num_rules; ++ i) {
-        const auto& rule = rules[i];
-        const bool in_lower_bound = rule.include_min
-            ? expected_tokens_per_expert >= rule.min_tokens_per_expert
-            : expected_tokens_per_expert > rule.min_tokens_per_expert;
-        if (!in_lower_bound or expected_tokens_per_expert >= rule.max_tokens_per_expert)
-            continue;
-
-        if (rule.num_experts_per_wave == 0) {
-            if (num_experts_per_rank <= 0)
-                continue;
-            num_experts_per_wave = num_experts_per_rank;
-            return true;
-        }
-        DG_HOST_ASSERT(rule.required_expert_divisor > 0);
-        if (num_experts_per_rank % rule.required_expert_divisor == 0) {
-            num_experts_per_wave = rule.num_experts_per_wave;
-            return true;
-        }
-    }
-    return false;
-}
-
-static bool fp4_sm90_stage_shape_matches(
-    const FP4SM90StageShape& shape, const bool& fp4_flash_shape, const bool& fp4_pro_shape) {
-    switch (shape) {
-        case FP4SM90StageShape::Any:
-            return true;
-        case FP4SM90StageShape::Flash:
-            return fp4_flash_shape;
-        case FP4SM90StageShape::Pro:
-            return fp4_pro_shape;
-        case FP4SM90StageShape::NotPro:
-            return !fp4_pro_shape;
-    }
-    DG_HOST_ASSERT(false);
-    return false;
-}
-
-static int get_default_num_stages_cap_for_mega_moe_sm90_fp4(
-    const int& intermediate_hidden, const int& block_m, const int& block_n,
-    const float& expected_tokens_per_expert) {
-    if (!(block_m == 64 and block_n == 128)) {
-        return 0;
-    }
-
-    const bool fp4_flash_shape = intermediate_hidden <= 2048;
-    const bool fp4_pro_shape = intermediate_hidden >= 3072;
-    // Ordered first-match rules preserve the historical stage-cap priority.
-    static constexpr FP4SM90StageCapRule stage_cap_rules[] = {
-        {6.0f, 12.0f, true, false, FP4SM90StageShape::Flash, 4},
-        {3.0f, 6.0f, false, false, FP4SM90StageShape::Flash, 4},
-        {0.0f, 0.25f, false, false, FP4SM90StageShape::Pro, 5},
-        {0.375f, 0.75f, true, false, FP4SM90StageShape::Pro, 5},
-        {1.5f, 3.0f, true, false, FP4SM90StageShape::Pro, 5},
-        {1.0f, 1.5f, true, false, FP4SM90StageShape::Pro, 5},
-        {24.0f, 64.0f, true, false, FP4SM90StageShape::Pro, 5},
-        {0.375f, 0.75f, true, false, FP4SM90StageShape::Any, 6},
-        {3.0f, 6.0f, false, false, FP4SM90StageShape::Flash, 6},
-        {1.5f, 3.0f, true, false, FP4SM90StageShape::NotPro, 6},
-        {1.5f, 24.0f, true, true, FP4SM90StageShape::Any, 5},
-    };
-    for (const auto& rule: stage_cap_rules) {
-        const bool in_lower_bound = rule.include_min
-            ? expected_tokens_per_expert >= rule.min_tokens_per_expert
-            : expected_tokens_per_expert > rule.min_tokens_per_expert;
-        const bool in_upper_bound = rule.include_max
-            ? expected_tokens_per_expert <= rule.max_tokens_per_expert
-            : expected_tokens_per_expert < rule.max_tokens_per_expert;
-        if (in_lower_bound and in_upper_bound and
-            fp4_sm90_stage_shape_matches(rule.shape, fp4_flash_shape, fp4_pro_shape)) {
-            return rule.num_stages_cap;
-        }
-    }
-    return 0;
-}
-
 static int get_num_experts_per_wave_for_mega_moe_sm90_fp4(
     const int& num_experts_per_rank, const int& num_tokens, const int& num_topk,
     const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms,
     const int& num_ring_tokens, const int& num_max_tokens_per_rank, const int& num_ranks) {
-    const float expected_tokens_per_expert =
-        static_cast<float>(num_tokens) * num_topk / num_experts_per_rank;
-    const bool fp4_small_block_n_kernel =
-        block_m == 64 and block_n == 128;
-    const bool fp4_flash_shape = intermediate_hidden <= 2048;
-    const bool fp4_pro_shape = intermediate_hidden >= 3072;
-    int fp4_num_experts_per_wave = 0;
-    if (fp4_small_block_n_kernel and fp4_flash_shape) {
-        static constexpr FP4SM90WaveRule flash_wave_rules[] = {
-            {0.75f, 1.0f, true, 16, 16},
-            {1.5f, 2.0f, true, 16, 16},
-            {3.0f, 6.0f, true, 16, 16},
-            {6.0f, 12.0f, true, 32, 32},
-            {6.0f, 12.0f, true, 8, 8},
-            {24.0f, 32.0f, true, 16, 16},
-            {12.0f, 24.0f, true, 32, 32},
-            {12.0f, 24.0f, true, 8, 8},
-            {32.0f, 64.0f, true, 16, 16},
-        };
-        if (try_get_num_experts_per_wave_for_sm90_fp4(
-                flash_wave_rules,
-                static_cast<int>(sizeof(flash_wave_rules) / sizeof(flash_wave_rules[0])),
-                expected_tokens_per_expert, num_experts_per_rank,
-                fp4_num_experts_per_wave)) {
-            return fp4_num_experts_per_wave;
-        }
-    }
-    if (fp4_small_block_n_kernel and fp4_pro_shape) {
-        static constexpr FP4SM90WaveRule pro_wave_rules[] = {
-            {0.0f, 0.25f, false, 16, 16},
-            {0.25f, 0.375f, true, 16, 16},
-            {0.375f, 0.75f, true, 16, 16},
-            {0.25f, 1.0f, true, 24, 24},
-            {1.0f, 1.5f, true, 1, 0},
-            {1.5f, 3.0f, true, 16, 16},
-            {3.0f, 6.0f, true, 8, 8},
-            {6.0f, 12.0f, true, 16, 16},
-            {6.0f, 12.0f, true, 8, 8},
-            {12.0f, 24.0f, true, 24, 24},
-            {12.0f, 24.0f, true, 8, 8},
-            {24.0f, 64.0f, true, 8, 8},
-        };
-        if (try_get_num_experts_per_wave_for_sm90_fp4(
-                pro_wave_rules,
-                static_cast<int>(sizeof(pro_wave_rules) / sizeof(pro_wave_rules[0])),
-                expected_tokens_per_expert, num_experts_per_rank,
-                fp4_num_experts_per_wave)) {
-            return fp4_num_experts_per_wave;
-        }
-    }
-    if (expected_tokens_per_expert < 1.0f or expected_tokens_per_expert > 4.0f) {
-        return num_experts_per_rank;
-    }
-    return get_num_experts_per_wave_for_mega_moe(
+    // Simplified: schedule FP4 expert waves exactly like the FP8 path. The
+    // historical flash/pro wave tables (9 + 12 first-match rules) were tuned
+    // point-by-point on benchmark batches and are retired.
+    return get_num_experts_per_wave_for_mega_moe_sm90(
         num_experts_per_rank, num_tokens, num_topk,
         intermediate_hidden, block_m, block_n, num_sms,
         num_ring_tokens, num_max_tokens_per_rank, num_ranks);
@@ -319,7 +164,6 @@ static std::pair<int, int> get_pipeline_config_for_mega_moe_sm90_fp4(
     const int& num_dispatch_warps, const int& num_epilogue_warps,
     const bool& use_early_b_decode = false,
     const bool& use_decode_done_mbarrier = false,
-    const int& default_num_stages_cap = 0,
     const bool& use_swap_ab = false,
     const bool& use_swap_ab_fast_amax = false) {
     constexpr int kSmemAlignment = 1024;
@@ -380,12 +224,10 @@ static std::pair<int, int> get_pipeline_config_for_mega_moe_sm90_fp4(
     const int smem_fixed =
         smem_dispatch_size + smem_cd + smem_amax_scratch + smem_barriers_fixed;
 
-    const int max_num_stages = (smem_capacity - smem_fixed) /
-                               (smem_per_stage + smem_barriers_per_stage);
-    int num_stages = max_num_stages;
-    if (default_num_stages_cap > 0) {
-        num_stages = std::min(num_stages, default_num_stages_cap);
-    }
+    // No FP4 stage cap (FP8 parity): always use as many pipeline stages as
+    // SMEM allows. The historical 11-rule cap table is retired.
+    const int num_stages = (smem_capacity - smem_fixed) /
+                           (smem_per_stage + smem_barriers_per_stage);
     DG_HOST_ASSERT(num_stages >= 2);
     return {num_stages,
             smem_fixed + num_stages * (smem_per_stage + smem_barriers_per_stage)};
@@ -461,15 +303,12 @@ static MegaMoESM90Config get_mega_moe_config_sm90_fp4(
                    num_non_epilogue_threads % 64 == 0);
     DG_HOST_ASSERT((num_dispatch_threads + num_non_epilogue_threads) % 128 == 0);
 
-    const int default_num_stages_cap = get_default_num_stages_cap_for_mega_moe_sm90_fp4(
-        intermediate_hidden, block_m, block_n, expected_tokens_per_expert);
-
     const auto [num_stages, smem_size] = get_pipeline_config_for_mega_moe_sm90_fp4(
         SM90ArchSpec::smem_capacity,
         num_experts, hidden,
         block_m, block_n, block_k,
         num_dispatch_threads / 32, fp4_num_epilogue_threads / 32,
-        use_early_b_decode, use_decode_done_mbarrier, default_num_stages_cap,
+        use_early_b_decode, use_decode_done_mbarrier,
         use_swap_ab, use_swap_ab_fast_amax);
 
     const auto config = MegaMoESM90Config {