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

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

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

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

from typing import Literal

import pytest
import torch

from flashinfer import (
    ActivationType,
    RoutingMethodType,
    fp4_quantize,
    mxfp8_quantize,
    reorder_rows_for_gated_act_gemm,
    shuffle_matrix_a,
    shuffle_matrix_sf_a,
)
from flashinfer.fused_moe import (
    WeightLayout,
    convert_to_block_layout,
    trtllm_bf16_moe,
    trtllm_bf16_routed_moe,
    trtllm_fp4_block_scale_moe,
    trtllm_fp4_block_scale_routed_moe,
    trtllm_fp8_block_scale_moe,
    trtllm_fp8_block_scale_routed_moe,
    trtllm_fp8_per_tensor_scale_moe,
    trtllm_fp8_per_tensor_scale_routed_moe,
    trtllm_mxint4_block_scale_moe,
    trtllm_mxint4_block_scale_routed_moe,
)
from flashinfer.fused_moe.core import Fp8QuantizationType
from flashinfer.utils import device_support_pdl, get_compute_capability
from .trtllm_gen_fused_moe_utils import (
    FP8BlockScaleMoe,
    QuantMode,
    routing_reference_renormalize,
    routing_reference_renormalize_naive,
    routing_reference_topk,
)

pytestmark = pytest.mark.solo


def _run_trtllm_gen_routed_fused_moe_case(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    num_experts: int,
    routing_method_type: RoutingMethodType,
    quant_mode: Literal["NvFP4xNvFP4", "MxFP4xMxFP8", "MxFP4xBf16"],
    routing_format: Literal["packed", "unpacked", "unpacked_fp32"],
    activation_type: ActivationType = ActivationType.Swiglu,
    gemm1_alpha: torch.Tensor | None = None,
    gemm1_beta: torch.Tensor | None = None,
    gemm1_clamp_limit: torch.Tensor | None = None,
):
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")
    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)
    routing_logits = torch.rand(num_tokens, num_experts, device=device).to(
        torch.bfloat16
    )
    hidden_states = (
        torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) * 0.1
    )
    if quant_mode == "NvFP4xNvFP4":
        hidden_states, hidden_states_scale = fp4_quantize(
            hidden_states,
            torch.tensor([448.0 * 6.0], device=device),
            sf_vec_size=16,
            sf_use_ue8m0=False,
            is_sf_swizzled_layout=False,
        )
        hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape(
            num_tokens, -1
        )
        hidden_states_global_scale = 1.0 / 448.0 / 6.0
    elif quant_mode == "MxFP4xMxFP8":
        hidden_states, hidden_states_scale = mxfp8_quantize(hidden_states, False)
        hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape(
            num_tokens, -1
        )
        hidden_states_global_scale = 1.0
    else:  # MxFP4xBf16
        hidden_states_scale = None
        hidden_states_global_scale = 1.0

    w13 = (
        torch.randn(num_experts, intermediate_size * 2, hidden_size, device=device).to(
            torch.bfloat16
        )
        * 0.1
    )
    w2 = (
        torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
            torch.bfloat16
        )
        * 0.1
    )
    if quant_mode == "NvFP4xNvFP4":
        w13, w13_scale = fp4_quantize(
            w13,
            torch.tensor([448.0 * 6.0], device=device),
            sf_vec_size=16,
            sf_use_ue8m0=False,
        )
        w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
            num_experts, intermediate_size * 2, -1
        )
        w2, w2_scale = fp4_quantize(
            w2,
            torch.tensor([448.0 * 6.0], device=device),
            sf_vec_size=16,
            sf_use_ue8m0=False,
        )
        w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
            num_experts, hidden_size, -1
        )
        w13_global_scale = 1.0 / 448.0 / 6.0
        w2_global_scale = 1.0 / 448.0 / 6.0
    else:
        w13, w13_scale = fp4_quantize(
            w13, torch.tensor([1.0], device=device), sf_vec_size=32, sf_use_ue8m0=True
        )
        w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
            num_experts, intermediate_size * 2, -1
        )
        w2, w2_scale = fp4_quantize(
            w2, torch.tensor([1.0], device=device), sf_vec_size=32, sf_use_ue8m0=True
        )
        w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
            num_experts, hidden_size, -1
        )
        w13_global_scale = 1.0
        w2_global_scale = 1.0

    output1_scale_scalar = torch.tensor(
        [hidden_states_global_scale * w13_global_scale] * num_experts, device=device
    )
    output1_scale_gate_scalar = torch.tensor(
        [hidden_states_global_scale * w13_global_scale] * num_experts, device=device
    )
    output2_scale_scalar = torch.tensor(
        [hidden_states_global_scale * w2_global_scale] * num_experts, device=device
    )

    reference_output = trtllm_fp4_block_scale_moe(
        routing_logits,
        None,  # routing_bias
        hidden_states,
        hidden_states_scale,
        w13,
        w13_scale,
        None,  # w13_bias
        gemm1_alpha,  # gemm1_alpha
        gemm1_beta,  # gemm1_beta
        gemm1_clamp_limit,  # gemm1_clamp_limit
        w2,
        w2_scale,
        None,  # w2_bias
        output1_scale_scalar,
        output1_scale_gate_scalar,
        output2_scale_scalar,
        num_experts,
        top_k,
        None,  # n_group
        None,  # topk_group
        intermediate_size,
        0,  # local_expert_offset
        num_experts,
        None,  # routed_scaling_factor
        routing_method_type.value,
        True,  # do_finalize
        enable_pdl,
        activation_type.value,  # act_type
        None,
    )[0].to(torch.float)

    if routing_method_type == RoutingMethodType.Renormalize:
        permute_info, expert_weights = routing_reference_renormalize(
            routing_logits, top_k, num_experts, 8
        )
    elif routing_method_type == RoutingMethodType.RenormalizeNaive:
        permute_info, expert_weights = routing_reference_renormalize_naive(
            routing_logits, top_k, num_experts, 8
        )
    elif routing_method_type == RoutingMethodType.TopK:
        permute_info, expert_weights = routing_reference_topk(
            routing_logits, top_k, num_experts, 8
        )
    topk_ids = permute_info["topKIndices"].to(torch.int32)
    topk_weights = expert_weights.view(num_tokens, num_experts)[
        torch.arange(num_tokens).unsqueeze(1), topk_ids
    ]

    # Prepare routing input based on format
    if routing_format == "packed":
        # Packed format: (score << 16 | expert_id)
        topk_weights = topk_weights.to(torch.bfloat16)
        routing_input = (topk_ids.to(torch.int32) << 16) | topk_weights.view(
            torch.int16
        )
    elif routing_format == "unpacked_fp32":
        # Unpacked format: (topk_ids, topk_weights) tuple, float32 weights.
        routing_input = (topk_ids, topk_weights.to(torch.float32))
    else:
        # Unpacked format: (topk_ids, topk_weights) tuple, bfloat16 weights.
        routing_input = (topk_ids, topk_weights.to(torch.bfloat16))

    output = trtllm_fp4_block_scale_routed_moe(
        routing_input,
        None,  # routing_bias
        hidden_states,
        hidden_states_scale,
        w13,
        w13_scale,
        None,  # w13_bias
        gemm1_alpha,  # gemm1_alpha
        gemm1_beta,  # gemm1_beta
        gemm1_clamp_limit,  # gemm1_clamp_limit
        w2,
        w2_scale,
        None,  # w2_bias
        output1_scale_scalar,
        output1_scale_gate_scalar,
        output2_scale_scalar,
        num_experts,
        top_k,
        None,  # n_group
        None,  # topk_group
        intermediate_size,
        0,  # local_expert_offset
        num_experts,
        None,  # routed_scaling_factor
        routing_method_type.value,
        True,  # do_finalize
        enable_pdl,
        activation_type.value,  # act_type
        None,
    )[0].to(torch.float)

    mask = torch.isclose(output, reference_output, rtol=1e-3, atol=1e-3)

    # mismatch percentage
    mismatch_pct = (~mask).float().mean().item() * 100
    assert mismatch_pct < 6, f"Mismatch percentage is {mismatch_pct:.2f}"


