"""
Copyright (c) 2025 by FlashInfer team.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

"""
Numerical accuracy tests for CuteDSL Fused MoE NVFP4 on Blackwell and Rubin GPUs.

This test file covers both APIs:
1. Functional API: `cute_dsl_fused_moe_nvfp4`
2. Wrapper API: `CuteDslMoEWrapper`

Tests include:
- Numerical accuracy against reference implementation
- CUDA graph capture and replay
- Auto-tuning integration
- API consistency between functional and wrapper APIs
"""

import gc
import weakref

import pytest
import torch

from flashinfer.cute_dsl.utils import is_cute_dsl_arch_supported

_requires_dsl_arch = pytest.mark.skipif(
    torch.cuda.is_available()
    and not is_cute_dsl_arch_supported(
        *torch.cuda.get_device_capability(0), native_only=True
    ),
    reason="installed CuTe DSL does not support this GPU architecture",
)


from flashinfer.tllm_enums import ActivationType
from flashinfer.fused_moe.cute_dsl.moe_utils import (
    normalize_cute_dsl_moe_activation_type,
    validate_cute_dsl_moe_situ_config,
)
from flashinfer.cute_dsl import is_cute_dsl_available
from .utils import (
    check_accuracy,
    compute_reference_moe_fp4,
    create_moe_tensors,
)


def is_sm100_family():
    """Check for SM100 family (Blackwell: SM100, SM103).

    CuteDSL MoE NVFP4 does not target Rubin SM107.
    """
    if not torch.cuda.is_available():
        return False
    props = torch.cuda.get_device_properties(0)
    return (props.major, props.minor) in ((10, 0), (10, 3))


# Skip decorators
cute_dsl_available = pytest.mark.skipif(
    not is_cute_dsl_available(), reason="CuteDSL not available"
)
sm100_required = pytest.mark.skipif(
    not is_sm100_family(),
    reason="Requires CuteDSL MoE target SM100 or SM103",
)


_MOE_QUANT_MODES = ("w4a4", "w4a16")


_MOE_QUANT_MODE_CASES = (
    pytest.param("w4a4", False, id="w4a4-per-tensor"),
    pytest.param("w4a4", True, id="w4a4-per-token"),
    pytest.param("w4a16", False, id="w4a16"),
)


def _prepare_moe_quant_mode_inputs(
    tensors: dict,
    quant_mode: str,
) -> tuple[dict, dict]:
    """Select API and reference inputs for a shared MoE test case."""
    if quant_mode == "w4a4":
        return (
            {
                "x": tensors["x"],
                "x_sf": tensors["x_sf"],
                "fc2_input_scale": tensors["fc2_input_scale"],
                "per_token_scale": tensors["x_per_token_scale"],
            },
            {
                "hidden_states": tensors["x_ref"].float(),
                "gemm1_weights": tensors["w1_weight_bf16"].float(),
                "gemm2_weights": tensors["w2_weight_bf16"].float(),
                "gemm1_alpha": tensors["w1_alpha"],
                "gemm2_alpha": tensors["w2_alpha"],
                "fc2_input_scale": tensors["fc2_input_scale"],
                "use_per_token_activation": tensors["x_per_token_scale"] is not None,
            },
        )
    elif quant_mode == "w4a16":
        return (
            {
                "x": tensors["x_bf16"],
                "x_sf": None,
                "fc2_input_scale": None,
                "per_token_scale": None,
            },
            {
                "hidden_states": tensors["x_bf16"],
                "gemm1_weights": (
                    tensors["w1_weight_bf16"].float()
                    * tensors["w1_alpha"].view(-1, 1, 1)
                ).to(torch.bfloat16),
                "gemm2_weights": (
                    tensors["w2_weight_bf16"].float()
                    * tensors["w2_alpha"].view(-1, 1, 1)
                ).to(torch.bfloat16),
                "fc2_input_scale": None,
                "use_per_token_activation": False,
            },
        )
    else:
        raise ValueError(f"Unsupported test quant_mode {quant_mode!r}")


def test_geglu_tanh_activation_is_supported():
    activation_type, gated = normalize_cute_dsl_moe_activation_type(
        ActivationType.GegluTanh
    )
    assert activation_type == ActivationType.GegluTanh
    assert gated


@pytest.mark.parametrize(
    "activation_type,situ_beta,situ_linear_beta,error",
    [
        (ActivationType.Swiglu, None, 1.0, "requires situ_beta"),
        (ActivationType.Swiglu, 0.0, None, "positive and finite"),
        (ActivationType.GegluTanh, 1.0, None, "require ActivationType.Swiglu"),
    ],
)
def test_invalid_situ_config(activation_type, situ_beta, situ_linear_beta, error: str):
    with pytest.raises(ValueError, match=error):
        validate_cute_dsl_moe_situ_config(activation_type, situ_beta, situ_linear_beta)


def test_situ_changes_autotuner_cache_key():
    from flashinfer.fused_moe.cute_dsl.tuner import CuteDslFusedMoENvfp4Runner

    kwargs = {
        "forward_impl": lambda *args, **kwargs: None,
        "num_experts": 8,
        "top_k": 2,
        "num_local_experts": 8,
    }
    swiglu_runner = CuteDslFusedMoENvfp4Runner(**kwargs)
    situ_runner = CuteDslFusedMoENvfp4Runner(**kwargs, situ_beta=1.0)
    situ_beta_runner = CuteDslFusedMoENvfp4Runner(**kwargs, situ_beta=2.0)
    situ_linear_runner = CuteDslFusedMoENvfp4Runner(
        **kwargs, situ_beta=1.0, situ_linear_beta=1.5
    )
    runners = [swiglu_runner, situ_runner, situ_beta_runner, situ_linear_runner]
    assert len({hash(runner) for runner in runners}) == len(runners)
    assert len({runner.get_cache_key_extras([]) for runner in runners}) == len(runners)


# =============================================================================
# Test Class: GEMM input validation
# =============================================================================


@cute_dsl_available
@sm100_required
class TestKernelInputValidation:
    @staticmethod
    def _gather_kwargs():
        device = "cuda"
        return {
            "a": torch.empty((4, 64), dtype=torch.uint8, device=device),
            "b": torch.empty((2, 256, 64), dtype=torch.uint8, device=device),
            "a_scale": torch.empty(1, dtype=torch.uint8, device=device),
            "b_scale": torch.empty(1, dtype=torch.uint8, device=device),
            "alpha": torch.empty(2, dtype=torch.float32, device=device),
            "tile_idx_to_expert_idx": torch.empty(1, dtype=torch.int32, device=device),
            "tile_idx_to_mn_limit": torch.empty(1, dtype=torch.int32, device=device),
            "token_id_mapping": torch.empty(4, dtype=torch.int32, device=device),
            "num_non_exiting_tiles": torch.empty(1, dtype=torch.int32, device=device),
        }

    @staticmethod
    def _finalize_kwargs():
        device = "cuda"
        return {
            "a": torch.empty((4, 64), dtype=torch.uint8, device=device),
            "b": torch.empty((2, 128, 64), dtype=torch.uint8, device=device),
            "a_scale": torch.empty(1, dtype=torch.uint8, device=device),
            "b_scale": torch.empty(1, dtype=torch.uint8, device=device),
            "alpha": torch.empty(2, dtype=torch.float32, device=device),
            "tile_idx_to_expert_idx": torch.empty(1, dtype=torch.int32, device=device),
            "num_non_exiting_tiles": torch.empty(1, dtype=torch.int32, device=device),
            "tile_idx_to_mn_limit": torch.empty(1, dtype=torch.int32, device=device),
            "permuted_idx_to_expanded_idx": torch.empty(
                4, dtype=torch.int32, device=device
            ),
            "token_final_scales": torch.empty(
                (2, 2), dtype=torch.float32, device=device
            ),
        }

    @pytest.mark.parametrize("stage", ["gather", "finalize"])
    @pytest.mark.parametrize(
        ("invalid_kind", "match"),
        [
            ("device", "CUDA device"),
            ("dtype", "dtype torch.float32"),
            ("contiguous", "contiguous"),
            ("length", "must have shape"),
            ("rank", "must have shape"),
        ],
    )
    def test_rejects_invalid_per_token_scale(self, stage, invalid_kind, match):
        if stage == "gather":
            from flashinfer.fused_moe.cute_dsl.blockscaled_contiguous_gather_grouped_gemm_act_fusion import (
                blockscaled_contiguous_gather_grouped_gemm_act_fusion_nvfp4,
            )

            op = blockscaled_contiguous_gather_grouped_gemm_act_fusion_nvfp4
            kwargs = self._gather_kwargs()
        else:
            from flashinfer.fused_moe.cute_dsl.blockscaled_contiguous_grouped_gemm_finalize_fusion import (
                blockscaled_contiguous_grouped_gemm_finalize_fusion_nvfp4,
            )

            op = blockscaled_contiguous_grouped_gemm_finalize_fusion_nvfp4
            kwargs = self._finalize_kwargs()

        expected_rows = kwargs["a"].shape[0]
        if invalid_kind == "device":
            per_token_scale = torch.ones(expected_rows, dtype=torch.float32)
        elif invalid_kind == "dtype":
            per_token_scale = torch.ones(
                expected_rows, dtype=torch.float16, device="cuda"
            )
        elif invalid_kind == "contiguous":
            per_token_scale = torch.ones(
                expected_rows * 2, dtype=torch.float32, device="cuda"
            )[::2]
        elif invalid_kind == "length":
            per_token_scale = torch.ones(
                expected_rows + 1, dtype=torch.float32, device="cuda"
            )
        else:
            per_token_scale = torch.ones(
                (2, expected_rows // 2), dtype=torch.float32, device="cuda"
            )

        with pytest.raises(ValueError, match=match):
            op(**kwargs, a_per_token_scale=per_token_scale)

    @pytest.mark.parametrize("scale_name", ["out_scale", "global_scale"])
    def test_rejects_output_scales_for_non_fp4_output(self, scale_name):
        from flashinfer.fused_moe.cute_dsl.blockscaled_contiguous_gather_grouped_gemm_act_fusion import (
            blockscaled_contiguous_gather_grouped_gemm_act_fusion_nvfp4,
        )

        kwargs = self._gather_kwargs()
        kwargs[scale_name] = torch.ones(1, dtype=torch.float32, device="cuda")
        with pytest.raises(ValueError, match="only supported"):
            blockscaled_contiguous_gather_grouped_gemm_act_fusion_nvfp4(
                **kwargs, c_dtype="bfloat16"
            )


# =============================================================================
# Test Class: Tactic-enumeration structural invariants (no GPU required)
# =============================================================================


@cute_dsl_available
class TestTacticEnumeration:
    """Structural invariants for the tactic-enumeration helpers in
    flashinfer.fused_moe.cute_dsl.tuner.

    These tests run without a GPU. They exercise the enumeration
    functions directly to enforce invariants that the end-to-end
    accuracy tests can fail to detect when a tile size is gated out of
    ALL_MOE_TACTICS as a workaround.

    The MoE pipeline runs gemm1 (gather + SwiGLU) followed by gemm2
    (finalize fusion) back-to-back on the same padded token sequence.
    For the layouts to match, gemm1 and gemm2 must share the same
    mma_tiler M dimension and the same cluster_shape M dimension.
    """

    @pytest.mark.parametrize("tile_size", [128, 256])
    def test_gemm1_tactics_match_tile_size(self, tile_size):
        """Every gemm1 tactic must have mma_tiler[0] == tile_size and
        cluster_shape[0] == tile_size // 128 (1-CTA at tile=128, 2-CTA
        at tile=256)."""
        from flashinfer.fused_moe.cute_dsl.tuner import (
            get_gemm1_valid_tactics,
        )

        tactics = get_gemm1_valid_tactics(tile_size)
        assert len(tactics) > 0, f"no gemm1 tactics returned at tile_size={tile_size}"
        expected_cluster_m = tile_size // 128
        for mma_tiler_mn, cluster_shape_mn, _ in tactics:
            assert mma_tiler_mn[0] == tile_size, (
                f"gemm1 mma_tiler[0]={mma_tiler_mn[0]} does not match "
                f"tile_size={tile_size}; tactic={(mma_tiler_mn, cluster_shape_mn)}"
            )
            assert cluster_shape_mn[0] == expected_cluster_m, (
                f"gemm1 cluster_shape[0]={cluster_shape_mn[0]} does not "
                f"match tile_size//128={expected_cluster_m}; "
                f"tactic={(mma_tiler_mn, cluster_shape_mn)}"
            )

    @pytest.mark.parametrize("tile_size", [128, 256])
    def test_gemm2_tactics_match_tile_size(self, tile_size):
        """Every gemm2 tactic must have mma_tiler[0] == tile_size and
        cluster_shape[0] == tile_size // 128. The finalize kernel
        consumes the upstream gemm1 output layout — a 1-CTA gemm2
        tactic at tile_size=256 cannot consume a 2-CTA gemm1 output
        and produces incorrect results (regression for #3067)."""
        from flashinfer.fused_moe.cute_dsl.tuner import (
            get_gemm2_valid_tactics,
        )

        tactics = get_gemm2_valid_tactics(tile_size)
        assert len(tactics) > 0, f"no gemm2 tactics returned at tile_size={tile_size}"
        expected_cluster_m = tile_size // 128
        for mma_tiler_mn, cluster_shape_mn, _ in tactics:
            assert mma_tiler_mn[0] == tile_size, (
                f"gemm2 mma_tiler[0]={mma_tiler_mn[0]} does not match "
                f"tile_size={tile_size}; tactic={(mma_tiler_mn, cluster_shape_mn)}"
            )
            assert cluster_shape_mn[0] == expected_cluster_m, (
                f"gemm2 cluster_shape[0]={cluster_shape_mn[0]} does not "
                f"match tile_size//128={expected_cluster_m}; "
                f"tactic={(mma_tiler_mn, cluster_shape_mn)}"
            )

    def test_all_moe_tactics_pair_gemm1_and_gemm2_consistently(self):
        """Every (tile_size, gemm1_tactic, gemm2_tactic) tuple in
        ALL_MOE_TACTICS must have gemm1 and gemm2 share both
        mma_tiler[0] and cluster_shape[0] (the M dimensions). This
        catches a class of bug where the product loop in
        get_moe_valid_tactics accidentally pairs incompatible
        gemm1/gemm2 tactics, even if each individual enumeration is
        internally consistent."""
        from flashinfer.fused_moe.cute_dsl.tuner import ALL_MOE_TACTICS

        assert len(ALL_MOE_TACTICS) > 0
        for tile_size, gemm1_tactic, gemm2_tactic in ALL_MOE_TACTICS:
            gemm1_mma_m = gemm1_tactic[0][0]
            gemm1_cluster_m = gemm1_tactic[1][0]
            gemm2_mma_m = gemm2_tactic[0][0]
            gemm2_cluster_m = gemm2_tactic[1][0]
            assert gemm1_mma_m == gemm2_mma_m == tile_size, (
                f"gemm1/gemm2 mma_m mismatch in ALL_MOE_TACTICS at "
                f"tile_size={tile_size}: gemm1_mma_m={gemm1_mma_m}, "
                f"gemm2_mma_m={gemm2_mma_m}"
            )
            assert gemm1_cluster_m == gemm2_cluster_m == tile_size // 128, (
                f"gemm1/gemm2 cluster_m mismatch in ALL_MOE_TACTICS at "
                f"tile_size={tile_size}: gemm1_cluster_m={gemm1_cluster_m}, "
                f"gemm2_cluster_m={gemm2_cluster_m}"
            )

    def test_w4a16_tactic_enumeration_invariants(self):
        from flashinfer.fused_moe.cute_dsl.blackwell.moe_w4a16 import (
            DEFAULT_W4A16_MOE_TACTIC,
        )
        from flashinfer.fused_moe.cute_dsl.tuner import W4A16_MOE_TACTICS

        assert len(W4A16_MOE_TACTICS) == len(set(W4A16_MOE_TACTICS))
        assert DEFAULT_W4A16_MOE_TACTIC in W4A16_MOE_TACTICS

        valid_topologies = {
            (128, (1, 1)),
            (128, (2, 1)),
            (256, (2, 1)),
        }
        for gemm1_tactic, gemm2_tactic in W4A16_MOE_TACTICS:
            assert gemm1_tactic[0][1] == gemm2_tactic[0][1]
            for mma_tiler, cluster_shape, _ in (
                gemm1_tactic,
                gemm2_tactic,
            ):
                mma_m, route_tile, mma_k = mma_tiler
                assert (mma_m, cluster_shape) in valid_topologies
                assert mma_k % 16 == 0
                if route_tile < 16:
                    assert mma_m == 128
                elif route_tile == 192:
                    assert mma_m == 256


