perf(moe_ep): add folded K128/N64 tuning candidate for Hopper MXFP4 x FP8 MegaMoE

flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/TUNING.md

@@ -881,13 +881,22 @@ of a complete tactic does not isolate the speedup of any one field.
 
 ### Bounded candidates and cache replay
 
-The normal MXFP4 list starts from 17 deduplicated base tactics drawn from the
-block-permutation and published-exact manifests and two H20-derived anchors.
-For H7168/I3072/E384/EP4, optional N8/readiness choices expand these to 23;
+The normal MXFP4 list starts from 22 deduplicated base tactics drawn from the
+block-permutation and published-exact manifests, two H20-derived anchors,
+and five K128 candidates for intermediate widths such as 640. The folded
+M256/N64/K128 CGA1x1x1 group-330 candidate was measured on H100 and H200.
+An H200 sweep at H2560/I640/E512/EP2 adds an unfolded M256/N16/K128
+CGA1x1x1 group-330 candidate with two dispatch warps and FC1 store offload,
+folded N64 CGA1x1x1 candidates with groups 264 and 132, and a folded N64
+CGA1x2x1 group-132 tail-pair candidate. All five use one scheduler stage;
+the folded candidates use one dispatch warp, early FC1 publication and
+direct FC1 stores. For this shape, the complete eligible union grows from
+three to eight candidates without changing the heuristic defaults.
+For H7168/I3072/E384/EP4, optional N8/readiness choices expand these to 28;
 three large-token configurations and 12 distinct tail-pair neighbors bring
-the total to 38. Every model considers the same bounded extension catalog,
+the total to 43. Every model considers the same bounded extension catalog,
 then filters candidates through tile alignment and strategy eligibility.
-Other shapes can therefore have fewer than 38 candidates. Token capacity
+Other shapes can therefore have fewer than 43 candidates. Token capacity
 changes ordering, not the union; the untuned heuristic buckets are unchanged.
 
 Online `knobs="auto"` and offline `flashinfer.moe_ep.tune` use the same

flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/mxfp4_tuner.py

@@ -159,6 +159,57 @@ def _tactic(
         fold_producer_warps=True,
     ),
     _tactic((256, 16, 256), group_hint=512, stages=1),
+    # Keep a folded N64 option when intermediate is not divisible by K256.
+    # The complete producer/publication combination was measured on H100/H200.
+    _tactic(
+        (256, 64, 128),
+        group_hint=330,
+        stages=1,
+        cluster=(1, 1, 1),
+        active_dispatch_warps=1,
+        fc1_early_done_publish=True,
+        fold_producer_warps=True,
+    ),
+    # H200 token sweeps favor N16 at small loads and different folded N64
+    # scheduling groups at larger loads; keep these as autotune candidates.
+    _tactic(
+        (256, 16, 128),
+        group_hint=330,
+        stages=1,
+        cluster=(1, 1, 1),
+        active_dispatch_warps=2,
+        fc1_store_offload=True,
+    ),
+    _tactic(
+        (256, 64, 128),
+        group_hint=264,
+        stages=1,
+        cluster=(1, 1, 1),
+        active_dispatch_warps=1,
+        fc1_early_done_publish=True,
+        fold_producer_warps=True,
+    ),
+    _tactic(
+        (256, 64, 128),
+        group_hint=132,
+        stages=1,
+        cluster=(1, 1, 1),
+        active_dispatch_warps=1,
+        fc1_early_done_publish=True,
+        fold_producer_warps=True,
+    ),
+    dict(
+        _tactic(
+            (256, 64, 128),
+            group_hint=132,
+            stages=1,
+            cluster=(1, 2, 1),
+            active_dispatch_warps=1,
+            fc1_early_done_publish=True,
+            fold_producer_warps=True,
+        ),
+        tail_split_pairs=True,
+    ),
 )
 
 # Indices into _BASE_TACTICS for the token buckets above.

tests/moe_ep/test_moe_ep_sm90_pull_mxfp4_mega_multirank.py

@@ -632,6 +632,34 @@ def test_moe_ep_sm90_pull_mxfp4_mega_multirank_raw_oracle_and_workspace_reuse(
     _run_raw_oracle(routing_pattern)
 
 
+@pytest.mark.gpu_2
+@pytest.mark.arch_hopper
+@pytest.mark.parametrize("tail_pairs", (False, True), ids=("folded", "folded-tail"))
+def test_moe_ep_sm90_pull_mxfp4_m256_n64_k128_independent_oracle(
+    monkeypatch, tail_pairs
+):
+    """Folded K128 candidates match raw math across full and partial N64 tiles."""
+    import sys
+
+    module = sys.modules[__name__]
+    monkeypatch.setattr(module, "HIDDEN", 2560)
+    monkeypatch.setattr(module, "INTERMEDIATE", 640)
+    # EP2 routes 64/65 rows per expert, crossing an N64 tile boundary.
+    monkeypatch.setattr(module, "TOKENS_PER_RANK", 129)
+    tactic = dict(
+        _fused_tactic(),
+        mma_tiler_mnk=(256, 64, 128),
+        cluster_shape_mnk=(1, 2, 1) if tail_pairs else (1, 1, 1),
+        group_hint=132 if tail_pairs else 330,
+        num_sched_stages=1,
+        load_balance_mode="atomic_counter",
+        fc1_store_offload=False,
+        fc1_early_done_publish=True,
+        tail_split_pairs=tail_pairs,
+    )
+    _run_raw_oracle("cross_rank", tactic=tactic)
+
+
 @pytest.mark.gpu_4
 @pytest.mark.arch_hopper
 @pytest.mark.parametrize(