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

"""TraceTemplates for Mixture-of-Experts operations."""

import math
import inspect

import torch

from ...tllm_enums import (
    ActivationType,
    DEFAULT_SWIGLU_ALPHA,
    DEFAULT_SWIGLU_BETA,
    DEFAULT_SWIGLU_LIMIT,
    normalize_activation_type,
)
from ..template import Const, Scalar, Tensor, TraceTemplate, Var
from ._init_helpers import fp8_block_quant_1d, fp8_block_quant_2d
from .quantize import _fp4_quantize_reference

# ---------------------------------------------------------------------------
# SM90 mixed-input weight preparation
# ---------------------------------------------------------------------------

_SM90_MIXED_WEIGHT_PREP_AXES: dict[str, Var | Const] = {
    "num_experts": Var(description="Number of local experts."),
    "rows": Var(description="Logical weight rows per expert."),
    "packed_k": Var(description="Packed 4-bit K bytes per row."),
    "scale_groups": Var(description="Weight-scale groups per row."),
}

_SM90_MIXED_WEIGHT_PREP_INPUTS: dict[str, Tensor | Scalar] = {
    "weight": Tensor(
        ["num_experts", "rows", "packed_k"],
        description="Packed 4-bit expert weights.",
    ),
    "raw_scale": Tensor(
        ["num_experts", "rows", "scale_groups"],
        description="Logical E8M0 weight scales.",
    ),
    "max_range": Scalar("int32", description="Maximum fused E8M0 exponent range."),
    "interleave": Scalar(
        "bool", description="Whether to emit the SM90 physical layout."
    ),
}

sm90_mixed_gemm_humming_weight_preprocess_trace = TraceTemplate(
    op_type="moe_preprocess",
    name_prefix="sm90_mixed_gemm_humming_weight_preprocess",
    description="Prepare Humming-style MXFP4 weights and scales for SM90 mixed-input MoE.",
    axes={
        **_SM90_MIXED_WEIGHT_PREP_AXES,
        "m64_blocks": Var(description="Folded 64-row scale blocks."),
        "k128_blocks": Var(description="Folded K128 scale blocks."),
        "folded_m": Var(description="Physical folded-M scale extent."),
        "physical_cols": Var(description="Physical scale columns per folded block."),
    },
    inputs=_SM90_MIXED_WEIGHT_PREP_INPUTS,
    outputs={
        "processed_weight": Tensor(
            ["num_experts", "rows", "packed_k"], dtype_from="weight"
        ),
        "processed_scale": Tensor(
            ["num_experts", "m64_blocks", "k128_blocks", "folded_m", "physical_cols"],
            dtype_from="raw_scale",
        ),
        "residual": Tensor(["num_experts"], dtype="float32"),
    },
    constraints=["packed_k == scale_groups * 16"],
    tags=["moe:sm90", "quantization:mxfp4"],
)

sm90_mixed_gemm_humming_weight_preprocess_logical_trace = TraceTemplate(
    op_type="moe_preprocess",
    name_prefix="sm90_mixed_gemm_humming_weight_preprocess_logical",
    description="Prepare logical Humming-style MXFP4 payloads and exponent offsets.",
    axes=_SM90_MIXED_WEIGHT_PREP_AXES,
    inputs=_SM90_MIXED_WEIGHT_PREP_INPUTS,
    outputs={
        "processed_weight": Tensor(
            ["num_experts", "rows", "packed_k"], dtype_from="weight"
        ),
        "processed_scale": Tensor(
            ["num_experts", "rows", "scale_groups"], dtype_from="raw_scale"
        ),
        "residual": Tensor(["num_experts"], dtype="float32"),
    },
    constraints=["packed_k == scale_groups * 16"],
    tags=["moe:sm90", "quantization:mxfp4"],
)


def sm90_mixed_gemm_humming_weight_preprocess_trace_dispatch(**kwargs):
    return (
        sm90_mixed_gemm_humming_weight_preprocess_trace
        if kwargs.get("interleave", True)
        else sm90_mixed_gemm_humming_weight_preprocess_logical_trace
    )


sm90_mixed_gemm_humming_weight_preprocess_trace_dispatch.templates = [  # type: ignore[attr-defined]
    sm90_mixed_gemm_humming_weight_preprocess_trace,
    sm90_mixed_gemm_humming_weight_preprocess_logical_trace,
]

sm90_mixed_gemm_scale_interleave_trace = TraceTemplate(
    op_type="moe_preprocess",
    name_prefix="sm90_mixed_gemm_scale_interleave",
    description="Fold logical weight scales into the SM90 mixed-input layout.",
    axes={
        "num_experts": Var(description="Number of local experts."),
        "rows": Var(description="Logical scale rows per expert."),
        "scale_groups": Var(description="Logical scale groups per row."),
        "m64_blocks": Var(description="Folded 64-row scale blocks."),
        "k128_blocks": Var(description="Folded K128 scale blocks."),
        "folded_m": Var(description="Physical folded-M scale extent."),
        "physical_cols": Var(description="Physical scale columns per folded block."),
    },
    inputs={
        "scales": Tensor(["num_experts", "rows", "scale_groups"]),
        "group_size": Scalar("int32", description="Weight quantization group size."),
    },
    outputs={
        "interleaved_scales": Tensor(
            ["num_experts", "m64_blocks", "k128_blocks", "folded_m", "physical_cols"],
            dtype_from="scales",
        )
    },
    tags=["moe:sm90", "layout:folded_scale"],
)

sm90_mixed_gemm_weight_interleave_trace = TraceTemplate(
    op_type="moe_preprocess",
    name_prefix="sm90_mixed_gemm_weight_interleave",
    description="Interleave packed 4-bit weights for the SM90 mixed-input MoE GEMM.",
    axes={
        "num_experts": Var(description="Number of local experts."),
        "rows": Var(description="Logical weight rows per expert."),
        "packed_k": Var(description="Packed 4-bit K bytes per row."),
    },
    inputs={
        "weight": Tensor(["num_experts", "rows", "packed_k"]),
        "quant_type": Scalar("str", description="4-bit payload/interleave variant."),
    },
    outputs={
        "interleaved_weight": Tensor(
            ["num_experts", "rows", "packed_k"], dtype_from="weight"
        )
    },
    tags=["moe:sm90", "layout:mixed_input"],
)

# ---------------------------------------------------------------------------
# Shared GEMM computation helper
# ---------------------------------------------------------------------------


@torch.no_grad()
def _fp8_moe_run_experts(
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    weights,
    topk_idx,
    local_expert_offset,
    E_global,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
):
    """FP8 block-scale dequantization + SwiGLU + GEMM for all routing types.

    ``weights``   : [T, TOP_K] float32 — per-token expert weights (already normalised)
    ``topk_idx``  : [T, TOP_K] int64   — selected global expert indices
    """
    T, H = hidden_states.shape
    E_local, gemm1_out_size, _ = gemm1_weights.shape
    I = gemm1_out_size // 2
    BLOCK = 128
    if gemm1_out_size != 2 * I:
        raise ValueError(
            f"gemm1_weights.shape[1]={gemm1_out_size} is not 2*intermediate_size; "
            "SwiGLU requires gemm1_out_size == 2 * intermediate_size."
        )
    device = hidden_states.device

    A_fp32 = hidden_states.to(torch.float32)
    A_scale = hidden_states_scale.to(torch.float32)  # [H/128, T]
    A_scale_TH = A_scale.permute(1, 0).contiguous()  # [T, H/128]
    A_scale_expanded = (
        A_scale_TH.unsqueeze(-1).repeat(1, 1, BLOCK).reshape(T, H).contiguous()
    )
    A = A_fp32 * A_scale_expanded

    W13_fp32 = gemm1_weights.to(torch.float32)
    S13 = gemm1_weights_scale.to(torch.float32)
    S13_expanded = torch.repeat_interleave(S13, BLOCK, dim=1)
    S13_expanded = torch.repeat_interleave(S13_expanded, BLOCK, dim=2)
    W13 = W13_fp32 * S13_expanded

    W2_fp32 = gemm2_weights.to(torch.float32)
    S2 = gemm2_weights_scale.to(torch.float32)
    S2_expanded = torch.repeat_interleave(S2, BLOCK, dim=1)
    S2_expanded = torch.repeat_interleave(S2_expanded, BLOCK, dim=2)
    W2 = W2_fp32 * S2_expanded

    output = torch.zeros((T, H), dtype=torch.float32, device=device)
    local_start = int(local_expert_offset)

    for le in range(E_local):
        ge = local_start + le
        if ge < 0 or ge >= E_global:
            continue
        # tokens that selected this expert
        sel_mask = (topk_idx == ge).any(dim=1)
        if not sel_mask.any():
            continue
        token_idx = torch.nonzero(sel_mask, as_tuple=False).squeeze(1)
        A_e = A.index_select(0, token_idx)
        G1 = A_e.matmul(W13[le].t())
        X1, X2 = G1[:, :I], G1[:, I:]
        if gemm1_clamp_limit is not None:
            limit = gemm1_clamp_limit[le].to(device=X1.device, dtype=torch.float32)
            X1 = torch.clamp(X1, min=-limit, max=limit)
            X2 = torch.clamp(X2, max=limit)
        alpha = (
            1.0
            if gemm1_alpha is None
            else gemm1_alpha[le].to(device=X2.device, dtype=torch.float32)
        )
        beta = (
            0.0
            if gemm1_beta is None
            else gemm1_beta[le].to(device=X1.device, dtype=torch.float32)
        )
        silu_X2 = X2 * torch.sigmoid(alpha * X2)
        O = (silu_X2 * (X1 + beta)).matmul(W2[le].t())
        # per-expert contribution weight for each token
        w_tok = weights.index_select(0, token_idx)
        # find which slot in topk_idx[token_idx] corresponds to ge
        match = (topk_idx.index_select(0, token_idx) == ge).float()
        w_e = (w_tok * match).sum(dim=1)
        output.index_add_(0, token_idx, O * w_e.unsqueeze(1))

    return output.to(torch.bfloat16)


# ---------------------------------------------------------------------------
# Per-routing-type reference implementations
# ---------------------------------------------------------------------------


@torch.no_grad()
def _trtllm_fp8_block_scale_moe_ds_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    top_k,
    n_group,
    topk_group,
    local_expert_offset,
    routed_scaling_factor,
    *,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
):
    """
    FP8 block-scale MoE with DeepSeek-V3 routing:
        s = sigmoid(logits)
        s_with_bias = s + bias
        group by n_group; per group take top-2 sum → pick topk_group groups
        on the kept groups, take global top_k experts
        combine with weights derived from s (without bias), normalised and
        scaled by routed_scaling_factor
    """
    E_global = routing_logits.shape[1]
    T = routing_logits.shape[0]
    TOP_K = int(top_k)
    N_GROUP = int(n_group)
    TOPK_GROUP = int(topk_group)

    logits = routing_logits.to(torch.float32)
    bias = routing_bias.to(torch.float32).reshape(-1)

    s = 1.0 / (1.0 + torch.exp(-logits))
    s_with_bias = s + bias

    group_size = E_global // N_GROUP
    s_wb_grouped = s_with_bias.view(T, N_GROUP, group_size)
    top2_vals, _ = torch.topk(s_wb_grouped, k=2, dim=2, largest=True, sorted=False)
    group_scores = top2_vals.sum(dim=2)

    _, group_idx = torch.topk(
        group_scores, k=TOPK_GROUP, dim=1, largest=True, sorted=False
    )
    group_mask = torch.zeros_like(group_scores)
    group_mask.scatter_(1, group_idx, 1.0)
    score_mask = (
        group_mask.unsqueeze(2).expand(T, N_GROUP, group_size).reshape(T, E_global)
    )

    neg_inf = torch.finfo(torch.float32).min
    scores_pruned = s_with_bias.masked_fill(score_mask == 0, neg_inf)
    _, topk_idx = torch.topk(scores_pruned, k=TOP_K, dim=1, largest=True, sorted=False)

    M = torch.zeros_like(s)
    M.scatter_(1, topk_idx, 1.0)
    raw_w = s * M
    weights_sum = raw_w.sum(dim=1, keepdim=True) + 1e-20
    weights = (raw_w / weights_sum) * routed_scaling_factor

    # Gather per-row weights into [T, TOP_K] for the shared GEMM helper
    w_topk = weights.gather(1, topk_idx)

    return _fp8_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        w_topk,
        topk_idx,
        local_expert_offset,
        E_global,
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_fp8_block_scale_moe_default_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    *,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
):
    """
    FP8 block-scale MoE with Default routing: Softmax → TopK.
    routing_bias is added to logits before softmax when provided.
    """
    TOP_K = int(top_k)
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    s = torch.softmax(logits, dim=-1)
    _, topk_idx = torch.topk(s, k=TOP_K, dim=1, largest=True, sorted=False)
    weights = s.gather(1, topk_idx) * routed_scaling_factor
    return _fp8_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        weights,
        topk_idx,
        local_expert_offset,
        E_global,
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_fp8_block_scale_moe_renormalize_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    *,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
):
    """
    FP8 block-scale MoE with Renormalize routing: TopK → Softmax.
    TopK is applied on raw logits; weights are then derived by softmax
    over the selected logits.
    """
    TOP_K = int(top_k)
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    _, topk_idx = torch.topk(logits, k=TOP_K, dim=1, largest=True, sorted=False)
    gathered = logits.gather(1, topk_idx)
    weights = torch.softmax(gathered, dim=-1) * routed_scaling_factor
    return _fp8_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        weights,
        topk_idx,
        local_expert_offset,
        E_global,
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_fp8_block_scale_moe_llama4_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    *,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
):
    """
    FP8 block-scale MoE with Llama4 routing: Top1 → Sigmoid.
    Single expert selected per token; weight derived from sigmoid of its logit.
    By definition Llama4 routing uses top_k=1; the parameter is accepted for
    schema consistency with the other routing methods.
    """
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    topk_idx = logits.argmax(dim=-1, keepdim=True)  # [T, 1]
    top1_logit = logits.gather(1, topk_idx)
    weights = (1.0 / (1.0 + torch.exp(-top1_logit))) * routed_scaling_factor
    return _fp8_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        weights,
        topk_idx,
        local_expert_offset,
        E_global,
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_fp8_block_scale_moe_renormalize_naive_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    *,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
):
    """
    FP8 block-scale MoE with RenormalizeNaive routing: Softmax → TopK → Renormalize.
    Same as Default but the selected weights are re-normalised to sum to 1.
    """
    TOP_K = int(top_k)
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    s = torch.softmax(logits, dim=-1)
    _, topk_idx = torch.topk(s, k=TOP_K, dim=1, largest=True, sorted=False)
    gathered = s.gather(1, topk_idx)
    weights = gathered / (gathered.sum(dim=1, keepdim=True) + 1e-20)
    weights = weights * routed_scaling_factor
    return _fp8_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        weights,
        topk_idx,
        local_expert_offset,
        E_global,
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_fp8_block_scale_moe_topk_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    *,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
):
    """
    FP8 block-scale MoE with TopK-only routing: TopK, uniform weights.
    No softmax or sigmoid; all selected experts receive equal weight.
    """
    TOP_K = int(top_k)
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    _, topk_idx = torch.topk(logits, k=TOP_K, dim=1, largest=True, sorted=False)
    T = logits.shape[0]
    weights = torch.full(
        (T, TOP_K),
        routed_scaling_factor / TOP_K,
        dtype=torch.float32,
        device=logits.device,
    )
    return _fp8_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        weights,
        topk_idx,
        local_expert_offset,
        E_global,
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


# ---------------------------------------------------------------------------
# Template factory: shared axes/inputs/outputs for all routing types
# ---------------------------------------------------------------------------

_STANDARD_AXES = {
    "seq_len": Var(description="Sequence length (number of tokens)"),
    "num_experts": Const(description="Total number of experts.", abbrev=""),
    "top_k": Const(
        description="Number of experts to route to per token.", abbrev="topk"
    ),
    "num_local_experts": Const(description="Number of local experts.", abbrev="e"),
    "hidden_size": Const(description="Hidden dimension size.", abbrev="h"),
    "intermediate_size": Const(description="MoE intermediate layer size.", abbrev="i"),
    "gemm1_out_size": Const(
        description="Output size of the first GEMM (W13). Should be 2 * intermediate_size.",
        abbrev="",
    ),
    "num_hidden_blocks": Const(
        description="Number of quantized blocks along the hidden_size dimension (block_size=128).",
        abbrev="",
    ),
    "num_intermediate_blocks": Const(
        description="Number of quantized blocks along the intermediate_size dimension (block_size=128).",
        abbrev="",
    ),
    "num_gemm1_out_blocks": Const(
        description="Number of quantized blocks along the gemm1_out_size dimension (block_size=128).",
        abbrev="",
    ),
}