# =============================================================================
# Test Class: CuteDslMoEInputsHelper.inputs_pre_hook layout contract
# (no GPU required)
# =============================================================================


@cute_dsl_available
class TestInputsHelperContract:
    """Structural invariants for ``CuteDslMoEInputsHelper.inputs_pre_hook``.

    Tests run without a GPU. They exercise the cross-file contract that
    the hook's unpacking pattern (``x, x_sf, tse, *rest = inputs``) must
    match the input list layout produced by ``CuteDslMoEWrapper.run`` so
    that autotune profile inputs are not silently corrupted by a
    refactor that reorders the wrapper's inputs list.
    """

    def _build_synthetic_inputs(
        self,
        num_tokens: int,
        num_local_experts: int,
        use_per_token_activation: bool = False,
    ):
        """Mirror ``CuteDslMoEWrapper.run``'s inputs-list layout with
        small-but-shape-faithful tensors so the test runs in <1s on CPU."""
        n = num_tokens
        # Small dimensions for fast CPU allocation; sizes only matter for
        # shape checks, not numerical results.
        hidden = 128
        intermediate = 64
        top_k = 8
        sf_vec = 16
        inputs = [
            torch.zeros(n, hidden // 2, dtype=torch.uint8),  # 0: x
            torch.zeros(n, hidden // sf_vec, dtype=torch.uint8),  # 1: x_sf
            torch.zeros(n, top_k, dtype=torch.int32),  # 2: token_selected_experts
            torch.zeros(n, top_k, dtype=torch.float32),  # 3: token_final_scales
            torch.zeros(
                num_local_experts, 2 * intermediate, hidden // 2, dtype=torch.uint8
            ),  # 4: w1_weight
            torch.zeros(
                num_local_experts, 2 * intermediate, hidden // sf_vec, dtype=torch.uint8
            ),  # 5: w1_weight_sf
            torch.zeros(num_local_experts, dtype=torch.float32),  # 6: w1_alpha
            torch.zeros(num_local_experts, dtype=torch.float32),  # 7: fc2_input_scale
            torch.zeros(
                num_local_experts, hidden, intermediate // 2, dtype=torch.uint8
            ),  # 8: w2_weight
            torch.zeros(
                num_local_experts, hidden, intermediate // sf_vec, dtype=torch.uint8
            ),  # 9: w2_weight_sf
            torch.zeros(num_local_experts, dtype=torch.float32),  # 10: w2_alpha
        ]
        if use_per_token_activation:
            inputs.append(torch.ones(n, dtype=torch.float32))  # per_token_scale
        inputs.append(torch.zeros(n, hidden, dtype=torch.bfloat16))  # moe_output
        return inputs

    @pytest.mark.parametrize("use_per_token_activation", [False, True])
    def test_hook_replaces_input_2_and_passes_through_rest(
        self, use_per_token_activation: bool
    ):
        """``inputs_pre_hook`` must replace ``inputs[2]``
        (token_selected_experts) and pass through every other input
        unchanged. Pins the contract with ``CuteDslMoEWrapper.run`` —
        if someone reorders the inputs list (e.g. moves x_sf or a
        weight tensor) without updating the helper's unpacking, the
        autotune profile silently corrupts a different tensor."""
        from flashinfer.fused_moe.cute_dsl._inputs_helper import (
            CuteDslMoEInputsHelper,
        )

        num_local_experts = 16
        helper = CuteDslMoEInputsHelper(
            num_experts=256,
            top_k=8,
            num_local_experts=num_local_experts,
            local_expert_offset=0,
        )

        inputs = self._build_synthetic_inputs(
            num_tokens=64,
            num_local_experts=num_local_experts,
            use_per_token_activation=use_per_token_activation,
        )
        original_tse = inputs[2]

        output = helper.inputs_pre_hook(inputs)

        expected_len = 13 if use_per_token_activation else 12
        assert len(output) == expected_len, (
            f"Expected {expected_len} outputs, got {len(output)}"
        )
        # Index 2 must be replaced (different object identity), with
        # the same shape and dtype.
        assert output[2] is not original_tse, (
            "inputs[2] (token_selected_experts) must be REPLACED by the hook, "
            "not passed through. Hook implementation in _inputs_helper.py "
            "must build a fresh tensor for input #2."
        )
        assert output[2].shape == original_tse.shape, (
            f"Replaced tse has shape {output[2].shape}, expected {original_tse.shape}"
        )
        assert output[2].dtype == original_tse.dtype, (
            f"Replaced tse has dtype {output[2].dtype}, expected {original_tse.dtype}"
        )
        # Every other input MUST pass through with object identity preserved.
        # If this breaks, the hook is mutating something it shouldn't, OR the
        # wrapper's inputs-list ordering has drifted from the hook's unpacking.
        for i in range(len(inputs)):
            if i == 2:
                continue
            assert output[i] is inputs[i], (
                f"inputs[{i}] must pass through the hook unchanged (object identity). "
                f"This typically indicates the inputs-list ordering in "
                f"CuteDslMoEWrapper.run has drifted from the hook's unpacking pattern "
                f"in CuteDslMoEInputsHelper.inputs_pre_hook."
            )

    def test_hook_is_deterministic_across_helper_instances(self):
        """Two ``CuteDslMoEInputsHelper`` instances with the same seed
        must produce tensor-equal replacement ``token_selected_experts``
        for the same input. This is the core determinism property the
        helper provides — cross-process autotune-pick variance is
        eliminated only if seeded sampling is deterministic."""
        from flashinfer.fused_moe.cute_dsl._inputs_helper import (
            CuteDslMoEInputsHelper,
        )

        helper_a = CuteDslMoEInputsHelper(
            num_experts=256, top_k=8, num_local_experts=16, local_expert_offset=0
        )
        helper_b = CuteDslMoEInputsHelper(
            num_experts=256, top_k=8, num_local_experts=16, local_expert_offset=0
        )

        inputs_a = self._build_synthetic_inputs(num_tokens=64, num_local_experts=16)
        inputs_b = self._build_synthetic_inputs(num_tokens=64, num_local_experts=16)

        out_a = helper_a.inputs_pre_hook(inputs_a)
        out_b = helper_b.inputs_pre_hook(inputs_b)

        assert torch.equal(out_a[2], out_b[2]), (
            "Two helpers with the same seed produced different "
            "token_selected_experts tensors. The seeded "
            "torch.random.fork_rng + manual_seed pattern in "
            "generate_token_selected_experts is broken."
        )


# =============================================================================
# Test Class: autotune replay stream contract (no GPU required)
# =============================================================================


@cute_dsl_available
class TestAutotuneReplayMemsetContract:
    @pytest.mark.parametrize("api", ["functional", "wrapper"])
    @pytest.mark.parametrize("is_tuning_mode", [False, True])
    @pytest.mark.parametrize("quant_mode", ["w4a4", "w4a16"])
    def test_selected_tactic_memset_stream_contract(
        self, monkeypatch, api, is_tuning_mode, quant_mode
    ):
        from flashinfer.fused_moe.cute_dsl import fused_moe

        # CPU-only contract test (see class header): stub out the functional
        # API's arch guard, which otherwise calls torch.cuda.get_device_capability()
        # on the CPU input tensors and raises "Expected a cuda device, but got: cpu".
        monkeypatch.setattr(
            fused_moe, "_require_cute_dsl_arch_for", lambda *a, **k: None
        )

        calls = []

        class RecordingRunner:
            def __init__(self, *args, **kwargs):
                self.tuning_config = object()

            def __call__(self, inputs, **kwargs):
                calls.append(kwargs)
                return inputs[-1]

        class StubTuner:
            def __init__(self):
                self.is_tuning_mode = is_tuning_mode

            def choose_one(self, custom_op, runners, tuning_config, inputs, **kwargs):
                assert "use_async_memset" not in kwargs
                return runners[0], ("selected",)

        monkeypatch.setattr(
            fused_moe.AutoTuner, "get", staticmethod(lambda: StubTuner())
        )
        monkeypatch.setattr(
            fused_moe,
            "_require_cute_dsl_arch_for",
            lambda *_args, **_kwargs: None,
        )

        tensors = {
            "x": torch.empty((2, 8), dtype=torch.uint8),
            "x_sf": torch.empty((2, 1), dtype=torch.uint8),
            "token_selected_experts": torch.zeros((2, 1), dtype=torch.int32),
            "token_final_scales": torch.ones((2, 1), dtype=torch.float32),
            "w1_weight": torch.empty((1, 32, 8), dtype=torch.uint8),
            "w1_weight_sf": torch.empty((1, 32, 1), dtype=torch.uint8),
            "w1_alpha": torch.ones(1, dtype=torch.float32),
            "fc2_input_scale": torch.ones(1, dtype=torch.float32),
            "w2_weight": torch.empty((1, 16, 8), dtype=torch.uint8),
            "w2_weight_sf": torch.empty((1, 16, 1), dtype=torch.uint8),
            "w2_alpha": torch.ones(1, dtype=torch.float32),
        }
        if quant_mode == "w4a16":
            tensors["x"] = torch.empty((2, 16), dtype=torch.bfloat16)
            tensors["x_sf"] = None
            tensors["fc2_input_scale"] = None

        if api == "functional":
            runner_name = (
                "CuteDslFusedMoENvfp4Runner"
                if quant_mode == "w4a4"
                else "CuteDslFusedMoEW4A16Runner"
            )
            monkeypatch.setattr(fused_moe, runner_name, RecordingRunner)
            result = fused_moe.cute_dsl_fused_moe_nvfp4(
                **tensors,
                num_experts=1,
                top_k=1,
                quant_mode=quant_mode,
            )
        else:
            wrapper = fused_moe.CuteDslMoEWrapper(
                num_experts=1,
                top_k=1,
                hidden_size=16,
                intermediate_size=16,
                use_cuda_graph=False,
                quant_mode=quant_mode,
            )
            if quant_mode == "w4a4":
                wrapper._runner = RecordingRunner()
                wrapper._per_token_runner = RecordingRunner()
            else:
                wrapper._w4a16_runner = RecordingRunner()
            result = wrapper.run(**tensors)

        assert result.shape == (2, 16)
        if quant_mode == "w4a4":
            assert calls[-1]["use_async_memset"] is not is_tuning_mode
        else:
            assert "use_async_memset" not in calls[-1]


# =============================================================================
# Test Class: get_max_num_tiles / get_max_num_permuted_tokens (no GPU required)
# =============================================================================


@cute_dsl_available
class TestGetMaxNumTiles:
    """Worst-case upper-bound tests for
    ``flashinfer.fused_moe.cute_dsl.moe_utils.get_max_num_tiles``.

    The function must return the tight upper bound on
    ``sum_e ceil(K_e / tile_size)``, where ``K_e`` is the per-local-expert
    token count and ``sum_e K_e = num_tokens * top_k``. The moe_sort
    routing kernel writes exactly that many entries to the
    ``tile_idx_to_expert_idx`` and ``tile_idx_to_mn_limit`` buffers (see
    ``include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh``,
    ``ExclusiveSum`` of ``ceil(count[e] / tile_size)``). An under-tight
    bound would buffer-overflow at runtime; an over-loose bound wastes
    memory and slows autotune profiling.
    """

    @staticmethod
    def _trtllm_compact_formula(
        num_tokens: int, top_k: int, num_local_experts: int, tile_size: int
    ) -> int:
        """Reference implementation matching TRT-LLM's compact form
        in ``tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py``."""
        num_expanded_tokens = num_tokens * top_k
        if num_expanded_tokens <= num_local_experts:
            return num_expanded_tokens
        return (num_expanded_tokens + (tile_size - 1) * num_local_experts) // tile_size

    @staticmethod
    def _worst_case_actual_tile_count(
        num_tokens: int, top_k: int, num_local_experts: int, tile_size: int
    ) -> int:
        """Construct the worst-case routing distribution (one expert
        receives ``E - L + 1`` tokens, the others each receive 1 token)
        and compute the actual ``sum_e ceil(K_e / tile_size)``. This is
        the value the moe_sort kernel writes to ``num_non_exiting_tiles``
        in the worst case."""
        num_expanded_tokens = num_tokens * top_k
        L = num_local_experts
        T = tile_size
        if num_expanded_tokens <= L:
            return num_expanded_tokens
        per_expert = [1] * (L - 1) + [num_expanded_tokens - (L - 1)]
        assert sum(per_expert) == num_expanded_tokens
        return sum((k + T - 1) // T for k in per_expert)

    # DeepSeek-V3-style production shapes: num_experts=256, top_k=8,
    # num_tokens=16384 (prefill) — full EP × tile_size matrix.
    # EP=1 → num_local_experts=256, EP=8 → 32, EP=16 → 16, EP=32 → 8.
    _DEEPSEEK_V3_SHAPES = [
        pytest.param(16384, 8, 256, 128, id="dsv3-ep1-tile128"),
        pytest.param(16384, 8, 256, 256, id="dsv3-ep1-tile256"),
        pytest.param(16384, 8, 32, 128, id="dsv3-ep8-tile128"),
        pytest.param(16384, 8, 32, 256, id="dsv3-ep8-tile256"),
        pytest.param(16384, 8, 16, 128, id="dsv3-ep16-tile128"),
        pytest.param(16384, 8, 16, 256, id="dsv3-ep16-tile256"),
        pytest.param(16384, 8, 8, 128, id="dsv3-ep32-tile128"),
        pytest.param(16384, 8, 8, 256, id="dsv3-ep32-tile256"),
    ]

    # Generic MoE shapes covering diverse model configurations, so the
    # formula is exercised across model families (not only DeepSeek-V3).
    # Each entry is shaped (num_tokens, top_k, num_local_experts, tile_size).
    _GENERIC_MOE_SHAPES = [
        # Small MoEs: (num_experts=8, top_k=2)-style — Mixtral-8x7B family.
        pytest.param(4096, 2, 8, 128, id="generic-mixtral-style-ep1-tile128"),
        pytest.param(4096, 2, 8, 256, id="generic-mixtral-style-ep1-tile256"),
        # Mid-size MoEs: (num_experts=16, top_k=4)-style — DBRX family.
        pytest.param(4096, 4, 16, 128, id="generic-dbrx-style-ep1-tile128"),
        pytest.param(4096, 4, 16, 256, id="generic-dbrx-style-ep1-tile256"),
        pytest.param(2048, 4, 4, 128, id="generic-dbrx-style-ep4-tile128"),
        # Mid-large MoEs: (num_experts=64, top_k=6)-style.
        pytest.param(8192, 6, 64, 128, id="generic-64expert-ep1-tile128"),
        pytest.param(8192, 6, 8, 128, id="generic-64expert-ep8-tile128"),
        pytest.param(8192, 6, 8, 256, id="generic-64expert-ep8-tile256"),
        # Large MoEs with high top_k beyond DeepSeek-V3's 8 — guards
        # against future top_k regimes.
        pytest.param(4096, 16, 32, 128, id="generic-topk16-ep4-tile128"),
        pytest.param(4096, 16, 32, 256, id="generic-topk16-ep4-tile256"),
    ]

    # Edge cases: boundary values and divisibility cases that have
    # historically tripped off-by-one logic. Independent of any specific
    # model.
    _EDGE_CASE_SHAPES = [
        # num_expanded == num_local_experts (early-return boundary).
        pytest.param(4, 8, 32, 128, id="edge-expanded-equals-local"),
        # num_expanded just past num_local_experts — first non-trivial
        # branch path.
        pytest.param(5, 8, 32, 128, id="edge-expanded-just-past-local"),
        # (E - L) exactly divisible by tile_size — historically the
        # off-by-one case.
        pytest.param(20, 8, 32, 128, id="edge-remainder-zero"),
        # (E - L) one short of divisibility.
        pytest.param(19, 8, 32, 128, id="edge-remainder-near-zero"),
        # Smallest meaningful: num_tokens=1.
        pytest.param(1, 8, 32, 128, id="edge-num-tokens-1"),
        # Smallest meaningful: top_k=1.
        pytest.param(128, 1, 8, 128, id="edge-top-k-1"),
    ]

    @pytest.mark.parametrize(
        "num_tokens,top_k,num_local_experts,tile_size",
        _DEEPSEEK_V3_SHAPES + _GENERIC_MOE_SHAPES + _EDGE_CASE_SHAPES,
    )
    def test_matches_trtllm_compact_formula(
        self,
        num_tokens: int,
        top_k: int,
        num_local_experts: int,
        tile_size: int,
    ) -> None:
        """Pin ``get_max_num_tiles`` to TRT-LLM's compact closed-form
        across diverse MoE shape configurations.

        Coverage:
        - DeepSeek-V3 production shapes at every supported EP partition
          (EP=1, 8, 16, 32) × every tile_size (128, 256).
        - Generic MoE shapes representing other model families (small,
          medium, large; varied top_k and EP partitions).
        - Edge cases (boundary values, divisibility cases).
        """
        from flashinfer.fused_moe.cute_dsl.moe_utils import get_max_num_tiles

        actual = get_max_num_tiles(num_tokens, top_k, num_local_experts, tile_size)
        expected = self._trtllm_compact_formula(
            num_tokens, top_k, num_local_experts, tile_size
        )
        assert actual == expected, (
            f"get_max_num_tiles({num_tokens=}, {top_k=}, {num_local_experts=}, "
            f"{tile_size=}) returned {actual}; TRT-LLM compact formula "
            f"returns {expected}. A discrepancy may indicate drift away "
            f"from the upstream formula or an off-by-one in floor/ceil division."
        )

    @pytest.mark.parametrize(
        "num_tokens,top_k,num_local_experts,tile_size",
        _DEEPSEEK_V3_SHAPES + _GENERIC_MOE_SHAPES + _EDGE_CASE_SHAPES,
    )
    def test_worst_case_construction_is_tight(
        self,
        num_tokens: int,
        top_k: int,
        num_local_experts: int,
        tile_size: int,
    ) -> None:
        """The bound must equal the actual tile count produced by the
        worst-case routing distribution. If ``get_max_num_tiles`` exceeds
        the worst case, buffers are over-allocated; if it falls short,
        runtime will buffer-overflow."""
        from flashinfer.fused_moe.cute_dsl.moe_utils import get_max_num_tiles

        bound = get_max_num_tiles(num_tokens, top_k, num_local_experts, tile_size)
        worst_case = self._worst_case_actual_tile_count(
            num_tokens, top_k, num_local_experts, tile_size
        )
        assert bound == worst_case, (
            f"get_max_num_tiles({num_tokens=}, {top_k=}, {num_local_experts=}, "
            f"{tile_size=}) returned {bound}; worst-case routing "
            f"(one expert gets E - L + 1 tokens, others get 1) actually "
            f"produces {worst_case}. Mismatch implies the formula is "
            f"either over- or under-allocating."
        )

    def test_zero_tokens(self) -> None:
        """Zero tokens => zero tiles."""
        from flashinfer.fused_moe.cute_dsl.moe_utils import get_max_num_tiles

        assert get_max_num_tiles(0, 8, 32, 128) == 0
        assert get_max_num_tiles(0, 8, 32, 256) == 0

    def test_below_or_at_local_experts_threshold(self) -> None:
        """When ``num_expanded <= num_local_experts``, every token can
        go to a distinct expert and each contributes 1 fully-padded
        tile. Output equals num_expanded."""
        from flashinfer.fused_moe.cute_dsl.moe_utils import get_max_num_tiles

        assert get_max_num_tiles(2, 8, 32, 128) == 16  # E = 16 < L = 32
        assert get_max_num_tiles(4, 8, 32, 128) == 32  # E = 32 == L
        assert get_max_num_tiles(2, 8, 32, 256) == 16  # tile_size irrelevant here

    @pytest.mark.parametrize(
        "top_k,num_local_experts,tile_size",
        [
            # DeepSeek-V3 (num_experts=256, top_k=8) at every EP partition.
            pytest.param(8, 256, 128, id="dsv3-ep1-tile128"),
            pytest.param(8, 256, 256, id="dsv3-ep1-tile256"),
            pytest.param(8, 32, 128, id="dsv3-ep8-tile128"),
            pytest.param(8, 32, 256, id="dsv3-ep8-tile256"),
            pytest.param(8, 16, 128, id="dsv3-ep16-tile128"),
            pytest.param(8, 16, 256, id="dsv3-ep16-tile256"),
            pytest.param(8, 8, 128, id="dsv3-ep32-tile128"),
            pytest.param(8, 8, 256, id="dsv3-ep32-tile256"),
            # Generic shapes representing other model families.
            pytest.param(2, 8, 128, id="generic-mixtral-style"),
            pytest.param(4, 16, 128, id="generic-dbrx-style"),
            pytest.param(16, 32, 256, id="generic-topk16"),
        ],
    )
    def test_monotonic_in_num_tokens(
        self, top_k: int, num_local_experts: int, tile_size: int
    ) -> None:
        """Increasing ``num_tokens`` must never decrease the tile count,
        across the same DeepSeek-V3 EP × tile_size matrix and the
        same generic-model shapes covered by the formula tests above.
        Catches sign errors / off-by-one in future refactors."""
        from flashinfer.fused_moe.cute_dsl.moe_utils import get_max_num_tiles

        prev = -1
        for n in [1, 16, 256, 1024, 16384]:
            cur = get_max_num_tiles(n, top_k, num_local_experts, tile_size)
            assert cur >= prev, (
                f"non-monotonic at num_tokens={n}, top_k={top_k}, "
                f"num_local_experts={num_local_experts}, "
                f"tile_size={tile_size}: prev={prev}, cur={cur}"
            )
            prev = cur


@cute_dsl_available
class TestGetMaxNumPermutedTokens:
    """Tests for ``get_max_num_permuted_tokens``, which is defined as
    ``get_max_num_tiles * tile_size``."""

    @pytest.mark.parametrize(
        "num_tokens,top_k,num_local_experts,tile_size",
        TestGetMaxNumTiles._DEEPSEEK_V3_SHAPES
        + TestGetMaxNumTiles._GENERIC_MOE_SHAPES
        + TestGetMaxNumTiles._EDGE_CASE_SHAPES,
    )
    def test_consistent_with_get_max_num_tiles(
        self,
        num_tokens: int,
        top_k: int,
        num_local_experts: int,
        tile_size: int,
    ) -> None:
        """Result must equal ``get_max_num_tiles * tile_size`` exactly,
        across the same DeepSeek-V3, generic-model, and edge-case shape
        coverage as :class:`TestGetMaxNumTiles`."""
        from flashinfer.fused_moe.cute_dsl.moe_utils import (
            get_max_num_permuted_tokens,
            get_max_num_tiles,
        )

        permuted = get_max_num_permuted_tokens(
            num_tokens, top_k, num_local_experts, tile_size
        )
        tiles = get_max_num_tiles(num_tokens, top_k, num_local_experts, tile_size)
        assert permuted == tiles * tile_size


# =============================================================================
# Test Class: Autotuner bucket configuration (no GPU required)
# =============================================================================


@pytest.fixture(scope="module")
def bucket_spec():
    """The first ``DynamicTensorSpec`` of a default-configured
    ``CuteDslFusedMoENvfp4Runner`` — the spec that owns the
    ``gen_tuning_buckets`` / ``map_to_tuning_buckets`` callables under
    test. Module-scoped: the runner is stateless for these checks.
    """
    from flashinfer.fused_moe.cute_dsl.tuner import (
        CuteDslFusedMoENvfp4Runner,
    )

    runner = CuteDslFusedMoENvfp4Runner(
        forward_impl=lambda *a, **k: None,
        num_experts=256,
        top_k=8,
        num_local_experts=256,
    )
    return runner.tuning_config.dynamic_tensor_specs[0]


@cute_dsl_available
class TestAutotunerBucketConfig:
    """Structural tests for the ``gen_tuning_buckets`` /
    ``map_to_tuning_buckets`` configuration on
    ``CuteDslFusedMoENvfp4Runner.tuning_config``.

    These tests run without a GPU. They guard against bucket-config
    forms that bake a hardcoded cap into the autotuner's input-dim
    bucket logic. A capped form silently clamps the autotune to a
    fixed shape — at runtime any token count larger than the cap
    maps to the smaller cached bucket and uses a tactic profiled at
    the wrong workload size.

    The correct form passes the bucket generators as bare callables;
    the autotuner invokes them with the actual input dim at autotune
    time so the bucket set adapts to the workload.
    """

    def test_gen_tuning_buckets_is_callable_not_static_tuple(self, bucket_spec):
        """``gen_tuning_buckets`` must be a callable that adapts to the
        actual input dim at autotune time — not a pre-computed
        tuple/sequence that bakes in a hardcoded cap.
        """
        assert callable(bucket_spec.gen_tuning_buckets), (
            f"gen_tuning_buckets must be a callable that adapts to the "
            f"runtime input dim — got "
            f"{type(bucket_spec.gen_tuning_buckets).__name__}. A "
            f"pre-computed sequence (e.g., a tuple) likely indicates a "
            f"bucket set with a hardcoded cap; pass the bare function "
            f"reference instead."
        )

    def test_gen_tuning_buckets_responds_to_input_dim(self, bucket_spec):
        """Calling ``gen_tuning_buckets`` with successively larger input
        dims must produce bucket sets whose maximum grows with the
        input. A capped form would produce identical (capped) bucket
        sets regardless of input.
        """
        small = bucket_spec.gen_tuning_buckets(8192)
        medium = bucket_spec.gen_tuning_buckets(16384)
        large = bucket_spec.gen_tuning_buckets(32768)
        assert max(small) >= 8192, (
            f"gen_tuning_buckets(8192) max should reach 8192; got {max(small)}"
        )
        assert max(medium) >= 16384, (
            f"gen_tuning_buckets(16384) max should reach 16384; "
            f"got {max(medium)}. Likely a hardcoded cap below 16384."
        )
        assert max(large) >= 32768, (
            f"gen_tuning_buckets(32768) max should reach 32768; "
            f"got {max(large)}. Likely a hardcoded cap below 32768."
        )
        assert max(medium) > max(small), (
            f"larger input dim should produce larger bucket max, but "
            f"max(buckets@8192)={max(small)} >= "
            f"max(buckets@16384)={max(medium)} — likely a hardcoded cap."
        )

    @pytest.mark.parametrize("x", [16384, 32768, 65536])
    def test_map_to_tuning_buckets_responds_to_large_input(self, bucket_spec, x: int):
        """``map_to_tuning_buckets(x)`` for large x must return a value
        that scales with x, not collapse to a smaller constant. A
        capped form would silently return the cap value for any input
        above it.
        """
        result = bucket_spec.map_to_tuning_buckets(x)
        assert result >= x, (
            f"map_to_tuning_buckets({x}) = {result}; expected >= {x}. "
            f"Likely a hardcoded cap below {x}."
        )

    @pytest.mark.parametrize("x", [1, 2, 4, 8, 16, 32, 64, 128, 256])
    def test_map_to_tuning_buckets_matches_trtllm_at_small_powers_of_2(
        self, bucket_spec, x: int
    ):
        """At power-of-2 inputs in the small-N regime (≤ 256), fi's
        ``map_to_tuning_buckets(x)`` must equal ``x`` — matching
        TRT-LLM's ``last_positive_power_of_2(x)`` behavior. Locks in
        the fi/trt-llm parity that IS achievable in this regime.
        """
        from flashinfer.utils import last_positive_power_of_2

        result = bucket_spec.map_to_tuning_buckets(x)
        assert result == x == last_positive_power_of_2(x), (
            f"At x={x} (power of 2 in small-N regime), fi's "
            f"map_to_tuning_buckets should equal x and equal "
            f"last_positive_power_of_2(x) (TRT-LLM's pattern). Got "
            f"fi={result}, expected {x}."
        )

    def test_map_to_tuning_buckets_is_monotonic(self, bucket_spec):
        """``map_to_tuning_buckets`` must be monotonically non-decreasing
        in its input — a property TRT-LLM's mapper also satisfies.
        Catches a regression that would introduce non-monotonic
        bucket-mapping behavior.
        """
        test_xs = [
            1,
            2,
            8,
            100,
            256,
            257,
            512,
            768,
            1024,
            2048,
            2049,
            4096,
            4097,
            8192,
            16384,
            32768,
            65536,
        ]
        results = [bucket_spec.map_to_tuning_buckets(x) for x in test_xs]
        for prev_x, prev_y, curr_x, curr_y in zip(
            test_xs, results, test_xs[1:], results[1:], strict=False
        ):
            assert prev_y <= curr_y, (
                f"map_to_tuning_buckets must be monotonically "
                f"non-decreasing; got map({prev_x})={prev_y} > "
                f"map({curr_x})={curr_y}. Full mapping at probe "
                f"points: {list(zip(test_xs, results, strict=False))}."
            )

    @pytest.mark.parametrize("max_n", [256, 4096, 16384])
    def test_gen_tuning_buckets_covers_trtllm_power_of_2_points(
        self, bucket_spec, max_n: int
    ):
        """fi's bucket set must be a superset of TRT-LLM's power-of-2
        bucket set at every input dim tested, so the autotuner has at
        least the same coarse-grained coverage as TRT-LLM at every
        power-of-2 boundary up to the input dim.
        """
        from flashinfer.utils import last_positive_power_of_2

        fi_buckets = set(bucket_spec.gen_tuning_buckets(max_n))
        # Mirror TRT-LLM's get_last_power_of_2_num_tokens_buckets:
        # powers of 2 from 1 up to last_positive_power_of_2(max_n).
        trtllm_top = last_positive_power_of_2(max_n)
        trtllm_buckets = {1 << i for i in range(trtllm_top.bit_length())}
        missing = trtllm_buckets - fi_buckets
        assert not missing, (
            f"At max_n={max_n}, fi's bucket set is missing power-of-2 "
            f"values that TRT-LLM's bucket set would include: "
            f"{sorted(missing)}. fi: {sorted(fi_buckets)}; "
            f"TRT-LLM: {sorted(trtllm_buckets)}."
        )


# =============================================================================
# Test Class: W4A16-specific contracts
# =============================================================================


@cute_dsl_available
@sm100_required
class TestCuteDslMoeW4A16:
    @pytest.mark.parametrize("use_wrapper", [False, True])
    def test_weight_scale_update(self, use_wrapper: bool):
        from flashinfer import CuteDslMoEWrapper, cute_dsl_fused_moe_nvfp4
        from flashinfer.fused_moe.cute_dsl.blackwell.moe_w4a16 import (
            DEFAULT_W4A16_MOE_TACTIC,
        )

        num_tokens, hidden_size, intermediate_size = 17, 256, 512
        num_experts, top_k = 8, 2
        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
        )
        kwargs = {
            "x": tensors["x_bf16"],
            "x_sf": None,
            "token_selected_experts": tensors["token_selected_experts"],
            "token_final_scales": tensors["token_final_scales"],
            "w1_weight": tensors["w1_weight"],
            "w1_weight_sf": tensors["w1_weight_sf"],
            "w1_alpha": tensors["w1_alpha"],
            "fc2_input_scale": None,
            "w2_weight": tensors["w2_weight"],
            "w2_weight_sf": tensors["w2_weight_sf"],
            "w2_alpha": tensors["w2_alpha"],
        }
        if use_wrapper:
            moe = CuteDslMoEWrapper(
                num_experts=num_experts,
                top_k=top_k,
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                use_cuda_graph=False,
                quant_mode="w4a16",
            )
        else:
            moe = None

        def run():
            if moe is not None:
                return moe.run(**kwargs, tactic=DEFAULT_W4A16_MOE_TACTIC)
            return cute_dsl_fused_moe_nvfp4(
                **kwargs,
                num_experts=num_experts,
                top_k=top_k,
                quant_mode="w4a16",
            )

        # Weight scales are serving-owned tensors and may be updated in place.
        # Reuse the same allocation to ensure W4A16 reads the current contents.
        result = run()
        assert torch.count_nonzero(result).item() > 0
        tensors["w1_weight_sf"].zero_()
        updated = run()
        torch.testing.assert_close(updated, torch.zeros_like(updated), rtol=0, atol=0)

    @pytest.mark.parametrize("scale_axis", ["row", "k"])
    @pytest.mark.parametrize(
        "rows,route_tile,tactic",
        [
            pytest.param(128, 32, ((128, 32, 64), (1, 1), True), id="1cta"),
            pytest.param(256, 64, ((256, 64, 128), (2, 1), True), id="2cta"),
            pytest.param(256, 128, ((128, 128, 256), (2, 1), True), id="cluster2"),
        ],
    )
    def test_weight_scale_mapping(
        self,
        scale_axis: str,
        rows: int,
        route_tile: int,
        tactic: tuple,
    ):
        from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
        from flashinfer.fp4_quantization import block_scale_interleave
        from flashinfer.fused_moe.cute_dsl.blackwell.moe_w4a16 import (
            _run_grouped_gemm,
        )
        from flashinfer.tllm_enums import (
            DEFAULT_SWIGLU_ALPHA,
            DEFAULT_SWIGLU_BETA,
            DEFAULT_SWIGLU_LIMIT,
        )

        k = 256
        scale_k = k // 16
        weight = torch.full((1, rows, k // 2), 0x22, dtype=torch.uint8, device="cuda")
        # Cycle through positive normal E4M3 codes; 0x7F is NaN.
        if scale_axis == "row":
            scale_codes = 0x08 + torch.arange(rows, device="cuda") % (0x7F - 0x08)
            scale_codes = scale_codes[:, None].expand(rows, scale_k)
        else:
            scale_codes = 0x08 + torch.arange(scale_k, device="cuda")
            scale_codes = scale_codes[None, :].expand(rows, scale_k)
        scale_codes = scale_codes.to(torch.uint8).contiguous().unsqueeze(0)
        weight_sf = convert_sf_to_mma_layout(
            block_scale_interleave(scale_codes),
            m=rows,
            k=k,
            num_groups=1,
            sf_vec_size=16,
        )

        activations = torch.zeros((route_tile, k), dtype=torch.bfloat16, device="cuda")
        block_idx = torch.arange(scale_k, device="cuda")
        activations[block_idx, block_idx * 16] = 1
        output = torch.empty((route_tile, rows), dtype=torch.bfloat16, device="cuda")
        _run_grouped_gemm(
            weight=weight,
            weight_sf=weight_sf,
            activations=activations,
            tile_idx_to_expert_idx=torch.zeros(1, dtype=torch.int32, device="cuda"),
            tile_idx_to_mn_limit=torch.full(
                (1,), route_tile, dtype=torch.int32, device="cuda"
            ),
            num_non_exiting_tiles=torch.ones(1, dtype=torch.int32, device="cuda"),
            alpha=torch.ones(1, dtype=torch.float32, device="cuda"),
            output=output,
            num_local_experts=1,
            activation_type=None,
            swiglu_alpha=DEFAULT_SWIGLU_ALPHA,
            swiglu_beta=DEFAULT_SWIGLU_BETA,
            swiglu_limit=DEFAULT_SWIGLU_LIMIT,
            use_fused_finalize=False,
            permuted_idx_to_expanded_idx=None,
            token_final_scales=None,
            enable_pdl=False,
            tactic=tactic,
        )

        expected = torch.zeros_like(output)
        expected[:scale_k] = (
            scale_codes.squeeze(0).view(torch.float8_e4m3fn).to(torch.bfloat16).T
        )
        torch.testing.assert_close(output, expected, rtol=0, atol=0)

    @pytest.mark.parametrize(
        "route_tile,gemm1_tactic,gemm2_tactic",
        [
            pytest.param(
                8,
                ((128, 8, 256), (2, 1), True),
                ((128, 8, 256), (2, 1), True),
                id="route8-1cta",
            ),
            pytest.param(
                16,
                ((128, 16, 256), (1, 1), True),
                ((128, 16, 256), (1, 1), True),
                id="route16-1cta",
            ),
            pytest.param(
                32,
                ((128, 32, 256), (2, 1), True),
                ((256, 32, 256), (2, 1), True),
                id="route32-mixed",
            ),
            pytest.param(
                64,
                ((256, 64, 256), (2, 1), True),
                ((128, 64, 256), (2, 1), True),
                id="route64-mixed",
            ),
            pytest.param(
                128,
                ((256, 128, 256), (2, 1), True),
                ((256, 128, 256), (2, 1), True),
                id="route128-2cta",
            ),
            pytest.param(
                192,
                ((256, 192, 256), (2, 1), True),
                ((256, 192, 256), (2, 1), True),
                id="route192",
            ),
        ],
    )
    def test_route_tile_boundary_accuracy(
        self,
        route_tile: int,
        gemm1_tactic: tuple,
        gemm2_tactic: tuple,
    ):
        from flashinfer.fused_moe.cute_dsl.blackwell.moe_w4a16 import (
            launch_w4a16_moe,
        )
        from flashinfer.fused_moe.cute_dsl.tuner import W4A16_MOE_TACTICS

        num_tokens, hidden_size, intermediate_size = route_tile + 1, 256, 512
        num_experts, top_k = 8, 2
        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
        )
        # Give two experts one full route tile and one boundary tile each.
        tensors["token_selected_experts"][:] = torch.arange(
            top_k, device=tensors["token_selected_experts"].device
        )
        _, reference_inputs = _prepare_moe_quant_mode_inputs(tensors, "w4a16")
        tactic = (gemm1_tactic, gemm2_tactic)
        assert tactic in W4A16_MOE_TACTICS

        result = launch_w4a16_moe(
            x=tensors["x_bf16"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            num_experts=num_experts,
            num_local_experts=num_experts,
            local_expert_offset=0,
            moe_output=torch.empty(
                (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
            ),
            use_fused_finalize=False,
            enable_pdl=False,
            activation_type=ActivationType.Swiglu,
            tactic=tactic,
        )
        ref_output = compute_reference_moe_fp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            activation_type=ActivationType.Swiglu,
            **reference_inputs,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"Only {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )


# =============================================================================
# Test Class: Functional API (cute_dsl_fused_moe_nvfp4)
# =============================================================================


@cute_dsl_available
@sm100_required
class TestCuteDslFusedMoeFunctional:
    """Tests for the functional API: cute_dsl_fused_moe_nvfp4."""

    pytestmark = _requires_dsl_arch

    @pytest.mark.parametrize(
        "hidden_size,intermediate_size", [(256, 512), (1024, 2048)]
    )
    @pytest.mark.parametrize(
        "quant_mode,use_per_token_activation", _MOE_QUANT_MODE_CASES
    )
    @pytest.mark.parametrize("top_k", [1, 2, 8])
    @pytest.mark.parametrize("num_tokens", [128, 515, 1024])
    @pytest.mark.parametrize("num_experts", [256, 384])
    @pytest.mark.parametrize(
        "activation_type", [ActivationType.Swiglu, ActivationType.Relu2]
    )
    def test_numerical_accuracy(
        self,
        activation_type: ActivationType,
        num_tokens: int,
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
        num_experts: int,
        quant_mode: str,
        use_per_token_activation: bool,
    ):
        """Accuracy test for functional API across configurations."""
        self._run_numerical_accuracy(
            activation_type,
            num_tokens,
            top_k,
            hidden_size,
            intermediate_size,
            num_experts,
            quant_mode,
            use_per_token_activation,
            use_fused_finalize=True,
        )

    @pytest.mark.parametrize(
        "quant_mode,activation_type,num_tokens,top_k,hidden_size,intermediate_size,num_experts,use_per_token_activation",
        [
            pytest.param(
                "w4a4",
                ActivationType.Swiglu,
                128,
                1,
                256,
                512,
                256,
                False,
                id="swiglu-per-tensor",
            ),
            pytest.param(
                "w4a4",
                ActivationType.Relu2,
                515,
                8,
                1024,
                2048,
                384,
                True,
                id="relu2-per-token",
            ),
            pytest.param(
                "w4a16",
                ActivationType.Swiglu,
                128,
                1,
                256,
                512,
                256,
                False,
                id="w4a16-swiglu",
            ),
            pytest.param(
                "w4a16",
                ActivationType.Relu2,
                515,
                8,
                1024,
                2048,
                384,
                False,
                id="w4a16-relu2",
            ),
        ],
    )
    def test_deterministic_finalize_numerical_accuracy(
        self,
        quant_mode: str,
        activation_type: ActivationType,
        num_tokens: int,
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
        num_experts: int,
        use_per_token_activation: bool,
    ):
        self._run_numerical_accuracy(
            activation_type,
            num_tokens,
            top_k,
            hidden_size,
            intermediate_size,
            num_experts,
            quant_mode,
            use_per_token_activation,
            use_fused_finalize=False,
        )

    @pytest.mark.parametrize(
        "quant_mode, use_per_token_activation",
        _MOE_QUANT_MODE_CASES,
    )
    @pytest.mark.parametrize("hidden_size", [256, 384])
    def test_finalize_handles_cluster_padding_and_partial_tiles(
        self,
        quant_mode: str,
        use_per_token_activation: bool,
        hidden_size: int,
        monkeypatch: pytest.MonkeyPatch,
    ):
        from flashinfer.autotuner import AutoTuner

        if quant_mode == "w4a4":
            # 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),
            )
        elif quant_mode == "w4a16":
            # W4A16 clusters 128-wide M CTAs in pairs, so hidden=384 leaves a
            # padding peer.
            tail_config = (
                ((256, 128, 256), (2, 1), True),
                ((256, 128, 256), (2, 1), True),
            )
        else:
            raise ValueError(f"unsupported quant_mode {quant_mode!r}")

        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,
            quant_mode=quant_mode,
            use_per_token_activation=use_per_token_activation,
            use_fused_finalize=True,
        )

    def _run_numerical_accuracy(
        self,
        activation_type: ActivationType,
        num_tokens: int,
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
        num_experts: int,
        quant_mode: str,
        use_per_token_activation: bool,
        use_fused_finalize: bool,
    ):
        from flashinfer import cute_dsl_fused_moe_nvfp4

        _, gated = normalize_cute_dsl_moe_activation_type(activation_type)
        num_local_experts = num_experts

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            top_k=top_k,
            gated=gated,
            use_per_token_activation=use_per_token_activation,
        )
        api_inputs, reference_inputs = _prepare_moe_quant_mode_inputs(
            tensors, quant_mode
        )

        result = cute_dsl_fused_moe_nvfp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            num_experts=num_experts,
            top_k=top_k,
            num_local_experts=num_local_experts,
            activation_type=activation_type,
            use_fused_finalize=use_fused_finalize,
            quant_mode=quant_mode,
            **api_inputs,
        )

        assert result.shape == (num_tokens, hidden_size)
        assert result.dtype == torch.bfloat16
        assert not torch.isnan(result).any()
        assert not torch.isinf(result).any()

        ref_output = compute_reference_moe_fp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            activation_type=activation_type,
            **reference_inputs,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"Only {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )

    @pytest.mark.parametrize("use_per_token_activation", [False, True])
    def test_geglu_tanh_accuracy(self, use_per_token_activation: bool):
        """Accuracy test for tanh-approximate GeGLU in both scaling modes."""
        from flashinfer import cute_dsl_fused_moe_nvfp4

        activation_type = ActivationType.GegluTanh
        num_tokens, hidden_size, intermediate_size = 128, 256, 512
        num_experts, top_k = 256, 2

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
            gated=True,
            use_per_token_activation=use_per_token_activation,
        )

        result = cute_dsl_fused_moe_nvfp4(
            x=tensors["x"],
            x_sf=tensors["x_sf"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            fc2_input_scale=tensors["fc2_input_scale"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            num_experts=num_experts,
            top_k=top_k,
            activation_type=activation_type,
            per_token_scale=tensors["x_per_token_scale"],
        )

        ref_output = compute_reference_moe_fp4(
            hidden_states=tensors["x_ref"].float().cuda(),
            gemm1_weights=tensors["w1_weight_bf16"].float().cuda(),
            gemm2_weights=tensors["w2_weight_bf16"].float().cuda(),
            gemm1_alpha=tensors["w1_alpha"],
            gemm2_alpha=tensors["w2_alpha"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            fc2_input_scale=tensors["fc2_input_scale"],
            activation_type=activation_type,
            use_per_token_activation=use_per_token_activation,
        )

        swiglu_ref_output = compute_reference_moe_fp4(
            hidden_states=tensors["x_ref"].float().cuda(),
            gemm1_weights=tensors["w1_weight_bf16"].float().cuda(),
            gemm2_weights=tensors["w2_weight_bf16"].float().cuda(),
            gemm1_alpha=tensors["w1_alpha"],
            gemm2_alpha=tensors["w2_alpha"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            fc2_input_scale=tensors["fc2_input_scale"],
            activation_type=ActivationType.Swiglu,
            use_per_token_activation=use_per_token_activation,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"Only {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )
        geglu_error = torch.mean(torch.abs(result.float() - ref_output)).item()
        swiglu_error = torch.mean(torch.abs(result.float() - swiglu_ref_output)).item()
        assert geglu_error < swiglu_error, (
            f"GeGLU output is not closer to GeGLU reference ({geglu_error:.4f}) "
            f"than SwiGLU reference ({swiglu_error:.4f})"
        )

    @pytest.mark.parametrize("use_per_token_activation", [False, True])
    @pytest.mark.parametrize(
        "situ_beta,situ_linear_beta",
        [(1.75, None), (0.8, 1.5)],
        ids=["gate-clamp", "gate-and-up-clamp"],
    )
    def test_situ_accuracy(
        self,
        use_per_token_activation: bool,
        situ_beta: float,
        situ_linear_beta: float | None,
    ):
        """Accuracy test for SiTU with optional smooth up-branch clamping."""
        from flashinfer import cute_dsl_fused_moe_nvfp4

        num_tokens, hidden_size, intermediate_size = 128, 256, 512
        num_experts, top_k = 256, 2

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
            gated=True,
            use_per_token_activation=use_per_token_activation,
        )

        result = cute_dsl_fused_moe_nvfp4(
            x=tensors["x"],
            x_sf=tensors["x_sf"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            fc2_input_scale=tensors["fc2_input_scale"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            num_experts=num_experts,
            top_k=top_k,
            activation_type=ActivationType.Swiglu,
            situ_beta=situ_beta,
            situ_linear_beta=situ_linear_beta,
            per_token_scale=tensors["x_per_token_scale"],
        )

        ref_output = compute_reference_moe_fp4(
            hidden_states=tensors["x_ref"].float().cuda(),
            gemm1_weights=tensors["w1_weight_bf16"].float().cuda(),
            gemm2_weights=tensors["w2_weight_bf16"].float().cuda(),
            gemm1_alpha=tensors["w1_alpha"],
            gemm2_alpha=tensors["w2_alpha"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            fc2_input_scale=tensors["fc2_input_scale"],
            activation_type=ActivationType.Swiglu,
            situ_beta=situ_beta,
            situ_linear_beta=situ_linear_beta,
            use_per_token_activation=use_per_token_activation,
        )

        swiglu_ref_output = compute_reference_moe_fp4(
            hidden_states=tensors["x_ref"].float().cuda(),
            gemm1_weights=tensors["w1_weight_bf16"].float().cuda(),
            gemm2_weights=tensors["w2_weight_bf16"].float().cuda(),
            gemm1_alpha=tensors["w1_alpha"],
            gemm2_alpha=tensors["w2_alpha"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            fc2_input_scale=tensors["fc2_input_scale"],
            activation_type=ActivationType.Swiglu,
            use_per_token_activation=use_per_token_activation,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"Only {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )
        situ_error = torch.mean(torch.abs(result.float() - ref_output)).item()
        swiglu_error = torch.mean(torch.abs(result.float() - swiglu_ref_output)).item()
        assert situ_error < swiglu_error, (
            f"SiTU output is not closer to SiTU reference ({situ_error:.4f}) "
            f"than SwiGLU reference ({swiglu_error:.4f})"
        )

    @pytest.mark.parametrize("quant_mode", _MOE_QUANT_MODES)
    def test_with_autotune(self, quant_mode: str):
        """Test functional API with autotune context."""
        from flashinfer import autotune
        from flashinfer import cute_dsl_fused_moe_nvfp4

        num_tokens, hidden_size, intermediate_size = 256, 256, 512
        num_experts, top_k = 256, 2

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
        )
        api_inputs, _ = _prepare_moe_quant_mode_inputs(tensors, quant_mode)

        with autotune(True):
            result = cute_dsl_fused_moe_nvfp4(
                token_selected_experts=tensors["token_selected_experts"],
                token_final_scales=tensors["token_final_scales"],
                w1_weight=tensors["w1_weight"],
                w1_weight_sf=tensors["w1_weight_sf"],
                w1_alpha=tensors["w1_alpha"],
                w2_weight=tensors["w2_weight"],
                w2_weight_sf=tensors["w2_weight_sf"],
                w2_alpha=tensors["w2_alpha"],
                num_experts=num_experts,
                top_k=top_k,
                activation_type=ActivationType.Swiglu,
                swiglu_alpha=1.702,
                swiglu_beta=1.0,
                swiglu_limit=7.0,
                quant_mode=quant_mode,
                **api_inputs,
            )

        assert result.shape == (num_tokens, hidden_size)
        assert not torch.isnan(result).any()

    @pytest.mark.parametrize("quant_mode", _MOE_QUANT_MODES)
    def test_swiglu_oai_accuracy(self, quant_mode: str):
        """Accuracy test for the OAI SwiGLU epilogue variant."""
        from flashinfer import cute_dsl_fused_moe_nvfp4

        num_tokens, hidden_size, intermediate_size = 128, 256, 512
        num_experts, top_k = 256, 2

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
        )
        api_inputs, reference_inputs = _prepare_moe_quant_mode_inputs(
            tensors, quant_mode
        )

        result = cute_dsl_fused_moe_nvfp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            num_experts=num_experts,
            top_k=top_k,
            activation_type=ActivationType.Swiglu,
            swiglu_alpha=1.702,
            swiglu_beta=1.0,
            swiglu_limit=7.0,
            quant_mode=quant_mode,
            **api_inputs,
        )

        ref_output = compute_reference_moe_fp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            activation_type=ActivationType.Swiglu,
            swiglu_alpha=1.702,
            swiglu_beta=1.0,
            swiglu_limit=7.0,
            **reference_inputs,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"Only {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )


# =============================================================================
# Test Class: Wrapper API (CuteDslMoEWrapper)
# =============================================================================


@cute_dsl_available
@sm100_required
class TestCuteDslMoEWrapper:
    """Tests for the wrapper API: CuteDslMoEWrapper."""

    pytestmark = _requires_dsl_arch

    @pytest.mark.parametrize("num_tokens", [128, 256, 512])
    @pytest.mark.parametrize("use_fused_finalize", [False, True])
    @pytest.mark.parametrize(
        "quant_mode,use_per_token_activation", _MOE_QUANT_MODE_CASES
    )
    @pytest.mark.parametrize("top_k", [2, 8])
    @pytest.mark.parametrize("num_experts", [256, 384])
    def test_wrapper_accuracy(
        self,
        num_tokens: int,
        top_k: int,
        num_experts: int,
        use_fused_finalize: bool,
        quant_mode: str,
        use_per_token_activation: bool,
    ):
        """Accuracy test for wrapper API."""
        from flashinfer import CuteDslMoEWrapper

        hidden_size, intermediate_size = 256, 512

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
            use_per_token_activation=use_per_token_activation,
        )
        api_inputs, reference_inputs = _prepare_moe_quant_mode_inputs(
            tensors, quant_mode
        )

        # Create wrapper WITHOUT CUDA graph
        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=False,
            use_fused_finalize=use_fused_finalize,
            quant_mode=quant_mode,
        )

        result = moe.run(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            **api_inputs,
        )

        assert result.shape == (num_tokens, hidden_size)
        assert not torch.isnan(result).any()
        assert not torch.isinf(result).any()

        ref_output = compute_reference_moe_fp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            **reference_inputs,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"Only {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )

    @pytest.mark.parametrize("quant_mode", _MOE_QUANT_MODES)
    def test_wrapper_swiglu_oai_accuracy(self, quant_mode: str):
        """Accuracy test for wrapper API with OAI SwiGLU."""
        from flashinfer import CuteDslMoEWrapper

        num_tokens, hidden_size, intermediate_size = 128, 256, 512
        num_experts, top_k = 256, 2

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
        )
        api_inputs, reference_inputs = _prepare_moe_quant_mode_inputs(
            tensors, quant_mode
        )

        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=False,
            activation_type=ActivationType.Swiglu,
            swiglu_alpha=1.702,
            swiglu_beta=1.0,
            swiglu_limit=7.0,
            quant_mode=quant_mode,
        )

        result = moe.run(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            **api_inputs,
        )

        ref_output = compute_reference_moe_fp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            activation_type=ActivationType.Swiglu,
            swiglu_alpha=1.702,
            swiglu_beta=1.0,
            swiglu_limit=7.0,
            **reference_inputs,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"Only {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )

    @pytest.mark.parametrize("use_fused_finalize", [False, True])
    @pytest.mark.parametrize(
        "quant_mode,use_per_token_activation", _MOE_QUANT_MODE_CASES
    )
    @pytest.mark.parametrize("num_tokens", [64, 128, 256])
    @pytest.mark.parametrize("num_experts", [256, 384])
    def test_wrapper_cuda_graph(
        self,
        num_tokens: int,
        num_experts: int,
        quant_mode: str,
        use_per_token_activation: bool,
        use_fused_finalize: bool,
    ):
        """Test wrapper API with CUDA graph capture and replay."""
        from flashinfer import CuteDslMoEWrapper

        hidden_size, intermediate_size = 256, 512
        top_k = 2

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
            use_per_token_activation=use_per_token_activation,
        )
        api_inputs, reference_inputs = _prepare_moe_quant_mode_inputs(
            tensors, quant_mode
        )

        # Create wrapper WITH CUDA graph
        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=True,
            max_num_tokens=num_tokens,
            activation_type=ActivationType.Swiglu,
            swiglu_alpha=1.702,
            swiglu_beta=1.0,
            swiglu_limit=7.0,
            use_fused_finalize=use_fused_finalize,
            quant_mode=quant_mode,
        )

        # Warmup
        for _ in range(3):
            moe.run(
                token_selected_experts=tensors["token_selected_experts"],
                token_final_scales=tensors["token_final_scales"],
                w1_weight=tensors["w1_weight"],
                w1_weight_sf=tensors["w1_weight_sf"],
                w1_alpha=tensors["w1_alpha"],
                w2_weight=tensors["w2_weight"],
                w2_weight_sf=tensors["w2_weight_sf"],
                w2_alpha=tensors["w2_alpha"],
                **api_inputs,
            )
        torch.cuda.synchronize()

        # Capture CUDA graph
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            output = moe.run(
                token_selected_experts=tensors["token_selected_experts"],
                token_final_scales=tensors["token_final_scales"],
                w1_weight=tensors["w1_weight"],
                w1_weight_sf=tensors["w1_weight_sf"],
                w1_alpha=tensors["w1_alpha"],
                w2_weight=tensors["w2_weight"],
                w2_weight_sf=tensors["w2_weight_sf"],
                w2_alpha=tensors["w2_alpha"],
                **api_inputs,
            )
        torch.cuda.synchronize()

        # Note: CUDA graph capture doesn't execute - output may be zeros here
        # Actual execution happens during replay
        assert output.shape == (num_tokens, hidden_size)

        # First replay to get actual output
        g.replay()
        torch.cuda.synchronize()

        # Verify output is valid after first replay
        assert not torch.isnan(output).any(), "NaN after first replay"
        assert not (output == 0).all(), "All zeros after first replay"

        results = []
        for _ in range(3):
            g.replay()
            torch.cuda.synchronize()
            results.append(output.clone())

        for i in range(1, len(results)):
            if use_fused_finalize:
                max_diff = (results[0] - results[i]).abs().max().item()
                assert max_diff < 0.5, (
                    f"Replay {i} differs too much: max_diff={max_diff}"
                )
            else:
                assert torch.equal(results[0], results[i]), f"Replay {i} differs"

        # Verify accuracy
        ref_output = compute_reference_moe_fp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            activation_type=ActivationType.Swiglu,
            swiglu_alpha=1.702,
            swiglu_beta=1.0,
            swiglu_limit=7.0,
            **reference_inputs,
        )

        passed, percent_within, atol = check_accuracy(results[0], ref_output)
        assert passed, (
            f"CUDA graph accuracy: {percent_within * 100:.2f}% (atol={atol:.4f})"
        )

    @pytest.mark.parametrize(
        "activation_type,situ_beta,situ_linear_beta,quant_mode",
        [
            (ActivationType.Swiglu, None, None, "w4a4"),
            (ActivationType.Swiglu, 1.0, 1.5, "w4a4"),
            (ActivationType.GegluTanh, None, None, "w4a4"),
            (ActivationType.Relu2, None, None, "w4a4"),
            (ActivationType.Swiglu, None, None, "w4a16"),
            (ActivationType.Relu2, None, None, "w4a16"),
        ],
        ids=[
            "w4a4-swiglu",
            "w4a4-situ",
            "w4a4-geglu-tanh",
            "w4a4-relu2",
            "w4a16-swiglu",
            "w4a16-relu2",
        ],
    )
    def test_wrapper_with_autotune(
        self,
        activation_type: ActivationType,
        situ_beta: float | None,
        situ_linear_beta: float | None,
        quant_mode: str,
    ):
        """Test wrapper API with autotune context."""
        from flashinfer import autotune
        from flashinfer import CuteDslMoEWrapper

        _, gated = normalize_cute_dsl_moe_activation_type(activation_type)
        num_tokens, hidden_size, intermediate_size = 256, 256, 512
        num_experts, top_k = 256, 2

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
            gated=gated,
        )
        api_inputs, reference_inputs = _prepare_moe_quant_mode_inputs(
            tensors, quant_mode
        )

        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=False,
            activation_type=activation_type,
            quant_mode=quant_mode,
            situ_beta=situ_beta,
            situ_linear_beta=situ_linear_beta,
        )

        with autotune(True):
            result = moe.run(
                token_selected_experts=tensors["token_selected_experts"],
                token_final_scales=tensors["token_final_scales"],
                w1_weight=tensors["w1_weight"],
                w1_weight_sf=tensors["w1_weight_sf"],
                w1_alpha=tensors["w1_alpha"],
                w2_weight=tensors["w2_weight"],
                w2_weight_sf=tensors["w2_weight_sf"],
                w2_alpha=tensors["w2_alpha"],
                **api_inputs,
            )

        assert result.shape == (num_tokens, hidden_size)
        assert not torch.isnan(result).any()

        ref_output = compute_reference_moe_fp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            activation_type=activation_type,
            **reference_inputs,
            situ_beta=situ_beta,
            situ_linear_beta=situ_linear_beta,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"Only {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )

    def test_cuda_graph_wrapper_lifetime_before_autotune(self):
        """Dropped CUDA graph wrappers should not wait for cyclic GC."""
        from flashinfer import autotune
        from flashinfer import CuteDslMoEWrapper

        hidden_size, intermediate_size = 256, 512
        num_experts, top_k = 256, 2
        target_num_tokens = 256

        def run_wrapper(moe, tensors):
            api_inputs, _ = _prepare_moe_quant_mode_inputs(tensors, "w4a4")
            return moe.run(
                token_selected_experts=tensors["token_selected_experts"],
                token_final_scales=tensors["token_final_scales"],
                w1_weight=tensors["w1_weight"],
                w1_weight_sf=tensors["w1_weight_sf"],
                w1_alpha=tensors["w1_alpha"],
                w2_weight=tensors["w2_weight"],
                w2_weight_sf=tensors["w2_weight_sf"],
                w2_alpha=tensors["w2_alpha"],
                **api_inputs,
            )

        def warmup_and_drop_cuda_graph_wrapper():
            tensors = create_moe_tensors(
                num_tokens=64,
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                num_experts=num_experts,
                num_local_experts=num_experts,
                top_k=top_k,
            )
            moe = CuteDslMoEWrapper(
                num_experts=num_experts,
                top_k=top_k,
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                use_cuda_graph=True,
                max_num_tokens=64,
            )
            ref = weakref.ref(moe)
            finalized = []
            weakref.finalize(moe, lambda: finalized.append(True))

            for _ in range(3):
                result = run_wrapper(moe, tensors)
            torch.cuda.synchronize()
            assert not torch.isnan(result).any()
            return ref, finalized

        gc_was_enabled = gc.isenabled()
        gc.disable()
        try:
            wrapper_ref, finalized = warmup_and_drop_cuda_graph_wrapper()
            assert wrapper_ref() is None
            assert finalized == [True]
        finally:
            if gc_was_enabled:
                gc.enable()

        tensors = create_moe_tensors(
            num_tokens=target_num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
        )
        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=False,
        )

        with autotune(True):
            result = run_wrapper(moe, tensors)
        torch.cuda.synchronize()

        assert result.shape == (target_num_tokens, hidden_size)
        assert not torch.isnan(result).any()

    def test_cuda_graph_wrapper_lifetime_after_autotune(self):
        """Dropped CUDA graph wrappers should not wait for cyclic GC,
        even after autotune profiling has populated the autotuner cache."""
        from flashinfer import autotune
        from flashinfer import CuteDslMoEWrapper
        from flashinfer.autotuner import AutoTuner

        hidden_size, intermediate_size = 256, 512
        num_experts, top_k = 256, 2

        def run_wrapper(moe, tensors):
            api_inputs, _ = _prepare_moe_quant_mode_inputs(tensors, "w4a4")
            return moe.run(
                token_selected_experts=tensors["token_selected_experts"],
                token_final_scales=tensors["token_final_scales"],
                w1_weight=tensors["w1_weight"],
                w1_weight_sf=tensors["w1_weight_sf"],
                w1_alpha=tensors["w1_alpha"],
                w2_weight=tensors["w2_weight"],
                w2_weight_sf=tensors["w2_weight_sf"],
                w2_alpha=tensors["w2_alpha"],
                **api_inputs,
            )

        def autotune_and_drop_cuda_graph_wrapper():
            tensors = create_moe_tensors(
                num_tokens=64,
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                num_experts=num_experts,
                num_local_experts=num_experts,
                top_k=top_k,
            )
            moe = CuteDslMoEWrapper(
                num_experts=num_experts,
                top_k=top_k,
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                use_cuda_graph=True,
                max_num_tokens=64,
            )
            ref = weakref.ref(moe)
            finalized = []
            weakref.finalize(moe, lambda: finalized.append(True))

            # Clear the autotuner cache so this test forces a profile pass.
            # Otherwise a prior test in the same process may have populated
            # a matching cache entry and `autotune(True)` would take the
            # cache-hit path, skipping the runner-wrapper profiling
            # interaction this test is meant to exercise.
            autotuner = AutoTuner.get()
            autotuner.clear_cache()

            with autotune(True):
                result = run_wrapper(moe, tensors)
            torch.cuda.synchronize()
            assert not torch.isnan(result).any()
            # Confirm profiling actually ran for this custom op. Cache keys
            # are ProfilingCacheKey instances; see AutoTuner._get_cache_key
            # in flashinfer/autotuner/autotuner.py.
            assert any(
                k.custom_op == "CuteDslMoEWrapper::run::Swiglu"
                for k in autotuner.profiling_cache
            ), "autotune(True) did not populate a CuteDslMoEWrapper::run cache entry"
            return ref, finalized

        gc_was_enabled = gc.isenabled()
        gc.disable()
        try:
            wrapper_ref, finalized = autotune_and_drop_cuda_graph_wrapper()
            assert wrapper_ref() is None
            assert finalized == [True]
        finally:
            if gc_was_enabled:
                gc.enable()


