submission 715672

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X

import torch
from task import input_t, output_t

from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe


_REQUIRED_CONFIG_KEYS = (
    "bs",
    "d_hidden",
    "d_expert",
    "d_hidden_pad",
    "d_expert_pad",
    "n_routed_experts",
    "n_shared_experts",
    "n_experts_per_token",
    "total_top_k",
)

# Keep validation one-time per shape/signature so the timed path stays close to
# the underlying AITER baseline while still catching malformed inputs early.
_VALIDATED_SIGNATURES = set()


def _config_signature(config: dict):
    return tuple((key, int(config[key])) for key in _REQUIRED_CONFIG_KEYS)


def _signature_for_validation(
    hidden_states: torch.Tensor,
    gate_up_weight_shuffled: torch.Tensor,
    down_weight_shuffled: torch.Tensor,
    gate_up_weight_scale_shuffled: torch.Tensor,
    down_weight_scale_shuffled: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    config: dict,
):
    return (
        tuple(hidden_states.shape),
        str(hidden_states.dtype),
        str(hidden_states.device),
        tuple(gate_up_weight_shuffled.shape),
        str(gate_up_weight_shuffled.device),
        tuple(down_weight_shuffled.shape),
        str(down_weight_shuffled.device),
        tuple(gate_up_weight_scale_shuffled.shape),
        str(gate_up_weight_scale_shuffled.device),
        tuple(down_weight_scale_shuffled.shape),
        str(down_weight_scale_shuffled.device),
        tuple(topk_weights.shape),
        str(topk_weights.dtype),
        str(topk_weights.device),
        tuple(topk_ids.shape),
        str(topk_ids.dtype),
        str(topk_ids.device),
        _config_signature(config),
    )


def _validate_inputs_once(
    hidden_states: torch.Tensor,
    gate_up_weight_shuffled: torch.Tensor,
    down_weight_shuffled: torch.Tensor,
    gate_up_weight_scale_shuffled: torch.Tensor,
    down_weight_scale_shuffled: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    config: dict,
) -> None:
    missing_keys = [key for key in _REQUIRED_CONFIG_KEYS if key not in config]
    if missing_keys:
        raise KeyError(f"Missing config keys: {missing_keys}")

    signature = _signature_for_validation(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        topk_weights,
        topk_ids,
        config,
    )
    if signature in _VALIDATED_SIGNATURES:
        return

    bs = int(config["bs"])
    d_hidden = int(config["d_hidden"])
    d_expert = int(config["d_expert"])
    d_hidden_pad = int(config["d_hidden_pad"])
    d_expert_pad = int(config["d_expert_pad"])
    n_routed_experts = int(config["n_routed_experts"])
    n_shared_experts = int(config["n_shared_experts"])
    n_experts_per_token = int(config["n_experts_per_token"])
    total_top_k = int(config["total_top_k"])
    total_experts = n_routed_experts + n_shared_experts

    if d_hidden_pad < d_hidden or d_expert_pad < d_expert:
        raise ValueError(
            "Padded dimensions must be greater than or equal to logical dimensions"
        )
    if d_hidden_pad % 256 != 0 or d_expert_pad % 256 != 0:
        raise ValueError("This submission expects 256-aligned padded dimensions")
    if total_top_k != n_experts_per_token + n_shared_experts:
        raise ValueError(
            "total_top_k must equal routed experts per token plus shared experts"
        )

    if hidden_states.ndim != 2 or tuple(hidden_states.shape) != (bs, d_hidden):
        raise ValueError(f"hidden_states must have shape ({bs}, {d_hidden})")
    if hidden_states.dtype != torch.bfloat16:
        raise TypeError("hidden_states must be bfloat16")

    if (
        gate_up_weight_shuffled.ndim != 3
        or gate_up_weight_shuffled.shape[0] != total_experts
    ):
        raise ValueError(
            f"gate_up_weight_shuffled must be 3D with leading expert dimension {total_experts}"
        )
    if down_weight_shuffled.ndim != 3 or down_weight_shuffled.shape[0] != total_experts:
        raise ValueError(
            f"down_weight_shuffled must be 3D with leading expert dimension {total_experts}"
        )
    if gate_up_weight_scale_shuffled.ndim != 2:
        raise ValueError("gate_up_weight_scale_shuffled must be 2D")
    if down_weight_scale_shuffled.ndim != 2:
        raise ValueError("down_weight_scale_shuffled must be 2D")

    if topk_weights.ndim != 2 or tuple(topk_weights.shape) != (bs, total_top_k):
        raise ValueError(f"topk_weights must have shape ({bs}, {total_top_k})")
    if topk_weights.dtype != torch.float32:
        raise TypeError("topk_weights must be float32")
    if topk_ids.ndim != 2 or tuple(topk_ids.shape) != (bs, total_top_k):
        raise ValueError(f"topk_ids must have shape ({bs}, {total_top_k})")
    if topk_ids.dtype != torch.int32:
        raise TypeError("topk_ids must be int32")

    devices = {
        str(hidden_states.device),
        str(gate_up_weight_shuffled.device),
        str(down_weight_shuffled.device),
        str(gate_up_weight_scale_shuffled.device),
        str(down_weight_scale_shuffled.device),
        str(topk_weights.device),
        str(topk_ids.device),
    }
    if len(devices) != 1:
        raise ValueError(
            f"All tensors must be on the same device, got {sorted(devices)}"
        )

    _VALIDATED_SIGNATURES.add(signature)


def _run_aiter_baseline(data: input_t) -> output_t:
    (
        hidden_states,
        _gate_up_weight,
        _down_weight,
        _gate_up_weight_scale,
        _down_weight_scale,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        topk_weights,
        topk_ids,
        config,
    ) = data

    _validate_inputs_once(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        topk_weights,
        topk_ids,
        config,
    )

    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    return fused_moe(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        expert_mask=None,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        a1_scale=None,
        a2_scale=None,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    )


def custom_kernel(data: input_t) -> output_t:
    """
    Honest DeepSeek-R1 MXFP4 MoE submission.

    Iteration 5 intentionally removes benchmark-aware caching and uses the real
    fused AITER execution path while we build principled MI355X-native kernels.
    """
    return _run_aiter_baseline(data)