Raise SM90 FP4 MegaMoE decode->prefill boundary to e=80

The decode config (BLOCK_M=64) and prefill config (BLOCK_M=128) flipped
at a hardcoded e>=64. Measured on 8xH20 both modes across the boundary:
decode wins for e in [64, 80) (its first m-block is exactly full while
prefill's 128-row block runs half empty, -13% at GLM5.2 b256); prefill
wins from e=80 up (decode's second m-block is mostly empty, +2~9%).

Parameterize the four boundary checks (block config auto_split_mn,
split-N band, 2-WG thread band, API prefill band) through a single
DG_SM90_FP4_PREFILL_E knob (default 80) plus a PREFILL_E mirror
constant in the test.

GLM5.2 b256: 1277 -> 1123 us; vs FP8 low-latency 0.98x -> 1.11x (cliff
fixed, b245->b320 now monotone). DSV4-Flash e=72 also improves -1.8%.
Accuracy: forced-decode reference checks at e=64/96/128 pass
(diff ~ 6e-4, tol 0.1).

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

Author: yinding

csrc/apis/sm90_mega.hpp

@@ -68,7 +68,11 @@ static FP4SM90APIDefaults get_fp4_sm90_api_defaults(
     (void)intermediate_hidden;
     const float expected_tokens_per_expert =
         static_cast<float>(num_tokens) * num_topk / num_experts_per_rank;
-    const bool prefill_band = expected_tokens_per_expert >= 64.0f;
+    // Decode -> prefill boundary; keep in sync with the JIT heuristics'
+    // get_fp4_sm90_prefill_threshold (DG_SM90_FP4_PREFILL_E, default 80).
+    const float prefill_threshold =
+        static_cast<float>(get_env<int>("DG_SM90_FP4_PREFILL_E", 80));
+    const bool prefill_band = expected_tokens_per_expert >= prefill_threshold;
     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

csrc/jit_kernels/heuristics/sm90_mega_moe.hpp

@@ -119,6 +119,17 @@ static std::pair<int, int> get_pipeline_config_for_mega_moe_sm90(
     return {num_stages, smem_size};
 }
 
+// Decode -> prefill boundary for the FP4 MegaMoE path, in expected tokens per
+// expert. At the boundary the kernel flips from the decode config (BLOCK_M=64,
+// split-N epilogue) to the prefill config (BLOCK_M=128, 2-WG decode offload).
+// 80 = measured on H20: decode wins for e in [64, 80) (its first m-block is
+// exactly full while prefill's 128-row block runs half empty); prefill wins
+// from e=80 up (decode's second m-block is mostly empty). Overridable via
+// DG_SM90_FP4_PREFILL_E for boundary A/B tuning.
+static float get_fp4_sm90_prefill_threshold() {
+    return static_cast<float>(get_env<int>("DG_SM90_FP4_PREFILL_E", 80));
+}
+
 static std::tuple<int, int> get_block_config_for_mega_moe_sm90_fp4(
     const int& num_ranks, const int& num_experts,
     const int& num_max_tokens_per_rank, const int& num_topk,
@@ -127,7 +138,8 @@ static std::tuple<int, int> get_block_config_for_mega_moe_sm90_fp4(
 
     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 =
+        expected_tokens_per_expert >= get_fp4_sm90_prefill_threshold();
     const bool ultra_small_split_n =
         expected_tokens_per_expert > 0.0f and expected_tokens_per_expert < 0.375f;
     int block_m = auto_split_mn ? 128 : 64;
@@ -257,7 +269,8 @@ static MegaMoESM90Config get_mega_moe_config_sm90_fp4(
         block_m == 64 and block_n % 128 == 0;
     const bool fp4_split_n_shape_band =
         fp4_flash_or_pro_shape and
-        expected_tokens_per_expert > 0.0f and expected_tokens_per_expert < 64.0f;
+        expected_tokens_per_expert > 0.0f and
+        expected_tokens_per_expert < get_fp4_sm90_prefill_threshold();
     if (fp4_split_n_eligible and fp4_split_n_shape_band) {
         fp4_num_epilogue_warpgroups = 2;
     }
@@ -285,7 +298,8 @@ static MegaMoESM90Config get_mega_moe_config_sm90_fp4(
         fp4_small_block_n_kernel and fp4_split_n_shape_band;
     const bool fp4_2wg_decode_offload_kernel_band =
         block_m == 128 and block_n == 128 and
-        fp4_num_epilogue_threads == 256 and expected_tokens_per_expert >= 64.0f;
+        fp4_num_epilogue_threads == 256 and
+        expected_tokens_per_expert >= get_fp4_sm90_prefill_threshold();
     const bool fp4_decode_assist_thread_kernel_band =
         fp4_2wg_decode_offload_kernel_band or
         (fp4_small_block_n_kernel and

sgl_deep_gemm/tests/test_mega_moe_fp4_hopper.py

@@ -33,6 +33,12 @@
 
 GENERIC_FALLBACK = "generic"
 
+# Decode -> prefill boundary (expected tokens per expert). Mirrors the C++
+# default of DG_SM90_FP4_PREFILL_E (the CPU mirror assumes the env is unset).
+# 80 = measured: BLOCK_M=64 decode wins for e in [64, 80) (full first m-block),
+# BLOCK_M=128 prefill wins from e=80 up (decode's second m-block mostly empty).
+PREFILL_E = 80.0
+
 
 def table_wave(
     num_experts_per_rank: int,
@@ -65,11 +71,11 @@ def legacy_threads(
     small_block_n = block_m == 64 and block_n == 128
     e = expected_tokens_per_expert
     decode_heavy_small_batch = small_block_n and 0.0 < e <= 24.0
-    pro_large_decode_assist_batch = small_block_n and intermediate_hidden >= 3072 and 24.0 <= e < 64.0
-    pro_split_n_decode_threads = small_block_n and intermediate_hidden >= 3072 and 0.0 < e < 64.0
-    flash_split_n_decode_threads = small_block_n and intermediate_hidden <= 2048 and 0.0 < e < 64.0
+    pro_large_decode_assist_batch = small_block_n and intermediate_hidden >= 3072 and 24.0 <= e < PREFILL_E
+    pro_split_n_decode_threads = small_block_n and intermediate_hidden >= 3072 and 0.0 < e < PREFILL_E
+    flash_split_n_decode_threads = small_block_n and intermediate_hidden <= 2048 and 0.0 < e < PREFILL_E
     two_wg_decode_offload = (
-        block_m == 128 and block_n == 128 and num_epilogue_threads == 256 and e >= 64.0
+        block_m == 128 and block_n == 128 and num_epilogue_threads == 256 and e >= PREFILL_E
     )
     dispatch = (
         64
@@ -101,9 +107,9 @@ def legacy_epilogue_threads(
     e = expected_tokens_per_expert
     epilogue_warpgroups = num_epilogue_threads // 128
     split_n_eligible = block_m == 64 and block_n % 128 == 0
-    split_n_band = 32.0 <= e < 64.0
-    pro_split_n_band = intermediate_hidden >= 3072 and 0.0 < e < 64.0
-    flash_split_n_band = intermediate_hidden <= 2048 and 0.0 < e < 64.0
+    split_n_band = 32.0 <= e < PREFILL_E
+    pro_split_n_band = intermediate_hidden >= 3072 and 0.0 < e < PREFILL_E
+    flash_split_n_band = intermediate_hidden <= 2048 and 0.0 < e < PREFILL_E
     small_split_n_band = flash_split_n_band or pro_split_n_band
     default_split_n = (
         split_n_eligible
@@ -126,7 +132,7 @@ def table_epilogue_threads(
     epilogue_warpgroups = num_epilogue_threads // 128
     split_n_eligible = block_m == 64 and block_n % 128 == 0
     split_n_shape_band = (
-        (intermediate_hidden <= 2048 or intermediate_hidden >= 3072) and 0.0 < e < 64.0
+        (intermediate_hidden <= 2048 or intermediate_hidden >= 3072) and 0.0 < e < PREFILL_E
     )
     if split_n_eligible and split_n_shape_band:
         epilogue_warpgroups = 2
@@ -143,11 +149,11 @@ def table_threads(
     small_block_n_kernel = block_m == 64 and block_n == 128
     e = expected_tokens_per_expert
     split_n_shape_band = (
-        (intermediate_hidden <= 2048 or intermediate_hidden >= 3072) and 0.0 < e < 64.0
+        (intermediate_hidden <= 2048 or intermediate_hidden >= 3072) and 0.0 < e < PREFILL_E
     )
     split_n_decode_thread_kernel_band = small_block_n_kernel and split_n_shape_band
     two_wg_decode_offload_kernel_band = (
-        block_m == 128 and block_n == 128 and num_epilogue_threads == 256 and e >= 64.0
+        block_m == 128 and block_n == 128 and num_epilogue_threads == 256 and e >= PREFILL_E
     )
     decode_assist_thread_kernel_band = two_wg_decode_offload_kernel_band or (
         small_block_n_kernel and 0.0 < e <= 24.0
@@ -246,8 +252,8 @@ def table_api_features(
     threshold, shape-agnostic (no per-(shape x e-band) table)."""
     del num_experts_per_rank, intermediate_hidden
     e = expected_tokens_per_expert
-    prefill = e >= 64.0
-    decode = 0.0 < e < 64.0
+    prefill = e >= PREFILL_E
+    decode = 0.0 < e < PREFILL_E
     return {
         "math_wg_participates": False,
         "first_decode_assist_warp": 2,