# =============================================================================
# Test Class: API Consistency
# =============================================================================


@cute_dsl_available
@sm100_required
class TestApiConsistency:
    """Tests verifying consistency between functional and wrapper APIs."""

    pytestmark = _requires_dsl_arch

    @pytest.mark.parametrize("quant_mode", _MOE_QUANT_MODES)
    def test_functional_vs_wrapper_output(self, quant_mode: str):
        """Verify functional and wrapper APIs produce the same output."""
        from flashinfer import CuteDslMoEWrapper, cute_dsl_fused_moe_nvfp4

        num_tokens, hidden_size, intermediate_size = 128, 256, 512
        num_experts, top_k = 256, 2

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_experts,
            top_k=top_k,
        )
        api_inputs, _ = _prepare_moe_quant_mode_inputs(tensors, quant_mode)

        # Functional API
        result_functional = cute_dsl_fused_moe_nvfp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            num_experts=num_experts,
            top_k=top_k,
            quant_mode=quant_mode,
            **api_inputs,
        )

        # Wrapper API
        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=False,
            quant_mode=quant_mode,
        )

        result_wrapper = moe.run(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            **api_inputs,
        )

        # Both should produce valid outputs
        assert result_functional.shape == result_wrapper.shape
        assert not torch.isnan(result_functional).any()
        assert not torch.isnan(result_wrapper).any()

        # Outputs should be very close (may not be exactly equal due to different
        # tuning paths, but should be within FP4 tolerance)
        diff = (result_functional - result_wrapper).abs()
        max_diff = diff.max().item()
        # Allow small differences from autotuner path differences
        assert max_diff < 1e-3, f"Max diff between APIs: {max_diff}"


