import inspect
from collections import defaultdict
from typing import Optional

import numpy as np
import torch

import flashinfer

try:
    from flashinfer import ActivationType
except ImportError:
    # ActivationType was not exported from the top-level package until 0.6.3
    from flashinfer.fused_moe.core import ActivationType

from flashinfer.autotuner import autotune
from flashinfer.fused_moe import (
    trtllm_fp4_block_scale_moe,
    trtllm_fp8_block_scale_moe,
    trtllm_fp8_per_tensor_scale_moe,
    cutlass_fused_moe,
    fused_topk_deepseek,
)
from flashinfer.tllm_enums import RoutingMethodType
from flashinfer import fp4_quantize, mxfp8_quantize
from flashinfer.testing.utils import (
    bench_gpu_time,
)

from .flashinfer_benchmark_utils import (
    dtype_str_to_torch_dtype,
    enum_type,
    get_device,
    print_perf_metrics,
    filter_backends_by_compute_capability,
    warn_if_pdl_unsupported,
)
from .moe_utils import (
    calculate_fp4_global_scale,
    quantize_fp4,
    quantize_fp4_batched,
    quantize_fp8,
    calculate_moe_tflops,
    calculate_moe_kernel_bandwidth,
    compute_routing,
    generate_moe_weights,
    add_common_moe_args,
    process_fp8_weight_layout,
    create_moe_output_scale_scalars,
    FLOAT8_E4M3_MAX,
    FLOAT4_E2M1_MAX,
    SF_VEC_SIZE,
)

# Before 0.6.3, MoE APIs used "gated_act_type" (SwiGlu=0, GeGlu=1) instead of
# "activation_type" (Swiglu=3, Geglu=4). Some prior APIs omit the parameter entirely.
_ACTIVATION_TO_GATED_ACT = {
    ActivationType.Swiglu: 0,
    ActivationType.Geglu: 1,
}


def _activation_kwarg(fn, activation_type: ActivationType) -> dict:
    """Return the correct activation keyword argument for *fn* in the installed version."""
    sig = inspect.signature(fn)
    if "activation_type" in sig.parameters:
        return {"activation_type": activation_type.value}
    if "gated_act_type" in sig.parameters:
        if activation_type not in _ACTIVATION_TO_GATED_ACT:
            raise ValueError(
                f"Activation type {activation_type.name} is not supported by the "
                f"installed flashinfer version (pre-0.6.3 only supports "
                f"{[k.name for k in _ACTIVATION_TO_GATED_ACT]})"
            )
        return {"gated_act_type": _ACTIVATION_TO_GATED_ACT[activation_type]}
    return {}


def is_gated_activation(activation_type: ActivationType) -> bool:
    """Whether the activation splits FC1 output into gate/up halves (FC1 weight has 2*intermediate
    rows). SwigluStep's clamp limit defaults to 7.0 (the Step-3 model value) in the kernel, so no
    swiglu_limit tensor needs to be passed.
    """
    return activation_type in (
        ActivationType.Swiglu,
        ActivationType.Geglu,
        ActivationType.SwigluBias,
        ActivationType.SwigluStep,
    )


