"""
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.
"""

"""
Auto-tuner for CuteDSL NVFP4 MoE kernels.

This module provides a TunableRunner implementation for the CuteDSL NVFP4 MoE
kernels, enabling automatic performance tuning across different GEMM tactics.

Tactic format follows TRT-LLM's style:
- GEMM1 (Gather + FC1 activation): (mma_tiler_mn, cluster_shape_mn, raster_along_m)
- GEMM2 (Finalize): (mma_tiler_mn, cluster_shape_mn, raster_along_m)

Reference: TensorRT-LLM/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
- Sm100BlockScaledContiguousGatherGroupedGemmSwigluFusionRunner.get_valid_tactics (line 1867)
- Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner.get_valid_tactics (line 1163)
"""

import itertools
import logging
from typing import Any, Callable, Dict, List, Optional, Tuple

import torch

from ...autotuner import (
    DynamicTensorSpec,
    OptimizationProfile,
    TunableRunner,
    TuningConfig,
)
from ...tllm_enums import (
    ActivationType,
    DEFAULT_SWIGLU_ALPHA,
    DEFAULT_SWIGLU_BETA,
    DEFAULT_SWIGLU_LIMIT,
)
from ..utils import (
    get_hybrid_num_tokens_buckets,
    map_to_hybrid_bucket_uncapped,
)
from ._inputs_helper import CuteDslMoEInputsHelper
from .blackwell.moe_w4a16 import launch_w4a16_moe
from .blackwell.moe_w4a16_kernel import Sm100W4A16GroupedGemmKernel
from .moe_utils import (
    normalize_cute_dsl_moe_activation_type,
    validate_cute_dsl_moe_situ_config,
)

logger = logging.getLogger(__name__)


# =============================================================================
# GEMM1 Tactics (Gather + FC1 Activation Fusion)
# =============================================================================
# Reference: TRT-LLM cute_dsl_custom_ops.py line 1867-1897
# Sm100BlockScaledContiguousGatherGroupedGemmSwigluFusionRunner.get_valid_tactics
#
# Format: (mma_tiler_mn, cluster_shape_mn, raster_along_m)
# - mma_tiler_mn: (tile_size, N_tile) where tile_size is 128 or 256, N_tile is 128 or 256
# - cluster_shape_mn: (tile_size // 128, cluster_n) where cluster_n is fixed to 1 for Gather kernel
# - raster_along_m: False (fixed)