_STANDARD_INPUTS = {
    "routing_logits": Tensor(
        ["seq_len", "num_experts"],
        description="Routing logits for expert selection.",
    ),
    "routing_bias": Tensor(
        ["num_experts"],
        description="Bias added to logits before routing. Pass None for no bias.",
        optional=True,
    ),
    "hidden_states": Tensor(
        ["seq_len", "hidden_size"],
        description="Input hidden states tensor (FP8 quantized).",
    ),
    "hidden_states_scale": Tensor(
        ["num_hidden_blocks", "seq_len"],
        description="Block-wise scaling factors for hidden states.",
    ),
    "gemm1_weights": Tensor(
        ["num_local_experts", "gemm1_out_size", "hidden_size"],
        description="First GEMM weights for all local experts (gate and up projections).",
    ),
    "gemm1_weights_scale": Tensor(
        ["num_local_experts", "num_gemm1_out_blocks", "num_hidden_blocks"],
        description="Block-wise scaling factors for first GEMM weights.",
    ),
    "gemm1_alpha": Tensor(
        ["num_local_experts"],
        dtype="float32",
        description="Optional MxFp8-only per-expert SwiGLU OA alpha.",
        optional=True,
    ),
    "gemm1_beta": Tensor(
        ["num_local_experts"],
        dtype="float32",
        description="Optional MxFp8-only per-expert SwiGLU OA beta.",
        optional=True,
    ),
    "gemm1_clamp_limit": Tensor(
        ["num_local_experts"],
        dtype="float32",
        description="Optional MxFp8-only per-expert SwiGLU OA clamp limit.",
        optional=True,
    ),
    "gemm2_weights": Tensor(
        ["num_local_experts", "hidden_size", "intermediate_size"],
        description="Second GEMM weights for all local experts (down projection).",
    ),
    "gemm2_weights_scale": Tensor(
        ["num_local_experts", "num_hidden_blocks", "num_intermediate_blocks"],
        description="Block-wise scaling factors for second GEMM weights.",
    ),
    "top_k": Scalar(
        "int32",
        description="Number of experts to route to per token.",
    ),
    "local_expert_offset": Scalar(
        "int32",
        description="Offset of local experts in global expert space.",
    ),
    "routed_scaling_factor": Scalar(
        "float32",
        description="Scaling factor applied to routing weights.",
    ),
}

_STANDARD_OUTPUTS = {
    "output": Tensor(
        ["seq_len", "hidden_size"],
        dtype="bfloat16",
        description="Final MoE output tensor.",
    ),
}

_STANDARD_TAGS = ["status:verified", "quantization:float8_e4m3fn"]


def _moe_fp8_block_scale_standard_init(
    *,
    seq_len: int,
    routing_method_type: int = 0,
    num_experts: int = 256,
    top_k: int = 8,
    num_local_experts: int = 32,
    hidden_size: int = 7168,
    intermediate_size: int = 2048,
    gemm1_out_size: int = 0,
    num_hidden_blocks: int = 0,
    num_intermediate_blocks: int = 0,
    num_gemm1_out_blocks: int = 0,
    device: str = "cuda",
    seed: int = 0,
):
    """Build inputs for ``trtllm_fp8_block_scale_moe`` (non-DS routing).

    Sourced from ``tests/moe/test_dpsk_fused_moe_fp8.py`` (the
    ``_make_inputs_*`` fixture): activations and weights are sampled
    from ``randn`` (with a ``2.0×`` scale on activations matching the
    test) and quantized via the ``_fp8_block_quant_*`` helpers
    (block=128). ``hidden_states_scale`` is transposed to ``[H/128, T]``
    to match the kernel's expected layout.
    """
    del gemm1_out_size, num_hidden_blocks, num_intermediate_blocks
    del num_gemm1_out_blocks
    torch.manual_seed(seed)
    BS = 128
    routing_logits = torch.randn(
        seq_len, num_experts, dtype=torch.float32, device=device
    )
    routing_bias = torch.zeros(num_experts, dtype=torch.bfloat16, device=device)

    a_bf16 = 2.0 * torch.randn(
        seq_len, hidden_size, dtype=torch.bfloat16, device=device
    )
    hs, hs_scale_TxNb = fp8_block_quant_1d(a_bf16, block=BS)
    hs_scale = hs_scale_TxNb.transpose(0, 1).contiguous()  # [H/128, T]

    w13_bf16 = torch.randn(
        num_local_experts,
        2 * intermediate_size,
        hidden_size,
        dtype=torch.bfloat16,
        device=device,
    )
    w2_bf16 = torch.randn(
        num_local_experts,
        hidden_size,
        intermediate_size,
        dtype=torch.bfloat16,
        device=device,
    )
    w1, w1s = fp8_block_quant_2d(w13_bf16, block=BS)
    w2, w2s = fp8_block_quant_2d(w2_bf16, block=BS)

    return {
        "routing_logits": routing_logits,
        "routing_bias": routing_bias,
        "hidden_states": hs,
        "hidden_states_scale": hs_scale,
        "gemm1_weights": w1,
        "gemm1_weights_scale": w1s,
        "gemm1_alpha": None,
        "gemm1_beta": None,
        "gemm1_clamp_limit": None,
        "gemm2_weights": w2,
        "gemm2_weights_scale": w2s,
        "top_k": int(top_k),
        "num_experts": int(num_experts),
        "intermediate_size": int(intermediate_size),
        "local_expert_offset": 0,
        "local_num_experts": int(num_local_experts),
        "routed_scaling_factor": 2.5,
        "routing_method_type": int(routing_method_type),
    }


def _bind_init_dependency(wrapper, dependency):
    wrapper._trace_init_dependencies = (dependency,)
    wrapper.__signature__ = inspect.signature(dependency)
    return wrapper


def _moe_fp8_block_scale_renormalize_init(**kwargs):
    kwargs["routing_method_type"] = 1
    return _moe_fp8_block_scale_standard_init(**kwargs)


def _moe_fp8_block_scale_llama4_init(**kwargs):
    kwargs["routing_method_type"] = 3
    kwargs["top_k"] = 1
    return _moe_fp8_block_scale_standard_init(**kwargs)


def _moe_fp8_block_scale_renormalize_naive_init(**kwargs):
    kwargs["routing_method_type"] = 4
    return _moe_fp8_block_scale_standard_init(**kwargs)


def _moe_fp8_block_scale_topk_init(**kwargs):
    kwargs["routing_method_type"] = 5
    return _moe_fp8_block_scale_standard_init(**kwargs)


_moe_fp8_block_scale_renormalize_init = _bind_init_dependency(
    _moe_fp8_block_scale_renormalize_init, _moe_fp8_block_scale_standard_init
)
_moe_fp8_block_scale_llama4_init = _bind_init_dependency(
    _moe_fp8_block_scale_llama4_init, _moe_fp8_block_scale_standard_init
)
_moe_fp8_block_scale_renormalize_naive_init = _bind_init_dependency(
    _moe_fp8_block_scale_renormalize_naive_init, _moe_fp8_block_scale_standard_init
)
_moe_fp8_block_scale_topk_init = _bind_init_dependency(
    _moe_fp8_block_scale_topk_init, _moe_fp8_block_scale_standard_init
)


def _make_standard_moe_trace(name_prefix, description, reference, init):
    """Factory for standard (non-DS) routing templates (same inputs/axes)."""
    return TraceTemplate(
        op_type="moe",
        name_prefix=name_prefix,
        description=description,
        axes=dict(_STANDARD_AXES),
        inputs=dict(_STANDARD_INPUTS),
        outputs=dict(_STANDARD_OUTPUTS),
        tags=_STANDARD_TAGS,
        reference=reference,
        init=init,
    )


# ---------------------------------------------------------------------------
# Template instances — one per RoutingMethodType value
# ---------------------------------------------------------------------------


# RoutingMethodType.DeepSeekV3 = 2
# Uses additional n_group / topk_group axes and requires routing_bias.
def _moe_fp8_block_scale_ds_init(
    *,
    seq_len: int,
    num_experts: int = 256,
    top_k: int = 8,
    n_group: int = 8,
    topk_group: int = 4,
    num_local_experts: int = 32,
    hidden_size: int = 7168,
    intermediate_size: int = 2048,
    gemm1_out_size: int = 0,  # derived: 2 * intermediate_size
    num_hidden_blocks: int = 0,  # derived: hidden_size / 128
    num_intermediate_blocks: int = 0,  # derived
    num_gemm1_out_blocks: int = 0,  # derived
    device: str = "cuda",
    seed: int = 0,
):
    """Build inputs for ``trtllm_fp8_block_scale_moe`` (DeepSeek-V3 routing).

    Sourced from ``tests/moe/test_dpsk_fused_moe_fp8.py``: same fp8
    block-quantization pipeline as the standard variant
    (``_fp8_block_quant_*``, block=128, ``2.0×`` activation scale)
    plus DS routing kwargs (``n_group``, ``topk_group``).
    """
    del gemm1_out_size, num_hidden_blocks, num_intermediate_blocks
    del num_gemm1_out_blocks
    torch.manual_seed(seed)
    BS = 128
    routing_logits = torch.randn(
        seq_len, num_experts, dtype=torch.float32, device=device
    )
    routing_bias = torch.zeros(num_experts, dtype=torch.bfloat16, device=device)

    a_bf16 = 2.0 * torch.randn(
        seq_len, hidden_size, dtype=torch.bfloat16, device=device
    )
    hs, hs_scale_TxNb = fp8_block_quant_1d(a_bf16, block=BS)
    hs_scale = hs_scale_TxNb.transpose(0, 1).contiguous()  # [H/128, T]

    w13_bf16 = torch.randn(
        num_local_experts,
        2 * intermediate_size,
        hidden_size,
        dtype=torch.bfloat16,
        device=device,
    )
    w2_bf16 = torch.randn(
        num_local_experts,
        hidden_size,
        intermediate_size,
        dtype=torch.bfloat16,
        device=device,
    )
    w1, w1s = fp8_block_quant_2d(w13_bf16, block=BS)
    w2, w2s = fp8_block_quant_2d(w2_bf16, block=BS)

    return {
        "routing_logits": routing_logits,
        "routing_bias": routing_bias,
        "hidden_states": hs,
        "hidden_states_scale": hs_scale,
        "gemm1_weights": w1,
        "gemm1_weights_scale": w1s,
        "gemm1_alpha": None,
        "gemm1_beta": None,
        "gemm1_clamp_limit": None,
        "gemm2_weights": w2,
        "gemm2_weights_scale": w2s,
        "top_k": int(top_k),
        "n_group": int(n_group),
        "topk_group": int(topk_group),
        "num_experts": int(num_experts),
        "intermediate_size": int(intermediate_size),
        "local_expert_offset": 0,
        "local_num_experts": int(num_local_experts),
        "routed_scaling_factor": 2.5,
        "routing_method_type": 2,
    }


trtllm_fp8_block_scale_moe_ds_routing_trace = TraceTemplate(
    op_type="moe",
    name_prefix="moe_fp8_block_scale_ds_routing",
    description="FP8 block scale MoE with DeepSeek-V3 routing. Includes grouped sigmoid routing and two grouped-GEMM.",
    axes={
        "seq_len": Var(description="Sequence length (number of tokens)"),
        "num_experts": Const(description="Total number of experts.", abbrev=""),
        "top_k": Const(
            description="Number of experts to route to per token.", abbrev="topk"
        ),
        "n_group": Const(
            description="Number of expert groups for group routing.", abbrev="ng"
        ),
        "topk_group": Const(
            description="Number of groups to select for top-k routing.", abbrev="kg"
        ),
        "num_local_experts": Const(description="Number of local experts.", abbrev="e"),
        "hidden_size": Const(description="Hidden dimension size.", abbrev="h"),
        "intermediate_size": Const(
            description="MoE intermediate layer size.", abbrev="i"
        ),
        "gemm1_out_size": Const(
            description="Output size of the first GEMM (W13). Should be 2 * intermediate_size.",
            abbrev="",
        ),
        "num_hidden_blocks": Const(
            description="Number of quantized blocks along the hidden_size dimension (block_size=128).",
            abbrev="",
        ),
        "num_intermediate_blocks": Const(
            description="Number of quantized blocks along the intermediate_size dimension (block_size=128).",
            abbrev="",
        ),
        "num_gemm1_out_blocks": Const(
            description="Number of quantized blocks along the gemm1_out_size dimension (block_size=128).",
            abbrev="",
        ),
    },
    inputs={
        "routing_logits": Tensor(
            ["seq_len", "num_experts"],
            description="Routing logits for expert selection.",
        ),
        "routing_bias": Tensor(
            ["num_experts"],
            description="Bias tensor for routing. Pass all zeros for no bias.",
        ),
        "hidden_states": Tensor(
            ["seq_len", "hidden_size"],
            description="Input hidden states tensor (FP8 quantized).",
        ),
        "hidden_states_scale": Tensor(
            ["num_hidden_blocks", "seq_len"],
            description="Block-wise scaling factors for hidden states.",
        ),
        "gemm1_weights": Tensor(
            ["num_local_experts", "gemm1_out_size", "hidden_size"],
            description="First GEMM weights for all local experts (gate and up projections).",
        ),
        "gemm1_weights_scale": Tensor(
            ["num_local_experts", "num_gemm1_out_blocks", "num_hidden_blocks"],
            description="Block-wise scaling factors for first GEMM weights.",
        ),
        "gemm1_alpha": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Optional MxFp8-only per-expert SwiGLU OA alpha.",
            optional=True,
        ),
        "gemm1_beta": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Optional MxFp8-only per-expert SwiGLU OA beta.",
            optional=True,
        ),
        "gemm1_clamp_limit": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Optional MxFp8-only per-expert SwiGLU OA clamp limit.",
            optional=True,
        ),
        "gemm2_weights": Tensor(
            ["num_local_experts", "hidden_size", "intermediate_size"],
            description="Second GEMM weights for all local experts (down projection).",
        ),
        "gemm2_weights_scale": Tensor(
            ["num_local_experts", "num_hidden_blocks", "num_intermediate_blocks"],
            description="Block-wise scaling factors for second GEMM weights.",
        ),
        "top_k": Scalar(
            "int32",
            description="Number of experts to route to per token (DeepSeek-V3 uses 8).",
        ),
        "n_group": Scalar(
            "int32",
            description="Number of expert groups (DeepSeek-V3 uses 8).",
        ),
        "topk_group": Scalar(
            "int32",
            description="Number of groups to keep after group-level top-k (DeepSeek-V3 uses 4).",
        ),
        "local_expert_offset": Scalar(
            "int32",
            description="Offset of local experts in global expert space.",
        ),
        "routed_scaling_factor": Scalar(
            "float32",
            description="Scaling factor for routing weights.",
        ),
    },
    outputs={
        "output": Tensor(
            ["seq_len", "hidden_size"],
            dtype="bfloat16",
            description="Final MoE output tensor.",
        ),
    },
    tags=["status:verified", "quantization:float8_e4m3fn"],
    reference=_trtllm_fp8_block_scale_moe_ds_routing_reference,
    init=_moe_fp8_block_scale_ds_init,
)

# Backward-compatible alias (the original name used in fused_moe/core.py import).
trtllm_fp8_block_scale_moe_trace = trtllm_fp8_block_scale_moe_ds_routing_trace

# RoutingMethodType.Default = 0 — Softmax → TopK
trtllm_fp8_block_scale_moe_default_routing_trace = _make_standard_moe_trace(
    name_prefix="moe_fp8_block_scale_default_routing",
    description="FP8 block scale MoE with Default routing (Softmax → TopK).",
    reference=_trtllm_fp8_block_scale_moe_default_routing_reference,
    init=_moe_fp8_block_scale_standard_init,
)

# RoutingMethodType.Renormalize = 1 — TopK → Softmax
trtllm_fp8_block_scale_moe_renormalize_routing_trace = _make_standard_moe_trace(
    name_prefix="moe_fp8_block_scale_renormalize_routing",
    description="FP8 block scale MoE with Renormalize routing (TopK → Softmax).",
    reference=_trtllm_fp8_block_scale_moe_renormalize_routing_reference,
    init=_moe_fp8_block_scale_renormalize_init,
)

# RoutingMethodType.Llama4 = 3 — Top1 → Sigmoid
trtllm_fp8_block_scale_moe_llama4_routing_trace = _make_standard_moe_trace(
    name_prefix="moe_fp8_block_scale_llama4_routing",
    description="FP8 block scale MoE with Llama4 routing (Top1 → Sigmoid).",
    reference=_trtllm_fp8_block_scale_moe_llama4_routing_reference,
    init=_moe_fp8_block_scale_llama4_init,
)

# RoutingMethodType.RenormalizeNaive = 4 — Softmax → TopK → Renormalize
trtllm_fp8_block_scale_moe_renormalize_naive_routing_trace = _make_standard_moe_trace(
    name_prefix="moe_fp8_block_scale_renormalize_naive_routing",
    description="FP8 block scale MoE with RenormalizeNaive routing (Softmax → TopK → Renormalize).",
    reference=_trtllm_fp8_block_scale_moe_renormalize_naive_routing_reference,
    init=_moe_fp8_block_scale_renormalize_naive_init,
)

# RoutingMethodType.TopK = 5 — TopK only (no softmax), uniform weights
trtllm_fp8_block_scale_moe_topk_routing_trace = _make_standard_moe_trace(
    name_prefix="moe_fp8_block_scale_topk_routing",
    description="FP8 block scale MoE with TopK-only routing (no softmax, uniform weights).",
    reference=_trtllm_fp8_block_scale_moe_topk_routing_reference,
    init=_moe_fp8_block_scale_topk_init,
)

# ---------------------------------------------------------------------------
# Dispatch function — maps routing_method_type → TraceTemplate
# ---------------------------------------------------------------------------

_MOE_TRACE_BY_ROUTING_TYPE = {
    0: trtllm_fp8_block_scale_moe_default_routing_trace,  # Default
    1: trtllm_fp8_block_scale_moe_renormalize_routing_trace,  # Renormalize
    2: trtllm_fp8_block_scale_moe_ds_routing_trace,  # DeepSeekV3
    3: trtllm_fp8_block_scale_moe_llama4_routing_trace,  # Llama4
    4: trtllm_fp8_block_scale_moe_renormalize_naive_routing_trace,  # RenormalizeNaive
    5: trtllm_fp8_block_scale_moe_topk_routing_trace,  # TopK
    # 6 = Unspecified: no trace
}