# =============================================================================
# Test Class: Expert Parallelism
# =============================================================================


@cute_dsl_available
@sm100_required
class TestExpertParallelism:
    """Tests for expert parallelism (EP) configurations."""

    pytestmark = _requires_dsl_arch

    @pytest.mark.parametrize("ep_size", [1, 8, 32])
    @pytest.mark.parametrize("ep_rank", [0, -1])  # -1 means last rank
    @pytest.mark.parametrize("quant_mode", _MOE_QUANT_MODES)
    def test_wrapper_with_ep(self, ep_size: int, ep_rank: int, quant_mode: str):
        """Test wrapper API with expert parallelism and numerical accuracy.

        Tests different EP ranks to ensure local_expert_offset handling is correct.
        ep_rank=-1 is converted to the last rank (ep_size-1) to test non-zero offsets.
        """
        from flashinfer import CuteDslMoEWrapper

        # Convert -1 to last rank
        if ep_rank == -1:
            ep_rank = ep_size - 1

        num_tokens, hidden_size, intermediate_size = 256, 256, 512
        num_experts, top_k = 256, 8
        num_local_experts = num_experts // ep_size
        local_expert_offset = ep_rank * num_local_experts

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            top_k=top_k,
        )
        api_inputs, reference_inputs = _prepare_moe_quant_mode_inputs(
            tensors, quant_mode
        )

        # Keep original routing - the kernel should handle filtering
        # based on local_expert_offset and num_local_experts
        token_selected_experts = tensors["token_selected_experts"].clone()

        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_local_experts=num_local_experts,
            local_expert_offset=local_expert_offset,
            quant_mode=quant_mode,
        )

        result = moe.run(
            token_selected_experts=token_selected_experts,
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            **api_inputs,
        )

        assert result.shape == (num_tokens, hidden_size)
        assert not torch.isnan(result).any()
        assert not torch.isinf(result).any()

        # Numerical accuracy verification against reference
        ref_output = compute_reference_moe_fp4(
            token_selected_experts=token_selected_experts,
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_local_experts=num_local_experts,
            local_expert_offset=local_expert_offset,
            **reference_inputs,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"EP accuracy test failed (ep_size={ep_size}, ep_rank={ep_rank}, "
            f"offset={local_expert_offset}): {percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )

    @pytest.mark.parametrize("use_fused_finalize", [False, True])
    @pytest.mark.parametrize(
        "quant_mode,use_per_token_activation", _MOE_QUANT_MODE_CASES
    )
    @pytest.mark.parametrize("ep_size", [8])
    def test_functional_with_ep(
        self,
        ep_size: int,
        quant_mode: str,
        use_per_token_activation: bool,
        use_fused_finalize: bool,
    ):
        """Test functional API with expert parallelism and numerical accuracy."""
        from flashinfer import cute_dsl_fused_moe_nvfp4

        # Test middle rank to ensure offset handling works
        ep_rank = ep_size // 2

        num_tokens, hidden_size, intermediate_size = 256, 256, 512
        num_experts, top_k = 256, 8
        num_local_experts = num_experts // ep_size
        local_expert_offset = ep_rank * num_local_experts

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            top_k=top_k,
            use_per_token_activation=use_per_token_activation,
        )
        api_inputs, reference_inputs = _prepare_moe_quant_mode_inputs(
            tensors, quant_mode
        )

        result = cute_dsl_fused_moe_nvfp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            num_experts=num_experts,
            top_k=top_k,
            num_local_experts=num_local_experts,
            local_expert_offset=local_expert_offset,
            use_fused_finalize=use_fused_finalize,
            quant_mode=quant_mode,
            **api_inputs,
        )

        assert result.shape == (num_tokens, hidden_size)
        assert not torch.isnan(result).any()
        assert not torch.isinf(result).any()

        # Numerical accuracy verification
        ref_output = compute_reference_moe_fp4(
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_local_experts=num_local_experts,
            local_expert_offset=local_expert_offset,
            **reference_inputs,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"EP functional API accuracy test failed (ep_size={ep_size}, ep_rank={ep_rank}): "
            f"{percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
        )