def get_gemm1_valid_tactics(tile_size: int) -> List[Tuple]:
    """Get valid tactics for GEMM1 (Gather + SwiGLU Fusion).

    Reference: TRT-LLM cute_dsl_custom_ops.py line 1879-1897

    Args:
        tile_size: MMA tile M dimension (128 or 256)

    Returns:
        List of (mma_tiler_mn, cluster_shape_mn, raster_along_m) tuples
    """
    # From TRT-LLM line 1879-1883:
    # mma_tiler_mn_candidates = [(self.tile_size, 128), (self.tile_size, 256)]
    # cluster_shape_mn_candidates = [(self.tile_size // 128, 1)]  # Note: Only 1, not 2!
    # raster_along_m_candidates = [False]

    mma_tiler_mn_candidates = [(tile_size, 128), (tile_size, 256)]
    cluster_shape_mn_candidates = [
        (tile_size // 128, 1)
    ]  # Gather kernel only supports cluster_n=1
    raster_along_m_candidates = [False]

    tactics = []
    for mma_tiler_mn, cluster_shape_mn, raster_along_m in itertools.product(
        mma_tiler_mn_candidates, cluster_shape_mn_candidates, raster_along_m_candidates
    ):
        tactics.append((mma_tiler_mn, cluster_shape_mn, raster_along_m))

    return tactics


# =============================================================================
# GEMM2 Tactics (Finalize Fusion)
# =============================================================================
# Reference: TRT-LLM cute_dsl_custom_ops.py line 1165-1202
# Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner.get_valid_tactics
#
# Format: (mma_tiler_mn, cluster_shape_mn, raster_along_m)
# - mma_tiler_mn: (tile_size, N_tile) where N_tile is 128 or 256.
#   At tile_size=256 use_2cta_instrs=True, so mma_m doubles to 256.
# - cluster_shape_mn: (tile_size // 128, cluster_n) where cluster_n is 1 or 2.
#   At tile_size=256 cluster_m=2 (2-CTA); at tile_size=128 cluster_m=1.
# - raster_along_m: False (fixed, theoretically more performant)


def get_gemm2_valid_tactics(tile_size: int) -> List[Tuple]:
    """Get valid tactics for GEMM2 (Finalize Fusion).

    Reference: TRT-LLM cute_dsl_custom_ops.py line 1165-1202

    The finalize kernel's MMA shape must match tile_size because the
    kernel consumes the upstream gemm1 output layout. At tile_size=128
    the kernel uses 1-CTA mma_m=128; at tile_size=256 it uses 2-CTA
    mma_m=256 (use_2cta_instrs=True). Returning a 1-CTA gemm2 tactic
    when tile_size=256 yields a layout mismatch and incorrect output.

    Args:
        tile_size: Tile size for moe_sort padding (128 or 256).
            Determines mma_tiler_mn[0] and cluster_shape_mn[0].

    Returns:
        List of (mma_tiler_mn, cluster_shape_mn, raster_along_m) tuples
    """
    mma_tiler_mn_candidates = [(tile_size, 128), (tile_size, 256)]
    cluster_shape_mn_candidates = [
        (tile_size // 128, 1),
        (tile_size // 128, 2),
    ]
    raster_along_m_candidates = [False]

    tactics = []
    for mma_tiler_mn, cluster_shape_mn, raster_along_m in itertools.product(
        mma_tiler_mn_candidates, cluster_shape_mn_candidates, raster_along_m_candidates
    ):
        tactics.append((mma_tiler_mn, cluster_shape_mn, raster_along_m))

    return tactics


# =============================================================================
# Combined MoE Tactics
# =============================================================================
# The MoE pipeline uses both GEMM1 and GEMM2, they must share the same tile_size
# (M dimension of mma_tiler_mn) because moe_sort uses tile_size for padding.
#
# Tactic format: (tile_size, gemm1_tactic, gemm2_tactic)
# - tile_size: 128 or 256 (shared by both GEMMs and moe_sort)
# - gemm1_tactic: (mma_tiler_mn, cluster_shape_mn, raster_along_m)
# - gemm2_tactic: (mma_tiler_mn, cluster_shape_mn, raster_along_m)


# Canonical list of tile_sizes the autotuner is allowed to pick.  Used by
# ``get_moe_valid_tactics`` for tactic enumeration.
VALID_TILE_SIZES: Tuple[int, ...] = (128, 256)


def get_moe_valid_tactics() -> List[Tuple]:
    """Get all valid MoE tactic combinations.

    Each tactic is a tuple: (tile_size, gemm1_tactic, gemm2_tactic)

    The tile_size must be shared between GEMM1 and GEMM2 because:
    1. moe_sort uses tile_size to pad tokens to tile boundaries
    2. Both GEMMs process the same padded token sequence

    Returns:
        List of (tile_size, gemm1_tactic, gemm2_tactic) tuples
    """
    tactics = []

    # Enable both 1-CTA (tile_size=128) and 2-CTA (tile_size=256,
    # use_2cta_instrs=True) variants; the autotuner picks per shape.
    # tile_size=256 typically wins at large batch where 2-CTA throughput
    # exceeds 1-CTA.
    for tile_size in VALID_TILE_SIZES:
        gemm1_tactics = get_gemm1_valid_tactics(tile_size)
        gemm2_tactics = get_gemm2_valid_tactics(tile_size)

        for gemm1_tactic, gemm2_tactic in itertools.product(
            gemm1_tactics, gemm2_tactics
        ):
            tactics.append((tile_size, gemm1_tactic, gemm2_tactic))

    return tactics


# Pre-generate all valid tactics
# tile_size=128: 2 GEMM1 tactics × 4 GEMM2 tactics = 8
# tile_size=256: 2 GEMM1 tactics × 4 GEMM2 tactics = 8
# Total: 16 tactics
ALL_MOE_TACTICS = get_moe_valid_tactics()

# Default tactic (tile_size=128, smallest MMA tiles, cluster_n=1)
DEFAULT_MOE_TACTIC = (
    128,  # tile_size
    ((128, 128), (1, 1), False),  # gemm1_tactic
    ((128, 128), (1, 1), False),  # gemm2_tactic
)


def _extract_tactic_params(tactic: Tuple) -> Dict[str, Any]:
    """Extract parameters from a MoE tactic tuple.

    Args:
        tactic: (tile_size, gemm1_tactic, gemm2_tactic)

    Returns:
        Dictionary with all tactic parameters
    """
    tile_size, gemm1_tactic, gemm2_tactic = tactic
    gemm1_mma_tiler_mn, gemm1_cluster_shape_mn, gemm1_raster_along_m = gemm1_tactic
    gemm2_mma_tiler_mn, gemm2_cluster_shape_mn, gemm2_raster_along_m = gemm2_tactic

    return {
        "tile_size": tile_size,
        "gemm1_mma_tiler_mn": gemm1_mma_tiler_mn,
        "gemm1_cluster_shape_mn": gemm1_cluster_shape_mn,
        "gemm1_raster_along_m": gemm1_raster_along_m,
        "gemm2_mma_tiler_mn": gemm2_mma_tiler_mn,
        "gemm2_cluster_shape_mn": gemm2_cluster_shape_mn,
        "gemm2_raster_along_m": gemm2_raster_along_m,
    }


class CuteDslFusedMoENvfp4Runner(TunableRunner):
    """TunableRunner for CuteDSL NVFP4 MoE kernels.

    This runner enables auto-tuning of the CuteDSL NVFP4 MoE pipeline by
    trying different combinations of GEMM tactics.

    Tactic format follows TRT-LLM style:
        (tile_size, gemm1_tactic, gemm2_tactic)
    where:
        - tile_size: 128 or 256
        - gemm1_tactic: (mma_tiler_mn, cluster_shape_mn, raster_along_m)
        - gemm2_tactic: (mma_tiler_mn, cluster_shape_mn, raster_along_m)

    Input tensor indices (for dynamic_tensor_specs):
        0: x (num_tokens, hidden_size//2) - FP4 packed input
        1: x_sf (num_tokens, hidden_size//sf_vec_size) - input scale factors
        2: token_selected_experts (num_tokens, top_k) - expert assignments
        3: token_final_scales (num_tokens, top_k) - routing weights
        4-10: weight tensors (fixed size, don't depend on num_tokens)
        11: moe_output, or per_token_scale when per-token activation is enabled
        12: moe_output when per-token activation is enabled

    Args:
        forward_impl: The actual MoE implementation function.
        num_experts: Total number of experts.
        top_k: Number of experts selected per token.
        num_local_experts: Number of local experts (for expert parallelism).
        local_expert_offset: Starting expert index for this partition.
        use_fused_finalize: Whether to use fused finalize (default: True).
        output_dtype: Output data type (default: torch.bfloat16).
        use_per_token_activation: Whether inputs include per-token row scales
            for GEMM1.
        situ_beta: When set with ActivationType.Swiglu, use the SiTU gate.
        situ_linear_beta: Optional SiTU tanh clamp for the up branch.
    """

    def __init__(
        self,
        forward_impl: Callable,
        num_experts: int,
        top_k: int,
        num_local_experts: int,
        local_expert_offset: int = 0,
        use_fused_finalize: bool = True,
        output_dtype: torch.dtype = torch.bfloat16,
        enable_pdl: bool = True,
        activation_type: int = ActivationType.Swiglu.value,
        swiglu_alpha: float = DEFAULT_SWIGLU_ALPHA,
        swiglu_beta: float = DEFAULT_SWIGLU_BETA,
        swiglu_limit: float = DEFAULT_SWIGLU_LIMIT,
        situ_beta: Optional[float] = None,
        situ_linear_beta: Optional[float] = None,
        use_per_token_activation: bool = False,
    ):
        activation_type, gated = normalize_cute_dsl_moe_activation_type(activation_type)
        validate_cute_dsl_moe_situ_config(activation_type, situ_beta, situ_linear_beta)
        self.forward_impl = forward_impl
        self.num_experts = num_experts
        self.top_k = top_k
        self.num_local_experts = num_local_experts
        self.local_expert_offset = local_expert_offset
        self.use_fused_finalize = use_fused_finalize
        self.output_dtype = output_dtype
        self.enable_pdl = enable_pdl
        self.activation_type = activation_type
        self.gated = gated
        self.swiglu_alpha = swiglu_alpha
        self.swiglu_beta = swiglu_beta
        self.swiglu_limit = swiglu_limit
        self.situ_beta = situ_beta
        self.situ_linear_beta = situ_linear_beta
        self.use_per_token_activation = use_per_token_activation

        # Helper that builds a deterministic balanced approx-max-load
        # assignment for token_selected_experts during autotune profiling.
        # See _inputs_helper.py for rationale -- the random tensor_initializer
        # for input #2 produces non-deterministic and unrealistic per-expert
        # load distributions, biasing autotune picks at marginal cells.
        self._inputs_helper = CuteDslMoEInputsHelper(
            num_experts, top_k, num_local_experts, local_expert_offset
        )

        # Instance-level so dummy expert IDs span all local experts
        # (randint(0, num_experts)) for realistic profiling.
        self.tuning_config = TuningConfig(
            dynamic_tensor_specs=(
                DynamicTensorSpec(
                    input_idx=(0, 1, 2, 3, 11)
                    + ((12,) if use_per_token_activation else ()),
                    dim_idx=(0,) * (6 if use_per_token_activation else 5),
                    # Bare callables: autotuner adapts the bucket set to
                    # the actual input dim (matches the
                    # _FP8_GEMM_SM100_TUNING_CONFIG pattern in
                    # `gemm/gemm_base.py`).
                    gen_tuning_buckets=get_hybrid_num_tokens_buckets,
                    map_to_tuning_buckets=map_to_hybrid_bucket_uncapped,
                ),
            ),
            # Per-input initializer closures, keyed by input index. Indices
            # match input_idx above: 11 is per_token_scale when per-token
            # activation is enabled (else moe_output), 12 is moe_output.
            tensor_initializers=(
                # 0: x — FP4 quantized input (uint8 packed). Seeded
                # for cross-process determinism of autotune picks
                # (matches trt-llm's seed=515 convention).
                (
                    0,
                    lambda shapes, dtype, device: torch.randint(
                        0,
                        256,
                        shapes,
                        dtype=torch.uint8,
                        device=device,
                        generator=torch.Generator(device=device).manual_seed(515),
                    ),
                ),
                # 1: x_sf — FP8 scale factors (uint8). Seeded.
                (
                    1,
                    lambda shapes, dtype, device: torch.randint(
                        1,
                        128,
                        shapes,
                        dtype=torch.uint8,
                        device=device,
                        generator=torch.Generator(device=device).manual_seed(515),
                    ),
                ),
                # 2: token_selected_experts — output is overwritten
                # by inputs_pre_hook (CuteDslMoEInputsHelper), but
                # seed the initializer too in case the hook is ever
                # disabled.
                (
                    2,
                    lambda shapes, dtype, device: torch.randint(
                        0,
                        max(num_experts, 1),
                        shapes,
                        dtype=torch.int32,
                        device=device,
                        generator=torch.Generator(device=device).manual_seed(515),
                    ),
                ),
                # 3: token_final_scales — softmax-normalized. Seeded.
                (
                    3,
                    lambda shapes, dtype, device: torch.softmax(
                        torch.randn(
                            shapes,
                            device=device,
                            generator=torch.Generator(device=device).manual_seed(515),
                        ),
                        dim=-1,
                    ).to(torch.float32),
                ),
                *(
                    [
                        # 11: per_token_scale — ones.
                        (
                            11,
                            lambda shapes, dtype, device: torch.ones(
                                shapes, dtype=torch.float32, device=device
                            ),
                        ),
                        # 12: moe_output — empty.
                        (
                            12,
                            lambda shapes, dtype, device: torch.empty(
                                shapes, dtype=dtype, device=device
                            ),
                        ),
                    ]
                    if use_per_token_activation
                    else [
                        # 11: moe_output — empty.
                        (
                            11,
                            lambda shapes, dtype, device: torch.empty(
                                shapes, dtype=dtype, device=device
                            ),
                        )
                    ]
                ),
            ),
            inputs_pre_hook=self._inputs_helper.inputs_pre_hook,
            # Cold-L2 measurement matches TRT-LLM's
            # CuteDslFusedMoENvfp4Runner.tuning_config; flushing L2
            # between profile iterations yields autotune timings
            # representative of production cold-cache conditions.
            use_cold_l2_cache=True,
        )

    def __hash__(self):
        return hash(
            (
                self.num_experts,
                self.top_k,
                self.num_local_experts,
                self.local_expert_offset,
                self.use_fused_finalize,
                self.output_dtype,
                int(self.activation_type),
                self.swiglu_alpha,
                self.swiglu_beta,
                self.swiglu_limit,
                self.situ_beta,
                self.situ_linear_beta,
                self.use_per_token_activation,
            )
        )

    def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple:
        return (
            int(self.activation_type),
            self.swiglu_alpha,
            self.swiglu_beta,
            self.swiglu_limit,
            self.situ_beta,
            self.situ_linear_beta,
        )

    def get_valid_tactics(  # type: ignore[override]
        self,
        inputs: List[torch.Tensor],
        profile: OptimizationProfile,
    ) -> List[Tuple[Any, ...]]:
        """Return list of valid tactics filtered by can_implement checks.

        Validates each candidate tactic against both GEMM1 and GEMM2 kernel
        can_implement methods using the actual problem dimensions from inputs.

        Returns tactics in TRT-LLM format:
            (tile_size, gemm1_tactic, gemm2_tactic)

        Args:
            inputs: List of input tensors used to derive problem dimensions.
            profile: Optimization profile (not used for tactic validation).

        Returns:
            List of valid tactic tuples.
        """
        import cutlass
        from .blackwell import (
            BlockScaledContiguousGatherGroupedGemmKernel,
            Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel,
        )
        from .moe_utils import get_max_num_permuted_tokens

        # Extract problem dimensions from inputs:
        #   0: x (num_tokens, hidden_size//2)
        #   4: w1_weight (num_local_experts, 2*intermediate_size, hidden_size//2)
        #   8: w2_weight (num_local_experts, hidden_size, intermediate_size//2)
        x = inputs[0]
        w1_weight = inputs[4]

        gated = self.gated
        num_tokens = x.shape[0]
        hidden_size = x.shape[1] * 2  # FP4 packed
        num_local_experts = w1_weight.shape[0]
        # Gated SwiGLU fuses gate+up (2*intermediate rows); non-gated ReLU^2
        # has a single intermediate-row projection.
        gemm1_n = w1_weight.shape[1]
        intermediate_size = gemm1_n // 2 if gated else gemm1_n

        # Fixed dtypes/layouts for NVFP4 MoE
        ab_dtype = cutlass.Float4E2M1FN
        sf_dtype = cutlass.Float8E4M3FN
        sf_vec_size = 16

        if self.use_per_token_activation:
            if self.output_dtype == torch.float16:
                gemm1_c_dtype = cutlass.Float16
            elif self.output_dtype == torch.bfloat16:
                gemm1_c_dtype = cutlass.BFloat16
            else:
                return []
        else:
            gemm1_c_dtype = cutlass.Float4E2M1FN
        gemm2_out_dtype = cutlass.BFloat16

        token_final_scales = inputs[3]
        if token_final_scales.dtype == torch.float32:
            final_scale_dtype = cutlass.Float32
        elif token_final_scales.dtype == torch.bfloat16:
            final_scale_dtype = cutlass.BFloat16
        else:
            final_scale_dtype = cutlass.Float16

        def _tactic_ok(tactic):
            tile_size, gemm1_tactic, gemm2_tactic = tactic
            gemm1_mma_tiler_mn = gemm1_tactic[0]
            gemm1_cluster_shape_mn = gemm1_tactic[1]
            gemm2_mma_tiler_mn = gemm2_tactic[0]
            gemm2_cluster_shape_mn = gemm2_tactic[1]

            permuted_m = get_max_num_permuted_tokens(
                num_tokens, self.top_k, self.num_local_experts, tile_size
            )

            gemm1_ok = BlockScaledContiguousGatherGroupedGemmKernel.can_implement(
                ab_dtype=ab_dtype,
                sf_dtype=sf_dtype,
                sf_vec_size=sf_vec_size,
                c_dtype=gemm1_c_dtype,
                mma_tiler_mn=gemm1_mma_tiler_mn,
                cluster_shape_mn=gemm1_cluster_shape_mn,
                m=permuted_m,
                n=gemm1_n,
                k=hidden_size,
                l=num_local_experts,
                a_major="k",
                b_major="k",
                c_major="n",
            )

            gemm2_ok = (
                Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel.can_implement(
                    ab_dtype=ab_dtype,
                    sf_dtype=sf_dtype,
                    sf_vec_size=sf_vec_size,
                    out_dtype=gemm2_out_dtype,
                    final_scale_dtype=final_scale_dtype,
                    mma_tiler_mn=gemm2_mma_tiler_mn,
                    cluster_shape_mn=gemm2_cluster_shape_mn,
                    m=permuted_m,
                    n=hidden_size,
                    k=intermediate_size,
                    l=num_local_experts,
                    a_major="k",
                    b_major="k",
                    out_major="n",
                )
            )

            return gemm1_ok and gemm2_ok

        valid_tactics = [t for t in ALL_MOE_TACTICS if _tactic_ok(t)]

        if not valid_tactics:
            # DEFAULT_MOE_TACTIC is a member of ALL_MOE_TACTICS, so an empty
            # list means even the default fails can_implement -- do not fall
            # back to it unvalidated (gh #3957). This early refusal is
            # diagnostics/defense-in-depth: the kernel wrappers re-validate
            # can_implement at launch and raise, so an unvalidated tactic
            # cannot reach the device -- but refusing here avoids pointless
            # profiling of a tactic that can only throw, and says why.
            logger.warning(
                "No valid tactics found for problem dims "
                "(tokens=%d, hidden=%d, intermediate=%d, experts=%d, top_k=%d). "
                "Falling back to default tactic.",
                num_tokens,
                hidden_size,
                intermediate_size,
                num_local_experts,
                self.top_k,
            )

        return valid_tactics

    def forward(  # type: ignore[override]
        self,
        inputs: List[torch.Tensor],
        tactic: Tuple[Any, ...] = None,  # type: ignore[assignment]
        do_preparation: bool = False,
        **kwargs: Any,
    ) -> torch.Tensor:
        """Execute the MoE forward pass with the specified tactic.

        Args:
            inputs: List of input tensors:
                [x, x_sf, token_selected_experts, token_final_scales,
                 w1_weight, w1_weight_sf, w1_alpha, fc2_input_scale,
                 w2_weight, w2_weight_sf, w2_alpha, per_token_scale (optional),
                 moe_output (optional)]
            tactic: Tactic tuple (tile_size, gemm1_tactic, gemm2_tactic) or None for default.
            do_preparation: If True, perform one-time setup (not used).
            **kwargs: Additional keyword arguments passed to forward_impl.

        Returns:
            Output tensor from the MoE computation.
        """
        if tactic is None or tactic == -1:
            tactic = DEFAULT_MOE_TACTIC

        # Extract parameters from tactic
        params = _extract_tactic_params(tactic)

        # Unpack inputs
        (
            x,
            x_sf,
            token_selected_experts,
            token_final_scales,
            w1_weight,
            w1_weight_sf,
            w1_alpha,
            fc2_input_scale,
            w2_weight,
            w2_weight_sf,
            w2_alpha,
            *optional_inputs,
        ) = inputs

        if self.use_per_token_activation:
            if not optional_inputs:
                raise ValueError(
                    "per_token_scale is required when use_per_token_activation=True"
                )
            per_token_scale = optional_inputs[0]
            moe_output = optional_inputs[1] if len(optional_inputs) > 1 else None
        else:
            per_token_scale = None
            moe_output = optional_inputs[0] if optional_inputs else None

        # Call the implementation with tactic parameters
        return self.forward_impl(
            x=x,
            x_sf=x_sf,
            token_selected_experts=token_selected_experts,
            token_final_scales=token_final_scales,
            w1_weight=w1_weight,
            w1_weight_sf=w1_weight_sf,
            w1_alpha=w1_alpha,
            fc2_input_scale=fc2_input_scale,
            w2_weight=w2_weight,
            w2_weight_sf=w2_weight_sf,
            w2_alpha=w2_alpha,
            num_experts=self.num_experts,
            top_k=self.top_k,
            num_local_experts=self.num_local_experts,
            local_expert_offset=self.local_expert_offset,
            tile_size=params["tile_size"],
            gemm1_mma_tiler_mn=params["gemm1_mma_tiler_mn"],
            gemm1_cluster_shape_mn=params["gemm1_cluster_shape_mn"],
            gemm2_mma_tiler_mn=params["gemm2_mma_tiler_mn"],
            gemm2_cluster_shape_mn=params["gemm2_cluster_shape_mn"],
            output_dtype=self.output_dtype,
            use_fused_finalize=self.use_fused_finalize,
            moe_output=moe_output,
            per_token_scale=per_token_scale,
            enable_pdl=self.enable_pdl,
            activation_type=int(self.activation_type),
            swiglu_alpha=self.swiglu_alpha,
            swiglu_beta=self.swiglu_beta,
            swiglu_limit=self.swiglu_limit,
            situ_beta=self.situ_beta,
            situ_linear_beta=self.situ_linear_beta,
            **kwargs,
        )


_W4A16_ROUTE_TILES = (8, 16, 32, 64, 128, 192)
_W4A16_K_TILE = 256
# W4A16 maps output channels to M and routed rows to N. Preserve the current
# M-major static scheduler order; CLC scheduling owns its traversal order.
_W4A16_RASTER_ALONG_M = True
# Grouped expert scheduling requires cluster N=1 when multiple routed rows
# target the same expert.
_W4A16_GEMM_TOPOLOGY_PAIRS = (
    ((128, (1, 1)), (128, (1, 1))),
    ((128, (2, 1)), (128, (2, 1))),
    ((256, (2, 1)), (256, (2, 1))),
    ((128, (2, 1)), (256, (2, 1))),
    ((256, (2, 1)), (128, (2, 1))),
)


def get_w4a16_moe_valid_tactics() -> List[Tuple]:
    """Get valid W4A16 GEMM1/GEMM2 tactic pairs."""
    tactics: List[Tuple] = []
    for route_tile in _W4A16_ROUTE_TILES:
        for gemm1_topology, gemm2_topology in _W4A16_GEMM_TOPOLOGY_PAIRS:
            gemm1_m, gemm1_cluster_shape = gemm1_topology
            gemm2_m, gemm2_cluster_shape = gemm2_topology
            if route_tile < 16 and (gemm1_m == 256 or gemm2_m == 256):
                continue
            # A 128x192x256 tile cannot retain the two load and transform
            # stages required by the warp-specialized pipeline.
            if route_tile == 192 and (gemm1_m == 128 or gemm2_m == 128):
                continue
            gemm1_tactic = (
                (gemm1_m, route_tile, _W4A16_K_TILE),
                gemm1_cluster_shape,
                _W4A16_RASTER_ALONG_M,
            )
            gemm2_tactic = (
                (gemm2_m, route_tile, _W4A16_K_TILE),
                gemm2_cluster_shape,
                _W4A16_RASTER_ALONG_M,
            )
            tactics.append((gemm1_tactic, gemm2_tactic))
    return tactics


W4A16_MOE_TACTICS = tuple(get_w4a16_moe_valid_tactics())


class CuteDslFusedMoEW4A16Runner(TunableRunner):
    """Tunable runner for the BF16-activation, NVFP4-weight MoE pipeline.

    Inputs:
        [x, token_selected_experts, token_final_scales,
         w1_weight, w1_weight_sf, w1_alpha,
         w2_weight, w2_weight_sf, w2_alpha, moe_output]
    """

    def __init__(
        self,
        num_experts: int,
        top_k: int,
        num_local_experts: int,
        local_expert_offset: int = 0,
        use_fused_finalize: bool = True,
        output_dtype: torch.dtype = torch.bfloat16,
        enable_pdl: bool = True,
        activation_type: int = ActivationType.Swiglu.value,
        swiglu_alpha: float = DEFAULT_SWIGLU_ALPHA,
        swiglu_beta: float = DEFAULT_SWIGLU_BETA,
        swiglu_limit: float = DEFAULT_SWIGLU_LIMIT,
    ):
        activation_type, _ = normalize_cute_dsl_moe_activation_type(activation_type)
        if output_dtype != torch.bfloat16:
            raise ValueError("W4A16 only supports BF16 output")
        self.num_experts = num_experts
        self.top_k = top_k
        self.num_local_experts = num_local_experts
        self.local_expert_offset = local_expert_offset
        self.use_fused_finalize = use_fused_finalize
        self.output_dtype = output_dtype
        self.enable_pdl = enable_pdl
        self.activation_type = activation_type
        self.swiglu_alpha = swiglu_alpha
        self.swiglu_beta = swiglu_beta
        self.swiglu_limit = swiglu_limit
        self._workspace_cache: Dict[Tuple, Any] = {}

        # Match production EP routing density while retaining seeded load
        # variance around route-tile boundaries.
        self.tuning_config = TuningConfig(
            dynamic_tensor_specs=(
                DynamicTensorSpec(
                    input_idx=(0, 1, 2, 9),
                    dim_idx=(0, 0, 0, 0),
                    gen_tuning_buckets=get_hybrid_num_tokens_buckets,
                    map_to_tuning_buckets=map_to_hybrid_bucket_uncapped,
                ),
            ),
            tensor_initializers=(
                (
                    0,
                    lambda shapes, dtype, device: torch.randn(
                        shapes,
                        dtype=dtype,
                        device=device,
                        generator=torch.Generator(device=device).manual_seed(515),
                    ),
                ),
                (
                    1,
                    lambda shapes, dtype, device: torch.randint(
                        0,
                        max(num_experts, 1),
                        shapes,
                        dtype=torch.int32,
                        device=device,
                        generator=torch.Generator(device=device).manual_seed(515),
                    ),
                ),
                (
                    2,
                    lambda shapes, dtype, device: torch.softmax(
                        torch.randn(
                            shapes,
                            device=device,
                            generator=torch.Generator(device=device).manual_seed(515),
                        ),
                        dim=-1,
                    ).to(torch.float32),
                ),
                (
                    9,
                    lambda shapes, dtype, device: torch.empty(
                        shapes, dtype=dtype, device=device
                    ),
                ),
            ),
            use_cold_l2_cache=True,
        )

    def __hash__(self):
        return hash(
            (
                self.num_experts,
                self.top_k,
                self.num_local_experts,
                self.local_expert_offset,
                self.use_fused_finalize,
                self.output_dtype,
                self.enable_pdl,
                int(self.activation_type),
                self.swiglu_alpha,
                self.swiglu_beta,
                self.swiglu_limit,
            )
        )

    def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple:
        return (
            self.num_experts,
            self.top_k,
            self.num_local_experts,
            self.local_expert_offset,
            self.use_fused_finalize,
            self.output_dtype,
            self.enable_pdl,
            int(self.activation_type),
            self.swiglu_alpha,
            self.swiglu_beta,
            self.swiglu_limit,
        )

    def get_valid_tactics(  # type: ignore[override]
        self,
        inputs: List[torch.Tensor],
        profile: OptimizationProfile,
    ) -> List[Tuple[Any, ...]]:
        import cutlass

        from .moe_utils import get_max_num_permuted_tokens

        w1_weight = inputs[3]
        w2_weight = inputs[6]
        num_tokens = inputs[0].shape[0]

        def can_implement(
            weight: torch.Tensor,
            route_slots: int,
            gemm_tactic: Tuple,
        ) -> bool:
            m = weight.shape[1]
            k = weight.shape[2] * 2
            mma_tiler_mnk, cluster_shape_mn, _ = gemm_tactic
            mma_tiler_m = mma_tiler_mnk[0]
            return Sm100W4A16GroupedGemmKernel.can_implement(
                mnkl=(m, route_slots, k, self.num_local_experts),
                a_dtype=cutlass.Float4E2M1FN,
                b_dtype=cutlass.BFloat16,
                c_dtype=cutlass.BFloat16,
                a_major="k",
                b_major="k",
                c_major="m",
                mma_tiler=mma_tiler_mnk,
                cluster_shape_mn=cluster_shape_mn,
                use_2cta_instrs=mma_tiler_m == 256,
            )

        valid_tactics = []
        for tactic in W4A16_MOE_TACTICS:
            gemm1_tactic, gemm2_tactic = tactic
            route_tile = gemm1_tactic[0][1]
            route_slots = get_max_num_permuted_tokens(
                num_tokens,
                self.top_k,
                self.num_local_experts,
                route_tile,
            )
            if can_implement(w1_weight, route_slots, gemm1_tactic) and can_implement(
                w2_weight, route_slots, gemm2_tactic
            ):
                valid_tactics.append(tactic)
        return valid_tactics

    def forward(  # type: ignore[override]
        self,
        inputs: List[torch.Tensor],
        tactic: Optional[Tuple] = None,
        do_preparation: bool = False,
        **kwargs: Any,
    ) -> torch.Tensor:
        (
            x,
            token_selected_experts,
            token_final_scales,
            w1_weight,
            w1_weight_sf,
            w1_alpha,
            w2_weight,
            w2_weight_sf,
            w2_alpha,
            moe_output,
        ) = inputs
        if x.dtype != torch.bfloat16:
            raise TypeError(f"W4A16 requires x.dtype=torch.bfloat16, got {x.dtype}")
        if token_final_scales.dtype != torch.float32:
            raise TypeError(
                "W4A16 requires token_final_scales.dtype=torch.float32, "
                f"got {token_final_scales.dtype}"
            )
        num_tokens = int(token_selected_experts.size(0))
        hidden_size = int(w2_weight.size(1))
        if tuple(x.shape) != (num_tokens, hidden_size):
            raise ValueError(
                f"x must have shape {(num_tokens, hidden_size)} for W4A16, "
                f"got {tuple(x.shape)}"
            )
        if tuple(moe_output.shape) != (num_tokens, hidden_size):
            raise ValueError(
                f"moe_output must have shape {(num_tokens, hidden_size)}, "
                f"got {tuple(moe_output.shape)}"
            )
        return launch_w4a16_moe(
            x=x,
            token_selected_experts=token_selected_experts,
            token_final_scales=token_final_scales,
            w1_weight=w1_weight,
            w1_weight_sf=w1_weight_sf,
            w1_alpha=w1_alpha,
            w2_weight=w2_weight,
            w2_weight_sf=w2_weight_sf,
            w2_alpha=w2_alpha,
            num_experts=self.num_experts,
            num_local_experts=self.num_local_experts,
            local_expert_offset=self.local_expert_offset,
            use_fused_finalize=self.use_fused_finalize,
            moe_output=moe_output,
            enable_pdl=self.enable_pdl,
            activation_type=self.activation_type,
            swiglu_alpha=self.swiglu_alpha,
            swiglu_beta=self.swiglu_beta,
            swiglu_limit=self.swiglu_limit,
            tactic=None if tactic is None or tactic == -1 else tactic,
            workspace_cache=self._workspace_cache,
        )


# =============================================================================
# Utility Functions
# =============================================================================


def print_all_tactics():
    """Print all valid MoE tactics for debugging."""
    logger.info("Total MoE tactics: %d", len(ALL_MOE_TACTICS))
    for i, tactic in enumerate(ALL_MOE_TACTICS):
        tile_size, gemm1_tactic, gemm2_tactic = tactic
        logger.info(
            "Tactic %d:\n  tile_size: %s\n  gemm1: mma_tiler_mn=%s, cluster_shape_mn=%s, raster_along_m=%s\n  gemm2: mma_tiler_mn=%s, cluster_shape_mn=%s, raster_along_m=%s",
            i,
            tile_size,
            gemm1_tactic[0],
            gemm1_tactic[1],
            gemm1_tactic[2],
            gemm2_tactic[0],
            gemm2_tactic[1],
            gemm2_tactic[2],
        )