@pytest.mark.parametrize("num_tokens", [1, 8, 1024])
@pytest.mark.parametrize("hidden_size", [1024, 2048, 3072, 4096])
@pytest.mark.parametrize("intermediate_size", [1024, 2048, 3072, 4096])
@pytest.mark.parametrize("num_experts", [128, 256])
@pytest.mark.parametrize("top_k", [4, 8])
@pytest.mark.parametrize(
    "routing_method_type",
    [
        RoutingMethodType.Renormalize,
        RoutingMethodType.RenormalizeNaive,
        RoutingMethodType.TopK,
    ],
)
@pytest.mark.parametrize("quant_mode", ["NvFP4xNvFP4", "MxFP4xMxFP8", "MxFP4xBf16"])
@pytest.mark.parametrize("routing_format", ["packed", "unpacked"])
def test_trtllm_gen_routed_fused_moe(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    num_experts: int,
    routing_method_type: RoutingMethodType,
    quant_mode: Literal["NvFP4xNvFP4", "MxFP4xMxFP8", "MxFP4xBf16"],
    routing_format: Literal["packed", "unpacked"],
):
    _run_trtllm_gen_routed_fused_moe_case(
        num_tokens,
        hidden_size,
        intermediate_size,
        top_k,
        num_experts,
        routing_method_type,
        quant_mode,
        routing_format,
    )


@pytest.mark.parametrize("num_tokens", [1, 8, 1024])
@pytest.mark.parametrize("hidden_size", [1024])
@pytest.mark.parametrize("intermediate_size", [1024])
@pytest.mark.parametrize("num_experts", [128, 256])
@pytest.mark.parametrize("top_k", [4])
@pytest.mark.parametrize(
    "routing_method_type",
    [
        RoutingMethodType.Renormalize,
        RoutingMethodType.RenormalizeNaive,
        RoutingMethodType.TopK,
    ],
)
@pytest.mark.parametrize("quant_mode", ["NvFP4xNvFP4", "MxFP4xMxFP8"])
@pytest.mark.parametrize("routing_format", ["packed", "unpacked"])
def test_trtllm_gen_routed_fused_moe_geglu(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    num_experts: int,
    routing_method_type: RoutingMethodType,
    quant_mode: Literal["NvFP4xNvFP4", "MxFP4xMxFP8"],
    routing_format: Literal["packed", "unpacked"],
):
    _run_trtllm_gen_routed_fused_moe_case(
        num_tokens,
        hidden_size,
        intermediate_size,
        top_k,
        num_experts,
        routing_method_type,
        quant_mode,
        routing_format,
        ActivationType.Geglu,
    )


def test_trtllm_gen_routed_fused_moe_unpacked_fp32():
    # All quantization modes share this finalize path.
    _run_trtllm_gen_routed_fused_moe_case(
        num_tokens=8,
        hidden_size=1024,
        intermediate_size=1024,
        top_k=4,
        num_experts=128,
        routing_method_type=RoutingMethodType.Renormalize,
        quant_mode="MxFP4xBf16",
        routing_format="unpacked_fp32",
    )


