fix(moe): handle CuTe DSL finalize output tails

flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py

@@ -3656,6 +3656,16 @@ def can_implement(
         # Skip unsupported A/B layout
         if not (a_major == "k" and b_major == "k"):
             can_implement = False
+        # N must not have a partial CTA tile (gh #3957 sibling): this kernel's
+        # SFC global store (autovec_copy -- see the epilogue TODO about the
+        # missing predicate) writes the full tile row with no column predicate,
+        # so a partial N-tile writes out of bounds past the intermediate
+        # buffer's row -- same class as the finalize kernel's bulk-reduce
+        # scatter OOB. Cluster padding along N cannot occur here: this kernel
+        # already requires cluster_shape_mn[1] == 1 above. M needs no
+        # analogue: rows are individually validity-guarded.
+        if n % mma_tiler_mn[1] != 0:
+            can_implement = False
         return can_implement
 
     @cute.jit

flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py

@@ -836,7 +836,6 @@ def __call__(
         epi_tile_size = epi_tile_m * epi_tile_n
         num_epilogue_threads = 32 * len(self.epilog_warp_id)
         self.ttr_racc_size = epi_tile_size // num_epilogue_threads
-        self.copy_size = self.cta_tile_shape_mnk[1] * (self.out_dtype.width // 8)
 
         if cutlass.const_expr(self.out_dtype == cutlass.BFloat16):
             # 8-element vectorization for BF16
@@ -2153,37 +2152,48 @@ def kernel(
                 reduce_permuted_row = tile_m_start + reduce_row
                 is_valid_reduce_row = reduce_permuted_row < tile_info[4]
                 if is_valid_reduce_row:
-                    reduce_token_idx = sMetaTokenIdx[
-                        (reduce_row, meta_consumer_state.index)
-                    ]
                     coord_n = tile_info[1] * self.cta_tile_shape_mnk[1]
-                    scatter_out_offset = cute.domain_offset(
-                        (reduce_token_idx, coord_n, 0), out
+                    valid_columns = cutlass.min(
+                        cutlass.Int64(out.shape[1]) - coord_n,
+                        cutlass.Int64(self.cta_tile_shape_mnk[1]),
                     )
-                    if cutlass.const_expr(not self.use_fused_finalize):
-                        blk_copy(
-                            scatter_out_offset,
-                            sC[reduce_row, None, 0],
-                            cutlass.Int32(self.copy_size),
-                        )
-                    elif cutlass.const_expr(self.out_dtype == cutlass.BFloat16):
-                        blk_reduce_bf16(
-                            scatter_out_offset,
-                            sC[reduce_row, None, 0],
-                            cutlass.Int32(self.copy_size),
+                    if valid_columns > 0:
+                        reduce_token_idx = sMetaTokenIdx[
+                            (reduce_row, meta_consumer_state.index)
+                        ]
+                        scatter_out_offset = cute.domain_offset(
+                            (reduce_token_idx, coord_n, 0), out
                         )
-                    elif cutlass.const_expr(self.out_dtype == cutlass.Float32):
-                        blk_reduce_fp32(
-                            scatter_out_offset,
-                            sC[reduce_row, None, 0],
-                            cutlass.Int32(self.copy_size),
-                        )
-                    elif cutlass.const_expr(self.out_dtype == cutlass.Float16):
-                        blk_reduce_fp16(
-                            scatter_out_offset,
-                            sC[reduce_row, None, 0],
-                            cutlass.Int32(self.copy_size),
+                        valid_copy_size = cutlass.Int32(
+                            valid_columns * (self.out_dtype.width // 8)
                         )
+                        # is_valid_tensor_alignment requires each output row to
+                        # end on a 16-byte boundary, matching the bulk-copy
+                        # instruction's size and address requirements.
+                        if cutlass.const_expr(not self.use_fused_finalize):
+                            blk_copy(
+                                scatter_out_offset,
+                                sC[reduce_row, None, 0],
+                                valid_copy_size,
+                            )
+                        elif cutlass.const_expr(self.out_dtype == cutlass.BFloat16):
+                            blk_reduce_bf16(
+                                scatter_out_offset,
+                                sC[reduce_row, None, 0],
+                                valid_copy_size,
+                            )
+                        elif cutlass.const_expr(self.out_dtype == cutlass.Float32):
+                            blk_reduce_fp32(
+                                scatter_out_offset,
+                                sC[reduce_row, None, 0],
+                                valid_copy_size,
+                            )
+                        elif cutlass.const_expr(self.out_dtype == cutlass.Float16):
+                            blk_reduce_fp16(
+                                scatter_out_offset,
+                                sC[reduce_row, None, 0],
+                                valid_copy_size,
+                            )
 
                 cute.arch.cp_async_bulk_commit_group()
                 cute.arch.cp_async_bulk_wait_group(0, read=True)

flashinfer/fused_moe/cute_dsl/tuner.py

@@ -481,8 +481,7 @@ def get_valid_tactics(  # type: ignore[override]
         else:
             final_scale_dtype = cutlass.Float16
 
-        valid_tactics = []
-        for tactic in ALL_MOE_TACTICS:
+        def _tactic_ok(tactic):
             tile_size, gemm1_tactic, gemm2_tactic = tactic
             gemm1_mma_tiler_mn = gemm1_tactic[0]
             gemm1_cluster_shape_mn = gemm1_tactic[1]
@@ -528,10 +527,18 @@ def get_valid_tactics(  # type: ignore[override]
                 )
             )
 
-            if gemm1_ok and gemm2_ok:
-                valid_tactics.append(tactic)
+            return gemm1_ok and gemm2_ok
+
+        valid_tactics = [t for t in ALL_MOE_TACTICS if _tactic_ok(t)]
 
         if not valid_tactics:
+            # DEFAULT_MOE_TACTIC is a member of ALL_MOE_TACTICS, so an empty
+            # list means even the default fails can_implement -- do not fall
+            # back to it unvalidated (gh #3957). This early refusal is
+            # diagnostics/defense-in-depth: the kernel wrappers re-validate
+            # can_implement at launch and raise, so an unvalidated tactic
+            # cannot reach the device -- but refusing here avoids pointless
+            # profiling of a tactic that can only throw, and says why.
             logger.warning(
                 "No valid tactics found for problem dims "
                 "(tokens=%d, hidden=%d, intermediate=%d, experts=%d, top_k=%d). "
@@ -542,7 +549,6 @@ def get_valid_tactics(  # type: ignore[override]
                 num_local_experts,
                 self.top_k,
             )
-            valid_tactics = [DEFAULT_MOE_TACTIC]
 
         return valid_tactics
 

tests/moe/test_cute_dsl_fused_moe.py

@@ -1036,6 +1036,41 @@ def test_deterministic_finalize_numerical_accuracy(
             use_fused_finalize=False,
         )
 
+    @pytest.mark.parametrize("hidden_size", [256, 384])
+    def test_finalize_handles_cluster_padding_and_partial_n_tiles(
+        self,
+        hidden_size: int,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        from flashinfer.autotuner import AutoTuner
+
+        # hidden=256 leaves one padding CTA in the N=256, cluster_n=2
+        # configuration. hidden=384 also gives the second CTA a partial tile.
+        # Force the 256-row route so both cases execute the kernel path that
+        # previously had to be filtered out.
+        tail_config = (
+            256,
+            ((256, 128), (2, 1), False),
+            ((256, 256), (2, 2), False),
+        )
+
+        def choose_tail_config(
+            _self, _custom_op, runners, _tuning_config, _inputs, **_kwargs
+        ):
+            return runners[0], tail_config
+
+        monkeypatch.setattr(AutoTuner, "choose_one", choose_tail_config)
+        self._run_numerical_accuracy(
+            activation_type=ActivationType.Relu2,
+            num_tokens=128,
+            top_k=2,
+            hidden_size=hidden_size,
+            intermediate_size=512,
+            num_experts=8,
+            use_per_token_activation=True,
+            use_fused_finalize=True,
+        )
+
     def _run_numerical_accuracy(
         self,
         activation_type: ActivationType,

tests/moe/test_cute_dsl_moe_can_implement.py

@@ -0,0 +1,89 @@
+"""Directed host-side checks for the gh #3957 N-tail handling.
+
+The finalize epilogue limits its bulk transfer to the remaining output columns,
+including for cluster-padding CTAs. The gemm1 SFC store is still unpredicated,
+so only gemm1 must reject configurations that leave a partial N tile. These are
+pure classmethod checks -- no GPU work.
+"""
+
+import pytest
+
+cutlass = pytest.importorskip("cutlass")
+
+from flashinfer.fused_moe.cute_dsl.blackwell.blockscaled_contiguous_gather_grouped_gemm_act_fusion import (  # noqa: E501
+    BlockScaledContiguousGatherGroupedGemmKernel,
+)
+from flashinfer.fused_moe.cute_dsl.blackwell.blockscaled_contiguous_grouped_gemm_finalize_fusion import (  # noqa: E501
+    Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel,
+)
+
+
+def _finalize_ok(n, mma_tiler_mn, cluster_shape_mn):
+    return Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel.can_implement(
+        ab_dtype=cutlass.Float4E2M1FN,
+        sf_dtype=cutlass.Float8E4M3FN,
+        sf_vec_size=16,
+        out_dtype=cutlass.BFloat16,
+        final_scale_dtype=cutlass.Float32,
+        mma_tiler_mn=mma_tiler_mn,
+        cluster_shape_mn=cluster_shape_mn,
+        m=1024,
+        n=n,
+        k=512,
+        l=8,
+        a_major="k",
+        b_major="k",
+        out_major="n",
+    )
+
+
+def _gemm1_ok(n, mma_tiler_mn, cluster_shape_mn):
+    return BlockScaledContiguousGatherGroupedGemmKernel.can_implement(
+        ab_dtype=cutlass.Float4E2M1FN,
+        sf_dtype=cutlass.Float8E4M3FN,
+        sf_vec_size=16,
+        c_dtype=cutlass.Float4E2M1FN,
+        mma_tiler_mn=mma_tiler_mn,
+        cluster_shape_mn=cluster_shape_mn,
+        m=1024,
+        n=n,
+        k=512,
+        l=8,
+        a_major="k",
+        b_major="k",
+        c_major="n",
+    )
+
+
+@pytest.mark.parametrize(
+    "n,mma,cluster,expect",
+    [
+        # A padding CTA has no remaining columns and skips its row transfer.
+        (256, (128, 256), (1, 2), True),
+        # 2 exact tiles / cluster_n=2 -> exact cluster tiling: fine.
+        (512, (128, 256), (1, 2), True),
+        # A partial tile transfers only columns 256..383.
+        (384, (128, 256), (1, 2), True),
+        # The same tail handling applies without an N cluster.
+        (384, (128, 256), (1, 1), True),
+        # 3 exact 128-tiles, no cluster: fine.
+        (384, (128, 128), (1, 1), True),
+    ],
+)
+def test_finalize_n_tiling(n, mma, cluster, expect):
+    assert _finalize_ok(n, mma, cluster) is expect
+
+
+@pytest.mark.parametrize(
+    "n,mma,expect",
+    [
+        # The scale-factor store requires exact tiling under mma_n=256.
+        (384, (128, 256), False),
+        # Exact tiling: fine.
+        (512, (128, 256), True),
+        (384, (128, 128), True),
+    ],
+)
+def test_gemm1_n_tiling_guard(n, mma, expect):
+    # gemm1 requires cluster_n == 1 (enforced independently).
+    assert _gemm1_ok(n, mma, (1, 1)) is expect