def trtllm_fp8_block_scale_moe_trace_dispatch(**kwargs):
    """Return the appropriate TraceTemplate for the given ``routing_method_type``.

    Pass this as ``trace=trtllm_fp8_block_scale_moe_trace_dispatch`` to
    ``@flashinfer_api`` so the correct template is selected at call time::

        @flashinfer_api(trace=trtllm_fp8_block_scale_moe_trace_dispatch)
        def trtllm_fp8_block_scale_moe(..., routing_method_type: int = 0, ...):
            ...

    Returns ``None`` for ``RoutingMethodType.Unspecified`` (6), which
    suppresses trace generation.
    """
    routing_method_type = int(kwargs.get("routing_method_type", 0))
    return _MOE_TRACE_BY_ROUTING_TYPE.get(routing_method_type)


# Expose all possible templates so _attach_fi_trace can auto-register them
# in _TRACE_REGISTRY for consistency testing.
trtllm_fp8_block_scale_moe_trace_dispatch.templates = list(  # type: ignore[attr-defined]
    _MOE_TRACE_BY_ROUTING_TYPE.values()
)


# ---------------------------------------------------------------------------
# FP4 block-scale MoE (trtllm_fp4_block_scale_moe)
# ---------------------------------------------------------------------------
# NvFP4: block_size=16, weights packed as uint8 (2 fp4 per byte).
#   hidden_states       : [seq_len, hidden_size // 2]   uint8
#   hidden_states_scale : [seq_len, hidden_size // 16]  float8  (optional for bf16 input)
#   gemm1_weights       : [E_loc, 2*I, hidden_size // 2]         uint8
#   gemm1_weights_scale : [E_loc, 2*I, hidden_size // 16]        float8
#   gemm2_weights       : [E_loc, hidden_size, I // 2]            uint8
#   gemm2_weights_scale : [E_loc, hidden_size, I // 16]           float8
# ---------------------------------------------------------------------------


# FP4 e2m1fn magnitudes. The 4-bit code is {sign(1), exponent(2), mantissa(1)};
# this table maps the 16 possible nibble values to the corresponding float32
# magnitude so dequantization is a single gather.
_E2M1_LUT_VALUES = [
    0.0,
    0.5,
    1.0,
    1.5,
    2.0,
    3.0,
    4.0,
    6.0,
    -0.0,
    -0.5,
    -1.0,
    -1.5,
    -2.0,
    -3.0,
    -4.0,
    -6.0,
]


@torch.no_grad()
def _unpack_fp4_e2m1(packed: torch.Tensor) -> torch.Tensor:
    """Unpack a uint8 tensor of packed e2m1fn FP4 values into float32.

    Each byte stores two 4-bit values (low nibble = first element along the
    last axis). The returned tensor has twice the last-dim size of *packed*.
    """
    lut = torch.tensor(_E2M1_LUT_VALUES, dtype=torch.float32, device=packed.device)
    p = packed.view(torch.uint8).to(torch.int64)
    lo = lut[p & 0x0F]
    hi = lut[(p >> 4) & 0x0F]
    stacked = torch.stack([lo, hi], dim=-1)  # pairs along a new last axis
    return stacked.reshape(*packed.shape[:-1], packed.shape[-1] * 2)


@torch.no_grad()
def _ue8m0_to_float32(scales: torch.Tensor) -> torch.Tensor:
    """Decode UE8M0 (uint8, unsigned exponent-only) scales to float32."""
    e = scales.view(torch.uint8).to(torch.int64)
    return torch.pow(torch.tensor(2.0, device=scales.device), (e - 127).float())


@torch.no_grad()
def _decode_block_scales(scales: torch.Tensor, is_ue8m0: bool) -> torch.Tensor:
    """Decode block scales: UE8M0 for MX formats, float8_e4m3fn otherwise."""
    if is_ue8m0:
        return _ue8m0_to_float32(scales)
    # fp8_e4m3fn (or already float): plain cast.
    return scales.to(torch.float32)


@torch.no_grad()
def _dequantize_fp4_tensor(
    packed: torch.Tensor,
    scales: torch.Tensor,
    is_ue8m0_scales: bool,
) -> torch.Tensor:
    """Unpack an FP4 tensor and apply its per-block scales along the last dim.

    The packed tensor has half the logical last-dim size of the output; the
    scale tensor has last-dim size = (output last dim) / block_size.
    block_size is inferred from the shape ratio.
    """
    unpacked = _unpack_fp4_e2m1(packed)  # float32, last dim = packed.last * 2
    block_size = unpacked.shape[-1] // scales.shape[-1]
    decoded_scales = _decode_block_scales(scales, is_ue8m0_scales)
    expanded = decoded_scales.repeat_interleave(block_size, dim=-1)
    return unpacked * expanded


@torch.no_grad()
def _dequantize_fp4_hidden_states(
    hidden_states: torch.Tensor,
    hidden_states_scale,
    is_weights_mxfp4: bool,
) -> torch.Tensor:
    """Dequantize hidden_states to float32.

    Three cases by dtype:
      * bfloat16 — pass-through (no scale).
      * float8_e4m3fn — MXFP8 activation with UE8M0 per-32 scales.
      * uint8 — NvFP4/MXFP4 packed activation with per-block scales (fp8_e4m3fn
        for NvFP4, UE8M0 for MXFP4; here both are treated as fp8_e4m3fn since
        the runtime FP4 path uses fp8_e4m3fn scales for activations).
    """
    if hidden_states.dtype == torch.bfloat16:
        return hidden_states.to(torch.float32)
    if hidden_states.dtype == torch.float8_e4m3fn:
        # MXFP8 hidden states: UE8M0 scales, block size 32.
        scales = _ue8m0_to_float32(hidden_states_scale)
        block_size = hidden_states.shape[-1] // scales.shape[-1]
        expanded = scales.repeat_interleave(block_size, dim=-1)
        return hidden_states.to(torch.float32) * expanded
    # uint8-packed FP4. For NvFP4 activation + NvFP4 weights the scales are
    # fp8_e4m3fn; for MXFP4 weights (and bf16-packed-as-fp4 corner cases) they
    # are UE8M0. Use the weight mode as the tiebreaker since activation scale
    # format tracks weight format in the trtllm-gen kernel.
    return _dequantize_fp4_tensor(
        hidden_states, hidden_states_scale, is_ue8m0_scales=is_weights_mxfp4
    )


@torch.no_grad()
def _fp4_moe_run_experts(
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    gemm1_bias,
    gemm2_bias,
    weights,
    topk_idx,
    local_expert_offset,
    E_global,
    activation_type=3,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
    **_unused,
):
    """FP4 dequantize + gated activation + GEMM for all routing types.

    ``weights``   : [T, TOP_K] float32 — per-token expert weights (normalised)
    ``topk_idx``  : [T, TOP_K] int64   — selected global expert indices

    Detects MXFP4 vs NvFP4 weight format from whether gemm1_weights_scale is
    fp8_e4m3fn (NvFP4) or uint8 (UE8M0, MXFP4). Block size is inferred from
    the ratio of unpacked K to scale K.
    """
    is_mxfp4 = gemm1_weights_scale.dtype == torch.uint8
    device = gemm1_weights.device

    # Dequantize both expert-weight tensors in one shot.
    W1 = _dequantize_fp4_tensor(
        gemm1_weights, gemm1_weights_scale, is_ue8m0_scales=is_mxfp4
    )  # [E_local, 2*I, H]
    W2 = _dequantize_fp4_tensor(
        gemm2_weights, gemm2_weights_scale, is_ue8m0_scales=is_mxfp4
    )  # [E_local, H, I]

    E_local, gemm1_out_size, H = W1.shape
    I = gemm1_out_size // 2
    if gemm1_out_size != 2 * I:
        raise ValueError(
            f"gemm1 output size {gemm1_out_size} is not 2*intermediate_size; "
            "FP4 MoE requires SwiGLU (gate + up)."
        )

    A = _dequantize_fp4_hidden_states(hidden_states, hidden_states_scale, is_mxfp4)
    T = A.shape[0]
    output = torch.zeros((T, H), dtype=torch.float32, device=device)
    local_start = int(local_expert_offset)
    activation_type = normalize_activation_type(activation_type)
    if activation_type not in (ActivationType.Swiglu, ActivationType.Situ):
        raise ValueError(
            "FP4 MoE trace reference supports "
            f"{ActivationType.Swiglu!r} and {ActivationType.Situ!r}, "
            f"got {activation_type!r}"
        )
    if activation_type == ActivationType.Situ:
        for name, param in (
            ("gemm1_alpha", gemm1_alpha),
            ("gemm1_beta", gemm1_beta),
            ("gemm1_clamp_limit", gemm1_clamp_limit),
        ):
            if param is None:
                continue
            if param.shape != (E_local,):
                raise ValueError(
                    f"{name} must have shape [{E_local}], got {tuple(param.shape)}"
                )
            if param.dtype != torch.float32:
                raise ValueError(f"{name} must have dtype float32, got {param.dtype}")
            if param.device != device:
                raise ValueError(f"{name} must be on {device}, got {param.device}")
            if not torch.isfinite(param).all() or not (param > 0).all():
                raise ValueError(f"{name} must contain only finite positive values")

    for le in range(E_local):
        ge = local_start + le
        if ge < 0 or ge >= E_global:
            continue
        sel_mask = (topk_idx == ge).any(dim=1)
        if not sel_mask.any():
            continue
        token_idx = torch.nonzero(sel_mask, as_tuple=False).squeeze(1)
        A_e = A.index_select(0, token_idx)  # [N, H]
        G1 = A_e.matmul(W1[le].t())  # [N, 2*I]
        if gemm1_bias is not None:
            G1 = G1 + gemm1_bias[le].to(torch.float32)
        # TRTLLM-Gen convention: x0 is linear/first; x1 is gate/second.
        x0, x1 = G1[:, :I], G1[:, I:]
        if activation_type == ActivationType.Situ:
            alpha = (
                torch.tensor(1.0, dtype=torch.float32, device=device)
                if gemm1_alpha is None
                else gemm1_alpha[le]
            )
            beta = (
                torch.tensor(1.0, dtype=torch.float32, device=device)
                if gemm1_beta is None
                else gemm1_beta[le]
            )
            if gemm1_clamp_limit is not None:
                limit = gemm1_clamp_limit[le]
                x0 = torch.clamp(x0, min=-limit, max=limit)
                x1 = torch.clamp(x1, max=limit)
            activated = (
                beta
                * torch.tanh(x0 / beta)
                * alpha
                * torch.tanh(x1 / alpha)
                * torch.sigmoid(x1)
            )
        else:
            activated = x1 * torch.sigmoid(x1) * x0
        O = activated.matmul(W2[le].t())  # [N, H]
        if gemm2_bias is not None:
            O = O + gemm2_bias[le].to(torch.float32)
        # Fold per-token expert weight.
        w_tok = weights.index_select(0, token_idx)
        match = (topk_idx.index_select(0, token_idx) == ge).float()
        w_e = (w_tok * match).sum(dim=1)
        output.index_add_(0, token_idx, O * w_e.unsqueeze(1))

    return output.to(torch.bfloat16)


@torch.no_grad()
def _trtllm_fp4_block_scale_moe_default_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm1_bias,
    gemm2_weights,
    gemm2_weights_scale,
    gemm2_bias,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    **activation_kwargs,
):
    """FP4 MoE with Default routing (Softmax → TopK)."""
    TOP_K = int(top_k)
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    s = torch.softmax(logits, dim=-1)
    _, topk_idx = torch.topk(s, k=TOP_K, dim=1, largest=True, sorted=False)
    scale = 1.0 if routed_scaling_factor is None else float(routed_scaling_factor)
    w_topk = s.gather(1, topk_idx) * scale
    return _fp4_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        gemm1_bias,
        gemm2_bias,
        w_topk,
        topk_idx,
        local_expert_offset,
        E_global,
        **activation_kwargs,
    )


@torch.no_grad()
def _trtllm_fp4_block_scale_moe_renormalize_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm1_bias,
    gemm2_weights,
    gemm2_weights_scale,
    gemm2_bias,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    **activation_kwargs,
):
    """FP4 MoE with Renormalize routing (TopK on logits → Softmax)."""
    TOP_K = int(top_k)
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    _, topk_idx = torch.topk(logits, k=TOP_K, dim=1, largest=True, sorted=False)
    gathered = logits.gather(1, topk_idx)
    scale = 1.0 if routed_scaling_factor is None else float(routed_scaling_factor)
    w_topk = torch.softmax(gathered, dim=-1) * scale
    return _fp4_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        gemm1_bias,
        gemm2_bias,
        w_topk,
        topk_idx,
        local_expert_offset,
        E_global,
        **activation_kwargs,
    )


@torch.no_grad()
def _trtllm_fp4_block_scale_moe_ds_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm1_bias,
    gemm2_weights,
    gemm2_weights_scale,
    gemm2_bias,
    top_k,
    n_group,
    topk_group,
    local_expert_offset,
    routed_scaling_factor,
    num_fused_shared_experts=0,
    **activation_kwargs,
):
    """FP4 MoE with DeepSeek-V3 routing: sigmoid + groups + top_k."""
    TOP_K = int(top_k)
    N_GROUP = int(n_group)
    TOPK_GROUP = int(topk_group)
    E_global = routing_logits.shape[1]
    T = routing_logits.shape[0]

    logits = routing_logits.to(torch.float32)
    bias = routing_bias.to(torch.float32).reshape(-1)
    s = 1.0 / (1.0 + torch.exp(-logits))
    s_with_bias = s + bias

    group_size = E_global // N_GROUP
    s_wb_grouped = s_with_bias.view(T, N_GROUP, group_size)
    top2_vals, _ = torch.topk(s_wb_grouped, k=2, dim=2, largest=True, sorted=False)
    group_scores = top2_vals.sum(dim=2)

    _, group_idx = torch.topk(
        group_scores, k=TOPK_GROUP, dim=1, largest=True, sorted=False
    )
    group_mask = torch.zeros_like(group_scores)
    group_mask.scatter_(1, group_idx, 1.0)
    score_mask = (
        group_mask.unsqueeze(2).expand(T, N_GROUP, group_size).reshape(T, E_global)
    )

    neg_inf = torch.finfo(torch.float32).min
    scores_pruned = s_with_bias.masked_fill(score_mask == 0, neg_inf)
    _, topk_idx = torch.topk(scores_pruned, k=TOP_K, dim=1, largest=True, sorted=False)

    M = torch.zeros_like(s)
    M.scatter_(1, topk_idx, 1.0)
    raw_w = s * M
    weights_sum = raw_w.sum(dim=1, keepdim=True) + 1e-20
    scale = 1.0 if routed_scaling_factor is None else float(routed_scaling_factor)
    full_weights = (raw_w / weights_sum) * scale
    w_topk = full_weights.gather(1, topk_idx)

    nsfe = int(num_fused_shared_experts or 0)
    if nsfe > 0:
        # Fused shared experts: ids num_experts + k with constant weight 1.0,
        # appended after the routed top-k slots.
        shared_idx = (
            torch.arange(
                E_global, E_global + nsfe, device=topk_idx.device, dtype=topk_idx.dtype
            )
            .unsqueeze(0)
            .expand(T, nsfe)
        )
        topk_idx = torch.cat([topk_idx, shared_idx], dim=1)
        w_topk = torch.cat(
            [w_topk, torch.ones(T, nsfe, dtype=w_topk.dtype, device=w_topk.device)],
            dim=1,
        )

    return _fp4_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        gemm1_bias,
        gemm2_bias,
        w_topk,
        topk_idx,
        local_expert_offset,
        E_global + nsfe,
        **activation_kwargs,
    )


@torch.no_grad()
def _trtllm_fp4_block_scale_moe_llama4_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm1_bias,
    gemm2_weights,
    gemm2_weights_scale,
    gemm2_bias,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    **activation_kwargs,
):
    """FP4 MoE with Llama4 routing (Top1 → Sigmoid). top_k is fixed at 1."""
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    topk_idx = logits.argmax(dim=-1, keepdim=True)
    top1_logit = logits.gather(1, topk_idx)
    scale = 1.0 if routed_scaling_factor is None else float(routed_scaling_factor)
    w_topk = (1.0 / (1.0 + torch.exp(-top1_logit))) * scale
    return _fp4_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        gemm1_bias,
        gemm2_bias,
        w_topk,
        topk_idx,
        local_expert_offset,
        E_global,
        **activation_kwargs,
    )


@torch.no_grad()
def _trtllm_fp4_block_scale_moe_renormalize_naive_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm1_bias,
    gemm2_weights,
    gemm2_weights_scale,
    gemm2_bias,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    **activation_kwargs,
):
    """FP4 MoE with RenormalizeNaive routing (Softmax → TopK → sum-to-1)."""
    TOP_K = int(top_k)
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    s = torch.softmax(logits, dim=-1)
    _, topk_idx = torch.topk(s, k=TOP_K, dim=1, largest=True, sorted=False)
    gathered = s.gather(1, topk_idx)
    w_topk = gathered / (gathered.sum(dim=1, keepdim=True) + 1e-20)
    scale = 1.0 if routed_scaling_factor is None else float(routed_scaling_factor)
    w_topk = w_topk * scale
    return _fp4_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        gemm1_bias,
        gemm2_bias,
        w_topk,
        topk_idx,
        local_expert_offset,
        E_global,
        **activation_kwargs,
    )