@pytest.mark.parametrize("num_tokens", [8, 64])
@pytest.mark.parametrize("hidden_size", [1024, 2048])
@pytest.mark.parametrize("intermediate_size", [1024, 2048])
@pytest.mark.parametrize("num_experts", [8, 16])
@pytest.mark.parametrize("top_k", [2, 4])
@pytest.mark.parametrize(
    "routing_method_type",
    [
        RoutingMethodType.Renormalize,
    ],
)
def test_trtllm_gen_fp8_routed_fused_moe(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    num_experts: int,
    routing_method_type: RoutingMethodType,
):
    """Test FP8 block scale routed MoE matches standard routing."""
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")
    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)

    # Generate random routing logits for reference
    routing_logits = torch.rand(num_tokens, num_experts, device=device).to(
        torch.bfloat16
    )

    # Generate random hidden states in FP8
    hidden_states_bf16 = (
        torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) * 0.1
    )
    hidden_states = hidden_states_bf16.to(torch.float8_e4m3fn)

    # Generate block scales for hidden states: [hidden_size // 128, num_tokens]
    hidden_states_scale = torch.ones(
        hidden_size // 128, num_tokens, device=device, dtype=torch.float32
    )

    # Generate FP8 weights
    gemm1_weights = torch.randn(
        num_experts, 2 * intermediate_size, hidden_size, device=device
    ).to(torch.float8_e4m3fn)
    gemm2_weights = torch.randn(
        num_experts, hidden_size, intermediate_size, device=device
    ).to(torch.float8_e4m3fn)

    # Generate block scales for weights
    gemm1_weights_scale = torch.ones(
        num_experts,
        2 * intermediate_size // 128,
        hidden_size // 128,
        device=device,
        dtype=torch.float32,
    )
    gemm2_weights_scale = torch.ones(
        num_experts,
        hidden_size // 128,
        intermediate_size // 128,
        device=device,
        dtype=torch.float32,
    )

    # Run reference with routing_logits
    reference_output = trtllm_fp8_block_scale_moe(
        routing_logits,
        None,  # routing_bias
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        num_experts,
        top_k,
        None,  # n_group
        None,  # topk_group
        intermediate_size,
        0,  # local_expert_offset
        num_experts,
        None,  # routed_scaling_factor
        routing_method_type.value,
        False,  # use_shuffled_weight
        0,  # weight_layout
        enable_pdl,
    ).to(torch.float)

    # Compute routing using reference implementation
    if routing_method_type == RoutingMethodType.Renormalize:
        permute_info, topk_weights_ref = routing_reference_renormalize(
            routing_logits, top_k, num_experts, 8
        )
    elif routing_method_type == RoutingMethodType.RenormalizeNaive:
        permute_info, topk_weights_ref = routing_reference_renormalize_naive(
            routing_logits, top_k, num_experts, 8
        )
    elif routing_method_type == RoutingMethodType.TopK:
        permute_info, topk_weights_ref = routing_reference_topk(
            routing_logits, top_k, num_experts, 8
        )
    topk_ids = permute_info["topKIndices"].to(torch.int32)
    topk_weights = topk_weights_ref.view(num_tokens, num_experts)[
        torch.arange(num_tokens, device=device).unsqueeze(1), topk_ids
    ].to(torch.bfloat16)

    # Pack topk_ids and topk_weights into single tensor
    # Format: (expert_id << 16) | (weight_bf16.view(int16))
    packed_topk_ids = (topk_ids << 16) | topk_weights.view(torch.int16).to(torch.int32)

    # Run with pre-computed routing (packed format)
    output = torch.empty(
        num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
    )
    trtllm_fp8_block_scale_routed_moe(
        topk_ids=packed_topk_ids,
        routing_bias=None,
        hidden_states=hidden_states,
        hidden_states_scale=hidden_states_scale,
        gemm1_weights=gemm1_weights,
        gemm1_weights_scale=gemm1_weights_scale,
        gemm2_weights=gemm2_weights,
        gemm2_weights_scale=gemm2_weights_scale,
        num_experts=num_experts,
        top_k=top_k,
        n_group=None,
        topk_group=None,
        intermediate_size=intermediate_size,
        local_expert_offset=0,
        local_num_experts=num_experts,
        routed_scaling_factor=None,
        routing_method_type=routing_method_type.value,
        use_shuffled_weight=False,
        weight_layout=0,
        enable_pdl=enable_pdl,
        output=output,
    )
    output = output.to(torch.float)

    mask = torch.isclose(output, reference_output, rtol=1e-2, atol=1e-2)

    # mismatch percentage
    mismatch_pct = (~mask).float().mean().item() * 100
    assert mismatch_pct < 10, f"Mismatch percentage is {mismatch_pct:.2f}%"


@pytest.mark.parametrize("num_tokens", [8, 64])
@pytest.mark.parametrize("hidden_size", [1024, 2048])
@pytest.mark.parametrize("intermediate_size", [1024, 2048])
@pytest.mark.parametrize("num_experts", [8, 16])
@pytest.mark.parametrize("top_k", [2, 4])
@pytest.mark.parametrize(
    "routing_method_type",
    [
        RoutingMethodType.Renormalize,
        RoutingMethodType.RenormalizeNaive,
    ],
)
def test_trtllm_gen_fp8_per_tensor_routed_fused_moe(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    num_experts: int,
    routing_method_type: RoutingMethodType,
):
    """Pre-routed FP8 per-tensor MoE matches the logits (self-routing) kernel.

    Feeds the packed (expert_id << 16 | weight) routing derived from the same
    logits to ``trtllm_fp8_per_tensor_scale_routed_moe`` and asserts it produces
    the same output as ``trtllm_fp8_per_tensor_scale_moe``. This is the numeric
    parity check for the pre-routed per-tensor path used by the vLLM modular
    (DP/EP/EPLB) kernel. The per-tensor scale scalars are shared verbatim by both
    calls, so the test isolates the routing path (weights need not be calibrated).
    """
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")
    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)

    # Random routing logits for the reference (self-routing) kernel.
    routing_logits = torch.rand(num_tokens, num_experts, device=device).to(
        torch.bfloat16
    )

    # Hidden states + weights in FP8 e4m3 (per-tensor: no block/hidden-states scale).
    hidden_states = (
        torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) * 0.1
    ).to(torch.float8_e4m3fn)
    gemm1_weights = (
        torch.randn(num_experts, 2 * intermediate_size, hidden_size, device=device)
        * 0.1
    ).to(torch.float8_e4m3fn)
    gemm2_weights = (
        torch.randn(num_experts, hidden_size, intermediate_size, device=device) * 0.1
    ).to(torch.float8_e4m3fn)

    # Per-expert per-tensor output scalars (folded activation/weight scales).
    # Shared verbatim by both kernels so only the routing path differs.
    output1_scales_scalar = torch.rand(num_experts, device=device) * 0.5 + 0.5
    output1_scales_gate_scalar = torch.rand(num_experts, device=device) * 0.5 + 0.5
    output2_scales_scalar = torch.rand(num_experts, device=device) * 0.5 + 0.5

    reference_output = trtllm_fp8_per_tensor_scale_moe(
        routing_logits,
        None,  # routing_bias
        hidden_states,
        gemm1_weights,
        output1_scales_scalar,
        output1_scales_gate_scalar,
        gemm2_weights,
        output2_scales_scalar,
        num_experts,
        top_k,
        None,  # n_group
        None,  # topk_group
        intermediate_size,
        0,  # local_expert_offset
        num_experts,
        None,  # routed_scaling_factor
        False,  # use_routing_scales_on_input
        routing_method_type.value,
    ).to(torch.float)

    # Compute the same routing with the reference implementation and pack it.
    if routing_method_type == RoutingMethodType.Renormalize:
        permute_info, topk_weights_ref = routing_reference_renormalize(
            routing_logits, top_k, num_experts, 8
        )
    else:  # RenormalizeNaive
        permute_info, topk_weights_ref = routing_reference_renormalize_naive(
            routing_logits, top_k, num_experts, 8
        )
    topk_ids = permute_info["topKIndices"].to(torch.int32)
    topk_weights = topk_weights_ref.view(num_tokens, num_experts)[
        torch.arange(num_tokens, device=device).unsqueeze(1), topk_ids
    ].to(torch.bfloat16)
    # Format: (expert_id << 16) | (weight_bf16.view(int16))
    packed_topk_ids = (topk_ids << 16) | topk_weights.view(torch.int16).to(torch.int32)

    output = torch.empty(
        num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
    )
    trtllm_fp8_per_tensor_scale_routed_moe(
        topk_ids=packed_topk_ids,
        routing_bias=None,
        hidden_states=hidden_states,
        gemm1_weights=gemm1_weights,
        output1_scales_scalar=output1_scales_scalar,
        output1_scales_gate_scalar=output1_scales_gate_scalar,
        gemm2_weights=gemm2_weights,
        output2_scales_scalar=output2_scales_scalar,
        num_experts=num_experts,
        top_k=top_k,
        n_group=None,
        topk_group=None,
        intermediate_size=intermediate_size,
        local_expert_offset=0,
        local_num_experts=num_experts,
        routed_scaling_factor=None,
        use_routing_scales_on_input=False,
        routing_method_type=routing_method_type.value,
        enable_pdl=enable_pdl,
        output=output,
    )
    output = output.to(torch.float)

    mask = torch.isclose(output, reference_output, rtol=1e-2, atol=1e-2)
    mismatch_pct = (~mask).float().mean().item() * 100
    assert mismatch_pct < 10, f"Mismatch percentage is {mismatch_pct:.2f}%"


