perf(moe): optimize sparse SM107 MegaMoE decode

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/SKILL.md

@@ -27,7 +27,9 @@ kernel_src/sm107/next_cutedsl_megamoe/
 └── TUNING.md               ← tuning and benchmark conventions
 ```
 
-**Keep `src/` byte-for-byte equal to the pinned upstream export.**
+**Keep `src/` byte-for-byte equal to the pinned upstream export, except for the
+explicit local patches inventoried in `VENDOR.md`.**
+Reapply and requalify those patches after an export until upstream absorbs them.
 Record the exporter revision, selected kernels and generated transformations in
 `VENDOR.md`. Fix device/exporter code upstream and re-export; put FlashInfer
 adaptation in `shim/` and the backend wrappers.
@@ -84,8 +86,9 @@ sha256sum next/export_src.py
 ```
 
 The current selection includes the generic and GenPhase inference kernels.
-Select GenPhase with `kernel_variant="genphase"` and preserve its fixed cluster
-and bulk-store requirements when adapting the launcher. Select additional
+Select GenPhase with `kernel_variant="genphase"` and preserve its admitted
+cluster/tile combinations and bulk-store requirements documented in `VENDOR.md`
+when adapting the launcher. Select additional
 kernels when their FlashInfer integration is included in the update.
 The export may include dependencies from other kernel families; record their
 source mappings rather than manually removing them.
@@ -132,8 +135,10 @@ layouts, ownership and reset semantics can change while names remain stable.
 
 Preserve the integration's correctness fixes when adapting: physical K-major
 weight concatenation, bounded preprocessing temporaries, padded physical router
-capacity, per-launch masked-slot clearing/output resets, routing validation,
-and normalization and scalar-override semantics. Recheck warmed graph bindings
+capacity, per-launch masked-slot handling/output resets, routing validation,
+and normalization and scalar-override semantics. Separate BF16 reduction masks
+invalid slots; quantized combine and in-kernel reduction still require their
+respective clears. Recheck warmed graph bindings
 across different weights/streams and safe workspace teardown if those contracts
 change.
 

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/TUNING.md

@@ -254,6 +254,24 @@ Use `--kernel-variant genphase` with the offline tuner. Its candidates and
 cache entries are separate from the generic inference kernel. A cache miss
 uses the fixed GenPhase cluster and bulk-store settings.
 
+GenPhase also accepts `mma_tiler_mnk=(256, 64, 128)` for MXFP8 E4M3
+activations and weights, or mixed MXFP4 weights with MXFP8 activations.
+The default candidate menu includes this tile with cluster `(4, 1)` and
+one FC2 TMA stage. MXFP8 E4M3 also admits an explicitly selected `(2, 1)`
+cluster; mixed MXFP4/MXFP8 requires `(4, 1)`. NVFP4 activations do not
+support this N64 tile. N64 is a tile dimension, not the number of live tokens.
+Gather4 issuance follows the live rows in each tile, including partially
+filled tiles, while retaining matching barrier transaction counts.
+
+Retune to compare the additional candidates for your workload. The default
+kernel family remains `inference`; these candidates do not automatically
+select GenPhase for existing callers.
+
+Both families mask invalid route slots during separate BF16 top-k reduction,
+so their production launch can omit clearing the combine payload first.
+This internal optimization applies without a caller API change. Quantized
+combine and in-kernel top-k reduction retain their existing clearing paths.
+
 ## Quantized combine selection
 
 Set `combine_dtype="nvfp4"` or `combine_dtype="mxfp8"` on a Rubin backend

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/VENDOR.md

@@ -1,6 +1,7 @@
 # Vendoring record: next_cutedsl_megamoe
 
-This directory contains one upstream kernel snapshot.
+This directory contains one pinned upstream kernel snapshot plus the explicit
+local patches recorded below.
 [SKILL.md](SKILL.md) describes updates; [TUNING.md](TUNING.md) covers measurements.
 
 ## Upstream
@@ -20,7 +21,7 @@ This directory contains one upstream kernel snapshot.
   Their dependency closure is exported to `src/sources/`: 32 implementation
   modules and nine generated package initializers, 41 files total.
 - **Verification**: two independent exports at the pinned commit produced
-  identical file lists and bytes. The local comment edit is listed below.
+  identical file lists and bytes before local patches. The exceptions are listed below.
   The exporter validates internal import closure and rejects
   unresolved imports and dependencies outside its external-module allowlist.
 
@@ -37,7 +38,8 @@ export manifest/log and file hashes with the integration evidence.
 
 ## Policy
 
-- `src/` is a **verbatim copy of the pinned upstream export output**. Do not
+- `src/` starts from a **verbatim copy of the pinned upstream export output**;
+  the explicit local exceptions below must be reapplied and qualified. Do not
   hand-edit, format or inject files into it. Compare it with a regenerated
   export, rather than expecting every generated file to match a raw upstream path.
 - All FlashInfer adaptation lives in `shim/`, re-exported through `__init__.py`;