@torch.no_grad()
def _trtllm_fp4_block_scale_moe_topk_routing_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm1_bias,
    gemm2_weights,
    gemm2_weights_scale,
    gemm2_bias,
    top_k,
    local_expert_offset,
    routed_scaling_factor,
    **activation_kwargs,
):
    """FP4 MoE with TopK-only routing (uniform weights)."""
    TOP_K = int(top_k)
    E_global = routing_logits.shape[1]
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    _, topk_idx = torch.topk(logits, k=TOP_K, dim=1, largest=True, sorted=False)
    T = logits.shape[0]
    scale = 1.0 if routed_scaling_factor is None else float(routed_scaling_factor)
    w_topk = torch.full(
        (T, TOP_K), scale / TOP_K, dtype=torch.float32, device=logits.device
    )
    return _fp4_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        gemm1_bias,
        gemm2_bias,
        w_topk,
        topk_idx,
        local_expert_offset,
        E_global,
        **activation_kwargs,
    )


_FP4_STANDARD_AXES: dict[str, Var | Const] = {
    "seq_len": Var(description="Number of tokens."),
    "num_experts": Const(description="Total number of experts.", abbrev=""),
    "top_k": Const(description="Number of experts selected per token.", abbrev="topk"),
    "num_local_experts": Const(description="Number of local experts.", abbrev="e"),
    "hidden_size": Const(description="Hidden dimension size.", abbrev="h"),
    "intermediate_size": Const(description="MoE intermediate layer size.", abbrev="i"),
    # Derived / block-count axes (abbrev="" → omitted from filename)
    "gemm1_out_size": Const(
        description="Output size of FC1 (2 × intermediate_size for SwiGLU).",
        abbrev="",
    ),
    "num_packed_hidden": Const(
        description="Packed hidden dimension (hidden_size // 2 for NvFP4).",
        abbrev="",
    ),
    "num_fp4_hidden_blocks": Const(
        description="Number of FP4 scale blocks along hidden_size (hidden_size // 16 for NvFP4).",
        abbrev="",
    ),
    "num_packed_intermediate": Const(
        description="Packed intermediate dimension (intermediate_size // 2 for NvFP4).",
        abbrev="",
    ),
    "num_fp4_intermediate_blocks": Const(
        description="Number of FP4 scale blocks along intermediate_size (intermediate_size // 16 for NvFP4).",
        abbrev="",
    ),
    "activation_type": Const(
        description="Fused activation type; 10 selects SiTU.", abbrev="act"
    ),
}

_FP4_STANDARD_INPUTS: dict[str, Tensor | Scalar] = {
    "routing_logits": Tensor(
        ["seq_len", "num_experts"],
        description="Routing logits for expert selection.",
    ),
    "routing_bias": Tensor(
        ["num_experts"],
        description="Bias added to routing logits. Pass None when not used.",
        optional=True,
    ),
    # Packed NvFP4 hidden states (2 values per uint8 byte).
    "hidden_states": Tensor(
        ["seq_len", "num_packed_hidden"],
        description="Input hidden states, NvFP4-packed (uint8, 2 fp4 per byte).",
    ),
    "hidden_states_scale": Tensor(
        ["seq_len", "num_fp4_hidden_blocks"],
        description="Block-wise scale factors for hidden_states (float8). None for bf16 input.",
        optional=True,
    ),
    "gemm1_weights": Tensor(
        ["num_local_experts", "gemm1_out_size", "num_packed_hidden"],
        description="FC1 weights, NvFP4-packed (uint8). Shape includes gate+up for SwiGLU.",
    ),
    "gemm1_weights_scale": Tensor(
        ["num_local_experts", "gemm1_out_size", "num_fp4_hidden_blocks"],
        description="Block-wise scale factors for gemm1_weights (float8).",
    ),
    "gemm1_bias": Tensor(
        ["num_local_experts", "gemm1_out_size"],
        description="FC1 bias (float32). Optional.",
        optional=True,
    ),
    "gemm1_alpha": Tensor(
        ["num_local_experts"],
        description="Per-expert SiTU alpha or SwiGLU alpha (float32). Optional.",
        optional=True,
    ),
    "gemm1_beta": Tensor(
        ["num_local_experts"],
        description="Per-expert SiTU beta or SwiGLU beta (float32). Optional.",
        optional=True,
    ),
    "gemm1_clamp_limit": Tensor(
        ["num_local_experts"],
        description="Per-expert gated-activation clamp limit (float32). Optional.",
        optional=True,
    ),
    "gemm2_weights": Tensor(
        ["num_local_experts", "hidden_size", "num_packed_intermediate"],
        description="FC2 weights, NvFP4-packed (uint8).",
    ),
    "gemm2_weights_scale": Tensor(
        ["num_local_experts", "hidden_size", "num_fp4_intermediate_blocks"],
        description="Block-wise scale factors for gemm2_weights (float8).",
    ),
    "gemm2_bias": Tensor(
        ["num_local_experts", "hidden_size"],
        description="FC2 bias (float32). Optional.",
        optional=True,
    ),
    "output1_scale_scalar": Tensor(
        ["num_local_experts"],
        description="Per-expert output scale for FC1 activation (float32). Optional.",
        optional=True,
    ),
    "output1_scale_gate_scalar": Tensor(
        ["num_local_experts"],
        description="Per-expert output scale for FC1 gate (float32). Optional.",
        optional=True,
    ),
    "output2_scale_scalar": Tensor(
        ["num_local_experts"],
        description="Per-expert output scale for FC2 (float32). Optional.",
        optional=True,
    ),
    "local_expert_offset": Scalar(
        "int32",
        description="Offset of local experts in the global expert array.",
    ),
    "routed_scaling_factor": Scalar(
        "float32",
        optional=True,
        description="Scaling factor applied to routing weights. None for some routing methods.",
    ),
    "activation_type": Scalar(
        "int32", description="Fused activation type; 10 selects SiTU."
    ),
}

_FP4_STANDARD_OUTPUTS = {
    "output": Tensor(
        ["seq_len", "hidden_size"],
        dtype="bfloat16",
        description="Final MoE output tensor.",
    ),
}

_FP4_STANDARD_TAGS = ["status:experimental", "quantization:nvfp4"]


def _moe_fp4_block_scale_init(
    *,
    seq_len: int,
    routing_method_type: int = 0,
    num_experts: int = 256,
    top_k: int = 8,
    n_group: int = 0,  # used only by DS routing
    topk_group: int = 0,  # used only by DS routing
    num_local_experts: int = 32,
    hidden_size: int = 7168,
    intermediate_size: int = 2048,
    gemm1_out_size: int = 0,  # derived
    num_packed_hidden: int = 0,  # derived: hidden_size // 2
    num_fp4_hidden_blocks: int = 0,  # derived: hidden_size // 16
    num_packed_intermediate: int = 0,  # derived
    num_fp4_intermediate_blocks: int = 0,  # derived
    device: str = "cuda",
    seed: int = 0,
    activation_type: int = 3,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
):
    """Build inputs for ``trtllm_fp4_block_scale_moe`` (any routing variant).

    Sourced from ``tests/moe/test_trtllm_gen_fused_moe.py``:
      hidden_states     = 2 * randn(...)                    (main fixture)
      hidden_states_sf  = 448*6 / amax(hs)                  (single global)
      weights[e]        = randn(...)
      weight_sf[e]      = 448*6 / amax(weight[e])           (**per-expert**,
        per ``quant_fp4_batches`` / ``calculate_fp4_global_scale_factor``)
      hidden_states_scale layout = swizzled (test default)
      routing_bias      = randn(...) (test parametrizes on/off; we use
        on by default for parity with the test).
    Requires SM100+ at runtime; CPU smoke tests skip.
    """
    activation_type = normalize_activation_type(activation_type)
    situ_alpha = gemm1_alpha
    situ_beta = gemm1_beta
    situ_clamp = gemm1_clamp_limit
    if activation_type == ActivationType.Situ:
        materialized: dict[str, torch.Tensor | None] = {}
        for name, value, default in (
            ("gemm1_alpha", gemm1_alpha, 1.0),
            ("gemm1_beta", gemm1_beta, 1.0),
            ("gemm1_clamp_limit", gemm1_clamp_limit, None),
        ):
            if value is None and default is None:
                materialized[name] = None
                continue
            scalar = float(default if value is None else value)
            if not math.isfinite(scalar) or scalar <= 0:
                raise ValueError(f"{name} must be finite and positive for SiTU")
            materialized[name] = torch.full(
                (num_local_experts,), scalar, dtype=torch.float32, device=device
            )
        situ_alpha = materialized["gemm1_alpha"]
        situ_beta = materialized["gemm1_beta"]
        situ_clamp = materialized["gemm1_clamp_limit"]

    del gemm1_out_size, num_packed_hidden, num_fp4_hidden_blocks
    del num_packed_intermediate, num_fp4_intermediate_blocks
    from flashinfer import fp4_quantize  # noqa: PLC0415

    SF_VEC = 16
    H = hidden_size
    I = intermediate_size  # noqa: E741
    torch.manual_seed(seed)
    routing_logits = torch.randn(
        seq_len, num_experts, dtype=torch.bfloat16, device=device
    )
    routing_bias = torch.randn(num_experts, dtype=torch.bfloat16, device=device)

    hs_bf16 = 2.0 * torch.randn(seq_len, H, dtype=torch.bfloat16, device=device)
    hs_amax = hs_bf16.float().abs().nan_to_num().max().clamp(min=1e-12)
    hs_global_sf = ((448.0 * 6.0) / hs_amax).reshape(1).contiguous()
    # Test default is swizzled hidden_states_scale layout.
    hidden_states, hidden_states_scale = fp4_quantize(
        hs_bf16,
        hs_global_sf,
        sf_vec_size=SF_VEC,
        sf_use_ue8m0=False,
        is_sf_swizzled_layout=True,
    )
    hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape(
        seq_len, -1
    )

    # Per-expert weight scales, matching ``quant_fp4_batches``.
    w13_bf16 = torch.randn(
        num_local_experts, 2 * I, H, dtype=torch.bfloat16, device=device
    )
    w13_amax_per_expert = (
        w13_bf16.reshape(num_local_experts, -1)
        .float()
        .abs()
        .nan_to_num()
        .amax(dim=1)
        .clamp(min=1e-12)
    )
    w13_global_sf_per_expert = (448.0 * 6.0) / w13_amax_per_expert  # [E]
    gemm1_weights_chunks, gemm1_scale_chunks = [], []
    for e in range(num_local_experts):
        wq, ws = fp4_quantize(
            w13_bf16[e],
            w13_global_sf_per_expert[e].reshape(1).contiguous(),
            sf_vec_size=SF_VEC,
            sf_use_ue8m0=False,
        )
        gemm1_weights_chunks.append(wq)
        gemm1_scale_chunks.append(ws.view(torch.float8_e4m3fn).reshape(2 * I, -1))
    gemm1_weights = torch.stack(gemm1_weights_chunks, dim=0)
    gemm1_weights_scale = torch.stack(gemm1_scale_chunks, dim=0)

    w2_bf16 = torch.randn(num_local_experts, H, I, dtype=torch.bfloat16, device=device)
    w2_amax_per_expert = (
        w2_bf16.reshape(num_local_experts, -1)
        .float()
        .abs()
        .nan_to_num()
        .amax(dim=1)
        .clamp(min=1e-12)
    )
    w2_global_sf_per_expert = (448.0 * 6.0) / w2_amax_per_expert  # [E]
    gemm2_weights_chunks, gemm2_scale_chunks = [], []
    for e in range(num_local_experts):
        wq, ws = fp4_quantize(
            w2_bf16[e],
            w2_global_sf_per_expert[e].reshape(1).contiguous(),
            sf_vec_size=SF_VEC,
            sf_use_ue8m0=False,
        )
        gemm2_weights_chunks.append(wq)
        gemm2_scale_chunks.append(ws.view(torch.float8_e4m3fn).reshape(H, -1))
    gemm2_weights = torch.stack(gemm2_weights_chunks, dim=0)
    gemm2_weights_scale = torch.stack(gemm2_scale_chunks, dim=0)

    # Per-expert FC1/FC2 output rescale = 1 / (hs_sf * w_sf[e]).
    hs_sf_v = float(hs_global_sf.item())
    out1_scale = (1.0 / (hs_sf_v * w13_global_sf_per_expert)).to(device)
    out1_gate_scale = out1_scale.clone()
    out2_scale = (1.0 / (hs_sf_v * w2_global_sf_per_expert)).to(device)
    result = {
        "routing_logits": routing_logits,
        "routing_bias": routing_bias,
        "hidden_states": hidden_states,
        "hidden_states_scale": hidden_states_scale,
        "gemm1_weights": gemm1_weights,
        "gemm1_weights_scale": gemm1_weights_scale,
        "gemm1_bias": None,
        "gemm1_alpha": situ_alpha,
        "gemm1_beta": situ_beta,
        "gemm1_clamp_limit": situ_clamp,
        "gemm2_weights": gemm2_weights,
        "gemm2_weights_scale": gemm2_weights_scale,
        "gemm2_bias": None,
        "output1_scale_scalar": out1_scale,
        "output1_scale_gate_scalar": out1_gate_scale,
        "output2_scale_scalar": out2_scale,
        "top_k": int(top_k),
        "num_experts": int(num_experts),
        "intermediate_size": int(intermediate_size),
        "local_expert_offset": 0,
        "local_num_experts": int(num_local_experts),
        "routed_scaling_factor": None,
        "routing_method_type": int(routing_method_type),
        "activation_type": int(activation_type),
    }
    if routing_method_type == 2:
        result["n_group"] = int(n_group) if n_group else 8
        result["topk_group"] = int(topk_group) if topk_group else 4
    return result


def _moe_fp4_block_scale_renormalize_init(**kwargs):
    kwargs["routing_method_type"] = 1
    return _moe_fp4_block_scale_init(**kwargs)


def _moe_fp4_block_scale_ds_init(**kwargs):
    # The generated benchmark covers the routed path
    # (num_fused_shared_experts == 0); the fused shared-expert slots only
    # append extra constant-weight columns and weight rows.
    kwargs["routing_method_type"] = 2
    out = _moe_fp4_block_scale_init(**kwargs)
    out["num_fused_shared_experts"] = 0
    return out


def _moe_fp4_block_scale_llama4_init(**kwargs):
    kwargs["routing_method_type"] = 3
    kwargs["top_k"] = 1
    return _moe_fp4_block_scale_init(**kwargs)


def _moe_fp4_block_scale_renormalize_naive_init(**kwargs):
    kwargs["routing_method_type"] = 4
    return _moe_fp4_block_scale_init(**kwargs)


def _moe_fp4_block_scale_topk_init(**kwargs):
    kwargs["routing_method_type"] = 5
    return _moe_fp4_block_scale_init(**kwargs)


_moe_fp4_block_scale_renormalize_init = _bind_init_dependency(
    _moe_fp4_block_scale_renormalize_init, _moe_fp4_block_scale_init
)
_moe_fp4_block_scale_ds_init = _bind_init_dependency(
    _moe_fp4_block_scale_ds_init, _moe_fp4_block_scale_init
)
_moe_fp4_block_scale_llama4_init = _bind_init_dependency(
    _moe_fp4_block_scale_llama4_init, _moe_fp4_block_scale_init
)
_moe_fp4_block_scale_renormalize_naive_init = _bind_init_dependency(
    _moe_fp4_block_scale_renormalize_naive_init, _moe_fp4_block_scale_init
)
_moe_fp4_block_scale_topk_init = _bind_init_dependency(
    _moe_fp4_block_scale_topk_init, _moe_fp4_block_scale_init
)


def _make_standard_fp4_moe_trace(name_prefix, description, reference=None, init=None):
    """Factory for FP4 MoE templates that share the standard (non-DS) axis set."""
    return TraceTemplate(
        op_type="moe",
        name_prefix=name_prefix,
        description=description,
        axes=dict(_FP4_STANDARD_AXES),
        inputs=dict(_FP4_STANDARD_INPUTS),
        outputs=dict(_FP4_STANDARD_OUTPUTS),
        tags=_FP4_STANDARD_TAGS,
        reference=reference,
        init=init or _moe_fp4_block_scale_init,
    )


# RoutingMethodType.Default = 0 — Softmax → TopK
trtllm_fp4_block_scale_moe_default_routing_trace = _make_standard_fp4_moe_trace(
    name_prefix="moe_fp4_block_scale_default_routing",
    description="NvFP4 block-scale MoE with Default routing (Softmax → TopK).",
    reference=_trtllm_fp4_block_scale_moe_default_routing_reference,
)

# RoutingMethodType.Renormalize = 1 — TopK → Softmax
trtllm_fp4_block_scale_moe_renormalize_routing_trace = _make_standard_fp4_moe_trace(
    name_prefix="moe_fp4_block_scale_renormalize_routing",
    description="NvFP4 block-scale MoE with Renormalize routing (TopK → Softmax).",
    reference=_trtllm_fp4_block_scale_moe_renormalize_routing_reference,
    init=_moe_fp4_block_scale_renormalize_init,
)