def test_trtllm_gen_fp8_per_tensor_routed_fused_moe_nonzero_expert_offset():
    """Packed global expert IDs map to the correct local weight rows."""
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")
    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)
    num_tokens = 8
    hidden_size = 1024
    intermediate_size = 1024
    num_experts = 8
    local_num_experts = 4
    top_k = 2

    hidden_states = (
        torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) * 0.1
    ).to(torch.float8_e4m3fn)
    gemm1_weights = (
        torch.randn(
            local_num_experts, 2 * intermediate_size, hidden_size, device=device
        )
        * 0.1
    ).to(torch.float8_e4m3fn)
    gemm2_weights = (
        torch.randn(local_num_experts, hidden_size, intermediate_size, device=device)
        * 0.1
    ).to(torch.float8_e4m3fn)
    output1_scales_scalar = torch.rand(local_num_experts, device=device) + 0.5
    output1_scales_gate_scalar = torch.rand(local_num_experts, device=device) + 0.5
    output2_scales_scalar = torch.rand(local_num_experts, device=device) + 0.5

    local_topk_ids = (
        torch.arange(num_tokens * top_k, dtype=torch.int32, device=device)
        .reshape(num_tokens, top_k)
        .remainder(local_num_experts)
    )
    topk_weights = torch.softmax(
        torch.rand(num_tokens, top_k, device=device), dim=-1
    ).to(torch.bfloat16)

    def run(local_expert_offset: int) -> torch.Tensor:
        global_topk_ids = local_topk_ids + local_expert_offset
        packed_topk_ids = (global_topk_ids << 16) | topk_weights.view(torch.int16).to(
            torch.int32
        )
        output = torch.empty(
            num_tokens, hidden_size, dtype=torch.bfloat16, device=device
        )
        trtllm_fp8_per_tensor_scale_routed_moe(
            topk_ids=packed_topk_ids,
            routing_bias=None,
            hidden_states=hidden_states,
            gemm1_weights=gemm1_weights,
            output1_scales_scalar=output1_scales_scalar,
            output1_scales_gate_scalar=output1_scales_gate_scalar,
            gemm2_weights=gemm2_weights,
            output2_scales_scalar=output2_scales_scalar,
            num_experts=num_experts,
            top_k=top_k,
            n_group=None,
            topk_group=None,
            intermediate_size=intermediate_size,
            local_expert_offset=local_expert_offset,
            local_num_experts=local_num_experts,
            routed_scaling_factor=None,
            use_routing_scales_on_input=False,
            routing_method_type=RoutingMethodType.Renormalize.value,
            enable_pdl=enable_pdl,
            output=output,
        )
        return output.float()

    baseline = run(0)
    sharded = run(num_experts - local_num_experts)

    assert torch.count_nonzero(sharded)
    torch.testing.assert_close(sharded, baseline, rtol=1e-2, atol=1e-2)