# =============================================================================
# Test Class: moe_sort buffer-init invariants (poisoning)
# =============================================================================


@cute_dsl_available
@sm100_required
class TestMoeSortBufferInitPoisoned:
    """Validate the invariant that the routing kernel writes every
    output entry that downstream code reads, by pre-poisoning the
    ``moe_sort`` output buffers with a sentinel value before invoking
    the MoE pipeline.

    The ``moe_sort`` wrapper in ``moe_utils.py`` allocates its output
    buffers via ``torch.empty(...)`` and relies on the routing kernel
    (``runPostTopKPipeline`` in ``trtllm_fused_moe_routing_common.cu``)
    to write every entry that any downstream kernel reads — including
    writing ``-1`` to masked slots of ``expanded_idx_to_permuted_idx``
    at EP > 1, and writing valid permuted indices to all unmasked
    slots. If the kernel is ever changed to skip an entry that
    downstream consumes, the uninitialized memory (or here, the
    sentinel) will leak through and produce dramatic numerical
    divergence — NaN/Inf via OOB index reads, or ``atomic-add`` into
    wildly wrong output rows.

    EP=32 specifically stresses the masked-position case for
    ``expanded_idx_to_permuted_idx``: with ``num_experts=256``,
    ``num_local_experts=8``, ``top_k=8`` only ~3.125% of expanded slots
    have their expert on this rank, so ~96% of the buffer must be
    written as ``-1`` by the kernel. If the kernel ever stops writing
    masked slots, this test catches it.

    ``_moe_core_impl`` accepts an external ``moe_sort_buffers`` dict
    for callers who want to manage their own routing-output buffers.
    The test allocates the dict via ``allocate_moe_sort_buffers``,
    fills it with the poison sentinel, and drives the full
    routing+gemm pipeline through ``_moe_core_impl`` directly.
    """

    pytestmark = _requires_dsl_arch

    @pytest.mark.parametrize(
        "ep_size,num_tokens",
        [
            (1, 256),
            (8, 256),
            (16, 256),
            (32, 256),
            # High-N case exercises the `num_tokens > 1024` branch in
            # ``moe_sort`` where ``expert_counts`` is allocated per-call
            # via ``torch.empty`` and the vendored kernel relies on
            # ``launchInitExpertCounts`` to zero it before reading. No
            # other test in this suite exercises that branch — without
            # this case a future regression in the kernel-side init
            # would slip past CI.
            (8, 1280),
        ],
    )
    def test_wrapper_with_poisoned_moe_sort_buffers(
        self, ep_size: int, num_tokens: int
    ):
        """Pre-poison all six moe_sort output buffers with a sentinel
        before invoking the MoE pipeline; verify output is well-formed
        and (at low N) matches the eager reference within tolerance.

        The high-N case (``num_tokens > 1024``) skips the
        ``compute_reference_moe_fp4`` comparison because that
        reference is ``O(num_tokens × top_k)`` Python iterations with
        ``.item()`` syncs and dominates CI runtime. The shape +
        NaN/Inf + non-zero-fraction assertions still run and are
        what the poisoning-detection logic actually relies on; the
        reference comparison is supplementary.
        """
        from flashinfer.fused_moe.cute_dsl.fused_moe import _moe_core_impl
        from flashinfer.fused_moe.cute_dsl.moe_utils import (
            allocate_moe_sort_buffers,
        )

        hidden_size, intermediate_size = 256, 512
        num_experts, top_k = 256, 8
        num_local_experts = num_experts // ep_size
        local_expert_offset = 0  # rank 0; offset > 0 covered by TestExpertParallelism

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            top_k=top_k,
        )

        # Allocate the moe_sort outputs externally so we can poison them
        # before invoking the kernel. ``_moe_core_impl`` accepts a
        # ``moe_sort_buffers`` dict that flows through to ``moe_sort`` via
        # **kwargs, taking the place of its internal ``torch.empty`` calls.
        tile_size = 128  # default tactic for _moe_core_impl
        moe_sort_buffers = allocate_moe_sort_buffers(
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            num_local_experts=num_local_experts,
            tile_tokens_dim=tile_size,
            device="cuda",
        )

        # Defensive guard: if a future refactor renames or restructures the
        # buffer dict returned by ``allocate_moe_sort_buffers``, the
        # poisoning loop below would silently iterate over zero items and
        # the test would pass without exercising the kernel-write
        # invariant. Fail loudly in that case.
        assert moe_sort_buffers and len(moe_sort_buffers) > 0, (
            "``allocate_moe_sort_buffers`` no longer returns a non-empty "
            "dict; the poisoning loop would be a no-op. Update this test "
            "to target the new buffer-allocation API."
        )

        # Sentinel: a non-zero, non-(-1), out-of-valid-index-range int32.
        # If the kernel writes every entry that downstream reads, none of
        # these sentinels survive into gemm2 finalize. If any do, gemm2's
        # atomic-add will scatter into wildly wrong output rows, producing
        # NaN/Inf or massive numerical divergence.
        POISON = 0x7FFFFFFE
        for buf in moe_sort_buffers.values():
            buf.fill_(POISON)

        result = _moe_core_impl(
            x=tensors["x"],
            x_sf=tensors["x_sf"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            w1_weight=tensors["w1_weight"],
            w1_weight_sf=tensors["w1_weight_sf"],
            w1_alpha=tensors["w1_alpha"],
            fc2_input_scale=tensors["fc2_input_scale"],
            w2_weight=tensors["w2_weight"],
            w2_weight_sf=tensors["w2_weight_sf"],
            w2_alpha=tensors["w2_alpha"],
            num_experts=num_experts,
            top_k=top_k,
            num_local_experts=num_local_experts,
            local_expert_offset=local_expert_offset,
            tile_size=tile_size,
            moe_sort_buffers=moe_sort_buffers,
            output_dtype=torch.bfloat16,
        )

        assert result.shape == (num_tokens, hidden_size)
        assert not torch.isnan(result).any(), (
            f"ep_size={ep_size}, num_tokens={num_tokens}: NaNs in output "
            f"after running with poisoned moe_sort buffers — strong "
            f"indication that the routing kernel left sentinel values "
            f"in positions that downstream gemm2 finalize reads, causing "
            f"OOB / garbage-index atomic-adds. The invariant that the "
            f"kernel writes every consumed entry has been violated."
        )
        assert not torch.isinf(result).any(), (
            f"ep_size={ep_size}, num_tokens={num_tokens}: Infs in output "
            f"after running with poisoned moe_sort buffers — same failure "
            f"mode as NaN; kernel left sentinels in consumed positions."
        )

        # Sanity: verify SOMETHING non-trivial actually got computed.
        # At EP=32 with 256 global experts and top_k=8, ~22% of tokens
        # have at least one local expert; the rest produce all-zero
        # output. If the kernel were silently broken and returned all
        # zeros, this check would catch it.
        nonzero_fraction = (result.abs() > 1e-3).float().mean().item()
        # Lower bound proportional to local-expert coverage at this EP
        # (probability that at least one of top_k selected experts is on
        # this rank). EP=1 → ~100%; EP=8 → most tokens covered;
        # EP=16 → ~40%+; EP=32 → ~20%+. Bounds set well below the real
        # fraction to absorb FP4 quantization producing legitimately
        # small (sub-threshold) outputs.
        min_expected_nonzero = {1: 0.5, 8: 0.3, 16: 0.15, 32: 0.05}[ep_size]
        assert nonzero_fraction >= min_expected_nonzero, (
            f"ep_size={ep_size}, num_tokens={num_tokens}: only "
            f"{nonzero_fraction * 100:.2f}% of output entries are "
            f"non-zero (expected at least "
            f"{min_expected_nonzero * 100:.0f}%) — kernel may not have "
            f"executed correctly with poisoned buffers."
        )

        # Reference comparison is supplementary to the NaN/Inf +
        # non-zero-fraction checks above. Skip at high N because
        # ``compute_reference_moe_fp4`` is ``O(num_tokens × top_k)``
        # Python iterations with ``.item()`` syncs, dominating CI
        # runtime per case. The poisoning-detection signal lives in
        # the assertions already executed.
        if num_tokens > 1024:
            return

        ref_output = compute_reference_moe_fp4(
            hidden_states=tensors["x_bf16"].float().cuda(),
            gemm1_weights=tensors["w1_weight_bf16"].float().cuda(),
            gemm2_weights=tensors["w2_weight_bf16"].float().cuda(),
            gemm1_alpha=tensors["w1_alpha"],
            gemm2_alpha=tensors["w2_alpha"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            fc2_input_scale=tensors["fc2_input_scale"],
            num_local_experts=num_local_experts,
            local_expert_offset=local_expert_offset,
        )

        passed, percent_within, atol = check_accuracy(result, ref_output)
        assert passed, (
            f"ep_size={ep_size}, num_tokens={num_tokens}: poisoned-buffer "
            f"test failed: only {percent_within * 100:.2f}% within "
            f"tolerance (atol={atol:.4f}). The kernel left poison "
            f"sentinels in positions that downstream reads — the "
            f"invariant that the routing kernel writes every consumed "
            f"entry is violated, so a Python-side ``.fill_()`` / "
            f"``.zero_()`` init step is load-bearing at this EP."
        )


# =============================================================================
# Test Class: All Valid Tactics
# =============================================================================


@cute_dsl_available
@sm100_required
class TestAllValidTactics:
    """Test that every tactic returned by get_valid_tactics produces correct output.

    For each problem configuration, gets the filtered list of valid tactics via
    can_implement checks, then runs CuteDslMoEWrapper with each tactic explicitly
    and verifies numerical accuracy against the reference implementation.
    """

    pytestmark = _requires_dsl_arch

    @pytest.mark.parametrize(
        "num_tokens,hidden_size,intermediate_size,num_experts,top_k",
        [
            (128, 256, 512, 256, 2),
            (256, 1024, 2048, 256, 8),
        ],
    )
    def test_all_tactics_accuracy(
        self,
        num_tokens: int,
        hidden_size: int,
        intermediate_size: int,
        num_experts: int,
        top_k: int,
    ):
        """Verify every valid tactic produces correct output."""
        from flashinfer import CuteDslMoEWrapper

        num_local_experts = num_experts

        tensors = create_moe_tensors(
            num_tokens=num_tokens,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            top_k=top_k,
        )

        ref_output = compute_reference_moe_fp4(
            hidden_states=tensors["x_bf16"].float().cuda(),
            gemm1_weights=tensors["w1_weight_bf16"].float().cuda(),
            gemm2_weights=tensors["w2_weight_bf16"].float().cuda(),
            gemm1_alpha=tensors["w1_alpha"],
            gemm2_alpha=tensors["w2_alpha"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            fc2_input_scale=tensors["fc2_input_scale"],
        )

        # Create wrapper without CUDA graph so we can freely try different tile_sizes
        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=False,
        )

        # Get the filtered list of valid tactics for this problem size
        inputs = [
            tensors["x"],
            tensors["x_sf"],
            tensors["token_selected_experts"],
            tensors["token_final_scales"],
            tensors["w1_weight"],
            tensors["w1_weight_sf"],
            tensors["w1_alpha"],
            tensors["fc2_input_scale"],
            tensors["w2_weight"],
            tensors["w2_weight_sf"],
            tensors["w2_alpha"],
        ]
        valid_tactics = moe._runner.get_valid_tactics(inputs, None)
        assert len(valid_tactics) > 0, "No valid tactics found"

        num_passed = 0
        num_failed = 0
        for tactic in valid_tactics:
            tile_size = tactic[0]
            result = moe.run(
                x=tensors["x"],
                x_sf=tensors["x_sf"],
                token_selected_experts=tensors["token_selected_experts"],
                token_final_scales=tensors["token_final_scales"],
                w1_weight=tensors["w1_weight"],
                w1_weight_sf=tensors["w1_weight_sf"],
                w1_alpha=tensors["w1_alpha"],
                fc2_input_scale=tensors["fc2_input_scale"],
                w2_weight=tensors["w2_weight"],
                w2_weight_sf=tensors["w2_weight_sf"],
                w2_alpha=tensors["w2_alpha"],
                tactic=tactic,
            )

            assert result.shape == (num_tokens, hidden_size)
            assert not torch.isnan(result).any(), f"NaN in output for tactic {tactic}"
            assert not torch.isinf(result).any(), f"Inf in output for tactic {tactic}"

            passed, percent_within, atol = check_accuracy(result, ref_output)
            if passed:
                num_passed += 1
            else:
                num_failed += 1
                # Don't fail immediately; report all failures at the end
                print(
                    f"[FAIL] tactic tile_size={tile_size} "
                    f"gemm1={tactic[1][0]},{tactic[1][1]} "
                    f"gemm2={tactic[2][0]},{tactic[2][1]}: "
                    f"{percent_within * 100:.2f}% within tolerance (atol={atol:.4f})"
                )

        total = len(valid_tactics)
        assert num_failed == 0, (
            f"{num_failed}/{total} tactics failed accuracy check "
            f"(tokens={num_tokens}, hidden={hidden_size}, "
            f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k})"
        )