# RoutingMethodType.DeepSeekV3 = 2 — Sigmoid → group selection → TopK
trtllm_fp4_block_scale_moe_ds_routing_trace = TraceTemplate(
    op_type="moe",
    name_prefix="moe_fp4_block_scale_ds_routing",
    description="NvFP4 block-scale MoE with DeepSeekV3 routing (Sigmoid → group selection → top_k).",
    axes={
        **_FP4_STANDARD_AXES,
        "n_group": Const(
            description="Number of expert groups for group routing.", abbrev="ng"
        ),
        "topk_group": Const(
            description="Number of groups selected in top-k routing.", abbrev="kg"
        ),
    },
    inputs={
        **_FP4_STANDARD_INPUTS,
        "num_fused_shared_experts": Scalar(
            "int32",
            description=(
                "Number of fused shared experts appended after the routed "
                "experts (weight rows num_experts + k, constant weight 1.0). "
                "Only supported with DeepSeekV3 routing."
            ),
        ),
    },
    outputs=dict(_FP4_STANDARD_OUTPUTS),
    tags=_FP4_STANDARD_TAGS,
    reference=_trtllm_fp4_block_scale_moe_ds_routing_reference,
    init=_moe_fp4_block_scale_ds_init,
)

# RoutingMethodType.Llama4 = 3 — Top1 → Sigmoid
trtllm_fp4_block_scale_moe_llama4_routing_trace = _make_standard_fp4_moe_trace(
    name_prefix="moe_fp4_block_scale_llama4_routing",
    description="NvFP4 block-scale MoE with Llama4 routing (Top1 → Sigmoid).",
    reference=_trtllm_fp4_block_scale_moe_llama4_routing_reference,
    init=_moe_fp4_block_scale_llama4_init,
)

# RoutingMethodType.RenormalizeNaive = 4 — Softmax → TopK → Renormalize
trtllm_fp4_block_scale_moe_renormalize_naive_routing_trace = _make_standard_fp4_moe_trace(
    name_prefix="moe_fp4_block_scale_renormalize_naive_routing",
    description="NvFP4 block-scale MoE with RenormalizeNaive routing (Softmax → TopK → Renormalize).",
    reference=_trtllm_fp4_block_scale_moe_renormalize_naive_routing_reference,
    init=_moe_fp4_block_scale_renormalize_naive_init,
)

# RoutingMethodType.TopK = 5 — plain TopK, uniform weights
trtllm_fp4_block_scale_moe_topk_routing_trace = _make_standard_fp4_moe_trace(
    name_prefix="moe_fp4_block_scale_topk_routing",
    description="NvFP4 block-scale MoE with TopK-only routing (no softmax, uniform weights).",
    reference=_trtllm_fp4_block_scale_moe_topk_routing_reference,
    init=_moe_fp4_block_scale_topk_init,
)

_FP4_MOE_TRACE_BY_ROUTING_TYPE = {
    0: trtllm_fp4_block_scale_moe_default_routing_trace,
    1: trtllm_fp4_block_scale_moe_renormalize_routing_trace,
    2: trtllm_fp4_block_scale_moe_ds_routing_trace,
    3: trtllm_fp4_block_scale_moe_llama4_routing_trace,
    4: trtllm_fp4_block_scale_moe_renormalize_naive_routing_trace,
    5: trtllm_fp4_block_scale_moe_topk_routing_trace,
    # 6 = Unspecified: no trace
}


def trtllm_fp4_block_scale_moe_trace_dispatch(**kwargs):
    """Return the FP4 TraceTemplate for the given ``routing_method_type``.

    Pass this as ``trace=trtllm_fp4_block_scale_moe_trace_dispatch`` to
    ``@flashinfer_api`` so the correct template is selected at call time::

        @flashinfer_api(trace=trtllm_fp4_block_scale_moe_trace_dispatch)
        def trtllm_fp4_block_scale_moe(..., routing_method_type: int = 0, ...):
            ...

    Returns ``None`` for ``RoutingMethodType.Unspecified`` (6).
    """
    routing_method_type = int(kwargs.get("routing_method_type", 0))
    return _FP4_MOE_TRACE_BY_ROUTING_TYPE.get(routing_method_type)


trtllm_fp4_block_scale_moe_trace_dispatch.templates = list(  # type: ignore[attr-defined]
    _FP4_MOE_TRACE_BY_ROUTING_TYPE.values()
)


# ---------------------------------------------------------------------------
# Additional MoE variants (CUTLASS fused MoE, bf16, routed, per-tensor, mxint4)
# ---------------------------------------------------------------------------

_MOE_COMMON_AXES: dict[str, Var | Const] = {
    "seq_len": Var(description="Number of input tokens."),
    "num_experts": Const(abbrev="", description="Total number of experts."),
    "top_k": Const(abbrev="topk"),
    "num_local_experts": Const(abbrev="e", description="Number of local experts."),
    "hidden_size": Const(abbrev="h"),
    "intermediate_size": Const(abbrev="i"),
}

# ---------------------------------------------------------------------------
# References for the additional MoE variants (bf16 / per-tensor FP8 / routed /
# mxint4). Each reference assumes inputs are already in their declared dtypes.
# ---------------------------------------------------------------------------


def _moe_expert_param(value, expert_idx, default, device):
    if value is None:
        return default
    if isinstance(value, torch.Tensor):
        value = value.reshape(-1)
        if value.numel() == 1:
            return value[0].to(device=device, dtype=torch.float32)
        return value[expert_idx].to(device=device, dtype=torch.float32)
    return float(value)


@torch.no_grad()
def _moe_bf16_run_experts(
    hidden_states,
    gemm1_weights,
    gemm2_weights,
    weights,
    topk_idx,
    local_expert_offset,
    E_global,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
    activation_type=ActivationType.Swiglu.value,
    situ_beta=None,
    situ_linear_beta=None,
):
    """Un-quantized (bf16) MoE expert computation."""
    activation_type = normalize_activation_type(activation_type)
    if situ_beta is not None or situ_linear_beta is not None:
        from ...fused_moe.cute_dsl.moe_utils import (
            validate_cute_dsl_moe_situ_config,
        )

        validate_cute_dsl_moe_situ_config(activation_type, situ_beta, situ_linear_beta)
    if activation_type not in (
        ActivationType.Swiglu,
        ActivationType.GegluTanh,
        ActivationType.Relu2,
    ):
        raise ValueError(
            f"Unsupported activation_type {activation_type!r}; "
            f"expected {ActivationType.Swiglu!r}, "
            f"{ActivationType.GegluTanh!r}, or {ActivationType.Relu2!r}"
        )
    T, H = hidden_states.shape
    E_local, gemm1_out, _ = gemm1_weights.shape
    I = gemm1_out // 2
    device = hidden_states.device
    A = hidden_states.to(torch.float32)
    W1 = gemm1_weights.to(torch.float32)
    W2 = gemm2_weights.to(torch.float32)
    output = torch.zeros((T, H), dtype=torch.float32, device=device)
    local_start = int(local_expert_offset)
    for le in range(E_local):
        ge = local_start + le
        if ge < 0 or ge >= E_global:
            continue
        sel_mask = (topk_idx == ge).any(dim=1)
        if not sel_mask.any():
            continue
        token_idx = torch.nonzero(sel_mask, as_tuple=False).squeeze(1)
        A_e = A.index_select(0, token_idx)
        G1 = A_e.matmul(W1[le].t())
        if activation_type == ActivationType.Relu2:
            act = torch.relu(G1) ** 2
        else:
            X1, X2 = G1[:, :I], G1[:, I:]
            if situ_beta is not None:
                gate = situ_beta * torch.tanh(X2 / situ_beta) * torch.sigmoid(X2)
                up = X1
                if situ_linear_beta is not None:
                    up = situ_linear_beta * torch.tanh(up / situ_linear_beta)
                act = gate * up
            elif activation_type == ActivationType.GegluTanh:
                act = torch.nn.functional.gelu(X2, approximate="tanh") * X1
            else:
                limit = _moe_expert_param(
                    gemm1_clamp_limit, le, DEFAULT_SWIGLU_LIMIT, X1.device
                )
                alpha = _moe_expert_param(
                    gemm1_alpha, le, DEFAULT_SWIGLU_ALPHA, X2.device
                )
                beta = _moe_expert_param(gemm1_beta, le, DEFAULT_SWIGLU_BETA, X1.device)
                up = torch.clamp(X1, min=-limit, max=limit)
                gate = torch.clamp(X2, max=limit)
                act = gate * torch.sigmoid(alpha * gate) * (up + beta)
        expert_out = act.matmul(W2[le].t())
        w_tok = weights.index_select(0, token_idx)
        match = (topk_idx.index_select(0, token_idx) == ge).float()
        w_e = (w_tok * match).sum(dim=1)
        output.index_add_(0, token_idx, expert_out * w_e.unsqueeze(1))
    return output.to(torch.bfloat16)


@torch.no_grad()
def _default_routing_weights(routing_logits, routing_bias, top_k, scale):
    logits = routing_logits.to(torch.float32)
    if routing_bias is not None:
        logits = logits + routing_bias.to(torch.float32).reshape(-1)
    s = torch.softmax(logits, dim=-1)
    _, topk_idx = torch.topk(s, k=int(top_k), dim=1, largest=True, sorted=False)
    return s.gather(1, topk_idx) * float(scale or 1.0), topk_idx


@torch.no_grad()
def _cutlass_fused_moe_reference(
    input,
    token_selected_experts,
    token_final_scales,
    fc1_expert_weights,
    fc2_expert_weights,
    **_unused,
):
    """Reference for CUTLASS fused MoE with precomputed routing."""
    E_global = fc1_expert_weights.shape[0]
    return _moe_bf16_run_experts(
        input,
        fc1_expert_weights,
        fc2_expert_weights,
        token_final_scales,
        token_selected_experts.to(torch.int64),
        local_expert_offset=0,
        E_global=E_global,
    )


@torch.no_grad()
def _trtllm_bf16_moe_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    gemm1_weights,
    gemm2_weights,
    num_experts,
    top_k,
    local_expert_offset,
    routed_scaling_factor=None,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
    **_unused,
):
    """Reference for TRT-LLM BF16 MoE (Default routing)."""
    w_topk, topk_idx = _default_routing_weights(
        routing_logits, routing_bias, top_k, routed_scaling_factor
    )
    return _moe_bf16_run_experts(
        hidden_states,
        gemm1_weights,
        gemm2_weights,
        w_topk,
        topk_idx,
        local_expert_offset,
        int(num_experts),
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_bf16_routed_moe_reference(
    topk_ids,
    hidden_states,
    gemm1_weights,
    gemm2_weights,
    num_experts,
    top_k,
    local_expert_offset,
    routed_scaling_factor=None,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
    **_unused,
):
    """Reference for TRT-LLM BF16 MoE with precomputed topk_ids."""
    T = topk_ids.shape[0]
    scale = float(routed_scaling_factor or 1.0)
    # Uniform weight per selected expert (real routing scales not available).
    w_topk = torch.full(
        (T, int(top_k)),
        scale / float(top_k),
        dtype=torch.float32,
        device=hidden_states.device,
    )
    return _moe_bf16_run_experts(
        hidden_states,
        gemm1_weights,
        gemm2_weights,
        w_topk,
        topk_ids.to(torch.int64),
        local_expert_offset,
        int(num_experts),
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_fp8_per_tensor_scale_moe_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    gemm1_weights,
    output1_scales_scalar,
    output1_scales_gate_scalar,
    gemm2_weights,
    output2_scales_scalar,
    num_experts,
    top_k,
    local_expert_offset,
    routed_scaling_factor=None,
    **_unused,
):
    """Reference for TRT-LLM FP8 per-tensor scale MoE. Dequantizes per-expert."""
    E_local = gemm1_weights.shape[0]
    w_topk, topk_idx = _default_routing_weights(
        routing_logits, routing_bias, top_k, routed_scaling_factor
    )
    # Per-expert dequant: each expert has its own scalar scale for FC1 gate,
    # FC1 up, and FC2. Scale broadcasts over the non-expert dims.
    W1 = gemm1_weights.to(torch.float32)
    W2 = gemm2_weights.to(torch.float32)
    s1 = output1_scales_scalar.to(torch.float32).view(E_local, 1, 1)
    s1g = output1_scales_gate_scalar.to(torch.float32).view(E_local, 1, 1)
    s2 = output2_scales_scalar.to(torch.float32).view(E_local, 1, 1)
    I = W1.shape[1] // 2
    # W1 is [E, 2I, H]: first half is gate, second half is up — apply scales.
    W1 = torch.cat([W1[:, :I] * s1g, W1[:, I:] * s1], dim=1)
    W2 = W2 * s2
    return _moe_bf16_run_experts(
        hidden_states.to(torch.float32),
        W1,
        W2,
        w_topk,
        topk_idx,
        local_expert_offset,
        int(num_experts),
    )


@torch.no_grad()
def _trtllm_fp8_per_tensor_scale_routed_moe_reference(
    topk_ids,
    hidden_states,
    gemm1_weights,
    output1_scales_scalar,
    output1_scales_gate_scalar,
    gemm2_weights,
    output2_scales_scalar,
    num_experts,
    local_expert_offset,
    **_unused,
):
    """Reference for routed TRT-LLM FP8 per-tensor scale MoE."""
    packed_topk = topk_ids.to(torch.int32)
    topk_idx = torch.bitwise_right_shift(packed_topk, 16).to(torch.int64)
    topk_weights = packed_topk.to(torch.int16).view(torch.bfloat16).to(torch.float32)

    T, H = hidden_states.shape
    E_local = gemm1_weights.shape[0]
    W1 = gemm1_weights.to(torch.float32)
    W2 = gemm2_weights.to(torch.float32)
    s1 = output1_scales_scalar.to(torch.float32).view(E_local, 1, 1)
    s1g = output1_scales_gate_scalar.to(torch.float32).view(E_local, 1, 1)
    s2 = output2_scales_scalar.to(torch.float32).view(E_local, 1, 1)
    I = W1.shape[1] // 2
    W1 = torch.cat([W1[:, :I] * s1g, W1[:, I:] * s1], dim=1)
    W2 = W2 * s2

    activations = hidden_states.to(torch.float32)
    output = torch.zeros((T, H), dtype=torch.float32, device=hidden_states.device)
    local_start = int(local_expert_offset)
    for local_idx in range(E_local):
        global_idx = local_start + local_idx
        if global_idx < 0 or global_idx >= int(num_experts):
            continue
        token_mask = (topk_idx == global_idx).any(dim=1)
        if not token_mask.any():
            continue
        token_idx = torch.nonzero(token_mask, as_tuple=False).squeeze(1)
        gemm1_out = activations.index_select(0, token_idx).matmul(W1[local_idx].t())
        gate, up = gemm1_out[:, :I], gemm1_out[:, I:]
        expert_out = (gate * torch.sigmoid(gate) * up).matmul(W2[local_idx].t())
        matches = (topk_idx.index_select(0, token_idx) == global_idx).to(torch.float32)
        weights = (topk_weights.index_select(0, token_idx) * matches).sum(dim=1)
        output.index_add_(0, token_idx, expert_out * weights.unsqueeze(1))
    return output.to(torch.bfloat16)


@torch.no_grad()
def _trtllm_fp8_block_scale_routed_moe_reference(
    topk_ids,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    num_experts,
    top_k,
    local_expert_offset,
    routed_scaling_factor=None,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
    **_unused,
):
    """Reference for TRT-LLM FP8 block-scale routed MoE (precomputed topk_ids).

    Reuses ``_fp8_moe_run_experts`` for the dequant + SwiGLU path, and builds
    a uniform per-token weight tensor (real routing scales are not available
    from topk_ids alone).
    """
    T = topk_ids.shape[0]
    TOP_K = int(top_k)
    scale = float(routed_scaling_factor or 1.0)
    w_topk = torch.full(
        (T, TOP_K),
        scale / TOP_K,
        dtype=torch.float32,
        device=hidden_states.device,
    )
    return _fp8_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        w_topk,
        topk_ids.to(torch.int64),
        local_expert_offset,
        int(num_experts),
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_fp4_block_scale_routed_moe_reference(
    topk_ids,
    hidden_states,
    hidden_states_scale,
    gemm1_weights,
    gemm1_weights_scale,
    gemm1_bias,
    gemm2_weights,
    gemm2_weights_scale,
    gemm2_bias,
    num_experts,
    top_k,
    local_expert_offset,
    routed_scaling_factor=None,
    activation_type=3,
    gemm1_alpha=None,
    gemm1_beta=None,
    gemm1_clamp_limit=None,
    **_unused,
):
    """Reference for TRT-LLM FP4 block-scale routed MoE (precomputed topk_ids)."""
    T = topk_ids.shape[0]
    TOP_K = int(top_k)
    scale = float(routed_scaling_factor or 1.0)
    w_topk = torch.full(
        (T, TOP_K),
        scale / TOP_K,
        dtype=torch.float32,
        device=hidden_states.device,
    )
    return _fp4_moe_run_experts(
        hidden_states,
        hidden_states_scale,
        gemm1_weights,
        gemm1_weights_scale,
        gemm2_weights,
        gemm2_weights_scale,
        gemm1_bias,
        gemm2_bias,
        w_topk,
        topk_ids.to(torch.int64),
        local_expert_offset,
        int(num_experts),
        activation_type=activation_type,
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
    )


@torch.no_grad()
def _trtllm_mxint4_block_scale_moe_reference(
    routing_logits,
    routing_bias,
    hidden_states,
    gemm1_weights,
    gemm1_weights_scale,
    gemm2_weights,
    gemm2_weights_scale,
    num_experts,
    top_k,
    local_expert_offset,
    routed_scaling_factor=None,
    **_unused,
):
    """Reference for TRT-LLM MxInt4 block-scale MoE.

    Weights are int4 packed as uint8 with bf16 per-32 block scales. Hidden
    states are bf16 (no activation quantization).
    """

    # Unpack int4: low nibble is first element, values are 4-bit signed (-8..7).
    def _unpack_int4(packed):
        lo = (packed & 0x0F).to(torch.int64)
        hi = ((packed >> 4) & 0x0F).to(torch.int64)
        # Sign-extend from 4-bit.
        lo = torch.where(lo >= 8, lo - 16, lo)
        hi = torch.where(hi >= 8, hi - 16, hi)
        stacked = torch.stack([lo, hi], dim=-1)
        return stacked.reshape(*packed.shape[:-1], packed.shape[-1] * 2).to(
            torch.float32
        )

    W1 = _unpack_int4(gemm1_weights)  # [E, 2I, H]
    W2 = _unpack_int4(gemm2_weights)  # [E, H, I]
    # Scales are bf16, broadcast per-32 along last axis.
    s1 = gemm1_weights_scale.to(torch.float32)
    s2 = gemm2_weights_scale.to(torch.float32)
    block1 = W1.shape[-1] // s1.shape[-1]
    block2 = W2.shape[-1] // s2.shape[-1]
    W1 = W1 * s1.repeat_interleave(block1, dim=-1)
    W2 = W2 * s2.repeat_interleave(block2, dim=-1)

    w_topk, topk_idx = _default_routing_weights(
        routing_logits, routing_bias, top_k, routed_scaling_factor
    )
    return _moe_bf16_run_experts(
        hidden_states,
        W1,
        W2,
        w_topk,
        topk_idx,
        local_expert_offset,
        int(num_experts),
    )


# CUTLASS fused MoE: precomputed token_selected_experts + token_final_scales
cutlass_fused_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="cutlass_fused_moe",
    description="CUTLASS fused MoE. Accepts precomputed per-token expert selections.",
    axes={
        "seq_len": Var(description="Number of input tokens."),
        "num_local_experts": Const(abbrev="e"),
        "hidden_size": Const(abbrev="h"),
        "intermediate_size": Const(abbrev="i"),
        "top_k": Const(abbrev="topk"),
    },
    inputs={
        "input": Tensor(
            ["seq_len", "hidden_size"],
            description="Input hidden states (bf16/fp8/fp4 depending on quant config).",
        ),
        "token_selected_experts": Tensor(
            ["seq_len", "top_k"],
            dtype="int32",
            description="Precomputed top-k expert ids per token.",
        ),
        "token_final_scales": Tensor(
            ["seq_len", "top_k"],
            dtype="float32",
            description="Precomputed per-token expert scales.",
        ),
        "fc1_expert_weights": Tensor(
            ["num_local_experts", "gemm1_out_size", "hidden_size"],
            description="FC1 weights per expert.",
        ),
        "fc2_expert_weights": Tensor(
            ["num_local_experts", "hidden_size", "intermediate_size"],
            description="FC2 weights per expert.",
        ),
    },
    outputs={
        "output": Tensor(["seq_len", "hidden_size"], dtype="bfloat16"),
    },
    tags=["status:verified", "backend:cutlass"],
    reference=_cutlass_fused_moe_reference,
)
cutlass_fused_moe_trace.axes["gemm1_out_size"] = Const(
    abbrev="", description="FC1 output size (typically 2 * intermediate_size)."
)


# ---------------------------------------------------------------------------
# mono_moe (monomoe) — single-kernel block-FP8 top-K MoE, Qwen3.5-35B shape
# ---------------------------------------------------------------------------
# Fixed-shape Hopper (SM90a) MonoMoe kernel: routing is FUSED in-kernel from
# router_logits, so unlike cutlass_fused_moe this template takes the logits
# (not precomputed selections).  Weights are block-FP8 (128x128):
#   activations_in     : [seq_len, hidden_size]                 bf16
#   router_logits      : [seq_len, num_experts]                 bf16
#   expert_weights_up  : [num_experts, gemm1_out_size, hidden_size]   fp8 (gate||up)
#   expert_scales_up   : [num_experts, num_gemm1_out_blocks, num_hidden_blocks] fp32
#   expert_weights_down: [num_experts, hidden_size, intermediate_size]    fp8
#   expert_scales_down : [num_experts, num_hidden_blocks, num_intermediate_blocks] fp32


@torch.no_grad()
def _mono_moe_reference(
    activations_in,
    router_logits,
    expert_weights_up,
    expert_scales_up,
    expert_weights_down,
    expert_scales_down,
    top_k,
    scoring_func="softmax",
    renormalize=True,
    **_unused,
):
    """Reference for mono_moe: in-kernel softmax/sigmoid routing + block-FP8 SwiGLU.

    Mirrors the monomoe math: score router_logits (softmax or sigmoid),
    take top-k, optionally renormalize the selected weights to sum to 1,
    then run the shared block-FP8 dequant + SwiGLU + GEMM helper.  The
    up-projection weights use the [gate || up] half ordering (monomoe's
    convention), which `_fp8_moe_run_experts` already assumes (X1=gate,
    X2=up with silu on X2).
    """
    E_global = router_logits.shape[1]
    TOP_K = int(top_k)
    logits = router_logits.to(torch.float32)
    if scoring_func == "sigmoid":
        scores = torch.sigmoid(logits)
    else:
        scores = torch.softmax(logits, dim=-1)
    topk_vals, topk_idx = torch.topk(scores, k=TOP_K, dim=1, largest=True, sorted=False)
    if renormalize:
        weights = topk_vals / (topk_vals.sum(dim=1, keepdim=True) + 1e-20)
    else:
        weights = topk_vals
    # hidden_states_scale: the helper expects [H/128, T] (block-wise act
    # scale).  monomoe quantizes activations internally, so for the
    # reference we treat the bf16 input as already-dequantized by passing a
    # unit scale of the right shape and an fp8 round-trip of the input.
    T, H = activations_in.shape
    BLOCK = 128
    hs_fp8, hs_scale_TxNb = fp8_block_quant_1d(
        activations_in.to(torch.bfloat16), block=BLOCK
    )
    hs_scale = hs_scale_TxNb.transpose(0, 1).contiguous()  # [H/128, T]
    return _fp8_moe_run_experts(
        hs_fp8,
        hs_scale,
        expert_weights_up,
        expert_scales_up,
        expert_weights_down,
        expert_scales_down,
        weights,
        topk_idx,
        local_expert_offset=0,
        E_global=E_global,
    )


def _mono_moe_init(
    *,
    seq_len: int,
    num_experts: int = 256,
    top_k: int = 8,
    hidden_size: int = 2048,
    intermediate_size: int = 512,
    gemm1_out_size: int = 0,  # derived: 2 * intermediate_size
    num_hidden_blocks: int = 0,  # derived: hidden_size / 128
    num_intermediate_blocks: int = 0,  # derived: intermediate_size / 128
    num_gemm1_out_blocks: int = 0,  # derived: gemm1_out_size / 128
    device: str = "cuda",
    seed: int = 0,
):
    """Build inputs for ``mono_moe`` (fixed Qwen3.5-35B block-FP8 shape).

    The monomoe kernel is hard-specialized to E=256, N=512, K=2048, BS<=8,
    so the defaults match that shape (``seq_len`` is the only Var, capped at
    8 by the kernel).  Weights are block-FP8 (128×128) quantized via the
    shared ``fp8_block_quant_2d`` helper; activations / router_logits are
    bf16 (the kernel quantizes activations internally).
    """
    del gemm1_out_size, num_hidden_blocks, num_intermediate_blocks
    del num_gemm1_out_blocks
    torch.manual_seed(seed)
    BLOCK = 128
    activations_in = torch.randn(
        seq_len, hidden_size, dtype=torch.bfloat16, device=device
    )
    router_logits = torch.randn(
        seq_len, num_experts, dtype=torch.bfloat16, device=device
    )
    w13_bf16 = 0.1 * torch.randn(
        num_experts,
        2 * intermediate_size,
        hidden_size,
        dtype=torch.bfloat16,
        device=device,
    )
    w2_bf16 = 0.1 * torch.randn(
        num_experts,
        hidden_size,
        intermediate_size,
        dtype=torch.bfloat16,
        device=device,
    )
    w13, w13s = fp8_block_quant_2d(w13_bf16, block=BLOCK)
    w2, w2s = fp8_block_quant_2d(w2_bf16, block=BLOCK)
    return {
        "activations_in": activations_in,
        "router_logits": router_logits,
        "expert_weights_up": w13,
        "expert_scales_up": w13s,
        "expert_weights_down": w2,
        "expert_scales_down": w2s,
        "top_k": int(top_k),
        "scoring_func": "softmax",
        "renormalize": True,
    }


mono_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="mono_moe",
    description=(
        "Single-kernel (monomoe) block-FP8 top-K MoE for the Qwen3.5-35B shape "
        "on Hopper (SM90a). Routing (softmax/sigmoid top-K + renormalize) is "
        "fused in-kernel from router_logits; weights are block-FP8 (128x128)."
    ),
    axes={
        "seq_len": Var(description="Sequence length (number of tokens, <= 8)."),
        "num_experts": Const(description="Total number of experts.", abbrev=""),
        "top_k": Const(
            description="Number of experts to route to per token.", abbrev="topk"
        ),
        "hidden_size": Const(description="Hidden dimension size (K).", abbrev="h"),
        "intermediate_size": Const(
            description="MoE intermediate layer size (N).", abbrev="i"
        ),
        "gemm1_out_size": Const(
            description="Up-projection output rows (2 * intermediate_size, gate||up).",
            abbrev="",
        ),
        "num_hidden_blocks": Const(
            description="Quantized blocks along hidden_size (block_size=128).",
            abbrev="",
        ),
        "num_intermediate_blocks": Const(
            description="Quantized blocks along intermediate_size (block_size=128).",
            abbrev="",
        ),
        "num_gemm1_out_blocks": Const(
            description="Quantized blocks along gemm1_out_size (block_size=128).",
            abbrev="",
        ),
    },
    inputs={
        "activations_in": Tensor(
            ["seq_len", "hidden_size"],
            dtype="bfloat16",
            description="Input activations (bf16; quantized to FP8 in-kernel).",
        ),
        "router_logits": Tensor(
            ["seq_len", "num_experts"],
            dtype="bfloat16",
            description="Router logits for in-kernel expert selection.",
        ),
        "expert_weights_up": Tensor(
            ["num_experts", "gemm1_out_size", "hidden_size"],
            description="Up/gate projection weights (FP8, [gate || up] rows).",
        ),
        "expert_scales_up": Tensor(
            ["num_experts", "num_gemm1_out_blocks", "num_hidden_blocks"],
            dtype="float32",
            description="Block-wise (128x128) FP8 scales for up/gate weights.",
        ),
        "expert_weights_down": Tensor(
            ["num_experts", "hidden_size", "intermediate_size"],
            description="Down projection weights (FP8, raw row-major).",
        ),
        "expert_scales_down": Tensor(
            ["num_experts", "num_hidden_blocks", "num_intermediate_blocks"],
            dtype="float32",
            description="Block-wise (128x128) FP8 scales for down weights.",
        ),
        "top_k": Scalar(
            "int32",
            description="Number of experts to route to per token (1..8).",
        ),
    },
    outputs={
        "output": Tensor(
            ["seq_len", "hidden_size"],
            dtype="bfloat16",
            description="Final MoE output tensor (also the return value).",
        ),
    },
    tags=["status:verified", "quantization:float8_e4m3fn", "arch:sm90a"],
    reference=_mono_moe_reference,
    init=_mono_moe_init,
)


# Shared factory for the remaining trtllm_* variants
_TRTLLM_MOE_COMMON_INPUTS: dict[str, Tensor | Scalar] = {
    "routing_logits": Tensor(
        ["seq_len", "num_experts"], description="Routing logits for expert selection."
    ),
    "routing_bias": Tensor(
        ["num_experts"], optional=True, description="Optional routing bias."
    ),
    "hidden_states": Tensor(
        ["seq_len", "hidden_size"],
        description="Input hidden states (dtype depends on variant).",
    ),
    "gemm1_weights": Tensor(
        ["num_local_experts", "gemm1_out_size", "hidden_size"],
        description="FC1 weights (gate+up).",
    ),
    "gemm2_weights": Tensor(
        ["num_local_experts", "hidden_size", "intermediate_size"],
        description="FC2 weights (down).",
    ),
    "top_k": Scalar("int32", description="Number of experts to route per token."),
    "n_group": Scalar(
        "int32", optional=True, description="Expert groups (DeepSeek-V3)."
    ),
    "topk_group": Scalar(
        "int32", optional=True, description="Groups to keep (DeepSeek-V3)."
    ),
    "local_expert_offset": Scalar(
        "int32", description="Offset of local experts in global expert space."
    ),
    "routed_scaling_factor": Scalar(
        "float32", optional=True, description="Scaling factor for routing weights."
    ),
    "routing_method_type": Scalar(
        "int32",
        optional=True,
        description="0=Default, 1=Renormalize, 2=DeepSeekV3, 3=Llama4, 4=RenormalizeNaive, 5=TopK.",
    ),
}

_TRTLLM_MOE_COMMON_AXES: dict[str, Var | Const] = {
    **_MOE_COMMON_AXES,
    "gemm1_out_size": Const(abbrev="", description="2 * intermediate_size."),
}

_TRTLLM_MOE_COMMON_OUTPUTS: dict[str, Tensor | Scalar] = {
    "output": Tensor(
        ["seq_len", "hidden_size"], dtype="bfloat16", description="MoE output."
    ),
}

_TRTLLM_BF16_SWIGLU_OA_INPUTS: dict[str, Tensor] = {
    "gemm1_alpha": Tensor(
        ["num_local_experts"],
        dtype="float32",
        optional=True,
        description="Optional per-expert SwiGLU OA alpha.",
    ),
    "gemm1_beta": Tensor(
        ["num_local_experts"],
        dtype="float32",
        optional=True,
        description="Optional per-expert SwiGLU OA beta.",
    ),
    "gemm1_clamp_limit": Tensor(
        ["num_local_experts"],
        dtype="float32",
        optional=True,
        description="Optional per-expert SwiGLU OA clamp limit.",
    ),
}

# BF16 MoE (no quantization)
trtllm_bf16_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="trtllm_bf16_moe",
    description="TRT-LLM BF16 MoE (no quantization).",
    axes=dict(_TRTLLM_MOE_COMMON_AXES),
    inputs={**_TRTLLM_MOE_COMMON_INPUTS, **_TRTLLM_BF16_SWIGLU_OA_INPUTS},
    outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS),
    tags=["status:verified", "backend:trtllm"],
    reference=_trtllm_bf16_moe_reference,
)

# BF16 routed MoE (accepts precomputed topk_ids instead of routing_logits)
# num_experts / intermediate_size become Var in routed variants because they
# are passed as scalar kwargs (no routing_logits tensor to resolve from).
_TRTLLM_MOE_ROUTED_AXES: dict[str, Var | Const] = {
    **_TRTLLM_MOE_COMMON_AXES,
    "num_experts": Var(description="Total number of experts (passed as kwarg)."),
    "intermediate_size": Var(
        description="MoE intermediate layer size (passed as kwarg)."
    ),
}
trtllm_bf16_routed_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="trtllm_bf16_routed_moe",
    description="TRT-LLM BF16 MoE with precomputed topk_ids.",
    axes=dict(_TRTLLM_MOE_ROUTED_AXES),
    inputs={
        "topk_ids": Tensor(
            ["seq_len", "top_k"],
            dtype="int32",
            description="Precomputed top-k expert ids per token.",
        ),
        "hidden_states": _TRTLLM_MOE_COMMON_INPUTS["hidden_states"],
        "gemm1_weights": _TRTLLM_MOE_COMMON_INPUTS["gemm1_weights"],
        "gemm2_weights": _TRTLLM_MOE_COMMON_INPUTS["gemm2_weights"],
        "num_experts": Scalar("int32", description="Total number of experts."),
        "top_k": _TRTLLM_MOE_COMMON_INPUTS["top_k"],
        "local_expert_offset": _TRTLLM_MOE_COMMON_INPUTS["local_expert_offset"],
        "routed_scaling_factor": _TRTLLM_MOE_COMMON_INPUTS["routed_scaling_factor"],
        **_TRTLLM_BF16_SWIGLU_OA_INPUTS,
    },
    outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS),
    tags=["status:verified", "backend:trtllm"],
    reference=_trtllm_bf16_routed_moe_reference,
)

# FP8 per-tensor scale MoE
trtllm_fp8_per_tensor_scale_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="trtllm_fp8_per_tensor_scale_moe",
    description="TRT-LLM FP8 MoE with per-tensor activation/weight scales.",
    axes=dict(_TRTLLM_MOE_COMMON_AXES),
    inputs={
        **_TRTLLM_MOE_COMMON_INPUTS,
        "output1_scales_scalar": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC1 output scale.",
        ),
        "output1_scales_gate_scalar": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC1 gate scale.",
        ),
        "output2_scales_scalar": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC2 output scale.",
        ),
    },
    outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS),
    tags=["status:verified", "backend:trtllm", "quantization:float8_e4m3fn"],
    reference=_trtllm_fp8_per_tensor_scale_moe_reference,
)


# FP8 per-tensor scale routed (packed precomputed topk_ids and weights)
trtllm_fp8_per_tensor_scale_routed_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="trtllm_fp8_per_tensor_scale_routed_moe",
    description="TRT-LLM FP8 per-tensor scale MoE with packed precomputed routing.",
    axes=dict(_TRTLLM_MOE_ROUTED_AXES),
    inputs={
        "topk_ids": Tensor(
            ["seq_len", "top_k"],
            dtype="int32",
            description="Packed BF16 routing weight and global expert id.",
        ),
        "routing_bias": Tensor(
            ["num_experts"], optional=True, description="Optional routing bias."
        ),
        "hidden_states": Tensor(
            ["seq_len", "hidden_size"], description="FP8-quantized hidden states."
        ),
        "gemm1_weights": Tensor(
            ["num_local_experts", "gemm1_out_size", "hidden_size"],
            description="FC1 FP8 weights.",
        ),
        "output1_scales_scalar": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC1 output scale.",
        ),
        "output1_scales_gate_scalar": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC1 gate scale.",
        ),
        "gemm2_weights": Tensor(
            ["num_local_experts", "hidden_size", "intermediate_size"],
            description="FC2 FP8 weights.",
        ),
        "output2_scales_scalar": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC2 output scale.",
        ),
        "num_experts": Scalar("int32", description="Total number of experts."),
        "top_k": Scalar("int32", description="Number of experts routed per token."),
        "local_expert_offset": Scalar(
            "int32", description="Offset of local experts in global expert space."
        ),
        "routed_scaling_factor": Scalar(
            "float32", optional=True, description="Scaling factor for routing weights."
        ),
    },
    outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS),
    tags=["status:verified", "backend:trtllm", "quantization:float8_e4m3fn"],
    reference=_trtllm_fp8_per_tensor_scale_routed_moe_reference,
)


# FP8 block-scale routed (precomputed topk_ids)
trtllm_fp8_block_scale_routed_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="trtllm_fp8_block_scale_routed_moe",
    description="TRT-LLM FP8 block-scale MoE with precomputed topk_ids.",
    axes={
        **_TRTLLM_MOE_ROUTED_AXES,
        "num_hidden_blocks": Const(abbrev=""),
        "num_intermediate_blocks": Const(abbrev=""),
        "num_gemm1_out_blocks": Const(abbrev=""),
    },
    inputs={
        "topk_ids": Tensor(
            ["seq_len", "top_k"], dtype="int32", description="Precomputed top-k."
        ),
        "routing_bias": Tensor(
            ["num_experts"], optional=True, description="Optional routing bias."
        ),
        "hidden_states": Tensor(
            ["seq_len", "hidden_size"],
            description="FP8-quantized hidden states.",
        ),
        "hidden_states_scale": Tensor(
            ["num_hidden_blocks", "seq_len"],
            description="Block-wise hidden_states scale.",
        ),
        "gemm1_weights": Tensor(
            ["num_local_experts", "gemm1_out_size", "hidden_size"],
            description="FC1 FP8 weights.",
        ),
        "gemm1_weights_scale": Tensor(
            ["num_local_experts", "num_gemm1_out_blocks", "num_hidden_blocks"],
            description="FC1 block-wise scale.",
        ),
        "gemm1_alpha": Tensor(
            ["num_local_experts"],
            dtype="float32",
            optional=True,
            description="Optional MxFp8-only per-expert SwiGLU OA alpha.",
        ),
        "gemm1_beta": Tensor(
            ["num_local_experts"],
            dtype="float32",
            optional=True,
            description="Optional MxFp8-only per-expert SwiGLU OA beta.",
        ),
        "gemm1_clamp_limit": Tensor(
            ["num_local_experts"],
            dtype="float32",
            optional=True,
            description="Optional MxFp8-only per-expert SwiGLU OA clamp limit.",
        ),
        "gemm2_weights": Tensor(
            ["num_local_experts", "hidden_size", "intermediate_size"],
            description="FC2 FP8 weights.",
        ),
        "gemm2_weights_scale": Tensor(
            ["num_local_experts", "num_hidden_blocks", "num_intermediate_blocks"],
            description="FC2 block-wise scale.",
        ),
        "num_experts": Scalar("int32", description="Total number of experts."),
        "top_k": Scalar("int32"),
        "local_expert_offset": Scalar("int32"),
        "routed_scaling_factor": Scalar("float32", optional=True),
    },
    outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS),
    tags=["status:verified", "backend:trtllm", "quantization:float8_e4m3fn"],
    reference=_trtllm_fp8_block_scale_routed_moe_reference,
)

# FP4 block-scale routed (precomputed topk_ids)
trtllm_fp4_block_scale_routed_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="trtllm_fp4_block_scale_routed_moe",
    description="TRT-LLM NvFP4 block-scale MoE with precomputed topk_ids.",
    axes={
        **_TRTLLM_MOE_ROUTED_AXES,
        "num_packed_hidden": Const(abbrev=""),
        # Var rather than Const because hidden_states_scale is optional and the
        # other tensors using this axis may have different shapes in routed mode.
        "num_fp4_hidden_blocks": Var(
            description="NvFP4 block count along hidden_size."
        ),
        "num_packed_intermediate": Const(abbrev=""),
        "num_fp4_intermediate_blocks": Const(abbrev=""),
        "activation_type": Const(
            description="Fused activation type; 10 selects SiTU.", abbrev="act"
        ),
    },
    inputs={
        "topk_ids": Tensor(
            ["seq_len", "top_k"], dtype="int32", description="Precomputed top-k."
        ),
        "routing_bias": Tensor(
            ["num_experts"], optional=True, description="Optional routing bias."
        ),
        "hidden_states": Tensor(
            ["seq_len", "num_packed_hidden"],
            description="NvFP4-packed hidden states.",
        ),
        "hidden_states_scale": Tensor(
            ["seq_len", "num_fp4_hidden_blocks"],
            optional=True,
            description="NvFP4 hidden_states scale.",
        ),
        "gemm1_weights": Tensor(
            ["num_local_experts", "gemm1_out_size", "num_packed_hidden"],
            description="FC1 NvFP4 weights.",
        ),
        "gemm1_weights_scale": Tensor(
            ["num_local_experts", "gemm1_out_size", "num_fp4_hidden_blocks"],
            description="FC1 NvFP4 scale.",
        ),
        "gemm1_alpha": Tensor(
            ["num_local_experts"],
            dtype="float32",
            optional=True,
            description="Per-expert SiTU alpha or SwiGLU alpha.",
        ),
        "gemm1_beta": Tensor(
            ["num_local_experts"],
            dtype="float32",
            optional=True,
            description="Per-expert SiTU beta or SwiGLU beta.",
        ),
        "gemm1_clamp_limit": Tensor(
            ["num_local_experts"],
            dtype="float32",
            optional=True,
            description="Per-expert gated-activation clamp limit.",
        ),
        "gemm2_weights": Tensor(
            ["num_local_experts", "hidden_size", "num_packed_intermediate"],
            description="FC2 NvFP4 weights.",
        ),
        "gemm2_weights_scale": Tensor(
            ["num_local_experts", "hidden_size", "num_fp4_intermediate_blocks"],
            description="FC2 NvFP4 scale.",
        ),
        "num_experts": Scalar("int32", description="Total number of experts."),
        "top_k": Scalar("int32"),
        "local_expert_offset": Scalar("int32"),
        "routed_scaling_factor": Scalar("float32", optional=True),
        "activation_type": Scalar(
            "int32", description="Fused activation type; 10 selects SiTU."
        ),
    },
    outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS),
    tags=["status:experimental", "backend:trtllm", "quantization:nvfp4"],
    reference=_trtllm_fp4_block_scale_routed_moe_reference,
)

# MxInt4 block-scale MoE
trtllm_mxint4_block_scale_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="trtllm_mxint4_block_scale_moe",
    description="TRT-LLM MxInt4 block-scale MoE.",
    axes={
        **_TRTLLM_MOE_COMMON_AXES,
        "intermediate_size": Var(description="MoE intermediate size (kwarg)."),
        "num_packed_hidden": Const(abbrev=""),
        "num_mxint4_hidden_blocks": Const(abbrev=""),
        "num_packed_intermediate": Const(abbrev=""),
        "num_mxint4_intermediate_blocks": Const(abbrev=""),
    },
    inputs={
        "routing_logits": Tensor(
            ["seq_len", "num_experts"], description="Routing logits."
        ),
        "routing_bias": Tensor(
            ["num_experts"], optional=True, description="Optional routing bias."
        ),
        "hidden_states": Tensor(
            ["seq_len", "hidden_size"],
            description="BF16/FP16 hidden states (quantized internally).",
        ),
        "gemm1_weights": Tensor(
            ["num_local_experts", "gemm1_out_size", "num_packed_hidden"],
            description="FC1 MxInt4-packed weights.",
        ),
        "gemm1_weights_scale": Tensor(
            ["num_local_experts", "gemm1_out_size", "num_mxint4_hidden_blocks"],
            description="FC1 MxInt4 scales.",
        ),
        "gemm2_weights": Tensor(
            ["num_local_experts", "hidden_size", "num_packed_intermediate"],
            description="FC2 MxInt4-packed weights.",
        ),
        "gemm2_weights_scale": Tensor(
            ["num_local_experts", "hidden_size", "num_mxint4_intermediate_blocks"],
            description="FC2 MxInt4 scales.",
        ),
        "top_k": Scalar("int32"),
        "n_group": Scalar("int32", optional=True),
        "topk_group": Scalar("int32", optional=True),
        "local_expert_offset": Scalar("int32"),
        "routed_scaling_factor": Scalar("float32", optional=True),
        "routing_method_type": Scalar("int32", optional=True),
    },
    outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS),
    tags=["status:experimental", "backend:trtllm", "quantization:mxint4"],
    reference=_trtllm_mxint4_block_scale_moe_reference,
)


# ---------------------------------------------------------------------------
# CuteDSL MoE variants (precomputed routing, NvFP4 weights on SM100+)
# ---------------------------------------------------------------------------

cute_dsl_fused_moe_nvfp4_trace = TraceTemplate(
    op_type="moe",
    name_prefix="cute_dsl_fused_moe_nvfp4",
    description=(
        "CuteDSL NVFP4 fused MoE (SM100/SM103). Runs explicit W4A4 or W4A16 "
        "compute with precomputed top-k routing and per-expert alpha scales."
    ),
    axes={
        "num_tokens": Var(description="Total tokens across the batch."),
        "num_experts": Const(abbrev="", description="Total number of experts."),
        "top_k": Const(abbrev="topk"),
        "num_local_experts": Const(abbrev="e"),
        "hidden_size": Const(abbrev="h"),
        "intermediate_size": Var(description="MoE intermediate size (kwarg)."),
        "input_width": Var(
            description="hidden_size // 2 for NVFP4 input or hidden_size for BF16."
        ),
        "num_packed_hidden": Var(description="hidden_size // 2 (NvFP4 packed)."),
        "num_packed_intermediate": Var(
            description="intermediate_size // 2 (NvFP4 packed)."
        ),
        "num_fp4_hidden_blocks": Var(
            description="NvFP4 scale-factor count along hidden_size."
        ),
        "num_fp4_intermediate_blocks": Var(
            description="NvFP4 scale-factor count along intermediate_size."
        ),
        "gemm1_out_size": Const(
            abbrev="",
            description=(
                "FC1 output rows: 2 * intermediate_size for gated SwiGLU, "
                "intermediate_size for non-gated ReLU^2."
            ),
        ),
    },
    inputs={
        "x": Tensor(
            ["num_tokens", "input_width"],
            description="Packed NVFP4 input for W4A4 or BF16 input for W4A16.",
        ),
        "x_sf": Tensor(
            ["num_tokens", "num_fp4_hidden_blocks"],
            optional=True,
            description="NVFP4 activation scales for W4A4; omitted for W4A16.",
        ),
        "token_selected_experts": Tensor(
            ["num_tokens", "top_k"],
            dtype="int32",
            description="Precomputed top-k expert ids per token.",
        ),
        "token_final_scales": Tensor(
            ["num_tokens", "top_k"],
            dtype="float32",
            description="Precomputed per-token routing scales.",
        ),
        "w1_weight": Tensor(
            ["num_local_experts", "gemm1_out_size", "num_packed_hidden"],
            description="FC1 weights, NvFP4-packed.",
        ),
        "w1_weight_sf": Tensor(
            ["num_local_experts", "gemm1_out_size", "num_fp4_hidden_blocks"],
            description="FC1 NvFP4 scales.",
        ),
        "w1_alpha": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC1 global scale.",
        ),
        "fc2_input_scale": Tensor(
            ["one"],
            dtype="float32",
            optional=True,
            description=(
                "Global scale for W4A4 FC2 input quantization; omitted for W4A16."
            ),
        ),
        "w2_weight": Tensor(
            ["num_local_experts", "hidden_size", "num_packed_intermediate"],
            description="FC2 weights, NvFP4-packed.",
        ),
        "w2_weight_sf": Tensor(
            ["num_local_experts", "hidden_size", "num_fp4_intermediate_blocks"],
            description="FC2 NvFP4 scales.",
        ),
        "w2_alpha": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC2 global scale.",
        ),
        "per_token_scale": Tensor(
            ["num_tokens"],
            dtype="float32",
            optional=True,
            description="Optional W4A4 per-token input row scale.",
        ),
        "quant_mode": Scalar(
            "string",
            optional=True,
            description="Compute mode: 'nvfp4'/'w4a4' or 'w4a16'.",
        ),
        "num_experts": Scalar("int32", description="Total number of experts."),
        "top_k": Scalar("int32", description="Number of experts per token."),
        "local_expert_offset": Scalar(
            "int32", optional=True, description="Offset of local experts."
        ),
        "activation_type": Scalar(
            "int32",
            optional=True,
            description=(
                "GEMM1 activation type: ActivationType.Swiglu for gated "
                "SwiGLU/OAI/SiTU, ActivationType.GegluTanh for "
                "tanh-approximate GeGLU, or ActivationType.Relu2 for "
                "non-gated ReLU^2. "
                "Determines gemm1_out_size."
            ),
        ),
        "swiglu_alpha": Scalar(
            "float32",
            optional=True,
            description="SwiGLU sigmoid multiplier.",
        ),
        "swiglu_beta": Scalar(
            "float32",
            optional=True,
            description="SwiGLU up-projection bias.",
        ),
        "swiglu_limit": Scalar(
            "float32",
            optional=True,
            description="SwiGLU clamp limit.",
        ),
        "situ_beta": Scalar(
            "float32",
            optional=True,
            description="SiTU gate tanh-clamp beta; enables SiTU when set.",
        ),
        "situ_linear_beta": Scalar(
            "float32",
            optional=True,
            description="Optional SiTU up-branch tanh-clamp beta.",
        ),
    },
    outputs={
        "output": Tensor(
            ["num_tokens", "hidden_size"],
            dtype="bfloat16",
            description="MoE output.",
        ),
    },
    tags=["status:experimental", "backend:cute-dsl", "quantization:nvfp4"],
)
cute_dsl_fused_moe_nvfp4_trace.axes["one"] = Var(
    description="Placeholder for shape [1] scalars."
)

_cute_dsl_wrapper_inputs = dict(cute_dsl_fused_moe_nvfp4_trace.inputs)
# num_experts / top_k live on the wrapper instance (set in __init__), not on run().
_cute_dsl_wrapper_inputs["num_experts"] = Scalar(
    "int32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)
_cute_dsl_wrapper_inputs["top_k"] = Scalar(
    "int32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)
_cute_dsl_wrapper_inputs["activation_type"] = Scalar(
    "int32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)
_cute_dsl_wrapper_inputs["swiglu_alpha"] = Scalar(
    "float32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)
_cute_dsl_wrapper_inputs["swiglu_beta"] = Scalar(
    "float32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)
_cute_dsl_wrapper_inputs["swiglu_limit"] = Scalar(
    "float32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)
_cute_dsl_wrapper_inputs["quant_mode"] = Scalar(
    "string",
    optional=True,
    description="Compute mode set at wrapper __init__, not passed to run().",
)
_cute_dsl_wrapper_inputs["situ_beta"] = Scalar(
    "float32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)
_cute_dsl_wrapper_inputs["situ_linear_beta"] = Scalar(
    "float32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)

_cute_dsl_wrapper_axes = dict(cute_dsl_fused_moe_nvfp4_trace.axes)
# num_experts / top_k are set at __init__ time — no tensor on run() has a
# num_experts dim, so the axis must be a Var here.
_cute_dsl_wrapper_axes["num_experts"] = Var(description="Total number of experts.")
_cute_dsl_wrapper_axes["top_k"] = Var(description="Experts per token.")

cute_dsl_moe_wrapper_run_trace = TraceTemplate(
    op_type="moe",
    name_prefix="cute_dsl_moe_wrapper",
    description=(
        "CuteDslMoEWrapper.run(): stateful version of cute_dsl_fused_moe_nvfp4 "
        "(same schema; wrapper persists autotuning state across calls)."
    ),
    axes=_cute_dsl_wrapper_axes,
    inputs=_cute_dsl_wrapper_inputs,
    outputs=dict(cute_dsl_fused_moe_nvfp4_trace.outputs),
    tags=cute_dsl_fused_moe_nvfp4_trace.tags,
)


# ---------------------------------------------------------------------------
# B12x MoE (SM120/SM121 CuTe-DSL, bf16 input + FP4 packed weights)
# ---------------------------------------------------------------------------

b12x_fused_moe_trace = TraceTemplate(
    op_type="moe",
    name_prefix="b12x_fused_moe",
    description=(
        "B12x CuTe-DSL fused MoE (SM120/SM121). BF16 input, FP4-packed "
        "weights, precomputed top-k routing; fuses quant + FC1 + activation + "
        "FC2 + scatter."
    ),
    axes={
        "num_tokens": Var(),
        "num_experts": Const(abbrev="", description="Total number of experts."),
        "top_k": Const(abbrev="topk"),
        "num_local_experts": Const(abbrev="e"),
        "hidden_size": Const(abbrev="h"),
        "intermediate_size": Var(description="MoE intermediate size (kwarg)."),
        "num_packed_hidden": Var(description="hidden_size // 2."),
        "num_packed_intermediate": Var(description="intermediate_size // 2."),
        "num_fp4_hidden_blocks": Var(),
        "num_fp4_intermediate_blocks": Var(),
        "gemm1_out_size": Const(
            abbrev="",
            description="2*I (SwiGLU) or I (ReLU2).",
        ),
        "input_global_scale_size": Var(
            description="1 (shared scale) or num_local_experts."
        ),
    },
    inputs={
        "x": Tensor(
            ["num_tokens", "hidden_size"], description="BF16 input activations."
        ),
        "w1_weight": Tensor(
            ["num_local_experts", "gemm1_out_size", "num_packed_hidden"],
            description="FC1 weights, FP4-packed.",
        ),
        "w1_weight_sf": Tensor(
            ["num_local_experts", "gemm1_out_size", "num_fp4_hidden_blocks"],
            description="FC1 FP4 scales.",
        ),
        "w2_weight": Tensor(
            ["num_local_experts", "hidden_size", "num_packed_intermediate"],
            description="FC2 weights, FP4-packed.",
        ),
        "w2_weight_sf": Tensor(
            ["num_local_experts", "hidden_size", "num_fp4_intermediate_blocks"],
            description="FC2 FP4 scales.",
        ),
        "token_selected_experts": Tensor(
            ["num_tokens", "top_k"],
            dtype="int32",
            description="Precomputed top-k expert ids per token.",
        ),
        "token_final_scales": Tensor(
            ["num_tokens", "top_k"],
            dtype="float32",
            description="Precomputed per-token routing scales.",
        ),
        "num_experts": Scalar("int32", description="Total experts."),
        "top_k": Scalar("int32"),
        "w1_alpha": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC1 global scale.",
        ),
        "w2_alpha": Tensor(
            ["num_local_experts"],
            dtype="float32",
            description="Per-expert FC2 global scale.",
        ),
        "fc2_input_scale": Tensor(
            ["one"],
            dtype="float32",
            optional=True,
            description=(
                "Global scale for FC2 input quantization. Required for "
                "activation_precision='fp4'; accepted but ignored for "
                "activation_precision='bf16'."
            ),
        ),
        "input_global_scale": Tensor(
            ["input_global_scale_size"],
            dtype="float32",
            optional=True,
            description=(
                "Global scale for FC1 input quantization (scalar or "
                "per-expert). Defaults to w1_alpha; ignored for "
                "activation_precision='bf16'."
            ),
        ),
        "activation_precision": Scalar(
            "string",
            optional=True,
            description=(
                "Backward-compatible alias: 'fp4' selects quant_mode='nvfp4'; "
                "'bf16' selects quant_mode='w4a16'."
            ),
        ),
        "quant_mode": Scalar(
            "string",
            optional=True,
            description="Quantization mode: 'nvfp4'/'w4a4' or 'w4a16'.",
        ),
        "source_format": Scalar(
            "string",
            optional=True,
            description=(
                "Source weight format for quant_mode='w4a16': 'modelopt' or "
                "'compressed_tensors'."
            ),
        ),
    },
    outputs={
        "output": Tensor(
            ["num_tokens", "hidden_size"],
            dtype="bfloat16",
            description="MoE output.",
        ),
    },
    tags=["status:experimental", "backend:cute-dsl", "quantization:fp4"],
)
b12x_fused_moe_trace.axes["one"] = Var(description="Placeholder for shape [1].")

_b12x_wrapper_inputs = dict(b12x_fused_moe_trace.inputs)
_b12x_wrapper_inputs.pop("activation_precision", None)
_b12x_wrapper_inputs["quant_mode"] = Scalar(
    "string",
    optional=True,
    description=(
        "Quantization mode set at wrapper __init__: 'nvfp4'/'w4a4' or 'w4a16'."
    ),
)
_b12x_wrapper_inputs["source_format"] = Scalar(
    "string",
    optional=True,
    description=(
        "Source weight format set at wrapper __init__ for quant_mode='w4a16'."
    ),
)
_b12x_wrapper_inputs["num_experts"] = Scalar(
    "int32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)
_b12x_wrapper_inputs["top_k"] = Scalar(
    "int32",
    optional=True,
    description="Set at wrapper __init__, not passed to run().",
)

_b12x_wrapper_axes = dict(b12x_fused_moe_trace.axes)
_b12x_wrapper_axes["num_experts"] = Var(description="Total number of experts.")
_b12x_wrapper_axes["top_k"] = Var(description="Experts per token.")


@torch.no_grad()
def _cute_dsl_fused_moe_nvfp4_reference(
    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,
    num_experts,
    top_k,
    activation_type=ActivationType.Swiglu.value,
    swiglu_alpha=DEFAULT_SWIGLU_ALPHA,
    swiglu_beta=DEFAULT_SWIGLU_BETA,
    swiglu_limit=DEFAULT_SWIGLU_LIMIT,
    quant_mode="w4a4",
    situ_beta=None,
    situ_linear_beta=None,
    per_token_scale=None,
    **_unused,
):
    """Reference for CuteDSL NvFP4 fused MoE — bridges to the FP4
    block-scale kernel with alpha scales folded into the dequantized
    weights."""
    E_local = w1_weight.shape[0]
    # Dequantize input and weights with alpha factors.
    quant_mode = quant_mode.lower()
    if quant_mode in ("nvfp4", "w4a4"):
        if x_sf is None:
            raise ValueError("x_sf is required when quant_mode='w4a4'")
        hs_deq = _dequantize_fp4_tensor(x, x_sf, is_ue8m0_scales=False)
    elif quant_mode == "w4a16":
        if x_sf is not None:
            raise ValueError("x_sf must be None when quant_mode='w4a16'")
        hs_deq = x.to(torch.float32)
    else:
        raise ValueError(f"Unsupported quant_mode {quant_mode!r}")
    W1 = _dequantize_fp4_tensor(w1_weight, w1_weight_sf, is_ue8m0_scales=False)
    W2 = _dequantize_fp4_tensor(w2_weight, w2_weight_sf, is_ue8m0_scales=False)
    if per_token_scale is not None:
        hs_deq = hs_deq * per_token_scale.to(torch.float32).view(-1, 1)
    W1 = W1 * w1_alpha.to(torch.float32).view(E_local, 1, 1)
    W2 = W2 * w2_alpha.to(torch.float32).view(E_local, 1, 1)
    return _moe_bf16_run_experts(
        hs_deq,
        W1,
        W2,
        token_final_scales,
        token_selected_experts.to(torch.int64),
        local_expert_offset=0,
        E_global=int(num_experts),
        activation_type=activation_type,
        gemm1_alpha=swiglu_alpha,
        gemm1_beta=swiglu_beta,
        gemm1_clamp_limit=swiglu_limit,
        situ_beta=situ_beta,
        situ_linear_beta=situ_linear_beta,
    )


@torch.no_grad()
def _b12x_activation_precision(activation_precision):
    normalized = str(activation_precision).lower()
    aliases = {
        "fp4": "fp4",
        "nvfp4": "fp4",
        "w4a4": "fp4",
        "bf16": "bf16",
        "w4a16": "bf16",
    }
    try:
        return aliases[normalized]
    except KeyError as exc:
        raise ValueError(
            "activation_precision must be 'fp4' or 'bf16' "
            f"(got {activation_precision!r})."
        ) from exc


@torch.no_grad()
def _b12x_quant_mode(quant_mode=None, activation_precision="fp4"):
    if quant_mode is None:
        return (
            "w4a16"
            if _b12x_activation_precision(activation_precision) == "bf16"
            else "nvfp4"
        )
    normalized = str(quant_mode).lower()
    aliases = {
        "fp4": "nvfp4",
        "nvfp4": "nvfp4",
        "w4a4": "nvfp4",
        "bf16": "w4a16",
        "w4a16": "w4a16",
    }
    try:
        return aliases[normalized]
    except KeyError as exc:
        raise ValueError(
            f"quant_mode must be 'nvfp4'/'w4a4' or 'w4a16' (got {quant_mode!r})."
        ) from exc


@torch.no_grad()
def _b12x_source_format(source_format="modelopt", quant_mode="nvfp4"):
    normalized = str(source_format).lower()
    aliases = {
        "modelopt": "modelopt",
        "compressed_tensors": "compressed_tensors",
        "compressed-tensors": "compressed_tensors",
        "ct": "compressed_tensors",
    }
    try:
        normalized = aliases[normalized]
    except KeyError as exc:
        raise ValueError(
            "source_format must be one of 'modelopt' or 'compressed_tensors', "
            f"got {source_format!r}"
        ) from exc
    if quant_mode == "nvfp4" and normalized == "compressed_tensors":
        raise ValueError(
            "source_format='compressed_tensors' requires quant_mode='w4a16'."
        )
    return normalized


@torch.no_grad()
def _b12x_moe_run_experts(
    hidden_states,
    gemm1_weights,
    gemm2_weights,
    weights,
    topk_idx,
    local_expert_offset,
    E_global,
    activation_precision,
    fc2_input_scale,
    quant_mode=None,
    source_format="modelopt",
):
    """B12x MoE expert computation with optional FP4 activation quantization."""
    quant_mode = _b12x_quant_mode(quant_mode, activation_precision)
    _b12x_source_format(source_format, quant_mode)
    activation_precision = "bf16" if quant_mode == "w4a16" else "fp4"
    if quant_mode == "nvfp4" and fc2_input_scale is None:
        raise ValueError("fc2_input_scale is required when quant_mode='nvfp4'.")

    T, H = hidden_states.shape
    E_local, gemm1_out, _ = gemm1_weights.shape
    I = gemm1_out // 2
    device = hidden_states.device
    A = hidden_states.to(torch.float32)
    W1 = gemm1_weights.to(torch.float32)
    W2 = gemm2_weights.to(torch.float32)
    output = torch.zeros((T, H), dtype=torch.float32, device=device)
    local_start = int(local_expert_offset)
    for le in range(E_local):
        ge = local_start + le
        if ge < 0 or ge >= E_global:
            continue
        sel_mask = (topk_idx == ge).any(dim=1)
        if not sel_mask.any():
            continue
        token_idx = torch.nonzero(sel_mask, as_tuple=False).squeeze(1)
        A_e = A.index_select(0, token_idx)
        G1 = A_e.matmul(W1[le].t())
        X1, X2 = G1[:, :I], G1[:, I:]
        silu_X2 = X2 / (1.0 + torch.exp(-X2))
        activated = silu_X2 * X1
        if activation_precision == "fp4":
            packed, scales = _fp4_quantize_reference(
                activated,
                global_scale=fc2_input_scale,
                sf_vec_size=16,
                is_sf_swizzled_layout=False,
            )
            activated = _dequantize_fp4_tensor(
                packed,
                scales,
                is_ue8m0_scales=False,
            )
        O = activated.matmul(W2[le].t())
        w_tok = weights.index_select(0, token_idx)
        match = (topk_idx.index_select(0, token_idx) == ge).float()
        w_e = (w_tok * match).sum(dim=1)
        output.index_add_(0, token_idx, O * w_e.unsqueeze(1))
    return output.to(torch.bfloat16)


@torch.no_grad()
def _b12x_fused_moe_reference(
    x,
    w1_weight,
    w1_weight_sf,
    w2_weight,
    w2_weight_sf,
    token_selected_experts,
    token_final_scales,
    num_experts,
    top_k,
    w1_alpha=None,
    w2_alpha=None,
    fc2_input_scale=None,
    activation_precision="fp4",
    quant_mode=None,
    source_format="modelopt",
    **_unused,
):
    """Reference for B12x CuTe-DSL fused MoE (bf16 input, FP4 weights)."""
    E_local = w1_weight.shape[0]
    W1 = _dequantize_fp4_tensor(w1_weight, w1_weight_sf, is_ue8m0_scales=False)
    W2 = _dequantize_fp4_tensor(w2_weight, w2_weight_sf, is_ue8m0_scales=False)
    if w1_alpha is not None:
        W1 = W1 * w1_alpha.to(torch.float32).view(E_local, 1, 1)
    if w2_alpha is not None:
        W2 = W2 * w2_alpha.to(torch.float32).view(E_local, 1, 1)
    return _b12x_moe_run_experts(
        x,
        W1,
        W2,
        token_final_scales,
        token_selected_experts.to(torch.int64),
        local_expert_offset=0,
        E_global=int(num_experts),
        activation_precision=activation_precision,
        fc2_input_scale=fc2_input_scale,
        quant_mode=quant_mode,
        source_format=source_format,
    )


cute_dsl_fused_moe_nvfp4_trace.reference = _cute_dsl_fused_moe_nvfp4_reference
cute_dsl_moe_wrapper_run_trace.reference = _cute_dsl_fused_moe_nvfp4_reference
b12x_fused_moe_trace.reference = _b12x_fused_moe_reference


b12x_moe_wrapper_run_trace = TraceTemplate(
    op_type="moe",
    name_prefix="b12x_moe_wrapper",
    description="B12xMoEWrapper.run(): wrapper form of b12x_fused_moe.",
    axes=_b12x_wrapper_axes,
    inputs=_b12x_wrapper_inputs,
    outputs=dict(b12x_fused_moe_trace.outputs),
    tags=b12x_fused_moe_trace.tags,
    reference=_b12x_fused_moe_reference,
)


# ---------------------------------------------------------------------------
# DSv4 hash-based MoE routing (MOE-01-HASH)
# ---------------------------------------------------------------------------


@torch.no_grad()
def _hash_topk_reference(
    router_logits: torch.Tensor,
    input_ids: torch.Tensor,
    tid2eid: torch.Tensor,
    num_fused_shared_experts: int = 0,
    routed_scaling_factor: float = 1.0,
    **_unused,
):
    """Reference for DSv4 hash-based MoE routing.

    1. Look up routed experts from the table: ``tid2eid[input_ids]``.
    2. Score the selected logits with ``sqrt(softplus(.))``.
    3. Normalize routed weights by their sum.
    4. Append a shared-expert slot (id ``num_routed_experts``, weight
       ``1 / routed_scaling_factor``) when ``num_fused_shared_experts == 1``.

    Returns ``(topk_weights, topk_ids)``.
    """
    num_tokens, num_routed_experts = router_logits.shape
    topk = tid2eid.shape[1]
    topk_fused = topk + int(num_fused_shared_experts)

    scores = torch.sqrt(torch.nn.functional.softplus(router_logits.to(torch.float32)))
    expert_ids = tid2eid[input_ids].to(torch.int64)
    gathered = torch.gather(scores, 1, expert_ids)
    routed_sum = gathered.sum(dim=-1, keepdim=True)

    weights = torch.empty(
        num_tokens, topk_fused, dtype=torch.float32, device=router_logits.device
    )
    ids = torch.empty(
        num_tokens, topk_fused, dtype=torch.int32, device=router_logits.device
    )
    weights[:, :topk] = gathered / routed_sum
    ids[:, :topk] = expert_ids.to(torch.int32)
    for s in range(int(num_fused_shared_experts)):
        ids[:, topk + s] = num_routed_experts + s
        weights[:, topk + s] = 1.0 / float(routed_scaling_factor)
    return weights, ids


def _hash_topk_init(
    *,
    num_tokens: int,
    num_routed_experts: int = 256,
    topk: int = 8,
    vocab: int = 1024,
    device: str = "cuda",
    seed: int = 0,
):
    """Build inputs for DSv4 hash-based MoE routing.

    The trace covers the routed path (``num_fused_shared_experts == 0``); the
    optional shared-expert slot only appends one constant column and does not
    change the kernel's per-token work.
    """
    torch.manual_seed(seed)
    router_logits = torch.randn(
        num_tokens, num_routed_experts, dtype=torch.float32, device=device
    )
    input_ids = torch.randint(0, vocab, (num_tokens,), dtype=torch.int64, device=device)
    tid2eid = torch.empty((vocab, topk), dtype=torch.int32, device=device)
    for v in range(vocab):
        tid2eid[v] = torch.randperm(num_routed_experts, device=device)[:topk].to(
            torch.int32
        )
    return {
        "router_logits": router_logits,
        "input_ids": input_ids,
        "tid2eid": tid2eid,
        "num_fused_shared_experts": 0,
        "routed_scaling_factor": 1.0,
    }


hash_topk_trace = TraceTemplate(
    op_type="moe_routing",
    name_prefix="hash_topk",
    description=(
        "DSv4 hash-based MoE routing: tid2eid[input_ids] table lookup -> "
        "sqrt(softplus) score -> normalize -> optional fused shared expert. "
        "Returns (topk_weights, topk_ids)."
    ),
    axes={
        "num_tokens": Var(),
        "num_routed_experts": Const(abbrev="e"),
        "topk": Const(abbrev="k"),
        "vocab": Const(abbrev="v"),
    },
    inputs={
        "router_logits": Tensor(["num_tokens", "num_routed_experts"]),
        "input_ids": Tensor(["num_tokens"], dtype="int64"),
        "tid2eid": Tensor(["vocab", "topk"], dtype="int32"),
        "num_fused_shared_experts": Scalar("int32"),
        "routed_scaling_factor": Scalar("float32"),
    },
    outputs={
        "topk_weights": Tensor(["num_tokens", "topk"], dtype="float32"),
        "topk_ids": Tensor(["num_tokens", "topk"], dtype="int32"),
    },
    tags=["status:verified", "moe"],
    reference=_hash_topk_reference,
    init=_hash_topk_init,
)