@pytest.mark.parametrize("num_tokens", [8, 64])
@pytest.mark.parametrize("hidden_size", [1024, 2048])
@pytest.mark.parametrize("intermediate_size", [1024, 2048])
@pytest.mark.parametrize("num_experts", [8, 16])
@pytest.mark.parametrize("top_k", [2, 4])
@pytest.mark.parametrize(
    "routing_method_type",
    [
        RoutingMethodType.Renormalize,
    ],
)
def test_trtllm_gen_bf16_routed_fused_moe(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    num_experts: int,
    routing_method_type: RoutingMethodType,
):
    """Test Bf16 scale routed MoE matches standard routing."""
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")
    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)

    # Generate random routing logits for reference
    routing_logits = torch.rand(num_tokens, num_experts, device=device).to(
        torch.bfloat16
    )

    # Generate random hidden states in FP8
    hidden_states = (
        torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) * 0.1
    )

    # Generate weights
    gemm1_weights = torch.randn(
        num_experts, 2 * intermediate_size, hidden_size, device=device
    ).to(torch.bfloat16)
    gemm2_weights = torch.randn(
        num_experts, hidden_size, intermediate_size, device=device
    ).to(torch.bfloat16)

    gemm1_weights_shuffled = []
    gemm2_weights_shuffled = []
    for i in range(num_experts):
        tmp_weights1 = shuffle_matrix_a(gemm1_weights[i].view(torch.uint8), 64)
        tmp_weights2 = shuffle_matrix_a(gemm2_weights[i].view(torch.uint8), 64)
        block_k = 128
        gemm1_weights_shuffled.append(convert_to_block_layout(tmp_weights1, block_k))
        gemm2_weights_shuffled.append(convert_to_block_layout(tmp_weights2, block_k))
    gemm1_weights = torch.stack(gemm1_weights_shuffled).view(torch.bfloat16)
    gemm2_weights = torch.stack(gemm2_weights_shuffled).view(torch.bfloat16)

    # Run reference with routing_logits
    reference_output = trtllm_bf16_moe(
        routing_logits=routing_logits,
        routing_bias=None,
        hidden_states=hidden_states,
        gemm1_weights=gemm1_weights,
        gemm2_weights=gemm2_weights,
        num_experts=num_experts,
        top_k=top_k,
        n_group=None,
        topk_group=None,
        intermediate_size=intermediate_size,
        local_expert_offset=0,
        local_num_experts=num_experts,
        routed_scaling_factor=None,
        routing_method_type=routing_method_type.value,
        use_shuffled_weight=True,
        weight_layout=WeightLayout.BlockMajorK,
        do_finalize=True,
        enable_pdl=enable_pdl,
    ).to(torch.float)

    # Compute routing using reference implementation
    if routing_method_type == RoutingMethodType.Renormalize:
        permute_info, expert_weights_ref = routing_reference_renormalize(
            routing_logits, top_k, num_experts, 8
        )
    elif routing_method_type == RoutingMethodType.RenormalizeNaive:
        permute_info, expert_weights_ref = routing_reference_renormalize_naive(
            routing_logits, top_k, num_experts, 8
        )
    elif routing_method_type == RoutingMethodType.TopK:
        permute_info, expert_weights_ref = routing_reference_topk(
            routing_logits, top_k, num_experts, 8
        )
    topk_ids = permute_info["topKIndices"].to(torch.int32)
    expert_weights = expert_weights_ref.view(num_tokens, num_experts)[
        torch.arange(num_tokens, device=device).unsqueeze(1), topk_ids
    ].to(torch.bfloat16)

    # Pack topk_ids and expert_weights into single tensor
    # Format: (expert_id << 16) | (weight_bf16.view(int16))
    packed_topk_ids = (topk_ids << 16) | expert_weights.view(torch.int16).to(
        torch.int32
    )

    # Run with pre-computed routing (packed format)
    output = trtllm_bf16_routed_moe(
        topk_ids=packed_topk_ids,
        hidden_states=hidden_states,
        gemm1_weights=gemm1_weights,
        gemm2_weights=gemm2_weights,
        num_experts=num_experts,
        top_k=top_k,
        n_group=None,
        topk_group=None,
        intermediate_size=intermediate_size,
        local_expert_offset=0,
        local_num_experts=num_experts,
        routed_scaling_factor=None,
        routing_method_type=routing_method_type.value,
        use_shuffled_weight=True,
        weight_layout=WeightLayout.BlockMajorK,
        do_finalize=True,
        enable_pdl=enable_pdl,
    ).to(torch.float)

    mask = torch.isclose(output, reference_output, rtol=1e-2, atol=1e-2)

    # mismatch percentage
    mismatch_pct = (~mask).float().mean().item() * 100
    assert mismatch_pct < 10, f"Mismatch percentage is {mismatch_pct:.2f}%"


@pytest.mark.parametrize("num_tokens", [8, 64])
@pytest.mark.parametrize("hidden_size", [1024])
@pytest.mark.parametrize("intermediate_size", [1024])
@pytest.mark.parametrize("num_experts", [16])
@pytest.mark.parametrize("top_k", [2, 4])
@pytest.mark.parametrize(
    "routing_method_type",
    [
        RoutingMethodType.Renormalize,
    ],
)
def test_trtllm_gen_mxint4_routed_fused_moe(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    num_experts: int,
    routing_method_type: RoutingMethodType,
):
    """Verify that the routing-logits and pre-computed-routing flavors of
    trtllm_mxint4_block_scale_moe produce numerically equivalent outputs, with
    and without a LoRA delta. Mirrors the BF16 routed parity test above."""
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")
    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)

    # Random routing logits used to generate both the reference output (via
    # routing_logits flavor) and the precomputed-topk for the routed flavor.
    routing_logits = torch.rand(num_tokens, num_experts, device=device).to(
        torch.float32
    )

    hidden_states = (
        torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) * 0.1
    )

    # Random mxint4 weights & scales. We only need both API calls to see the
    # same buffers, so we don't bother quantizing real bf16 weights — the test
    # measures parity between the two entry points, not absolute correctness.
    g1_shape = (num_experts, 2 * intermediate_size, hidden_size // 2)
    g2_shape = (num_experts, hidden_size, intermediate_size // 2)
    gemm1_weights = torch.randint(0, 256, g1_shape, dtype=torch.uint8, device=device)
    gemm2_weights = torch.randint(0, 256, g2_shape, dtype=torch.uint8, device=device)
    g1_scale_shape = (num_experts, 2 * intermediate_size, hidden_size // 32)
    g2_scale_shape = (num_experts, hidden_size, intermediate_size // 32)
    gemm1_weights_scale = torch.randn(
        g1_scale_shape, dtype=torch.bfloat16, device=device
    )
    gemm2_weights_scale = torch.randn(
        g2_scale_shape, dtype=torch.bfloat16, device=device
    )

    common_kwargs = dict(
        hidden_states=hidden_states,
        gemm1_weights=gemm1_weights,
        gemm1_weights_scale=gemm1_weights_scale,
        gemm1_alpha=None,
        gemm1_beta=None,
        gemm1_clamp_limit=None,
        gemm2_weights=gemm2_weights,
        gemm2_weights_scale=gemm2_weights_scale,
        num_experts=num_experts,
        top_k=top_k,
        n_group=None,
        topk_group=None,
        intermediate_size=intermediate_size,
        local_expert_offset=0,
        local_num_experts=num_experts,
        routed_scaling_factor=None,
        routing_method_type=routing_method_type.value,
        do_finalize=True,
        enable_pdl=enable_pdl,
    )

    # Reference: routing-logits flavor.
    reference = trtllm_mxint4_block_scale_moe(
        routing_logits=routing_logits,
        routing_bias=None,
        **common_kwargs,
    )
    reference_output = reference[0].to(torch.float)

    # Compute the same routing decisions on the host so we can hand them to the
    # routed flavor as packed int32 (expert_id << 16) | weight_bf16.view(int16).
    permute_info, expert_weights_ref = routing_reference_renormalize(
        routing_logits, top_k, num_experts, 8
    )
    topk_ids = permute_info["topKIndices"].to(torch.int32)
    expert_weights = expert_weights_ref.view(num_tokens, num_experts)[
        torch.arange(num_tokens, device=device).unsqueeze(1), topk_ids
    ].to(torch.bfloat16)
    packed_topk_ids = (topk_ids << 16) | expert_weights.view(torch.int16).to(
        torch.int32
    )

    routed = trtllm_mxint4_block_scale_routed_moe(
        topk_ids=packed_topk_ids,
        **common_kwargs,
    )
    routed_output = routed[0].to(torch.float)

    mask = torch.isclose(routed_output, reference_output, rtol=1e-2, atol=1e-2)
    mismatch_pct = (~mask).float().mean().item() * 100
    assert mismatch_pct < 10, f"Mismatch percentage is {mismatch_pct:.2f}%"


@pytest.mark.parametrize(
    "activation_type",
    [
        pytest.param(ActivationType.Swiglu.value, id="Swiglu"),
        pytest.param(ActivationType.Geglu.value, id="Geglu"),
        pytest.param(ActivationType.Relu2.value, id="Relu2"),
    ],
)
def test_trtllm_gen_fp8_mxfp8_routed_activation_parity(activation_type: int):
    """MXFP8 routed path should match non-routed reference for gated and non-gated activations."""
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")

    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)

    num_tokens = 32
    hidden_size = 512
    intermediate_size = 512
    num_experts = 64
    top_k = 8
    routing_method_type = RoutingMethodType.Renormalize
    is_gated = activation_type in [
        ActivationType.Swiglu.value,
        ActivationType.Geglu.value,
    ]

    routing_logits = torch.randn((num_tokens, num_experts), device=device).to(
        torch.bfloat16
    )
    hidden_states = torch.randn((num_tokens, hidden_size), device=device).to(
        torch.bfloat16
    )
    gemm1_weights = torch.randn(
        (
            num_experts,
            (2 if is_gated else 1) * intermediate_size,
            hidden_size,
        ),
        device=device,
        dtype=torch.bfloat16,
    )
    gemm2_weights = torch.randn(
        (num_experts, hidden_size, intermediate_size),
        device=device,
        dtype=torch.bfloat16,
    )

    quant_impl = FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_MXFP8)
    quant_weights = quant_impl.quantize_weights(
        gemm1_weights, gemm2_weights, hidden_states
    )
    quant_inputs = quant_impl.quantize_inputs(hidden_states)
    epilogue_tile_m = 128
    w13_rows = (2 if is_gated else 1) * intermediate_size
    gemm1_weights_shuffled = []
    gemm1_scales_shuffled = []
    gemm2_weights_shuffled = []
    gemm2_scales_shuffled = []
    for i in range(num_experts):
        w1_interleaved = quant_weights["gemm1_weights"][i].clone().reshape(w13_rows, -1)
        s1_interleaved = quant_weights["gemm1_scales"][i].clone().reshape(w13_rows, -1)
        if is_gated:
            w1_interleaved = reorder_rows_for_gated_act_gemm(w1_interleaved)
            s1_interleaved = reorder_rows_for_gated_act_gemm(s1_interleaved)
        gemm1_weights_shuffled.append(
            shuffle_matrix_a(w1_interleaved.view(torch.uint8), epilogue_tile_m)
            .contiguous()
            .view(quant_weights["gemm1_weights"].dtype)
        )
        gemm2_weights_shuffled.append(
            shuffle_matrix_a(
                quant_weights["gemm2_weights"][i].view(torch.uint8), epilogue_tile_m
            )
            .contiguous()
            .view(quant_weights["gemm2_weights"].dtype)
        )
        gemm1_scales_shuffled.append(
            shuffle_matrix_sf_a(
                s1_interleaved.view(torch.uint8).reshape(w13_rows, -1),
                epilogue_tile_m,
            )
            .contiguous()
            .view(quant_weights["gemm1_scales"].dtype)
        )
        gemm2_scales_shuffled.append(
            shuffle_matrix_sf_a(
                quant_weights["gemm2_scales"][i]
                .view(torch.uint8)
                .reshape(hidden_size, -1),
                epilogue_tile_m,
            )
            .contiguous()
            .view(quant_weights["gemm2_scales"].dtype)
        )
    gemm1_weights_kernel = torch.stack(gemm1_weights_shuffled)
    gemm1_scales_kernel = torch.stack(gemm1_scales_shuffled)
    gemm2_weights_kernel = torch.stack(gemm2_weights_shuffled)
    gemm2_scales_kernel = torch.stack(gemm2_scales_shuffled)

    output_ref = trtllm_fp8_block_scale_moe(
        routing_logits=routing_logits,
        routing_bias=None,
        hidden_states=quant_inputs["hidden_states"],
        hidden_states_scale=quant_inputs["hidden_states_scale"],
        gemm1_weights=gemm1_weights_kernel,
        gemm1_weights_scale=gemm1_scales_kernel,
        gemm2_weights=gemm2_weights_kernel,
        gemm2_weights_scale=gemm2_scales_kernel,
        num_experts=num_experts,
        top_k=top_k,
        n_group=None,
        topk_group=None,
        intermediate_size=intermediate_size,
        local_expert_offset=0,
        local_num_experts=num_experts,
        routed_scaling_factor=None,
        routing_method_type=routing_method_type.value,
        use_shuffled_weight=True,
        weight_layout=WeightLayout.MajorK.value,
        enable_pdl=enable_pdl,
        fp8_quantization_type=Fp8QuantizationType.MxFp8,
        activation_type=activation_type,
    ).to(torch.float)

    permute_info, expert_weights_full = routing_reference_renormalize(
        routing_logits, top_k, num_experts, 8
    )
    topk_ids = permute_info["topKIndices"].to(torch.int32)
    expert_weights = expert_weights_full.view(num_tokens, num_experts)[
        torch.arange(num_tokens, device=device).unsqueeze(1), topk_ids
    ].to(torch.bfloat16)
    packed_topk_ids = (topk_ids << 16) | expert_weights.view(torch.int16).to(
        torch.int32
    )

    output_routed = trtllm_fp8_block_scale_routed_moe(
        topk_ids=packed_topk_ids,
        routing_bias=None,
        hidden_states=quant_inputs["hidden_states"],
        hidden_states_scale=quant_inputs["hidden_states_scale"],
        gemm1_weights=gemm1_weights_kernel,
        gemm1_weights_scale=gemm1_scales_kernel,
        gemm2_weights=gemm2_weights_kernel,
        gemm2_weights_scale=gemm2_scales_kernel,
        num_experts=num_experts,
        top_k=top_k,
        n_group=None,
        topk_group=None,
        intermediate_size=intermediate_size,
        local_expert_offset=0,
        local_num_experts=num_experts,
        routed_scaling_factor=None,
        routing_method_type=routing_method_type.value,
        use_shuffled_weight=True,
        weight_layout=WeightLayout.MajorK.value,
        enable_pdl=enable_pdl,
        fp8_quantization_type=Fp8QuantizationType.MxFp8,
        activation_type=activation_type,
    ).to(torch.float)

    close = torch.isclose(output_ref, output_routed, atol=1e-2, rtol=1e-2)
    mismatch_pct = (~close).float().mean().item() * 100
    assert mismatch_pct < 10, f"Mismatch percentage is {mismatch_pct:.2f}%"


@pytest.mark.parametrize("num_tokens", [1, 7, 32])
@pytest.mark.parametrize("hidden_size", [1024])
@pytest.mark.parametrize("intermediate_size", [1024, 2048])
@pytest.mark.parametrize("num_experts", [16])
@pytest.mark.parametrize("top_k", [2, 4])
def test_fp8_block_scale_moe_routing_replay(
    num_tokens: int,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    num_experts: int,
):
    """Test that routing_replay_out in trtllm_fp8_block_scale_moe records correct expert IDs.

    Uses DeepSeekV3 routing (the only routing method with replay support).
    Runs the full MoE kernel twice with the same inputs: once with routing_replay_out
    and once without. Verifies that:
    1. The MoE output is identical (replay has no side effects).
    2. The replay buffer matches the reference routing result (sorted set equality).
    3. Tail rows beyond num_tokens remain sentinel (CUDA graph pre-alloc contract).
    """
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")
    n_group = 4
    topk_group = 2
    if topk_group * n_group < top_k or topk_group > n_group:
        pytest.skip("Invalid DeepSeek routing configuration")
    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)

    routing_logits = torch.rand(
        num_tokens, num_experts, device=device, dtype=torch.float32
    )
    routing_bias = torch.randn(num_experts, device=device, dtype=torch.bfloat16)

    hidden_states_bf16 = (
        torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) * 0.1
    )
    hidden_states = hidden_states_bf16.to(torch.float8_e4m3fn)

    hidden_states_scale = torch.ones(
        hidden_size // 128, num_tokens, device=device, dtype=torch.float32
    )

    gemm1_weights = torch.randn(
        num_experts, 2 * intermediate_size, hidden_size, device=device
    ).to(torch.float8_e4m3fn)
    gemm2_weights = torch.randn(
        num_experts, hidden_size, intermediate_size, device=device
    ).to(torch.float8_e4m3fn)

    gemm1_weights_scale = torch.ones(
        num_experts,
        2 * intermediate_size // 128,
        hidden_size // 128,
        device=device,
        dtype=torch.float32,
    )
    gemm2_weights_scale = torch.ones(
        num_experts,
        hidden_size // 128,
        intermediate_size // 128,
        device=device,
        dtype=torch.float32,
    )

    # Allocate oversized buffer to validate CUDA graph pre-allocation contract:
    # kernel should only write to [0, num_tokens) and leave tail rows as sentinel.
    replay_capacity = num_tokens + 5
    routing_replay_out = torch.full(
        (replay_capacity, top_k), -1, device=device, dtype=torch.int16
    )

    output_with_replay = trtllm_fp8_block_scale_moe(
        routing_logits,
        routing_bias,
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        num_experts,
        top_k,
        n_group,
        topk_group,
        intermediate_size,
        0,  # local_expert_offset
        num_experts,
        1.0,  # routed_scaling_factor
        RoutingMethodType.DeepSeekV3.value,
        False,  # use_shuffled_weight
        0,  # weight_layout
        enable_pdl,
        routing_replay_out=routing_replay_out,
    )

    output_without_replay = trtllm_fp8_block_scale_moe(
        routing_logits,
        routing_bias,
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        num_experts,
        top_k,
        n_group,
        topk_group,
        intermediate_size,
        0,
        num_experts,
        1.0,
        RoutingMethodType.DeepSeekV3.value,
        False,
        0,
        enable_pdl,
        routing_replay_out=None,
    )

    # MoE output should be identical regardless of replay
    torch.testing.assert_close(
        output_with_replay.to(torch.float),
        output_without_replay.to(torch.float),
        rtol=0,
        atol=0,
    )

    # Compare replay against reference routing — verify active rows only
    active_replay = routing_replay_out[:num_tokens]
    # Verify replay IDs are valid expert indices
    assert (active_replay >= 0).all() and (active_replay < num_experts).all(), (
        "Replay contains out-of-range expert IDs"
    )
    # Each token should have top_k unique experts
    for t in range(num_tokens):
        unique_experts = active_replay[t].unique()
        assert unique_experts.numel() == top_k, (
            f"Token {t}: expected {top_k} unique experts, got {unique_experts.numel()}"
        )

    # Tail rows beyond num_tokens should remain sentinel (-1)
    assert (routing_replay_out[num_tokens:] == -1).all(), (
        "Kernel should not write beyond active token rows"
    )


# Each (num_tokens, num_experts) entry is chosen to land in exactly one of the
# five top-K kernels in `routing_custom.cu`. See the dispatch logic comment in
# `trtllm_fused_moe_routing_custom.cu` (around the `useSplitTopKPath` block):
#
#   * BlockKernel             : num_tokens <= 4
#   * DynBlockKernel          : 5 <= num_tokens <= 16  (num_experts <= 512)
#   * ClusterKernel           : 17 <= num_tokens <= 256, num_experts < 160 (SM90+)
#   * BlockScoresKernel       : 17 <= num_tokens <= 256, num_experts >= 160
#                               (split-topK path; also requires policy opt-in)
#   * HistogramScoresKernel   : num_tokens > 256  (warp-per-token fallback)
_REPLAY_KERNEL_TIERS = [
    (4, 16, "BlockKernel"),
    (16, 16, "DynBlockKernel"),
    (32, 16, "ClusterKernel"),
    (32, 256, "BlockScoresKernel"),
    (512, 16, "HistogramScoresKernel"),
]


@pytest.mark.parametrize(
    "num_tokens, num_experts, kernel_tier",
    _REPLAY_KERNEL_TIERS,
    ids=[t[2] for t in _REPLAY_KERNEL_TIERS],
)
@pytest.mark.parametrize("hidden_size", [1024])
@pytest.mark.parametrize("intermediate_size", [1024])
@pytest.mark.parametrize("top_k", [2, 4])
@pytest.mark.parametrize(
    "routing_method_type",
    [
        RoutingMethodType.Renormalize,
        RoutingMethodType.RenormalizeNaive,
        RoutingMethodType.TopK,
    ],
)
def test_fp8_block_scale_moe_routing_replay_custom_routing(
    num_tokens: int,
    num_experts: int,
    kernel_tier: str,
    hidden_size: int,
    intermediate_size: int,
    top_k: int,
    routing_method_type: RoutingMethodType,
):
    """Test that ``routing_replay_out`` in ``trtllm_fp8_block_scale_moe`` records
    correct expert IDs for the routing methods that dispatch through
    ``routingCustom`` (Renormalize / RenormalizeNaive / TopK).

    Companion to ``test_fp8_block_scale_moe_routing_replay`` (which only
    exercises the DeepSeekV3 routing kernel). The parametrize grid is chosen
    so that each ``(num_tokens, num_experts)`` row lands in exactly one of
    the five custom-routing top-K kernels — together they cover every kernel
    that writes ``mPtrRoutingReplayOut`` on the non-DSV3 path.

    Mirrors the DSV3 test's verification structure:
      1. MoE output is identical with/without replay (replay is a side
         channel with no observable effect on the GEMM result).
      2. Replay buffer contains valid expert IDs in ``[0, num_experts)``.
      3. Each active row picks ``top_k`` distinct experts (TopK over raw
         scores cannot select the same expert twice within one token).
      4. Tail rows beyond ``num_tokens`` remain at the sentinel ``-1``
         (CUDA-graph pre-allocation contract).
    """
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if compute_capability[0] not in [10]:
        pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.")
    torch.manual_seed(42)
    device = torch.device("cuda:0")
    enable_pdl = device_support_pdl(device)

    routing_logits = torch.rand(
        num_tokens, num_experts, device=device, dtype=torch.float32
    )

    hidden_states_bf16 = (
        torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) * 0.1
    )
    hidden_states = hidden_states_bf16.to(torch.float8_e4m3fn)

    hidden_states_scale = torch.ones(
        hidden_size // 128, num_tokens, device=device, dtype=torch.float32
    )

    gemm1_weights = torch.randn(
        num_experts, 2 * intermediate_size, hidden_size, device=device
    ).to(torch.float8_e4m3fn)
    gemm2_weights = torch.randn(
        num_experts, hidden_size, intermediate_size, device=device
    ).to(torch.float8_e4m3fn)

    gemm1_weights_scale = torch.ones(
        num_experts,
        2 * intermediate_size // 128,
        hidden_size // 128,
        device=device,
        dtype=torch.float32,
    )
    gemm2_weights_scale = torch.ones(
        num_experts,
        hidden_size // 128,
        intermediate_size // 128,
        device=device,
        dtype=torch.float32,
    )

    # Oversized buffer + sentinel: validates the CUDA-graph pre-allocation
    # contract that the kernel writes only the [0, num_tokens) prefix.
    replay_capacity = num_tokens + 5
    routing_replay_out = torch.full(
        (replay_capacity, top_k), -1, device=device, dtype=torch.int16
    )

    output_with_replay = trtllm_fp8_block_scale_moe(
        routing_logits,
        None,  # routing_bias (DSV3-only)
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        num_experts,
        top_k,
        None,  # n_group (DSV3-only)
        None,  # topk_group (DSV3-only)
        intermediate_size,
        0,  # local_expert_offset
        num_experts,
        None,  # routed_scaling_factor
        routing_method_type.value,
        False,  # use_shuffled_weight
        0,  # weight_layout
        enable_pdl,
        routing_replay_out=routing_replay_out,
    )

    output_without_replay = trtllm_fp8_block_scale_moe(
        routing_logits,
        None,
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        num_experts,
        top_k,
        None,
        None,
        intermediate_size,
        0,
        num_experts,
        None,
        routing_method_type.value,
        False,
        0,
        enable_pdl,
        routing_replay_out=None,
    )

    # 1. MoE output is unaffected by the replay side channel.
    torch.testing.assert_close(
        output_with_replay.to(torch.float),
        output_without_replay.to(torch.float),
        rtol=0,
        atol=0,
    )

    active_replay = routing_replay_out[:num_tokens]

    # 2. All recorded IDs are valid expert indices.
    assert (active_replay >= 0).all() and (active_replay < num_experts).all(), (
        f"Replay contains out-of-range expert IDs (kernel={kernel_tier}, "
        f"routing={routing_method_type.name}): "
        f"min={active_replay.min().item()}, max={active_replay.max().item()}"
    )

    # 3. Each token's top_k selections are distinct.
    for t in range(num_tokens):
        unique_experts = active_replay[t].unique()
        assert unique_experts.numel() == top_k, (
            f"Token {t}: expected {top_k} unique experts, got "
            f"{unique_experts.numel()} ({active_replay[t].tolist()}) "
            f"(kernel={kernel_tier}, routing={routing_method_type.name})"
        )

    # 4. Tail rows untouched.
    assert (routing_replay_out[num_tokens:] == -1).all(), (
        f"Kernel wrote beyond active token rows "
        f"(kernel={kernel_tier}, routing={routing_method_type.name})"
    )