@@ -58,8 +60,12 @@ Shared helpers, schedulers, communication and workspace code come from the
 same pinned snapshot.
 
 `RubinInferenceGenphaseMegaMoE` (`BlockScaledSwapAbGenphaseMoeKernel`) is selected
-with `kernel_variant="genphase"`. It requires uniform `(4, 1)` clusters,
-`fc2_use_bulk=True`, and at most 1024 tokens per rank. Combine remains BF16.
+with `kernel_variant="genphase"`. It requires `fc2_use_bulk=True`, at most
+1024 tokens per rank, and uniform clusters (no mixed-CGA fallback). The
+inherited tile-N 128/256 configurations use `(4, 1)` clusters. The local
+`(256, 64, 128)` decode tile additionally admits MXFP8 E4M3 with `(2, 1)` or
+`(4, 1)` clusters, and MXFP4 weights/MXFP8 E4M3 activations with `(4, 1)`
+clusters. Other N64 combinations remain rejected. Combine remains BF16.
 Tile K is 256 for NVFP4 activations and 128 for MXFP8 activations.
 SiTU requires both beta parameters, matching the exported kernel.
 
@@ -106,11 +112,38 @@ transformations, not manual FlashInfer patches.
 
 ## Local differences
 
-The comment above `pusher_cta_count` in
-`src/sources/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py`
-describes the residency requirement without device-specific counts. Kernel
-logic and configuration values are unchanged. All other exported files match
-the pinned upstream export.
+The existing comment above `pusher_cta_count` in the GenPhase kernel describes
+residency without device-specific counts. In addition, FlashInfer PR
+[#6447](https://github.com/flashinfer-ai/flashinfer/pull/6447) carries the following
+local device patches. They are not part of upstream `1667b47a` and must not be
+silently dropped by a fresh export. Paths below are relative to `src/sources/`.
+
+| File | Local exception |
+| --- | --- |
+| `helpers/ptx_helpers.py` | Counted, unaligned CTA barrier helper for the actual participating warps. |
+| `kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py` | Use that barrier for the Generic MMA/accumulator TMEM allocation rendezvous. |
+| `kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py` | N64 SFB/TMEM parity, valid-row TMA-gather extent and transaction accounting, and two gather warps for CGA4. |
+| `kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py` | Pass routing IDs to separate BF16 reduction and advertise masked-slot handling to the shim. |
+| `kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py` | Propagate masked reduction and admitted N64 geometry; support a one-row allocation in the gather descriptor. |
+| `kernel_src/rubin/inference/mega/token_comm_gen_specialized.py` | Bound route metadata by the live extent, skip invalid histogram/scatter work, and reuse the existing expert prefix for scale-factor offsets. |
+| `kernel_src/rubin/inference/mega/topk_reduce.py` | Predicate BF16 combine loads by routing validity, yielding exact zero for rows with no valid routes. |
+
+The Generic router `communication/nvlink_domain/token_comm.py` remains the
+pinned implementation. A proposed invalid-route histogram/scatter shortcut was
+withdrawn after an EP4 captured-graph sparse-to-dense regression; it is not an
+exception to carry forward. The regression is retained in
+`tests/moe_ep/test_sm107_routing_graph_multirank.py`: every capacity row is checked
+after dense, empty, tail-only, sparse-hole, asymmetric-rank, and dense-restoration
+states, on both the first and repeated replays, with early/late routing weights.
+
+FlashInfer-owned `shim/block_scaled.py` admits the geometries listed above and
+omits BF16 masked-slot preclear only when the kernel advertises masked reduction.
+Quantized combine and in-kernel reduction retain their existing clearing rules.
+This shim change is outside the exported payload. The local patch inventory
+above is separate from the historical upstream qualification evidence; run the
+registered strict single/multirank suites when updating or rebasing these patches.
+All remaining exported implementation files retain the pinned export, apart
+from the previously documented comment.
 
 ## Related trees
 

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/shim/autotune.py

@@ -67,7 +67,9 @@ def sm107_candidates(
     - epi flag batches (1,4) vs (2,4),
     - FC2 bulk TMA (2 stages) on/off.
 
-    GenPhase fixes the cluster to (4, 1) and requires FC2 bulk stores.
+    Default GenPhase candidates use cluster (4, 1) and require FC2 bulk stores.
+    MXFP8 E4M3 and mixed FP4/FP8 additionally include the N64 decode tile,
+    using one FC2 stage and deterministic separate top-k reduction.
 
     ``token_in_flag_batch`` and the phase-interleave hint are skew-sensitive —
     sweep them with :func:`sm107_schedule_candidates` instead.  An ikr winner
@@ -111,6 +113,32 @@ def sm107_candidates(
                                 "reduce_topk_in_kernel": ikr,
                             }
                         )
+    if kernel_variant == "genphase" and quant_kind in (
+        "mxfp8_e4m3",
+        "mxfp4_mxfp8",
+    ):
+        # These are additional candidates, not a replacement for wider tiles:
+        # routing, cache residency and token count change the best pipeline.
+        for schedule, work_id in (
+            (("grouped", None), "grid_stride"),
+            (("grouped", None), "atomic_counter"),
+            (("phase_interleave", None), "atomic_counter"),
+        ):
+            out.append(
+                {
+                    "mma_tiler_mnk": (256, 64, 128),
+                    "cluster_shape_mn": (4, 1),
+                    "fallback_cluster_shape_mn": None,
+                    "schedule_policy": schedule,
+                    "work_id_mode": work_id,
+                    "fc2_use_bulk": True,
+                    "fc2_tma_stages": 1,
+                    "epi_flag_batches": (1, 4),
+                    "token_in_flag_batch": 1,
+                    "token_back_mode": "epi_warps",
+                    "reduce_topk_in_kernel": False,
+                }
+            )
     return out
 
 

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/shim/block_scaled.py

@@ -364,15 +364,29 @@ def __post_init__(self) -> None:
         if self.kernel_variant == "genphase":
             if self.padded_tokens_per_rank > 1024:
                 raise ValueError("GenPhase supports at most 1024 tokens per rank.")
+            n64_decode = (
+                self.cluster_shape_mn in ((2, 1), (4, 1))
+                and tiler == (256, 64, 128)
+                and (
+                    self.quant_kind == "mxfp8_e4m3"
+                    or (
+                        self.quant_kind == "mxfp4_mxfp8"
+                        and self.cluster_shape_mn == (4, 1)
+                    )
+                )
+            )
             if (
-                self.cluster_shape_mn != (4, 1)
-                or self.fallback_cluster_shape_mn is not None
-            ):
-                raise ValueError("GenPhase requires uniform cluster_shape_mn=(4, 1).")
+                self.cluster_shape_mn != (4, 1) and not n64_decode
+            ) or self.fallback_cluster_shape_mn is not None:
+                raise ValueError(
+                    "GenPhase requires uniform cluster (4, 1), or an admitted N64 decode geometry at (2, 1) or (4, 1)."
+                )
             if not self.fc2_use_bulk:
                 raise ValueError("GenPhase requires fc2_use_bulk=True.")