def run_moe_test(args):
    """
    Run a MOE test.

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.routine == "trtllm_fp4_block_scale_moe":
        return testTrtllmFp4BlockScaleMoe(args)
    elif args.routine == "trtllm_fp8_block_scale_moe":
        return testTrtllmFp8BlockScaleMoe(args)
    elif args.routine == "trtllm_fp8_per_tensor_scale_moe":
        return testTrtllmFp8PerTensorScaleMoe(args)
    elif args.routine == "cutlass_fused_moe":
        return testCutlassFusedMoe(args)
    elif args.routine == "cute_dsl_fp4_block_scale_moe":
        return testCuteDslFp4BlockScaleMoe(args)
    elif args.routine == "b12x_fused_moe":
        return testB12xFusedMoe(args)
    elif args.routine == "bgmv_moe":
        return testBgmvMoe(args)
    elif args.routine == "unified_nvfp4_moe":
        return testUnifiedNvfp4Moe(args)
    else:
        raise ValueError(f"Unsupported routine: {args.routine}")


def parse_moe_args(line, parser):
    """
    Parse command line arguments for MOE test configuration.

    Args:
        line: Command line arguments
        parser: ArgumentParser object already populated with shared arguments

    Returns:
        Parsed argument namespace
    """
    add_common_moe_args(parser)
    # Note: num_tokens/hidden_size is added by add_common_moe_args
    parser.add_argument(
        "--intermediate_size",
        type=int,
        required=True,
        help="Intermediate dimension size (not used for bgmv_moe).",
    )
    # Note: num_experts/top_k is added by add_common_moe_args
    parser.add_argument(
        "--n_group",
        type=int,
        required=False,
        default=None,
        help="Number of expert groups (for DeepSeek routing). Only used with DeepSeekV3 routing method.",
    )
    parser.add_argument(
        "--topk_group",
        type=int,
        required=False,
        default=None,
        help="Number of groups to consider for top-k routing. Only used with DeepSeekV3 routing method.",
    )
    parser.add_argument(
        "--routed_scaling_factor",
        type=float,
        required=False,
        default=2.5,
        help="Scaling factor for routing.",
    )
    parser.add_argument(
        "--local_expert_offset",
        type=int,
        required=False,
        default=0,
        help="Offset of local experts in global expert space.",
    )
    parser.add_argument(
        "--local_num_experts",
        type=int,
        required=False,
        default=None,
        help="Number of experts handled by this device. Defaults to num_experts.",
    )
    parser.add_argument(
        "--routing_method",
        type=str,
        required=False,
        default="deepseek_v3",
        choices=[
            "renormalize",
            "deepseek_v3",
            "llama4",
            "renormalize_naive",
            "topk",
        ],
        help=(
            "Routing method: renormalize | deepseek_v3 | llama4 | renormalize_naive | topk."
        ),
    )
    parser.add_argument(
        "--use_shuffled_weight",
        action="store_true",
        default=False,
        help="Whether to use shuffled weight layout.",
    )
    parser.add_argument(
        "--weight_layout",
        type=int,
        required=False,
        default=0,
        choices=[0, 1, 2],
        help="Weight layout: 0=MajorK, 1=MajorMn, 2=BlockMajorK.",
    )
    parser.add_argument(
        "--use_routing_bias",
        action="store_true",
        default=False,
        help="Whether to use routing bias.",
    )
    parser.add_argument(
        "--use_routing_scales_on_input",
        action="store_true",
        default=False,
        help="Whether to use routing scales on input (for Llama4 routing).",
    )
    # Note: input_dtype is added by add_common_moe_args
    parser.add_argument(
        "--weight_dtype",
        type=str,
        required=False,
        default="bfloat16",
        help="Data type of the weights (before quantization).",
    )
    parser.add_argument(
        "--activation-type",
        type=enum_type(ActivationType),
        metavar=str([e.name for e in ActivationType]),
        required=False,
        default=ActivationType.Swiglu,
        help=f"Type of activation function: {[e.name for e in ActivationType]}",
    )
    parser.add_argument(
        "--autotune",
        action="store_true",
        default=False,
        help=(
            "Enable autotuner warmup for supported routines (trtllm_fp4_block_scale_moe and cutlass_fused_moe)."
        ),
    )
    parser.add_argument(
        "--fp4_mode",
        type=str,
        required=False,
        default="nvfp4",
        choices=["nvfp4", "mxfp4_mxfp8", "mxfp4_bf16"],
        help=(
            "FP4 quantization mode for trtllm_fp4_block_scale_moe: "
            "nvfp4 (NvFP4 weights + NvFP4 hidden states, block_size=16), "
            "mxfp4_mxfp8 (MXFP4 weights + MXFP8 hidden states, block_size=32), "
            "mxfp4_bf16 (MXFP4 weights + BF16 hidden states, block_size=32). "
            "Default: nvfp4"
        ),
    )

    # CuTe DSL / b12x MoE specific
    parser.add_argument(
        "--use_functional_api",
        action="store_true",
        default=False,
        help=(
            "Use the functional MoE API instead of the wrapper class: "
            "cute_dsl_fused_moe_nvfp4 vs CuteDslMoEWrapper for "
            "cute_dsl_fp4_block_scale_moe, and b12x_fused_moe vs B12xMoEWrapper for "
            "b12x_fused_moe. Useful for verifying that the wrapper's workspace cache "
            "eliminates per-call allocation overhead."
        ),
    )
    parser.add_argument(
        "--b12x_quant_mode",
        type=str,
        required=False,
        default="nvfp4",
        choices=["nvfp4", "w4a16"],
        help=(
            "Quantization mode for b12x_fused_moe: nvfp4 (W4A4, FP4 weights + "
            "FP4 activations) or w4a16 (FP4 weights + BF16 activations). "
            "Default: nvfp4"
        ),
    )

    # CUTLASS fused MoE specific
    parser.add_argument(
        "--cutlass_variant",
        type=str,
        required=False,
        default="base",
        choices=["base", "fp8", "nvfp4"],
        help="Variant for cutlass_fused_moe benchmark: base (no quant), fp8 (per-tensor), nvfp4 (fp4 blockscale)",
    )
    parser.add_argument(
        "--quantized_input",
        action="store_true",
        default=False,
        help="Quantize input activations (only used for nvfp4).",
    )
    parser.add_argument(
        "--tp_size",
        type=int,
        required=False,
        default=1,
        help="Tensor parallel size for cutlass_fused_moe.",
    )
    parser.add_argument(
        "--tp_rank",
        type=int,
        required=False,
        default=0,
        help="Tensor parallel rank for cutlass_fused_moe.",
    )
    parser.add_argument(
        "--ep_size",
        type=int,
        required=False,
        default=1,
        help="Expert parallel size for cutlass_fused_moe.",
    )
    parser.add_argument(
        "--ep_rank",
        type=int,
        required=False,
        default=0,
        help="Expert parallel rank for cutlass_fused_moe.",
    )

    # BGMV MoE specific arguments
    parser.add_argument(
        "--rank",
        type=int,
        required=False,
        default=32,
        help="LoRA rank for bgmv_moe benchmark.",
    )
    parser.add_argument(
        "--num_loras",
        type=int,
        required=False,
        default=8,
        help="Number of LoRA adapters for bgmv_moe benchmark.",
    )
    parser.add_argument(
        "--num_slices",
        type=int,
        required=False,
        default=1,
        help="Number of weight slices for bgmv_moe benchmark (e.g., 2 for gate+up).",
    )

    args = parser.parse_args(line)

    # Normalize routing method (map string to internal int expected by kernels)
    routing_method_name_to_type = {
        "renormalize": 1,
        "deepseek_v3": 2,
        "llama4": 3,
        "renormalize_naive": 4,
        "topk": 5,
    }
    args.routing_method_type = routing_method_name_to_type[args.routing_method]

    if args.verbose >= 1:
        print(f"[INFO] {args = }")
    return args


def create_trtllm_moe_test_data(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    num_experts: int,
    routing_method_type: int,
    use_routing_bias: bool,
    input_dtype: torch.dtype,
    weight_dtype: torch.dtype,
    device: torch.device,
    moe_kernel_type: str = "fp8_per_tensor",
):
    """
    Create test data for TensorRT-LLM fused MoE benchmarking (trtllm_*_moe APIs).

    This helper prepares inputs for the TensorRT-LLM fused MoE kernels exposed via
    flashinfer.fused_moe (e.g., trtllm_fp4_block_scale_moe, trtllm_fp8_block_scale_moe,
    trtllm_fp8_per_tensor_scale_moe). It is NOT used for CUTLASS MoE benchmarks,
    which construct their own inputs specific to the CUTLASS path.

    Returns:
        Tuple of tensors needed for trtllm fused MoE computation
    """
    # Create routing logits - dtype depends on both routing method AND MOE kernel type
    # Different MOE kernels have different routing_logits dtype requirements:

    if moe_kernel_type == "fp8_block_scale":
        # DeepSeekV3 routing uses float32, others use bfloat16
        routing_dtype = torch.float32 if routing_method_type == 2 else torch.bfloat16
        routing_logits = torch.randn(
            (num_tokens, num_experts), device=device, dtype=routing_dtype
        )
    elif moe_kernel_type == "fp8_per_tensor":
        # FP8 per-tensor MOE dtype depends on use_routing_scales_on_input parameter
        # For Llama4: use_routing_scales_on_input=True -> bfloat16
        # For others: use_routing_scales_on_input=False -> float32
        routing_dtype = torch.bfloat16 if routing_method_type == 3 else torch.float32
        routing_logits = torch.randn(
            (num_tokens, num_experts), device=device, dtype=routing_dtype
        )
    elif moe_kernel_type == "fp4_block_scale":
        # FP4 block scale MOE follows the test pattern: float32 for DeepSeekV3, bfloat16 for others
        routing_dtype = torch.float32 if routing_method_type == 2 else torch.bfloat16
        routing_logits = torch.randn(
            (num_tokens, num_experts), device=device, dtype=routing_dtype
        )
    else:
        raise ValueError(f"Unknown MOE kernel type: {moe_kernel_type}")

    # Create routing bias if needed - always bfloat16
    routing_bias = None
    if use_routing_bias:
        # Use uniform routing bias for less skewed expert distribution
        routing_bias = (
            torch.ones(num_experts, device=device, dtype=torch.bfloat16) * 0.1
        )

    # Create hidden states - always start with bfloat16 for proper quantization
    hidden_states = 2 * torch.randn(
        (num_tokens, hidden_size), device=device, dtype=torch.bfloat16
    )

    # Create weights - always start with bfloat16 for proper quantization
    gemm1_weights, gemm2_weights = generate_moe_weights(
        num_experts, hidden_size, intermediate_size, device, dtype=torch.bfloat16
    )

    return routing_logits, routing_bias, hidden_states, gemm1_weights, gemm2_weights


def _compute_routing_for_method(
    routing_logits: torch.Tensor,
    routing_bias: Optional[torch.Tensor],
    top_k: int,
    routing_method_type: int,
    n_group: Optional[int] = None,
    topk_group: Optional[int] = None,
    routed_scaling_factor: Optional[float] = None,
) -> torch.Tensor:
    """
    Compute selected experts based on routing method type.
    Returns only the selected expert indices tensor.

    Args:
        routing_logits: [num_tokens, num_experts] routing scores
        routing_bias: Optional [num_experts] routing bias
        top_k: Number of experts to select per token
        routing_method_type: Type of routing method (see RoutingMethodType enum)
        n_group: Number of expert groups (for DeepSeekV3)
        topk_group: Number of top groups (for DeepSeekV3)
        routed_scaling_factor: Scaling factor (for DeepSeekV3)

    Returns:
        selected_experts: [num_tokens, top_k] tensor of selected expert indices
    """
    num_tokens = routing_logits.shape[0]
    device = routing_logits.device

    if routing_method_type == RoutingMethodType.DeepSeekV3:
        # Use fused_topk_deepseek for accurate DeepSeekV3 routing
        if n_group is None or topk_group is None or routed_scaling_factor is None:
            raise ValueError(
                "DeepSeekV3 routing requires n_group, topk_group, and routed_scaling_factor"
            )
        if routing_bias is None:
            routing_bias = torch.zeros(
                routing_logits.shape[1], device=device, dtype=routing_logits.dtype
            )

        # Allocate output tensors
        topk_values = torch.empty(num_tokens, top_k, device=device, dtype=torch.float32)
        topk_indices = torch.empty(num_tokens, top_k, device=device, dtype=torch.int32)

        fused_topk_deepseek(
            scores=routing_logits.float(),
            bias=routing_bias.float(),
            n_group=n_group,
            topk_group=topk_group,
            topk=top_k,
            routed_scaling_factor=routed_scaling_factor,
            topk_values=topk_values,
            topk_indices=topk_indices,
        )
        return topk_indices
    else:
        # For other routing methods, use simple top-k as approximation
        # This is accurate for Default, Renormalize, RenormalizeNaive, TopK
        # and approximate for Llama4
        _, selected_experts = compute_routing(routing_logits.float(), top_k)
        return selected_experts


def testTrtllmFp4BlockScaleMoe(args):
    """
    Test trtllm_fp4_block_scale_moe API (TensorRT-LLM fused MoE).

    This test:
    1. Creates quantized FP4 weights and scales
    2. Runs FP4 block scale MOE
    3. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.verbose >= 1:
        print("[INFO] Running testTrtllmFp4BlockScaleMoe")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype)

    # Parse configuration
    num_tokens = args.num_tokens
    hidden_size = args.hidden_size
    intermediate_size = args.intermediate_size
    num_experts = args.num_experts
    top_k = args.top_k
    n_group = (
        args.n_group
        if hasattr(args, "n_group") and args.n_group is not None and args.n_group > 0
        else None
    )
    topk_group = (
        args.topk_group
        if hasattr(args, "topk_group")
        and args.topk_group is not None
        and args.topk_group > 0
        else None
    )
    routed_scaling_factor = (
        args.routed_scaling_factor
        if hasattr(args, "routed_scaling_factor")
        and args.routed_scaling_factor is not None
        else None
    )
    local_expert_offset = args.local_expert_offset
    local_num_experts = args.local_num_experts or num_experts
    routing_method_type = args.routing_method_type
    use_shuffled_weight = args.use_shuffled_weight
    weight_layout = args.weight_layout
    is_cuda_graph_compatible = not args.no_cuda_graph
    activation_type = args.activation_type
    res = []

    backends = ["trtllm"]
    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    if args.verbose >= 1:
        print(
            f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, "
            f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}"
        )

    # Create test data
    routing_logits, routing_bias, hidden_states, gemm1_weights, gemm2_weights = (
        create_trtllm_moe_test_data(
            num_tokens,
            hidden_size,
            intermediate_size,
            num_experts,
            routing_method_type,
            args.use_routing_bias,
            input_dtype,
            weight_dtype,
            device,
            moe_kernel_type="fp4_block_scale",
        )
    )

    # Compute selected experts for accurate bandwidth calculation
    # Use the actual routing method to get correct expert assignments
    selected_experts = _compute_routing_for_method(
        routing_logits=routing_logits,
        routing_bias=routing_bias,
        top_k=top_k,
        routing_method_type=routing_method_type,
        n_group=n_group,
        topk_group=topk_group,
        routed_scaling_factor=routed_scaling_factor,
    )

    # Determine FP4 quantization mode
    fp4_mode = getattr(args, "fp4_mode", "nvfp4")
    is_mxfp4 = fp4_mode in ("mxfp4_mxfp8", "mxfp4_bf16")
    use_ue8m0 = is_mxfp4
    sf_vec_size = SF_VEC_SIZE["mxfp4" if is_mxfp4 else "nvfp4"]

    if args.verbose >= 1:
        print(
            f"[INFO] FP4 mode: {fp4_mode} (use_ue8m0={use_ue8m0}, sf_vec_size={sf_vec_size})"
        )

    # Quantize weights using proper FP4 quantization
    gemm1_weights_fp4_bytes, gemm1_scales_fp4_bytes, gemm1_scales_global = (
        quantize_fp4_batched(gemm1_weights, num_experts, use_ue8m0, True)
    )
    gemm2_weights_fp4_bytes, gemm2_scales_fp4_bytes, gemm2_scales_global = (
        quantize_fp4_batched(gemm2_weights, num_experts, use_ue8m0, True)
    )

    # Prepare hidden states and scale based on fp4_mode
    if fp4_mode == "mxfp4_bf16":
        hidden_states_fp4 = hidden_states.to(torch.bfloat16)
        hidden_states_scale_linear_fp4 = None
    elif fp4_mode == "mxfp4_mxfp8":
        if num_tokens % 128 != 0:
            raise ValueError(
                f"mxfp4_mxfp8 mode requires num_tokens to be a multiple of 128 "
                f"(got {num_tokens}) because mxfp8_quantize with swizzled scale "
                f"layout pads rows to 128-element boundaries."
            )
        hs_quant, hs_scale = mxfp8_quantize(hidden_states, True)
        hidden_states_fp4 = hs_quant
        hidden_states_scale_linear_fp4 = hs_scale.view(torch.float8_e4m3fn).reshape(
            num_tokens, -1
        )
    else:
        hidden_states_scale_global = calculate_fp4_global_scale(hidden_states)
        hidden_states_fp4_bytes, hidden_states_scale_fp4_bytes, _ = quantize_fp4(
            hidden_states, hidden_states_scale_global, use_ue8m0, True
        )
        hidden_states_fp4 = hidden_states_fp4_bytes.view(torch.uint8).reshape(
            num_tokens, hidden_size // 2
        )
        hidden_states_scale_linear_fp4 = hidden_states_scale_fp4_bytes.view(
            torch.float8_e4m3fn
        )
        expected_scale_elems = (num_tokens * hidden_size) // sf_vec_size
        if hidden_states_scale_linear_fp4.numel() != expected_scale_elems:
            print(
                f"[WARNING] FP4 hidden_states_scale element count mismatch "
                f"({hidden_states_scale_linear_fp4.numel()} vs expected {expected_scale_elems}), "
                f"substituting all-ones scale tensor. This is likely caused by swizzled "
                f"layout padding when num_tokens ({num_tokens}) is not a multiple of 128."
            )
            hidden_states_scale_linear_fp4 = torch.ones(
                expected_scale_elems, device=device, dtype=torch.float8_e4m3fn
            )
        hidden_states_scale_linear_fp4 = hidden_states_scale_linear_fp4.reshape(
            num_tokens, hidden_size // sf_vec_size
        )

    # Prepare weights for kernel (packed uint8 format, scale as float8_e4m3fn)
    gemm1_weights_fp4 = gemm1_weights_fp4_bytes.view(torch.uint8).reshape(
        num_experts, 2 * intermediate_size, hidden_size // 2
    )
    gemm1_weights_scale = gemm1_scales_fp4_bytes.view(torch.float8_e4m3fn).reshape(
        num_experts, 2 * intermediate_size, hidden_size // sf_vec_size
    )

    gemm2_weights_fp4 = gemm2_weights_fp4_bytes.view(torch.uint8).reshape(
        num_experts, hidden_size, intermediate_size // 2
    )
    gemm2_weights_scale = gemm2_scales_fp4_bytes.view(torch.float8_e4m3fn).reshape(
        num_experts, hidden_size, intermediate_size // sf_vec_size
    )

    gemm1_bias = None
    gemm1_alpha = None
    gemm1_beta = None
    gemm1_clamp_limit = None
    gemm2_bias = None

    # Create scale scalars using shared utility
    output1_scale_scalar, output1_scale_gate_scalar, output2_scale_scalar = (
        create_moe_output_scale_scalars(local_num_experts, device)
    )

    if args.verbose >= 2:
        print(f"[VVERBOSE] routing_logits.shape = {routing_logits.shape}")
        print(f"[VVERBOSE] hidden_states.shape = {hidden_states.shape}")
        print(f"[VVERBOSE] gemm1_weights_fp4.shape = {gemm1_weights_fp4.shape}")
        print(f"[VVERBOSE] gemm2_weights_fp4.shape = {gemm2_weights_fp4.shape}")

    def run_fp4_moe(
        routing_logits,
        routing_bias,
        hidden_states_fp4,
        hidden_states_scale_linear_fp4,
        gemm1_weights_fp4,
        gemm1_weights_scale,
        gemm2_weights_fp4,
        gemm2_weights_scale,
        output1_scale_scalar,
        output1_scale_gate_scalar,
        output2_scale_scalar,
    ):
        return trtllm_fp4_block_scale_moe(
            routing_logits=routing_logits,
            routing_bias=routing_bias,
            hidden_states=hidden_states_fp4,
            hidden_states_scale=hidden_states_scale_linear_fp4,
            gemm1_weights=gemm1_weights_fp4,
            gemm1_weights_scale=gemm1_weights_scale,
            gemm1_bias=gemm1_bias,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=gemm1_clamp_limit,
            gemm2_weights=gemm2_weights_fp4,
            gemm2_weights_scale=gemm2_weights_scale,
            gemm2_bias=gemm2_bias,
            output1_scale_scalar=output1_scale_scalar,
            output1_scale_gate_scalar=output1_scale_gate_scalar,
            output2_scale_scalar=output2_scale_scalar,
            num_experts=num_experts,
            top_k=top_k,
            n_group=n_group,
            topk_group=topk_group,
            intermediate_size=intermediate_size,
            local_expert_offset=local_expert_offset,
            local_num_experts=local_num_experts,
            routed_scaling_factor=routed_scaling_factor,
            routing_method_type=routing_method_type,
            do_finalize=True,
            enable_pdl=args.enable_pdl,
            **_activation_kwarg(trtllm_fp4_block_scale_moe, activation_type),
        )

    backend = "trtllm"

    # Optional autotune warmup (supported for FP4 TRTLlm fused MoE)
    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        backend = "trtllm_autotune"
        if args.verbose >= 1:
            print(
                f"[INFO] Autotune warmup for FP4 block scale MoE: {warmup_iters} iters"
            )
        with autotune(True, cache=cache_path):
            for _ in range(warmup_iters):
                run_fp4_moe(
                    routing_logits,
                    routing_bias,
                    hidden_states_fp4,
                    hidden_states_scale_linear_fp4,
                    gemm1_weights_fp4,
                    gemm1_weights_scale,
                    gemm2_weights_fp4,
                    gemm2_weights_scale,
                    output1_scale_scalar,
                    output1_scale_gate_scalar,
                    output2_scale_scalar,
                )
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    # Benchmark timing
    times = bench_gpu_time(
        fn=run_fp4_moe,
        dry_run_iters=args.dry_run_iters,
        repeat_iters=args.num_iters,
        sleep_after_run=False,
        enable_cupti=args.use_cupti,
        use_cuda_graph=is_cuda_graph_compatible,
        cold_l2_cache=True,
        input_args=(
            routing_logits,
            routing_bias,
            hidden_states_fp4,
            hidden_states_scale_linear_fp4,
            gemm1_weights_fp4,
            gemm1_weights_scale,
            gemm2_weights_fp4,
            gemm2_weights_scale,
            output1_scale_scalar,
            output1_scale_gate_scalar,
            output2_scale_scalar,
        ),
    )

    # Compute performance metrics
    median_time = np.median(times)
    std_time = np.std(times)
    tflops = calculate_moe_tflops(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        is_gated=args.activation_type in (ActivationType.Swiglu, ActivationType.Geglu),
    )
    input_format_str = {"nvfp4": "nvfp4", "mxfp4_mxfp8": "mxfp8", "mxfp4_bf16": "bf16"}[
        fp4_mode
    ]
    weight_format_str = "mxfp4" if is_mxfp4 else "nvfp4"
    tb_per_sec = calculate_moe_kernel_bandwidth(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        input_dtype,
        weight_dtype,
        input_format=input_format_str,
        weight_format=weight_format_str,
        routing_logits_dtype=routing_logits.dtype,
        active_experts=int(selected_experts.unique().numel()),
        verbose=args.verbose,
        is_gated=args.activation_type in (ActivationType.Swiglu, ActivationType.Geglu),
    )

    print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

    if args.output_path is not None:
        cur_res = defaultdict(str)
        cur_res["routine"] = args.routine
        cur_res["median_time"] = median_time
        cur_res["std_time"] = std_time
        cur_res["tflops"] = tflops
        cur_res["tb_per_sec"] = tb_per_sec
        cur_res["backend"] = backend
        cur_res["num_tokens"] = num_tokens
        cur_res["hidden_size"] = hidden_size
        cur_res["intermediate_size"] = intermediate_size
        cur_res["num_experts"] = num_experts
        cur_res["top_k"] = top_k
        cur_res["n_group"] = n_group
        cur_res["topk_group"] = topk_group
        cur_res["routed_scaling_factor"] = routed_scaling_factor
        cur_res["local_expert_offset"] = local_expert_offset
        cur_res["local_num_experts"] = local_num_experts
        cur_res["routing_method"] = args.routing_method
        cur_res["use_shuffled_weight"] = use_shuffled_weight
        cur_res["weight_layout"] = weight_layout
        cur_res["use_routing_bias"] = args.use_routing_bias
        cur_res["use_routing_scales_on_input"] = args.use_routing_scales_on_input
        cur_res["input_dtype"] = input_dtype
        cur_res["weight_dtype"] = weight_dtype
        cur_res["activation_type"] = args.activation_type.name
        cur_res["fp4_mode"] = fp4_mode
        res.append(cur_res)

    return res


def testCutlassFusedMoe(args):
    """
    Benchmark cutlass_fused_moe (CUTLASS MoE) with variants mirroring tests in tests/moe/test_trtllm_cutlass_fused_moe.py
    Variants:
      - base: no quantization
      - fp8: per-tensor fp8 for weights and activation scale
      - nvfp4: FP4 block-scale weights, optional quantized input
    Supports TP/EP via tp_size/tp_rank and ep_size/ep_rank.
    """
    if args.verbose >= 1:
        print("[INFO] Running testCutlassFusedMoe")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)

    # Shapes
    num_tokens = args.num_tokens
    hidden_size = args.hidden_size
    intermediate_size = args.intermediate_size
    num_experts = args.num_experts
    top_k = args.top_k
    tp_size = getattr(args, "tp_size", 1)
    tp_rank = getattr(args, "tp_rank", 0)
    ep_size = getattr(args, "ep_size", 1)
    ep_rank = getattr(args, "ep_rank", 0)
    is_cuda_graph_compatible = not args.no_cuda_graph
    res = []
    backends = ["cutlass"]
    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    # Create base tensors
    activation_type = args.activation_type
    is_gated = is_gated_activation(activation_type)
    w1_rows = (2 if is_gated else 1) * intermediate_size

    torch.manual_seed(args.random_seed)
    x = torch.randn(num_tokens, hidden_size, dtype=input_dtype, device=device)
    w31_weight = (
        torch.randn(
            num_experts,
            w1_rows,
            hidden_size,
            dtype=input_dtype,
            device=device,
        )
        / 10
    )
    w2_weight = (
        torch.randn(
            num_experts,
            hidden_size,
            intermediate_size,
            dtype=input_dtype,
            device=device,
        )
        / 10
    )

    # Routing
    router_logits = torch.randn(
        num_tokens, num_experts, dtype=input_dtype, device=device
    )
    routing_weights, selected_experts = compute_routing(router_logits, top_k)

    if args.verbose >= 2:
        print(f"[VVERBOSE] x.shape = {x.shape}")
        print(f"[VVERBOSE] w31_weight.shape = {w31_weight.shape}")
        print(f"[VVERBOSE] w2_weight.shape = {w2_weight.shape}")

    # Build local weights per EP/TP like tests do
    experts_per_rank = num_experts // max(ep_size, 1)
    expert_start = ep_rank * experts_per_rank
    expert_end = expert_start + experts_per_rank
    w31_ep = w31_weight[expert_start:expert_end, :]
    w2_ep = w2_weight[expert_start:expert_end, :]

    def build_tp_shards(w31_ep_tensor: torch.Tensor, w2_ep_tensor: torch.Tensor):
        if tp_size <= 1:
            return w31_ep_tensor, w2_ep_tensor
        shard = intermediate_size // tp_size
        start = tp_rank * shard
        end = start + shard
        if is_gated:
            # Split w31 into w3 and w1 along intermediate dim
            w3_weight, w1_weight = torch.chunk(w31_ep_tensor, 2, dim=1)
            w3_local = w3_weight[:, start:end, :]
            w1_local = w1_weight[:, start:end, :]
            w31_local = torch.cat([w3_local, w1_local], dim=1)
        else:
            w31_local = w31_ep_tensor[:, start:end, :]
        w2_local = w2_ep_tensor[:, :, start:end]
        return w31_local.contiguous(), w2_local.contiguous()

    w31_local, w2_local = build_tp_shards(w31_ep, w2_ep)

    # Prepare variant-specific inputs (outside of the timed/captured region)
    variant = getattr(args, "cutlass_variant", "base")
    out = torch.empty_like(x)

    if variant == "base":

        def run_cutlass(x, selected_experts, routing_weights, w31_local, w2_local, out):
            return cutlass_fused_moe(
                x,
                selected_experts.to(torch.int),
                routing_weights,
                w31_local,
                w2_local,
                input_dtype,
                tp_size=tp_size,
                tp_rank=tp_rank,
                ep_size=ep_size,
                ep_rank=ep_rank,
                quant_scales=None,
                output=out,
                enable_pdl=args.enable_pdl,
                **_activation_kwarg(cutlass_fused_moe, activation_type),
            )

        input_args_for_bench = (
            x,
            selected_experts,
            routing_weights,
            w31_local,
            w2_local,
            out,
        )

    elif variant == "fp8":
        # Per-tensor FP8 for weights and activation scale
        w31_weight_fp8 = torch.empty_like(w31_local, dtype=torch.float8_e4m3fn)
        w2_weight_fp8 = torch.empty_like(w2_local, dtype=torch.float8_e4m3fn)
        local_num_experts = w31_local.shape[0]
        w31_scales = torch.empty(local_num_experts, 2, dtype=input_dtype, device=device)
        w2_scales = torch.empty(local_num_experts, 1, dtype=input_dtype, device=device)

        # Quantize weights per expert
        for expert_id in range(local_num_experts):
            w31_expert = w31_local[expert_id]
            w2_expert = w2_local[expert_id]
            w31_q, s31 = quantize_fp8(w31_expert)
            w2_q, s2 = quantize_fp8(w2_expert)
            w31_weight_fp8[expert_id].copy_(w31_q)
            w2_weight_fp8[expert_id].copy_(w2_q)
            # Store the same scalar twice to mimic test layout (avoid torch.tensor())
            w31_scales[expert_id, 0] = s31.to(dtype=input_dtype, device=device)
            w31_scales[expert_id, 1] = s31.to(dtype=input_dtype, device=device)
            w2_scales[expert_id, 0] = s2.to(dtype=input_dtype, device=device)

        x_quant, hidden_states_scale = quantize_fp8(x)
        hidden_states_scale_scalar = hidden_states_scale[0].to(device)

        # Note: follow tests quant_scales format
        # [w1_scales * hidden_states_scale, 1.0, 1.0 * w2_scales, hidden_states_scale]
        w1_scales = w31_scales[:, 1]
        one_const = torch.ones((), device=device)
        quant_scales = [
            (w1_scales * hidden_states_scale_scalar).float().squeeze(),
            one_const,
            w2_scales.squeeze().float(),
            hidden_states_scale_scalar,
        ]

        def run_cutlass(
            x_quant,
            selected_experts,
            routing_weights,
            w31_weight_fp8,
            w2_weight_fp8,
            out,
        ):
            return cutlass_fused_moe(
                x_quant,
                selected_experts.to(torch.int),
                routing_weights,
                w31_weight_fp8,
                w2_weight_fp8,
                input_dtype,
                tp_size=tp_size,
                tp_rank=tp_rank,
                ep_size=ep_size,
                ep_rank=ep_rank,
                quant_scales=quant_scales,
                output=out,
                enable_pdl=args.enable_pdl,
                **_activation_kwarg(cutlass_fused_moe, activation_type),
            )

        input_args_for_bench = (
            x_quant,
            selected_experts,
            routing_weights,
            w31_weight_fp8,
            w2_weight_fp8,
            out,
        )

    elif variant == "nvfp4":
        # NVFP4: FP4 block-scale weights, optional quantized input

        def round_up(x_val, y):
            return (x_val + y - 1) // y * y

        e = w31_local.shape[0]
        n = w2_local.shape[2]  # local intermediate size after TP
        k = hidden_size
        quant_blocksize = 16
        w1_n = w31_local.shape[1]  # 2*n for gated, n for non-gated

        # Weight quantization buffers
        w1_q = torch.empty((e, w1_n, k // 2), device=device, dtype=torch.uint8)
        w2_q = torch.empty((e, k, n // 2), device=device, dtype=torch.uint8)
        w1_blockscale = torch.empty(
            (e, round_up(w1_n, 128), round_up(k // quant_blocksize, 4)),
            device=device,
            dtype=torch.float8_e4m3fn,
        )
        w2_blockscale = torch.empty(
            (e, round_up(k, 128), round_up(n // quant_blocksize, 4)),
            device=device,
            dtype=torch.float8_e4m3fn,
        )
        w1_gs = torch.empty((e,), device=device, dtype=torch.float32)
        w2_gs = torch.empty((e,), device=device, dtype=torch.float32)

        # Quantize from local shards
        for expert in range(e):
            w1_src = w31_local[expert]
            # w31 layout is [2n, k]; w2 layout is [k, n]
            w2_src = w2_local[expert].contiguous()  # [hidden_size, n]
            w1_amax = torch.abs(w1_src).max().to(torch.float32)
            w2_amax = torch.abs(w2_src).max().to(torch.float32)
            w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
            w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
            w1_q[expert], w1_blockscale[expert] = fp4_quantize(w1_src, w1_gs[expert])
            w2_q[expert], w2_blockscale[expert] = fp4_quantize(w2_src, w2_gs[expert])

        a1_gs = torch.ones((), device=device, dtype=torch.float32)
        a2_gs = torch.ones((), device=device, dtype=torch.float32)

        hidden_states = x
        input_sf = None
        if getattr(args, "quantized_input", False):
            hidden_states, input_sf = fp4_quantize(x, a1_gs)

        quant_scales = [
            a1_gs,
            w1_blockscale.view(torch.int32),
            1.0 / (a1_gs * w1_gs),
            a2_gs,
            w2_blockscale.view(torch.int32),
            1.0 / (a2_gs * w2_gs),
        ]

        def run_cutlass(
            hidden_states, selected_experts, routing_weights, w1_q, w2_q, out
        ):
            return cutlass_fused_moe(
                hidden_states,
                selected_experts.to(torch.int),
                routing_weights,
                w1_q.contiguous().view(torch.long),
                w2_q.contiguous().view(torch.long),
                input_dtype,
                tp_size=tp_size,
                tp_rank=tp_rank,
                ep_size=ep_size,
                ep_rank=ep_rank,
                quant_scales=quant_scales,
                input_sf=input_sf,
                output=out,
                enable_pdl=args.enable_pdl,
                **_activation_kwarg(cutlass_fused_moe, activation_type),
            )

        input_args_for_bench = (
            hidden_states,
            selected_experts,
            routing_weights,
            w1_q,
            w2_q,
            out,
        )
    else:
        raise ValueError(f"Unknown cutlass_variant: {variant}")

    backend = "cutlass"

    # Optional autotune warmup (supported for CUTLASS fused MoE)
    cache_path = getattr(args, "autotune_cache", None)
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        backend = "cutlass_autotune"
        if args.verbose >= 1:
            print(f"[INFO] Autotune warmup for CUTLASS fused MoE: {warmup_iters} iters")
        with autotune(True, cache=cache_path):
            for _ in range(warmup_iters):
                run_cutlass(*input_args_for_bench)
    elif cache_path:
        with autotune(False, cache=cache_path):
            pass

    # Measure
    times = bench_gpu_time(
        fn=run_cutlass,
        dry_run_iters=args.dry_run_iters,
        repeat_iters=args.num_iters,
        sleep_after_run=False,
        enable_cupti=args.use_cupti,
        use_cuda_graph=is_cuda_graph_compatible,
        cold_l2_cache=True,
        input_args=input_args_for_bench,
    )

    median_time = np.median(times)
    std_time = np.std(times)
    tflops = calculate_moe_tflops(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        is_gated=is_gated,
    )
    tb_per_sec = calculate_moe_kernel_bandwidth(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        input_dtype,
        input_dtype,
        input_format=variant,
        weight_format=variant,
        routing_logits_dtype=router_logits.dtype,
        active_experts=int(selected_experts.unique().numel()),
        verbose=args.verbose,
        is_gated=is_gated,
    )

    print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

    if args.output_path is not None:
        cur_res = defaultdict(str)
        cur_res["routine"] = args.routine
        cur_res["median_time"] = median_time
        cur_res["std_time"] = std_time
        cur_res["tflops"] = tflops
        cur_res["tb_per_sec"] = tb_per_sec
        cur_res["backend"] = backend
        cur_res["num_tokens"] = num_tokens
        cur_res["hidden_size"] = hidden_size
        cur_res["intermediate_size"] = intermediate_size
        cur_res["num_experts"] = num_experts
        cur_res["top_k"] = top_k
        # Routing method/weight layout not applicable; leave defaults
        cur_res["use_shuffled_weight"] = False
        cur_res["weight_layout"] = 0
        cur_res["use_routing_scales_on_input"] = False
        cur_res["input_dtype"] = input_dtype
        cur_res["weight_dtype"] = input_dtype
        # CUTLASS fused MoE specific
        cur_res["cutlass_variant"] = variant
        cur_res["quantized_input"] = args.quantized_input
        cur_res["tp_size"] = tp_size
        cur_res["tp_rank"] = tp_rank
        cur_res["ep_size"] = ep_size
        cur_res["ep_rank"] = ep_rank
        res.append(cur_res)

    return res


def _interleave_linear_and_gate(
    x: torch.Tensor, group_size: int = 64, dim: int = -1
) -> torch.Tensor:
    """Interleave linear and gate weights for CuteDSL SwiGLU layout."""
    sizes = x.size()
    dim = dim % x.dim()
    assert sizes[dim] % (group_size * 2) == 0
    prev_sizes = sizes[:dim]
    post_sizes = sizes[dim + 1 :]
    x = x.view(*prev_sizes, 2, sizes[dim] // (group_size * 2), group_size, *post_sizes)
    x = x.transpose(dim, dim + 1).contiguous().view(*sizes)
    return x


def _create_nvfp4_moe_test_data(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    num_experts: int,
    num_local_experts: int,
    top_k: int,
    device: torch.device,
    backend: str,
    is_gated: bool = True,
):
    """Create NVFP4-quantized test data for CuTe-DSL-family MoE kernels.

    Supports two backends that share the weight-quantization recipe but differ
    in the FC1 weight layout and in whether input activations are pre-quantized:

    - ``backend="cute-dsl"`` (SM100/SM103 ``CuteDslMoEWrapper``): gated SwiGLU only,
      FC1 weights interleaved via ``_interleave_linear_and_gate``, fp4-quantized
      input activations + scale factors.
    - ``backend="b12x"`` (SM120/SM121 ``B12xMoEWrapper``): either gated (SwiGLU,
      ``[E, 2n, k]``) or non-gated (ReLU2, ``[E, n, k]``) FC1, no interleave,
      bf16 input activations (the kernel fuses quantization internally).

    Routing is computed externally via simple top-k (the wrappers take
    pre-computed ``token_selected_experts`` and ``token_final_scales``).
    """
    from flashinfer.fp4_quantization import fp4_quantize
    from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout

    if backend not in ("cute-dsl", "b12x"):
        raise ValueError(f"Unsupported backend: {backend!r}")
    if backend == "cute-dsl" and not is_gated:
        raise ValueError("cute-dsl backend only supports gated (SwiGLU) activation")

    sf_vec_size = 16

    # Input activations: fp4-quantized for cute-dsl, bf16 for b12x
    x_bf16 = (
        torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) / 10
    )
    a1_gs = torch.tensor([1.0], device=device, dtype=torch.float32)
    x_quantized, x_sf = fp4_quantize(
        x_bf16, global_scale=a1_gs, sf_vec_size=sf_vec_size, is_sf_swizzled_layout=False
    )
    x_sf = x_sf.unsqueeze(-1)

    # Routing (simple top-k; the wrapper takes pre-computed assignments)
    routing_logits = torch.randn(num_tokens, num_experts, device=device)
    routing_weights, selected_experts = compute_routing(routing_logits, top_k)
    selected_experts = selected_experts.to(torch.int32)

    # GEMM1 weights
    # Gated (SwiGLU): [E, 2*n, k] — gate + up fused
    # Non-gated (ReLU2, b12x only): [E, n, k] — single FC1 matrix
    w1_rows = (2 if is_gated else 1) * intermediate_size
    w1_bf16 = (
        torch.randn(
            num_local_experts,
            w1_rows,
            hidden_size,
            dtype=torch.bfloat16,
            device=device,
        )
        / 10
    )
    # cute-dsl (SM100) expects gate/up interleaved; b12x (SM120) does not.
    if backend == "cute-dsl":
        w1_bf16_prepared = _interleave_linear_and_gate(w1_bf16, group_size=64, dim=1)
    else:
        w1_bf16_prepared = w1_bf16
    w1_gs = torch.tensor([1.0], device=device, dtype=torch.float32)
    w1_flat = w1_bf16_prepared.view(num_local_experts * w1_rows, hidden_size)
    w1_q_flat, w1_sf_flat = fp4_quantize(
        w1_flat, global_scale=w1_gs, sf_vec_size=sf_vec_size, is_sf_swizzled_layout=True
    )
    w1_weight = w1_q_flat.view(num_local_experts, w1_rows, hidden_size // 2)
    w1_weight_sf = convert_sf_to_mma_layout(
        w1_sf_flat,
        m=w1_rows,
        k=hidden_size,
        num_groups=num_local_experts,
        sf_vec_size=sf_vec_size,
    )
    w1_alpha = torch.ones(num_local_experts, device=device, dtype=torch.float32)

    # GEMM2 weights (down projection)
    w2_bf16 = (
        torch.randn(
            num_local_experts,
            hidden_size,
            intermediate_size,
            dtype=torch.bfloat16,
            device=device,
        )
        / 10
    )
    w2_gs = torch.tensor([1.0], device=device, dtype=torch.float32)
    w2_flat = w2_bf16.view(num_local_experts * hidden_size, intermediate_size)
    w2_q_flat, w2_sf_flat = fp4_quantize(
        w2_flat, global_scale=w2_gs, sf_vec_size=sf_vec_size, is_sf_swizzled_layout=True
    )
    w2_weight = w2_q_flat.view(num_local_experts, hidden_size, intermediate_size // 2)
    w2_weight_sf = convert_sf_to_mma_layout(
        w2_sf_flat,
        m=hidden_size,
        k=intermediate_size,
        num_groups=num_local_experts,
        sf_vec_size=sf_vec_size,
    )
    w2_alpha = torch.ones(num_local_experts, device=device, dtype=torch.float32)

    fc2_input_scale = torch.tensor([1.0], device=device, dtype=torch.float32)

    return {
        "x": x_quantized,
        "x_bf16": x_bf16,
        "x_sf": x_sf,
        "token_selected_experts": selected_experts,
        "token_final_scales": routing_weights,
        "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,
    }


def testCuteDslFp4BlockScaleMoe(args):
    """
    Test cute_dsl_fp4_block_scale_moe (CuTe DSL NVFP4 MoE on SM100/SM103).

    This test:
    1. Creates NVFP4-quantized weights and fp4-quantized inputs for CuTe DSL kernels
    2. Runs MoE via CuteDslMoEWrapper (or cute_dsl_fused_moe_nvfp4 when
       ``--use_functional_api`` is set). SwiGLU only.
    3. Measures performance metrics (TFLOPS, TB/sec)

    Note: SM120/SM121 has moved to the dedicated ``b12x_fused_moe`` routine
    (see ``testB12xFusedMoe``).

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.verbose >= 1:
        print("[INFO] Running testCuteDslFp4BlockScaleMoe")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    from flashinfer import CuteDslMoEWrapper

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype)

    num_tokens = args.num_tokens
    hidden_size = args.hidden_size
    intermediate_size = args.intermediate_size
    num_experts = args.num_experts
    top_k = args.top_k
    local_expert_offset = args.local_expert_offset
    local_num_experts = args.local_num_experts or num_experts
    is_cuda_graph_compatible = not args.no_cuda_graph
    res = []

    backends = ["cute-dsl"]
    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    if args.verbose >= 1:
        print(
            f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, "
            f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}"
        )

    # CuteDslMoEWrapper is gated SwiGLU only.
    activation_type = args.activation_type
    if activation_type != ActivationType.Swiglu:
        raise ValueError(
            f"cute_dsl_fp4_block_scale_moe only supports Swiglu activation, "
            f"got {activation_type.name}. Use --routine b12x_fused_moe for ReLU2."
        )

    # Create CuteDSL-specific NVFP4 test data (gated, fp4 input)
    tensors = _create_nvfp4_moe_test_data(
        num_tokens=num_tokens,
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        num_experts=num_experts,
        num_local_experts=local_num_experts,
        top_k=top_k,
        device=device,
        backend="cute-dsl",
        is_gated=True,
    )

    if args.verbose >= 2:
        print(f"[VVERBOSE] x.shape = {tensors['x'].shape}")
        print(f"[VVERBOSE] w1_weight.shape = {tensors['w1_weight'].shape}")
        print(f"[VVERBOSE] w2_weight.shape = {tensors['w2_weight'].shape}")

    use_functional = getattr(args, "use_functional_api", False)

    if use_functional:
        from functools import partial
        from flashinfer import cute_dsl_fused_moe_nvfp4

        # Pre-allocate output buffer to avoid per-call allocation
        moe_output = torch.empty(
            num_tokens, hidden_size, dtype=torch.bfloat16, device=device
        )

        if args.verbose >= 1:
            print("[INFO] Using CuTe DSL functional API (cute_dsl_fused_moe_nvfp4)")
        runner = partial(
            cute_dsl_fused_moe_nvfp4,
            num_experts=num_experts,
            top_k=top_k,
            num_local_experts=local_num_experts,
            local_expert_offset=local_expert_offset,
            moe_output=moe_output,
            enable_pdl=args.enable_pdl,
        )

        # Warmup call to populate workspace cache before timed region
        runner(
            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"],
        )
    else:
        moe = CuteDslMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=is_cuda_graph_compatible,
            max_num_tokens=num_tokens,
            num_local_experts=local_num_experts,
            local_expert_offset=local_expert_offset,
            enable_pdl=args.enable_pdl,
        )
        runner = moe.run

    def run_cute_dsl_moe(
        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,
    ):
        return runner(
            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,
        )

    input_args = (
        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"],
    )

    # Snapshot active expert count before any kernel execution, since
    # autotune tactic exploration may corrupt input tensors.
    num_active_experts = int(tensors["token_selected_experts"].unique().numel())

    backend = "cute-dsl"

    # Optional autotune warmup.
    # Clone input_args so autotune tactic exploration doesn't corrupt the
    # original tensors used by the subsequent benchmark.
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        backend = "cute-dsl_autotune"
        if args.verbose >= 1:
            print(f"[INFO] Autotune warmup for CuteDSL NVFP4 MoE: {warmup_iters} iters")
        autotune_args = tuple(
            t.clone() if isinstance(t, torch.Tensor) else t for t in input_args
        )
        with autotune(True):
            for _ in range(warmup_iters):
                run_cute_dsl_moe(*autotune_args)
        del autotune_args

    # Benchmark timing
    times = bench_gpu_time(
        fn=run_cute_dsl_moe,
        dry_run_iters=args.dry_run_iters,
        repeat_iters=args.num_iters,
        sleep_after_run=False,
        enable_cupti=args.use_cupti,
        use_cuda_graph=is_cuda_graph_compatible,
        cold_l2_cache=True,
        input_args=input_args,
    )

    # Compute performance metrics
    median_time = np.median(times)
    std_time = np.std(times)
    tflops = calculate_moe_tflops(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        is_gated=True,
    )
    tb_per_sec = calculate_moe_kernel_bandwidth(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        input_dtype,
        weight_dtype,
        input_format="nvfp4",
        weight_format="nvfp4",
        routing_logits_dtype=None,
        active_experts=num_active_experts,
        verbose=args.verbose,
        is_gated=True,
    )

    print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

    if args.output_path is not None:
        cur_res = defaultdict(str)
        cur_res["routine"] = args.routine
        cur_res["median_time"] = median_time
        cur_res["std_time"] = std_time
        cur_res["tflops"] = tflops
        cur_res["tb_per_sec"] = tb_per_sec
        cur_res["backend"] = backend
        cur_res["num_tokens"] = num_tokens
        cur_res["hidden_size"] = hidden_size
        cur_res["intermediate_size"] = intermediate_size
        cur_res["num_experts"] = num_experts
        cur_res["top_k"] = top_k
        cur_res["local_expert_offset"] = local_expert_offset
        cur_res["local_num_experts"] = local_num_experts
        cur_res["input_dtype"] = input_dtype
        cur_res["weight_dtype"] = weight_dtype
        cur_res["fp4_mode"] = "nvfp4"
        cur_res["activation_type"] = activation_type.name
        res.append(cur_res)

    return res


def testB12xFusedMoe(args):
    """
    Test b12x_fused_moe (SM120/SM121 CuTe DSL FP4-weight MoE).

    The b12x MoE takes **bf16** hidden states and NVFP4-quantized weights.
    ``--b12x_quant_mode`` selects the activation precision: ``nvfp4``
    (W4A4, the kernel fuses activation quantization internally) or
    ``w4a16`` (BF16 activations). Supports both SwiGLU (gated) and ReLU2
    (non-gated) activations.

    This test:
    1. Creates NVFP4-quantized weights and bf16 inputs for b12x kernels
    2. Runs MoE via B12xMoEWrapper (or b12x_fused_moe when
       ``--use_functional_api`` is set)
    3. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    warn_if_pdl_unsupported(args, args.routine)
    if args.verbose >= 1:
        print("[INFO] Running testB12xFusedMoe")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    from flashinfer import B12xMoEWrapper

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype)

    num_tokens = args.num_tokens
    hidden_size = args.hidden_size
    intermediate_size = args.intermediate_size
    num_experts = args.num_experts
    top_k = args.top_k
    local_num_experts = args.local_num_experts or num_experts
    is_cuda_graph_compatible = not args.no_cuda_graph
    res = []

    backends = ["b12x"]
    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    if args.verbose >= 1:
        print(
            f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, "
            f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}"
        )

    # b12x supports SwiGLU / GeGLU (gated) and ReLU2 (non-gated)
    activation_type = args.activation_type
    _ACT_STR = {
        ActivationType.Swiglu: "silu",
        ActivationType.Geglu: "gelu_tanh",
        ActivationType.GegluTanh: "gelu_tanh",
        ActivationType.Relu2: "relu2",
    }
    if activation_type not in _ACT_STR:
        raise ValueError(
            f"b12x_fused_moe only supports Swiglu, Geglu, GegluTanh, and Relu2 "
            f"activations, got {activation_type.name}"
        )
    activation_str = _ACT_STR[activation_type]
    is_gated = activation_type.is_gated

    # Create b12x-specific NVFP4 test data (weights quantized, input stays bf16)
    tensors = _create_nvfp4_moe_test_data(
        num_tokens=num_tokens,
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        num_experts=num_experts,
        num_local_experts=local_num_experts,
        top_k=top_k,
        device=device,
        backend="b12x",
        is_gated=is_gated,
    )

    if args.verbose >= 2:
        print(f"[VVERBOSE] x_bf16.shape = {tensors['x_bf16'].shape}")
        print(f"[VVERBOSE] w1_weight.shape = {tensors['w1_weight'].shape}")
        print(f"[VVERBOSE] w2_weight.shape = {tensors['w2_weight'].shape}")

    use_functional = getattr(args, "use_functional_api", False)
    quant_mode = getattr(args, "b12x_quant_mode", "nvfp4")
    x_input = tensors["x_bf16"]

    if use_functional:
        from functools import partial
        from flashinfer import b12x_fused_moe

        # Pre-allocate output buffer to avoid per-call allocation
        moe_output = torch.empty(
            num_tokens, hidden_size, dtype=torch.bfloat16, device=device
        )

        if args.verbose >= 1:
            print("[INFO] Using b12x functional API (b12x_fused_moe)")
        runner = partial(
            b12x_fused_moe,
            num_experts=num_experts,
            top_k=top_k,
            num_local_experts=local_num_experts,
            output=moe_output,
            activation=activation_str,
            quant_mode=quant_mode,
        )

        # Warmup call to populate workspace cache before timed region
        runner(
            x=x_input,
            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"],
            token_selected_experts=tensors["token_selected_experts"],
            token_final_scales=tensors["token_final_scales"],
        )
    else:
        moe = B12xMoEWrapper(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            use_cuda_graph=is_cuda_graph_compatible,
            max_num_tokens=num_tokens,
            num_local_experts=local_num_experts,
            activation=activation_str,
            quant_mode=quant_mode,
        )
        runner = moe.run

    def run_b12x_moe(
        x,
        w1_weight,
        w1_weight_sf,
        w1_alpha,
        fc2_input_scale,
        w2_weight,
        w2_weight_sf,
        w2_alpha,
        token_selected_experts,
        token_final_scales,
    ):
        return runner(
            x=x,
            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,
            token_selected_experts=token_selected_experts,
            token_final_scales=token_final_scales,
        )

    input_args = (
        x_input,
        tensors["w1_weight"],
        tensors["w1_weight_sf"],
        tensors["w1_alpha"],
        tensors["fc2_input_scale"],
        tensors["w2_weight"],
        tensors["w2_weight_sf"],
        tensors["w2_alpha"],
        tensors["token_selected_experts"],
        tensors["token_final_scales"],
    )

    # Snapshot active expert count before any kernel execution, since
    # autotune tactic exploration may corrupt input tensors.
    num_active_experts = int(tensors["token_selected_experts"].unique().numel())

    backend = "b12x"

    # Optional autotune warmup.
    if getattr(args, "autotune", False):
        warmup_iters = (
            args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10
        )
        backend = "b12x_autotune"
        if args.verbose >= 1:
            print(f"[INFO] Autotune warmup for b12x NVFP4 MoE: {warmup_iters} iters")
        autotune_args = tuple(
            t.clone() if isinstance(t, torch.Tensor) else t for t in input_args
        )
        with autotune(True):
            for _ in range(warmup_iters):
                run_b12x_moe(*autotune_args)
        del autotune_args

    # Cold-L2 rotation clones the weights inside CUDA-graph capture, but
    # W4A16 packs weights on first use and cannot pack during capture.
    # Keep L2 warm for that combination.
    cold_l2_cache = not (quant_mode == "w4a16" and is_cuda_graph_compatible)
    if not cold_l2_cache and args.verbose >= 1:
        print(
            "[INFO] w4a16 + CUDA graph: cold-L2 buffer rotation disabled "
            "(W4A16 packed-weight cache cannot be populated during capture)"
        )

    # Benchmark timing
    times = bench_gpu_time(
        fn=run_b12x_moe,
        dry_run_iters=args.dry_run_iters,
        repeat_iters=args.num_iters,
        sleep_after_run=False,
        enable_cupti=args.use_cupti,
        use_cuda_graph=is_cuda_graph_compatible,
        cold_l2_cache=cold_l2_cache,
        input_args=input_args,
    )

    # Compute performance metrics
    median_time = np.median(times)
    std_time = np.std(times)
    tflops = calculate_moe_tflops(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        is_gated=is_gated,
    )
    # Input format is bf16 for b12x (kernel fuses quantization), weights are nvfp4.
    tb_per_sec = calculate_moe_kernel_bandwidth(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        input_dtype,
        weight_dtype,
        input_format="bf16",
        weight_format="nvfp4",
        routing_logits_dtype=None,
        active_experts=num_active_experts,
        verbose=args.verbose,
        is_gated=is_gated,
    )

    print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

    if args.output_path is not None:
        cur_res = defaultdict(str)
        cur_res["routine"] = args.routine
        cur_res["median_time"] = median_time
        cur_res["std_time"] = std_time
        cur_res["tflops"] = tflops
        cur_res["tb_per_sec"] = tb_per_sec
        cur_res["backend"] = backend
        cur_res["num_tokens"] = num_tokens
        cur_res["hidden_size"] = hidden_size
        cur_res["intermediate_size"] = intermediate_size
        cur_res["num_experts"] = num_experts
        cur_res["top_k"] = top_k
        cur_res["local_num_experts"] = local_num_experts
        cur_res["input_dtype"] = input_dtype
        cur_res["weight_dtype"] = weight_dtype
        cur_res["fp4_mode"] = quant_mode
        cur_res["cold_l2_cache"] = cold_l2_cache
        cur_res["activation_type"] = activation_type.name
        res.append(cur_res)

    return res


def testTrtllmFp8BlockScaleMoe(args):
    """
    Test trtllm_fp8_block_scale_moe API (TensorRT-LLM fused MoE).

    This test:
    1. Creates quantized FP8 weights and block scales
    2. Runs FP8 block scale MOE
    3. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.verbose >= 1:
        print("[INFO] Running testTrtllmFp8BlockScaleMoe")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype)

    # Parse configuration
    num_tokens = args.num_tokens
    hidden_size = args.hidden_size
    intermediate_size = args.intermediate_size
    num_experts = args.num_experts
    top_k = args.top_k
    n_group = (
        args.n_group
        if hasattr(args, "n_group") and args.n_group is not None and args.n_group > 0
        else None
    )
    topk_group = (
        args.topk_group
        if hasattr(args, "topk_group")
        and args.topk_group is not None
        and args.topk_group > 0
        else None
    )
    routed_scaling_factor = (
        args.routed_scaling_factor
        if hasattr(args, "routed_scaling_factor")
        and args.routed_scaling_factor is not None
        else None
    )
    local_expert_offset = args.local_expert_offset
    local_num_experts = args.local_num_experts or num_experts
    routing_method_type = args.routing_method_type
    use_shuffled_weight = args.use_shuffled_weight
    weight_layout = args.weight_layout
    is_cuda_graph_compatible = not args.no_cuda_graph
    res = []
    backends = ["trtllm"]
    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    if args.verbose >= 1:
        print(
            f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, "
            f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}"
        )

    # Create test data
    routing_logits, routing_bias, hidden_states, gemm1_weights, gemm2_weights = (
        create_trtllm_moe_test_data(
            num_tokens,
            hidden_size,
            intermediate_size,
            num_experts,
            routing_method_type,
            args.use_routing_bias,
            input_dtype,
            weight_dtype,
            device,
            moe_kernel_type="fp8_block_scale",
        )
    )

    # Compute selected experts for accurate bandwidth calculation
    # Use the actual routing method to get correct expert assignments
    selected_experts = _compute_routing_for_method(
        routing_logits=routing_logits,
        routing_bias=routing_bias,
        top_k=top_k,
        routing_method_type=routing_method_type,
        n_group=n_group,
        topk_group=topk_group,
        routed_scaling_factor=routed_scaling_factor,
    )

    # For FP8 block scale, create quantized weights and block scales
    # Quantize to FP8
    gemm1_weights_fp8 = gemm1_weights.to(torch.float8_e4m3fn)
    gemm2_weights_fp8 = gemm2_weights.to(torch.float8_e4m3fn)

    # Optionally shuffle weights and convert to BlockMajorK layout to match kernel expectation
    if use_shuffled_weight:
        gemm1_weights_fp8_shuffled = []
        gemm2_weights_fp8_shuffled = []
        for i in range(num_experts):
            tmp_w1 = process_fp8_weight_layout(
                gemm1_weights_fp8[i], use_shuffled_weight, weight_layout
            )
            tmp_w2 = process_fp8_weight_layout(
                gemm2_weights_fp8[i], use_shuffled_weight, weight_layout
            )
            gemm1_weights_fp8_shuffled.append(tmp_w1.view(torch.uint8))
            gemm2_weights_fp8_shuffled.append(tmp_w2.view(torch.uint8))

        kernel_gemm1_weights = torch.stack(gemm1_weights_fp8_shuffled).view(
            torch.float8_e4m3fn
        )
        kernel_gemm2_weights = torch.stack(gemm2_weights_fp8_shuffled).view(
            torch.float8_e4m3fn
        )
    else:
        kernel_gemm1_weights = gemm1_weights_fp8
        kernel_gemm2_weights = gemm2_weights_fp8

    # Create block scale tensors for hidden states and weights (use float32 for scales)
    # TensorRT-LLM FP8 block-scale expects hidden_states_scale shape [hidden_size // 128, num_tokens]
    hidden_states_scale = 2.0 * torch.ones(
        (hidden_size // 128, num_tokens), device=device, dtype=torch.float32
    )
    gemm1_weights_scale = 2.0 * torch.ones(
        (num_experts, 2 * intermediate_size // 128, hidden_size // 128),
        device=device,
        dtype=torch.float32,
    )
    gemm2_weights_scale = 2.0 * torch.ones(
        (num_experts, hidden_size // 128, intermediate_size // 128),
        device=device,
        dtype=torch.float32,
    )

    if args.verbose >= 2:
        print(f"[VVERBOSE] routing_logits.shape = {routing_logits.shape}")
        print(f"[VVERBOSE] hidden_states.shape = {hidden_states.shape}")
        print(f"[VVERBOSE] gemm1_weights_fp8.shape = {gemm1_weights_fp8.shape}")
        print(f"[VVERBOSE] gemm2_weights_fp8.shape = {gemm2_weights_fp8.shape}")

    def run_fp8_block_moe(
        routing_logits,
        routing_bias,
        hidden_states,
        hidden_states_scale,
        kernel_gemm1_weights,
        gemm1_weights_scale,
        kernel_gemm2_weights,
        gemm2_weights_scale,
    ):
        # Quantize hidden states to FP8 for block scale MOE
        hidden_states_fp8 = hidden_states.to(torch.float8_e4m3fn)
        # Note: FP8 block scale MOE expects int64_t for n_group/topk_group, not Optional[int64_t]
        # So we convert None to 0 to indicate "no groups" mode
        return trtllm_fp8_block_scale_moe(
            routing_logits=routing_logits,
            routing_bias=routing_bias,
            hidden_states=hidden_states_fp8,
            hidden_states_scale=hidden_states_scale,
            gemm1_weights=kernel_gemm1_weights,
            gemm1_weights_scale=gemm1_weights_scale,
            gemm2_weights=kernel_gemm2_weights,
            gemm2_weights_scale=gemm2_weights_scale,
            num_experts=num_experts,
            top_k=top_k,
            n_group=n_group if n_group is not None else 0,
            topk_group=topk_group if topk_group is not None else 0,
            intermediate_size=intermediate_size,
            local_expert_offset=local_expert_offset,
            local_num_experts=local_num_experts,
            routed_scaling_factor=routed_scaling_factor,
            routing_method_type=routing_method_type,
            use_shuffled_weight=use_shuffled_weight,
            weight_layout=weight_layout,
            enable_pdl=args.enable_pdl,
        )

    # Benchmark timing
    times = bench_gpu_time(
        fn=run_fp8_block_moe,
        dry_run_iters=args.dry_run_iters,
        repeat_iters=args.num_iters,
        sleep_after_run=False,
        enable_cupti=args.use_cupti,
        use_cuda_graph=is_cuda_graph_compatible,
        cold_l2_cache=True,
        input_args=(
            routing_logits,
            routing_bias,
            hidden_states,
            hidden_states_scale,
            kernel_gemm1_weights,
            gemm1_weights_scale,
            kernel_gemm2_weights,
            gemm2_weights_scale,
        ),
    )

    # Compute performance metrics
    median_time = np.median(times)
    std_time = np.std(times)
    tflops = calculate_moe_tflops(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        is_gated=args.activation_type in (ActivationType.Swiglu, ActivationType.Geglu),
    )
    tb_per_sec = calculate_moe_kernel_bandwidth(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        input_dtype,
        weight_dtype,
        input_format="fp8",
        weight_format="fp8",
        routing_logits_dtype=routing_logits.dtype,
        active_experts=int(selected_experts.unique().numel()),
        verbose=args.verbose,
    )

    backend = "trtllm"
    print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

    if args.output_path is not None:
        cur_res = defaultdict(str)
        cur_res["routine"] = args.routine
        cur_res["median_time"] = median_time
        cur_res["std_time"] = std_time
        cur_res["tflops"] = tflops
        cur_res["tb_per_sec"] = tb_per_sec
        cur_res["backend"] = backend
        cur_res["num_tokens"] = num_tokens
        cur_res["hidden_size"] = hidden_size
        cur_res["intermediate_size"] = intermediate_size
        cur_res["num_experts"] = num_experts
        cur_res["top_k"] = top_k
        cur_res["n_group"] = n_group
        cur_res["topk_group"] = topk_group
        cur_res["routed_scaling_factor"] = routed_scaling_factor
        cur_res["local_expert_offset"] = local_expert_offset
        cur_res["local_num_experts"] = local_num_experts
        cur_res["routing_method"] = args.routing_method
        cur_res["use_shuffled_weight"] = use_shuffled_weight
        cur_res["weight_layout"] = weight_layout
        cur_res["use_routing_bias"] = args.use_routing_bias
        cur_res["use_routing_scales_on_input"] = args.use_routing_scales_on_input
        cur_res["input_dtype"] = input_dtype
        cur_res["weight_dtype"] = weight_dtype
        res.append(cur_res)

    return res


def testTrtllmFp8PerTensorScaleMoe(args):
    """
    Test trtllm_fp8_per_tensor_scale_moe API (TensorRT-LLM fused MoE).

    This test:
    1. Creates quantized FP8 weights and per-tensor scales
    2. Runs FP8 per-tensor scale MOE
    3. Measures performance metrics (TFLOPS, TB/sec)

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        dict: List of dictionaries containing performance results
    """
    if args.verbose >= 1:
        print("[INFO] Running testTrtllmFp8PerTensorScaleMoe")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype)

    # Parse configuration
    num_tokens = args.num_tokens
    hidden_size = args.hidden_size
    intermediate_size = args.intermediate_size
    num_experts = args.num_experts
    top_k = args.top_k
    n_group = (
        args.n_group
        if hasattr(args, "n_group") and args.n_group is not None and args.n_group > 0
        else None
    )
    topk_group = (
        args.topk_group
        if hasattr(args, "topk_group")
        and args.topk_group is not None
        and args.topk_group > 0
        else None
    )
    routed_scaling_factor = (
        args.routed_scaling_factor
        if hasattr(args, "routed_scaling_factor")
        and args.routed_scaling_factor is not None
        else None
    )
    local_expert_offset = args.local_expert_offset
    local_num_experts = args.local_num_experts or num_experts
    routing_method_type = args.routing_method_type
    use_routing_scales_on_input = args.use_routing_scales_on_input
    is_cuda_graph_compatible = not args.no_cuda_graph
    res = []
    backends = ["trtllm"]
    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    if args.verbose >= 1:
        print(
            f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, "
            f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}"
        )

    # Create test data
    routing_logits, routing_bias, hidden_states, gemm1_weights, gemm2_weights = (
        create_trtllm_moe_test_data(
            num_tokens,
            hidden_size,
            intermediate_size,
            num_experts,
            routing_method_type,
            args.use_routing_bias,
            input_dtype,
            weight_dtype,
            device,
            moe_kernel_type="fp8_per_tensor",
        )
    )

    # Compute selected experts for accurate bandwidth calculation
    # Use the actual routing method to get correct expert assignments
    selected_experts = _compute_routing_for_method(
        routing_logits=routing_logits,
        routing_bias=routing_bias,
        top_k=top_k,
        routing_method_type=routing_method_type,
        n_group=n_group,
        topk_group=topk_group,
        routed_scaling_factor=routed_scaling_factor,
    )

    # For FP8 per-tensor scale, create quantized weights and per-tensor scales
    # Quantize to FP8
    gemm1_weights_fp8 = gemm1_weights.to(torch.float8_e4m3fn)
    gemm2_weights_fp8 = gemm2_weights.to(torch.float8_e4m3fn)

    # Quantize hidden states to FP8 for per-tensor scale
    hidden_states_fp8 = hidden_states.to(torch.float8_e4m3fn)

    # Create per-tensor scale scalars using shared utility
    output1_scales_scalar, output1_scales_gate_scalar, output2_scales_scalar = (
        create_moe_output_scale_scalars(local_num_experts, device)
    )

    if args.verbose >= 2:
        print(f"[VVERBOSE] routing_logits.shape = {routing_logits.shape}")
        print(f"[VVERBOSE] hidden_states.shape = {hidden_states.shape}")
        print(f"[VVERBOSE] gemm1_weights_fp8.shape = {gemm1_weights_fp8.shape}")
        print(f"[VVERBOSE] gemm2_weights_fp8.shape = {gemm2_weights_fp8.shape}")

    def run_fp8_per_tensor_moe(
        routing_logits,
        routing_bias,
        hidden_states_fp8,
        gemm1_weights_fp8,
        output1_scales_scalar,
        output1_scales_gate_scalar,
        gemm2_weights_fp8,
        output2_scales_scalar,
        activation_type,
    ):
        # Note: FP8 per-tensor MOE expects int64_t for n_group/topk_group, not Optional[int64_t]
        # So we convert None to 0 to indicate "no groups" mode
        return trtllm_fp8_per_tensor_scale_moe(
            routing_logits=routing_logits,
            routing_bias=routing_bias,
            hidden_states=hidden_states_fp8,
            gemm1_weights=gemm1_weights_fp8,
            output1_scales_scalar=output1_scales_scalar,
            output1_scales_gate_scalar=output1_scales_gate_scalar,
            gemm2_weights=gemm2_weights_fp8,
            output2_scales_scalar=output2_scales_scalar,
            num_experts=num_experts,
            top_k=top_k,
            n_group=n_group if n_group is not None else 0,
            topk_group=topk_group if topk_group is not None else 0,
            intermediate_size=intermediate_size,
            local_expert_offset=local_expert_offset,
            local_num_experts=local_num_experts,
            routed_scaling_factor=routed_scaling_factor,
            use_routing_scales_on_input=use_routing_scales_on_input,
            routing_method_type=routing_method_type,
            enable_pdl=args.enable_pdl,
            **_activation_kwarg(trtllm_fp8_per_tensor_scale_moe, activation_type),
        )

    # Benchmark timing
    times = bench_gpu_time(
        fn=run_fp8_per_tensor_moe,
        dry_run_iters=args.dry_run_iters,
        repeat_iters=args.num_iters,
        sleep_after_run=False,
        enable_cupti=args.use_cupti,
        use_cuda_graph=is_cuda_graph_compatible,
        cold_l2_cache=True,
        input_args=(
            routing_logits,
            routing_bias,
            hidden_states_fp8,
            gemm1_weights_fp8,
            output1_scales_scalar,
            output1_scales_gate_scalar,
            gemm2_weights_fp8,
            output2_scales_scalar,
            args.activation_type,
        ),
    )

    # Compute performance metrics
    median_time = np.median(times)
    std_time = np.std(times)
    tflops = calculate_moe_tflops(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        is_gated=args.activation_type in (ActivationType.Swiglu, ActivationType.Geglu),
    )
    tb_per_sec = calculate_moe_kernel_bandwidth(
        num_tokens,
        hidden_size,
        intermediate_size,
        num_experts,
        top_k,
        median_time,
        input_dtype,
        weight_dtype,
        input_format="fp8",
        weight_format="fp8",
        routing_logits_dtype=routing_logits.dtype,
        active_experts=int(selected_experts.unique().numel()),
        verbose=args.verbose,
    )

    backend = "trtllm"
    print_perf_metrics(backend, median_time, std_time, tflops, tb_per_sec)

    if args.output_path is not None:
        cur_res = defaultdict(str)
        cur_res["routine"] = args.routine
        cur_res["median_time"] = median_time
        cur_res["std_time"] = std_time
        cur_res["tflops"] = tflops
        cur_res["tb_per_sec"] = tb_per_sec
        cur_res["backend"] = backend
        cur_res["num_tokens"] = num_tokens
        cur_res["hidden_size"] = hidden_size
        cur_res["intermediate_size"] = intermediate_size
        cur_res["num_experts"] = num_experts
        cur_res["top_k"] = top_k
        cur_res["n_group"] = n_group
        cur_res["topk_group"] = topk_group
        cur_res["routed_scaling_factor"] = routed_scaling_factor
        cur_res["local_expert_offset"] = local_expert_offset
        cur_res["local_num_experts"] = local_num_experts
        cur_res["routing_method"] = args.routing_method
        cur_res["use_routing_bias"] = args.use_routing_bias
        cur_res["use_routing_scales_on_input"] = use_routing_scales_on_input
        cur_res["input_dtype"] = input_dtype
        cur_res["weight_dtype"] = weight_dtype
        cur_res["activation_type"] = args.activation_type.name
        res.append(cur_res)

    return res


# =============================================================================
# Unified NVFP4 MoE — MoELayer-based cross-backend autotune
# =============================================================================


def testUnifiedNvfp4Moe(args):
    """MoELayer cross-backend autotune for NVFP4.

    Exercises the unified MoE API: per-shape, builds MoELayer over
    {CuteDslConfig, TrtllmFp4Config}, registers both backend-native weight
    views on a single MoEWeightPack, dispatches once, reports the winner
    and both backends' best-tactic latencies.

    Emits one result row per backend candidate so CSV output carries the
    cross-backend comparison directly.
    """
    from flashinfer.autotuner import AutoTuner
    from flashinfer.fused_moe import (
        ActivationConfig,
        CuteDslConfig,
        ExecutionConfig,
        ExpertConfig,
        MoEActivationPack,
        MoEConfig,
        MoELayer,
        MoEWeightPack,
        QuantConfig,
        QuantVariant,
        RoutingConfig,
        TrtllmFp4Config,
    )
    from flashinfer.fused_moe.api import BackendOptions

    if args.verbose >= 1:
        print("[INFO] Running testUnifiedNvfp4Moe")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)
    if args.generate_repro_command:
        print(
            f"[INFO] To reproduce this test case, run the following command: {args.repro_command}"
        )

    num_tokens = args.num_tokens
    hidden_size = args.hidden_size
    intermediate_size = args.intermediate_size
    num_experts = args.num_experts
    top_k = args.top_k
    local_expert_offset = args.local_expert_offset
    local_num_experts = args.local_num_experts or num_experts
    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)
    weight_dtype = dtype_str_to_torch_dtype(args.weight_dtype)

    res = []
    backends = ["unified"]
    backends = filter_backends_by_compute_capability(backends, args.routine, device)
    if len(backends) == 0:
        print("[ERROR] No backends to test. Exiting.")
        return res

    if args.verbose >= 1:
        print(
            f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, "
            f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}, "
            f"local_experts={local_num_experts}"
        )

    # ---- Build shared bf16 reference weights ------------------------------
    torch.manual_seed(0)
    w1_bf16 = (
        torch.randn(
            local_num_experts,
            2 * intermediate_size,
            hidden_size,
            dtype=torch.bfloat16,
            device=device,
        )
        / 10
    )
    w2_bf16 = (
        torch.randn(
            local_num_experts,
            hidden_size,
            intermediate_size,
            dtype=torch.bfloat16,
            device=device,
        )
        / 10
    )

    # ---- Backend-native weight views — first-class prepare helpers (CR2) --
    # Both views are built from the SAME canonical bf16 weights, so the two
    # backends are directly comparable (and a shared reference is meaningful).
    cute_dsl_view = CuteDslConfig.prepare_weights(
        w1_bf16,
        w2_bf16,
        num_local_experts=local_num_experts,
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        device=device,
    )
    trtllm_view = TrtllmFp4Config.prepare_weights(
        w1_bf16,
        w2_bf16,
        num_local_experts=local_num_experts,
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        device=device,
    )

    # ---- Activation pack --------------------------------------------------
    # The activation (NVFP4-quantized hidden states + pre-routed indices) still
    # comes from the canonical test data creator; only weight prep is
    # first-class today (activation-prep promotion tracked under CR2).
    import os
    import sys

    _repo_root = os.path.dirname(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    )
    if _repo_root not in sys.path:
        sys.path.insert(0, _repo_root)
    from tests.moe.test_cute_dsl_fused_moe import (
        check_accuracy,
        compute_reference_moe_fp4,
        create_moe_tensors,
    )

    # Wide-EP (MVP): model a single rank as a complete MoE over its
    # local_num_experts experts — route the activation WITHIN the local experts
    # (selected ids in [0, local_num_experts)) so every token is computed
    # locally. This is the "local-only" proxy: it avoids the unfaithful
    # global-routing-but-local-weights setup (most tokens skipped) and keeps the
    # derived FLOP/bandwidth metrics correct (all tokens computed, active
    # experts <= local). For EP=1 (local == global) this is unchanged.
    routing_num_experts = local_num_experts
    routing_top_k = min(top_k, local_num_experts)
    cute_dsl_data = create_moe_tensors(
        num_tokens=num_tokens,
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        num_experts=routing_num_experts,
        num_local_experts=local_num_experts,
        top_k=routing_top_k,
        device=device,
    )
    # cute_dsl_data["x_sf"] is already unsqueezed to [M, H//16, 1]; strip that
    # for the Pack (runner re-applies unsqueeze in pack_inputs).
    x_sf = cute_dsl_data["x_sf"].squeeze(-1)
    # The local-only proxy generates ids in [0, local_num_experts); the runners
    # expect GLOBAL ids (the kernel maps them to the local shard via the
    # separately passed local_expert_offset), so lift the pack ids to global.
    # The reference oracle below keeps the LOCAL ids: it indexes the local
    # weight slices directly.
    local_topk_ids = cute_dsl_data["token_selected_experts"]
    act_pack = MoEActivationPack(
        hidden_states_q=cute_dsl_data["x"],
        hidden_states_scale=x_sf,
        topk_ids=local_topk_ids + local_expert_offset,
        topk_weights=cute_dsl_data["token_final_scales"],
    )

    num_active_experts = int(local_topk_ids.unique().numel())

    weight_pack = MoEWeightPack()
    weight_pack.prepare_for("cute_dsl_nvfp4", cute_dsl_view)
    weight_pack.prepare_for("trtllm_fp4_routed", trtllm_view)

    # ---- MoELayer config --------------------------------------------------
    config = MoEConfig(
        routing=RoutingConfig(
            num_experts=num_experts,
            # top_k must match the pack's column count: the tensors above were
            # generated with routing_top_k (clamped to local_num_experts for
            # the local-only EP proxy), not the raw CLI top_k.
            top_k=routing_top_k,
            n_group=args.n_group,
            topk_group=args.topk_group,
            routed_scaling_factor=args.routed_scaling_factor,
        ),
        quant=QuantConfig(variant=QuantVariant.NVFP4),
        experts=ExpertConfig(
            intermediate_size=intermediate_size,
            local_expert_offset=local_expert_offset,
            local_num_experts=local_num_experts,
        ),
        activation=ActivationConfig(),
        backend=BackendOptions(candidates=(CuteDslConfig(), TrtllmFp4Config())),
        execution=ExecutionConfig(tune_max_num_tokens=max(num_tokens, 8192)),
    )

    # ---- Dispatch: trigger cross-backend selection ------------------------
    with autotune(True):
        layer = MoELayer(config, device=device)
        _ = layer(act_pack, weight_pack)

    winner_key = layer.winner_backend or "?"
    if args.verbose >= 1:
        print(
            f"[INFO] MoELayer winner: {winner_key} "
            f"(candidates: {[r.backend_key for r in layer.runners]})"
        )

    # ---- Optional shared-reference accuracy check (CR10/CR11) -------------
    # Both backend views derive from the same bf16 weights, so a single bf16
    # reference is valid for every candidate.  Catches shared-mode errors that
    # cross-backend agreement would miss.
    ref_output = None
    if args.refcheck:
        ref_output = compute_reference_moe_fp4(
            hidden_states=cute_dsl_data["x_bf16"].float().to(device),
            gemm1_weights=w1_bf16.float().to(device),
            gemm2_weights=w2_bf16.float().to(device),
            token_selected_experts=local_topk_ids,
            token_final_scales=act_pack.topk_weights,
            num_tokens=num_tokens,
            num_experts=num_experts,
            top_k=routing_top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            fc2_input_scale=cute_dsl_data["fc2_input_scale"],
        )

    # ---- Per-backend best-tactic timing (one result row per candidate) ----
    tuner = AutoTuner.get()
    for runner in layer.runners:
        inputs = runner.pack_inputs(act_pack, weight_pack)
        with autotune(True):
            _, tactic = tuner.choose_one(
                custom_op=f"moe_{runner.backend_key}",
                runners=[runner],
                tuning_config=runner.tuning_config,
                inputs=inputs,
            )

        def _call(r=runner, i=inputs, t=tactic):
            return r.forward(i, tactic=t)

        times = bench_gpu_time(
            fn=_call,
            dry_run_iters=args.dry_run_iters,
            repeat_iters=args.num_iters,
            sleep_after_run=False,
            enable_cupti=args.use_cupti,
            use_cuda_graph=not args.no_cuda_graph,
            cold_l2_cache=True,
        )
        median_time = float(np.median(times))
        std_time = float(np.std(times))
        # Use the actual routing geometry (local experts / clamped top_k), not
        # the global config, so the wide-EP local-only proxy reports honest
        # FLOP/byte counts. For EP=1 these equal num_experts/top_k.
        tflops = calculate_moe_tflops(
            num_tokens,
            hidden_size,
            intermediate_size,
            routing_num_experts,
            routing_top_k,
            median_time,
        )
        tb_per_sec = calculate_moe_kernel_bandwidth(
            num_tokens,
            hidden_size,
            intermediate_size,
            routing_num_experts,
            routing_top_k,
            median_time,
            input_dtype,
            weight_dtype,
            input_format="nvfp4",
            weight_format="nvfp4",
            routing_logits_dtype=None,
            active_experts=num_active_experts,
            verbose=args.verbose,
        )

        # Tag: "unified/<backend>" + star the winner
        is_winner = runner.backend_key == winner_key
        backend_label = f"unified/{runner.backend_key}" + ("*" if is_winner else "")
        print_perf_metrics(backend_label, median_time, std_time, tflops, tb_per_sec)

        # Shared-reference accuracy check for this candidate (CR10/CR11).
        refcheck_passed = None
        if ref_output is not None:
            out = runner.forward(inputs, tactic=tactic)
            refcheck_passed, pct, atol = check_accuracy(out, ref_output)
            status = "PASS" if refcheck_passed else "FAIL"
            print(
                f"[REFCHECK] {backend_label}: {status} "
                f"({pct * 100:.2f}% within atol={atol:.4f} vs bf16 reference)"
            )
            if not refcheck_passed and not args.allow_output_mismatch:
                print(
                    f"[ERROR] {backend_label} failed reference check. "
                    f"Re-run with --allow_output_mismatch to continue past mismatches."
                )

        if args.output_path is not None:
            cur_res = defaultdict(str)
            cur_res["routine"] = args.routine
            cur_res["median_time"] = median_time
            cur_res["std_time"] = std_time
            cur_res["tflops"] = tflops
            cur_res["tb_per_sec"] = tb_per_sec
            cur_res["backend"] = backend_label
            cur_res["refcheck_passed"] = refcheck_passed
            cur_res["num_tokens"] = num_tokens
            cur_res["hidden_size"] = hidden_size
            cur_res["intermediate_size"] = intermediate_size
            cur_res["num_experts"] = num_experts
            cur_res["top_k"] = top_k
            cur_res["local_expert_offset"] = local_expert_offset
            cur_res["local_num_experts"] = local_num_experts
            cur_res["input_dtype"] = input_dtype
            cur_res["weight_dtype"] = weight_dtype
            cur_res["fp4_mode"] = "nvfp4"
            res.append(cur_res)

    return res


def testBgmvMoe(args):
    """
    Benchmark BGMV MoE kernels for multi-LoRA inference.

    Measures the performance of the fused shrink + expand CUDA kernels
    that apply LoRA adapters in MoE models. Compares against FlashInfer's
    grouped_mm_bf16 as a baseline.

    Args:
        args: Parsed command line arguments containing test configuration

    Returns:
        list: List of dictionaries containing performance results
    """
    if args.verbose >= 1:
        print("[INFO] Running testBgmvMoe")
        print(f"[INFO] FlashInfer version: {flashinfer.__version__}")

    device = get_device(args)

    input_dtype = dtype_str_to_torch_dtype(args.input_dtype)

    num_tokens = args.num_tokens
    hidden_size = args.hidden_size
    num_experts = args.num_experts
    top_k = args.top_k
    rank = args.rank
    num_loras = args.num_loras
    num_slices = getattr(args, "num_slices", 1)
    res = []

    if args.verbose >= 1:
        print(
            f"[INFO] Configuration: tokens={num_tokens}, hidden={hidden_size}, "
            f"rank={rank}, experts={num_experts}, top_k={top_k}, "
            f"num_loras={num_loras}, num_slices={num_slices}"
        )

    # Import BGMV MoE kernels
    from flashinfer.fused_moe.bgmv_moe import (
        bgmv_moe_shrink,
        bgmv_moe_expand,
        fill_w_ptr,
    )

    # Generate test data
    num_pairs = num_tokens * top_k
    feat_out = hidden_size

    x = torch.randn(num_tokens, hidden_size, dtype=input_dtype, device=device) * 0.1
    lora_a_weights = [
        torch.randn(
            num_loras, num_experts, rank, hidden_size, dtype=input_dtype, device=device
        )
        * 0.01
        for _ in range(num_slices)
    ]
    lora_b_weights = [
        torch.randn(
            num_loras, num_experts, feat_out, rank, dtype=input_dtype, device=device
        )
        * 0.01
        for _ in range(num_slices)
    ]

    sorted_token_ids = torch.arange(
        num_tokens, device=device, dtype=torch.int64
    ).repeat_interleave(top_k)
    expert_ids = torch.randint(
        0, num_experts, (num_pairs,), device=device, dtype=torch.int64
    )
    topk_weights = (
        torch.softmax(torch.randn(num_tokens, top_k, device=device), dim=-1)
        .view(-1)
        .to(torch.float32)
    )
    lora_indices = torch.randint(
        0, num_loras, (num_tokens,), device=device, dtype=torch.int64
    )

    # Pre-allocate buffers for BGMV MoE
    w_ptr_a = torch.zeros(num_slices, num_experts, dtype=torch.int64, device=device)
    lora_stride_a = 0
    for s in range(num_slices):
        lora_stride_a = fill_w_ptr(w_ptr_a, lora_a_weights[s], num_experts, s)

    w_ptr_b = torch.zeros(num_slices, num_experts, dtype=torch.int64, device=device)
    lora_stride_b = 0
    for s in range(num_slices):
        lora_stride_b = fill_w_ptr(w_ptr_b, lora_b_weights[s], num_experts, s)

    shrink_out = torch.zeros(
        num_slices, num_pairs, rank, dtype=input_dtype, device=device
    )
    slice_start_loc = torch.zeros(num_slices, dtype=torch.int64, device=device)
    for s in range(num_slices):
        slice_start_loc[s] = s * feat_out
    output_slices = [feat_out] * num_slices

    y_accum = torch.zeros(
        num_tokens, feat_out * num_slices, dtype=torch.float32, device=device
    )

    # Define benchmark function (shrink + expand)
    def run_bgmv_moe(
        shrink_out,
        x,
        w_ptr_a,
        sorted_token_ids,
        expert_ids,
        lora_indices,
        lora_stride_a,
        y_accum,
        w_ptr_b,
        topk_weights,
        slice_start_loc,
        output_slices,
        lora_stride_b,
    ):
        shrink_out.zero_()
        y_accum.zero_()
        bgmv_moe_shrink(
            shrink_out,
            x,
            w_ptr_a,
            sorted_token_ids,
            expert_ids,
            lora_indices,
            lora_stride_a,
        )
        bgmv_moe_expand(
            y_accum,
            shrink_out,
            w_ptr_b,
            sorted_token_ids,
            expert_ids,
            topk_weights,
            lora_indices,
            slice_start_loc,
            output_slices,
            lora_stride_b,
        )

    is_cuda_graph_compatible = not args.no_cuda_graph

    # Benchmark BGMV MoE
    times = bench_gpu_time(
        fn=run_bgmv_moe,
        dry_run_iters=args.dry_run_iters,
        repeat_iters=args.num_iters,
        sleep_after_run=False,
        enable_cupti=args.use_cupti,
        use_cuda_graph=is_cuda_graph_compatible,
        cold_l2_cache=True,
        input_args=(
            shrink_out,
            x,
            w_ptr_a,
            sorted_token_ids,
            expert_ids,
            lora_indices,
            lora_stride_a,
            y_accum,
            w_ptr_b,
            topk_weights,
            slice_start_loc,
            output_slices,
            lora_stride_b,
        ),
    )

    bgmv_median = np.median(times)
    bgmv_std = np.std(times)

    # Benchmark grouped_mm_bf16 as baseline comparison (single-slice only)
    gg_kernel_median = float("nan")
    gg_full_median = float("nan")
    if num_slices != 1:
        print(
            "[INFO] grouped_mm_bf16 baseline skipped: only supported for num_slices=1"
        )
    else:
        try:
            from flashinfer.grouped_mm import grouped_mm_bf16

            # For grouped GEMM, sort tokens by (lora_id, expert_id) and build m_indptr
            num_groups = num_loras * num_experts

            # Assign each pair a group_id = lora_id * num_experts + expert_id
            lora_ids_expanded = lora_indices[sorted_token_ids]
            group_ids = lora_ids_expanded * num_experts + expert_ids
            group_ids[lora_ids_expanded < 0] = num_groups

            # Sort by group_id
            sorted_indices = torch.argsort(group_ids)
            sorted_group_ids = group_ids[sorted_indices]

            valid_mask = sorted_group_ids < num_groups
            num_valid = valid_mask.sum().item()

            sorted_token_indices = sorted_token_ids[sorted_indices[:num_valid]]
            g_input = x[sorted_token_indices]

            # Build m_indptr
            counts = torch.zeros(num_groups + 1, dtype=torch.int32, device=device)
            valid_groups = sorted_group_ids[:num_valid]
            for g in range(num_groups):
                counts[g + 1] = counts[g] + (valid_groups == g).sum().to(torch.int32)
            g_m_indptr = counts

            # Reshape LoRA weights for grouped GEMM
            g_lora_a = lora_a_weights[0].view(
                num_loras * num_experts, rank, hidden_size
            )
            g_lora_b = lora_b_weights[0].view(num_loras * num_experts, feat_out, rank)

            g_shrink_out = torch.zeros(
                num_valid, rank, dtype=input_dtype, device=device
            )
            g_expand_out = torch.zeros(
                num_valid, feat_out, dtype=input_dtype, device=device
            )

            # Warmup
            grouped_mm_bf16(g_input, g_lora_a, g_m_indptr, out=g_shrink_out)
            grouped_mm_bf16(g_shrink_out, g_lora_b, g_m_indptr, out=g_expand_out)
            torch.cuda.synchronize()

            # Benchmark: kernel only (pre-sorted, no sort overhead)
            def run_gg_kernel(
                g_input, g_lora_a, g_m_indptr, g_shrink_out, g_lora_b, g_expand_out
            ):
                g_shrink_out.zero_()
                g_expand_out.zero_()
                grouped_mm_bf16(g_input, g_lora_a, g_m_indptr, out=g_shrink_out)
                grouped_mm_bf16(g_shrink_out, g_lora_b, g_m_indptr, out=g_expand_out)

            gg_kernel_times = bench_gpu_time(
                fn=run_gg_kernel,
                dry_run_iters=args.dry_run_iters,
                repeat_iters=args.num_iters,
                sleep_after_run=False,
                enable_cupti=args.use_cupti,
                use_cuda_graph=is_cuda_graph_compatible,
                cold_l2_cache=True,
                input_args=(
                    g_input,
                    g_lora_a,
                    g_m_indptr,
                    g_shrink_out,
                    g_lora_b,
                    g_expand_out,
                ),
            )
            gg_kernel_median = np.median(gg_kernel_times)

            # Benchmark: sort + kernel (full end-to-end)
            def run_gg_full(
                x,
                sorted_token_ids,
                lora_indices,
                expert_ids,
                num_experts,
                num_groups,
                g_lora_a,
                g_lora_b,
                g_m_indptr,
                g_shrink_out,
                g_expand_out,
            ):
                _lora_ids = lora_indices[sorted_token_ids]
                _group_ids = _lora_ids * num_experts + expert_ids
                _group_ids[_lora_ids < 0] = num_groups
                _sorted_indices = torch.argsort(_group_ids)
                _sorted_token_indices = sorted_token_ids[_sorted_indices[:num_valid]]
                _g_input = x[_sorted_token_indices]
                g_shrink_out.zero_()
                g_expand_out.zero_()
                grouped_mm_bf16(_g_input, g_lora_a, g_m_indptr, out=g_shrink_out)
                grouped_mm_bf16(g_shrink_out, g_lora_b, g_m_indptr, out=g_expand_out)

            gg_full_times = bench_gpu_time(
                fn=run_gg_full,
                dry_run_iters=args.dry_run_iters,
                repeat_iters=args.num_iters,
                sleep_after_run=False,
                enable_cupti=args.use_cupti,
                use_cuda_graph=False,  # sort is not cuda-graph compatible
                cold_l2_cache=True,
                input_args=(
                    x,
                    sorted_token_ids,
                    lora_indices,
                    expert_ids,
                    num_experts,
                    num_groups,
                    g_lora_a,
                    g_lora_b,
                    g_m_indptr,
                    g_shrink_out,
                    g_expand_out,
                ),
            )
            gg_full_median = np.median(gg_full_times)

        except (ImportError, RuntimeError) as e:
            print(f"[INFO] grouped_mm_bf16 baseline skipped: {e}")

    # Print comparison results
    def _fmt(v):
        return f"{v:.3f}" if v == v else "N/A"

    def _speedup(baseline, kernel):
        if baseline != baseline or kernel != kernel or kernel == 0:
            return "N/A"
        return f"{baseline / kernel:.2f}x"

    print(f"\n{'=' * 70}")
    print(
        f"  BGMV MoE Benchmark: tokens={num_tokens}, hidden={hidden_size}, "
        f"rank={rank}, experts={num_experts}, top_k={top_k}"
    )
    print(f"{'=' * 70}")
    print(f"  {'Method':<25} {'Median (ms)':>12} {'Speedup vs BGMV MoE':>20}")
    print(f"  {'-' * 25} {'-' * 12} {'-' * 20}")
    print(f"  {'BGMV MoE (this PR)':<25} {_fmt(bgmv_median):>12} {'—':>20}")
    print(
        f"  {'grouped_mm (kernel)':<25} {_fmt(gg_kernel_median):>12} {_speedup(gg_kernel_median, bgmv_median):>20}"
    )
    print(
        f"  {'grouped_mm (sort+kern)':<25} {_fmt(gg_full_median):>12} {_speedup(gg_full_median, bgmv_median):>20}"
    )
    print(f"{'=' * 70}\n")

    # Also print in standard format
    flops = num_slices * num_pairs * (hidden_size * rank + rank * feat_out) * 2
    tflops = flops / (bgmv_median * 1e-3) / 1e12

    backend = "bgmv_moe"
    print_perf_metrics(backend, bgmv_median, bgmv_std, tflops, 0.0)

    if args.output_path is not None:
        cur_res = defaultdict(str)
        cur_res["routine"] = args.routine
        cur_res["median_time"] = bgmv_median
        cur_res["std_time"] = bgmv_std
        cur_res["tflops"] = tflops
        cur_res["backend"] = backend
        cur_res["num_tokens"] = num_tokens
        cur_res["hidden_size"] = hidden_size
        cur_res["num_experts"] = num_experts
        cur_res["top_k"] = top_k
        cur_res["rank"] = rank
        cur_res["num_loras"] = num_loras
        cur_res["num_slices"] = num_slices
        cur_res["input_dtype"] = str(input_dtype)
        cur_res["gg_kernel_median"] = gg_kernel_median
        cur_res["gg_full_median"] = gg_full_median
        res.append(cur_res)

    return res