@pytest.mark.parametrize(
    "alpha_value,beta_value,clamp_value",
    [
        pytest.param(4.0, 25.0, None, id="Situ_Alpha4Beta25"),
        pytest.param(1.7, 1.0, 7.0, id="Situ_Alpha1p7Beta1Clamp7"),
    ],
)
def test_situ_mxfp4_mxfp8_logits_match_pre_routed(alpha_value, beta_value, clamp_value):
    num_experts = 8
    device = torch.device("cuda:0")
    _run_trtllm_gen_routed_fused_moe_case(
        num_tokens=32,
        hidden_size=1024,
        intermediate_size=512,
        top_k=2,
        num_experts=num_experts,
        routing_method_type=RoutingMethodType.Renormalize,
        quant_mode="MxFP4xMxFP8",
        routing_format="unpacked",
        activation_type=ActivationType.Situ,
        gemm1_alpha=torch.full(
            (num_experts,), alpha_value, device=device, dtype=torch.float32
        ),
        gemm1_beta=torch.full(
            (num_experts,), beta_value, device=device, dtype=torch.float32
        ),
        gemm1_clamp_limit=(
            None
            if clamp_value is None
            else torch.full(
                (num_experts,), clamp_value, device=device, dtype=torch.float32
            )
        ),
    )