-            if tiler[1] not in (128, 256):
-                raise ValueError("GenPhase requires tile N 128 or 256.")
+            if tiler[1] not in (128, 256) and not n64_decode:
+                raise ValueError(
+                    "GenPhase requires tile N 128 or 256, or an admitted N64 decode geometry."
+                )
             if tiler[2] != 2 * instruction_k:
                 raise ValueError(
                     f"GenPhase {self.quant_kind} requires tile K {2 * instruction_k} "
@@ -733,7 +747,7 @@ def launch(
         if self.config.reduce_topk_in_kernel:
             # red.add accumulation base: zero the output every launch.
             self.output_activation.zero_()
-        else:
+        elif not getattr(self.kernel, "masks_invalid_combine_slots", False):
             # Invalid routes do not write a combine slot. Clear those slots on
             # every launch so a masked route cannot reuse an earlier result.
             invalid = (self.topk_idx < 0) | (

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/src/sources/helpers/ptx_helpers.py

@@ -749,3 +749,18 @@ def stg_e8m0x8_from_f32(
     "stg_f32",
     "tma_load_1d",
 ]
+
+
+@dsl_user_op
+def barrier_sync_unaligned(barrier_id, num_threads, *, loc=None, ip=None):
+    llvm.inline_asm(
+        res=None,
+        operands_=[Int32(barrier_id).ir_value(loc=loc, ip=ip), Int32(num_threads).ir_value(loc=loc, ip=ip)],
+        asm_string="barrier.cta.sync $0, $1;",
+        constraints="r,r,~{memory}",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/src/sources/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py

@@ -3,6 +3,7 @@
 from typing import ClassVar, Optional, Tuple
 
 import cutlass
+from .....helpers.ptx_helpers import barrier_sync_unaligned
 import cutlass.cute as cute
 import cutlass.pipeline as pipeline
 import cutlass.utils as utils
@@ -1061,7 +1062,7 @@ def run_mma(
         fragment_a = tiled_mma.make_fragment_A(self.a_smem_tensor)
         fragment_b = tiled_mma.make_fragment_B(self.b_smem_tensor)
 
-        tmem_allocator.wait_for_alloc()
+        barrier_sync_unaligned(self.tmem_allocation_barrier_id, 32 * (1 + self.num_accumulator_consumer_warps_per_cta))
         accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype)
         accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride)
         accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout)