diff --git a/flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py b/flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py index a3c2a0a938f..ac6b80a0080 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py +++ b/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 diff --git a/flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py b/flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py index 21a33b23453..229bd7a4105 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py +++ b/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) diff --git a/flashinfer/fused_moe/cute_dsl/tuner.py b/flashinfer/fused_moe/cute_dsl/tuner.py index 1bd1fc3c8ae..e8a32576d4f 100644 --- a/flashinfer/fused_moe/cute_dsl/tuner.py +++ b/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 diff --git a/tests/moe/test_cute_dsl_fused_moe.py b/tests/moe/test_cute_dsl_fused_moe.py index f516ecbf620..89b0d6dbe02 100644 --- a/tests/moe/test_cute_dsl_fused_moe.py +++ b/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, diff --git a/tests/moe/test_cute_dsl_moe_can_implement.py b/tests/moe/test_cute_dsl_moe_can_implement.py new file mode 100644 index 00000000000..386eaa173e5 --- /dev/null +++ b/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