# ============================================================================
# moe_output_memset_inplace (dense Path A) — unit tests
# ============================================================================
#
# Tests for the dense cudaMemsetAsync wrapper that mirrors TRT-LLM's
# `moe_output_memset_inplace` Path A. Used in `_moe_core_impl` before GEMM2
# finalize to zero the active output slice for the atomic scatter-add.
# See flashinfer/fused_moe/cute_dsl/moe_utils.py:moe_output_memset_inplace
# for the function and the audit doc follow-up #11 for the scope decision.


@cute_dsl_available
@sm100_required
class TestMoeOutputMemsetInplace:
    """Correctness + stream-handling tests for the dense memset wrapper."""

    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
    @pytest.mark.parametrize(
        "shape",
        [
            (1, 7168),  # decode N=1, DSv3 hidden
            (64, 1024),  # generic small
        ],
    )
    def test_zeros_buffer(self, dtype, shape):
        """Buffer is fully zeroed regardless of starting contents."""
        x = torch.randn(*shape, device="cuda", dtype=dtype)
        assert not (x == 0).all(), (
            "test setup invariant: a random tensor must not be all-zeros"
        )

        from flashinfer.fused_moe.cute_dsl.moe_utils import (
            moe_output_memset_inplace,
        )

        moe_output_memset_inplace(x)
        torch.cuda.synchronize()

        assert (x == 0).all(), (
            f"moe_output_memset_inplace did not zero the buffer "
            f"(shape={shape}, dtype={dtype})"
        )

    def test_unsupported_dtype_raises(self):
        """Reject dtypes we don't bind before native dispatch."""
        from flashinfer.fused_moe.cute_dsl.moe_utils import (
            moe_output_memset_inplace,
        )

        x = torch.randn(16, 32, device="cuda", dtype=torch.float32)
        with pytest.raises(ValueError, match="only supports"):
            moe_output_memset_inplace(x)

    def test_cuda_graph_capture(self):
        """
        Verify the wrapper captures cleanly into a CUDA graph and replays.

        CUDA-graph capture requires every queued op to land on the capture
        stream. If the wrapper's memset ended up on a *different* stream
        (e.g. TVM FFI's env stream when the Python current stream is the
        capture stream), capture would either error or produce an
        inconsistent graph whose replay doesn't zero the buffer. A clean
        capture + correct replay is strong evidence that the explicit
        stream pointer is being honored by the C++ binding.
        """
        from flashinfer.fused_moe.cute_dsl.moe_utils import (
            moe_output_memset_inplace,
        )

        x = torch.randn(128, 2048, device="cuda", dtype=torch.bfloat16)

        # Warm up on the capture stream so JIT compile doesn't race capture.
        capture_stream = torch.cuda.Stream()
        capture_stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(capture_stream):
            moe_output_memset_inplace(x)
        torch.cuda.current_stream().wait_stream(capture_stream)
        torch.cuda.synchronize()

        graph = torch.cuda.CUDAGraph()
        x.fill_(7.0)
        with torch.cuda.graph(graph, stream=capture_stream):
            moe_output_memset_inplace(x)

        # Replay: refill with non-zero so a successful replay is observable.
        x.fill_(7.0)
        torch.cuda.synchronize()
        graph.replay()
        torch.cuda.synchronize()
        assert (x == 0).all(), (
            "CUDA graph replay did not zero the buffer — wrapper's memset "
            "is not being captured onto the capture stream"
        )


class TestMoeOutputMemsetInplaceContract:
    """
    Python-level contract test for ``moe_output_memset_inplace``.

    Intentionally NOT decorated with ``@cute_dsl_available`` /
    ``@sm100_required`` — this test does not touch the GPU at all
    (CPU tensor, monkeypatched FFI dispatch, monkeypatched stream
    getter), so it runs everywhere and catches the highest-priority
    Python-side regressions deterministically.
    """

    def test_passes_current_torch_stream_ptr_to_binding(self, monkeypatch):
        """
        Verify the wrapper passes the FFI binding:
          - the right ``data_ptr`` (positional 1)
          - ``num_tokens`` and ``hidden_size`` (positional 2/3)
          - the result of ``_get_cuda_stream_ptr()`` (positional 4)

        Catches Python-side regressions:
          - dropping the ``_get_cuda_stream_ptr()`` argument
          - passing 0 (which would fall back to TVM FFI's env stream)
          - changing the order of FFI arguments
          - calling the wrong dtype-suffixed entry point

        Monkeypatches both ``_get_moe_utils_module`` (to capture the
        FFI call) and ``_get_cuda_stream_ptr`` (to return a known
        sentinel), so the test is fully deterministic and CPU-only.
        """
        from flashinfer.fused_moe.cute_dsl import moe_utils

        captured = {}

        def fake_memset(ptr, num_tokens, hidden_size, stream_ptr):
            captured["ptr"] = ptr
            captured["num_tokens"] = num_tokens
            captured["hidden_size"] = hidden_size
            captured["stream_ptr"] = stream_ptr

        SENTINEL_STREAM_PTR = 0x123456789ABCDEF0
        monkeypatch.setattr(
            moe_utils,
            "_get_moe_utils_module",
            lambda: {"flashinfer_moe_output_memset_inplace_bf16": fake_memset},
        )
        monkeypatch.setattr(
            moe_utils, "_get_cuda_stream_ptr", lambda: SENTINEL_STREAM_PTR
        )

        x = torch.empty((2, 3), device="cpu", dtype=torch.bfloat16)
        moe_utils.moe_output_memset_inplace(x)

        assert captured["ptr"] == x.data_ptr(), (
            "wrapper passed wrong data_ptr to FFI binding"
        )
        assert captured["num_tokens"] == 2, "wrapper passed wrong num_tokens"
        assert captured["hidden_size"] == 3, "wrapper passed wrong hidden_size"
        assert captured["stream_ptr"] == SENTINEL_STREAM_PTR, (
            "wrapper did not pass `_get_cuda_stream_ptr()`'s return value "
            "as the 4th FFI argument — the active PyTorch stream would be "
            "lost on the C++ side. Likely cause: `_get_cuda_stream_ptr()` "
            "was dropped, replaced with 0, or the FFI argument order "
            "changed."
        )

    @pytest.mark.parametrize(
        ("tensor", "match"),
        [
            (torch.empty((2, 3), dtype=torch.float32), "only supports"),
            (torch.empty((2, 3, 4), dtype=torch.bfloat16), "2D tensor"),
            (torch.empty((2, 3), dtype=torch.bfloat16).t(), "contiguous"),
        ],
    )
    def test_rejects_invalid_inputs_before_native_dispatch(
        self, monkeypatch, tensor, match
    ):
        """Validate dangerous inputs before resolving the native binding."""
        from flashinfer.fused_moe.cute_dsl import moe_utils

        def fail_module_load():
            raise AssertionError("native binding should not be loaded")

        monkeypatch.setattr(moe_utils, "_get_moe_utils_module", fail_module_load)
        with pytest.raises(ValueError, match=match):
            moe_utils.moe_output_memset_inplace(tensor)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
