submission 721746

# v_bc: Three stacked optimizations on top of v_ab:
#  1. Remove AITER_USE_NT=0 → let heuristic enable non-temporal loads for cfg3
#     (17 tokens/expert < 64 threshold → NT=True bypasses L2 for weight loads)
#  2. cfg6 GEMM1: wider 256×32×128×128 tile (DSv3-tuner, 2% gain from v_bb)
#  3. cfg6 GEMM2: FlyDSL t32×128×256_reduce (same kernel boosting cfg7)
#     This avoids the small 64-thread CK 64×32 tile and uses 128-wide N tiles.
# Hand-maintained single-file variant. Edit this file directly.
from __future__ import annotations

import os
from typing import Optional

import torch

from task import input_t, output_t

# Let the heuristic decide non-temporal loads per config:
# NT=True when tokens_per_expert < 64 (cfg1-5), NT=False for dense configs (cfg6/7).
# Removing the global AITER_USE_NT=0 unlocks NT for cfg3 (17 tok/expert).
# os.environ["AITER_USE_NT"] = "0"
os.environ["AITER_USE_OPUS_MOE_SORTING"] = "0"

from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import (
    fused_moe,
    fused_moe_2stages,
    get_2stage_cfgs,
    get_inter_dim,
    get_padded_M,
    moe_sorting,
)
import aiter.fused_moe as _fmoe


CFG1_SHAPE = (16, 256, 257)
CFG2_SHAPE = (128, 256, 257)
CFG3_SHAPE = (512, 256, 257)
CFG4_SHAPE = (16, 512, 33)
CFG5_SHAPE = (128, 512, 33)
CFG6_SHAPE = (512, 512, 33)
CFG7_SHAPE = (512, 2048, 33)

_SUFFIX = (
    "ActivationType.Silu",
    "torch.bfloat16",
    "torch.float4_e2m1fn_x2",
    "torch.float4_e2m1fn_x2",
    "QuantType.per_1x32",
    True,
    False,
)

_CFG_OVERRIDES_INSTALLED: set[tuple[object, ...]] = set()
_ROUTE_LOGGED: set[tuple[tuple[int, int, int], str]] = set()


def _log_route_once(shape: tuple[int, int, int], route: str, detail: str) -> None:
    key = (shape, route)
    if key in _ROUTE_LOGGED:
        return
    _ROUTE_LOGGED.add(key)
    print(f"[route] shape={shape} route={route} detail={detail}", flush=True)


def _route_detail(spec) -> str:
    backend = "custom-fast-ck-2stage" if spec.USE_FAST_PATH else "custom-aiter-fused-moe"
    return (
        f"cfg={spec.NAME} backend={backend} "
        f"sort_policy={spec.SORT_POLICY} use_opus_sort={int(spec.USE_OPUS_SORT)}"
    )


def _ensure_cfg_dict() -> None:
    if _fmoe.cfg_2stages is not None:
        return
    try:
        import pandas as pd
        from aiter.jit.core import AITER_CONFIGS

        tune_file = AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
        index_cols = [
            "cu_num",
            "token",
            "model_dim",
            "inter_dim",
            "expert",
            "topk",
            "act_type",
            "dtype",
            "q_dtype_a",
            "q_dtype_w",
            "q_type",
            "use_g1u1",
            "doweight_stage1",
        ]
        df = pd.read_csv(tune_file)
        if "_tag" in df.columns:
            df = df[df["_tag"].fillna("") == ""]
        _fmoe.cfg_2stages = df.set_index(index_cols).to_dict("index")
    except Exception:
        return


def _install_override(spec) -> None:
    cfg_override = spec.CFG_OVERRIDE
    if cfg_override is None or spec.CFG_KEY in _CFG_OVERRIDES_INSTALLED:
        return
    _ensure_cfg_dict()
    if _fmoe.cfg_2stages is None:
        return
    _fmoe.cfg_2stages[spec.CFG_KEY] = dict(cfg_override)
    get_2stage_cfgs.cache_clear()
    _CFG_OVERRIDES_INSTALLED.add(spec.CFG_KEY)


def _run_fast_ck_two_stage(
    spec,
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    w1_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    hidden_pad: int,
    intermediate_pad: int,
) -> torch.Tensor:
    token_num, topk = topk_ids.shape
    num_experts, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape)
    dtype = hidden_states.dtype
    q_dtype_w = w1.dtype
    q_dtype_a = dtypes.fp4x2
    is_shuffled = getattr(w1, "is_shuffled", False)

    metadata = get_2stage_cfgs(
        get_padded_M(token_num),
        model_dim,
        inter_dim,
        num_experts,
        topk,
        dtype,
        q_dtype_a,
        q_dtype_w,
        QuantType.per_1x32,
        True,
        ActivationType.Silu,
        False,
        hidden_pad,
        intermediate_pad,
        is_shuffled,
    )

    block_size_m = metadata.block_m if spec.BLOCK_M_OVERRIDE is None else int(spec.BLOCK_M_OVERRIDE)
    if metadata.run_1stage or metadata.ksplit > 0:
        return fused_moe(
            hidden_states,
            w1,
            w2,
            topk_weights,
            topk_ids,
            expert_mask=None,
            activation=ActivationType.Silu,
            quant_type=QuantType.per_1x32,
            doweight_stage1=False,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            a1_scale=None,
            a2_scale=None,
            block_size_M=block_size_m,
            hidden_pad=hidden_pad,
            intermediate_pad=intermediate_pad,
            moe_sorting_dispatch_policy=spec.SORT_POLICY,
        )

    sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_out = moe_sorting(
        topk_ids,
        topk_weights,
        num_experts,
        model_dim,
        dtype,
        block_size_m,
        None,
        None,
        spec.SORT_POLICY,
    )

    return fused_moe_2stages(
        hidden_states,
        w1,
        w2,
        topk,
        sorted_ids,
        sorted_weights,
        sorted_expert_ids,
        num_valid_ids,
        moe_out,
        True,
        block_size_m,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        q_dtype_a=q_dtype_a,
        q_dtype_w=q_dtype_w,
        w1_scale=w1_scale,
        w2_scale=w2_scale,
        a1_scale=None,
        a2_scale=None,
        num_local_tokens=None,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
        bias1=None,
        bias2=None,
    )


def _run_cfg(spec, data: input_t) -> output_t:
    _install_override(spec)
    _log_route_once(spec.SHAPE, "custom", _route_detail(spec))

    (
        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

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

    prev_use_opus = bool(getattr(_fmoe, "_USE_OPUS_MOE_SORTING", False))
    _fmoe._USE_OPUS_MOE_SORTING = spec.USE_OPUS_SORT
    try:
        if spec.USE_FAST_PATH:
            return _run_fast_ck_two_stage(
                spec,
                hidden_states,
                gate_up_weight_shuffled,
                down_weight_shuffled,
                gate_up_weight_scale_shuffled,
                down_weight_scale_shuffled,
                topk_weights,
                topk_ids,
                hidden_pad,
                intermediate_pad,
            )
        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,
            block_size_M=spec.BLOCK_M_OVERRIDE,
            hidden_pad=hidden_pad,
            intermediate_pad=intermediate_pad,
            moe_sorting_dispatch_policy=spec.SORT_POLICY,
        )
    finally:
        _fmoe._USE_OPUS_MOE_SORTING = prev_use_opus


def baseline_custom_kernel(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

    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,
        block_size_M=None,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
        moe_sorting_dispatch_policy=0,
    )


class _Cfg1Spec:
    """cfg1: bs=16, d_expert=256, experts=257."""

    NAME = "moe_cfg1"
    SHAPE = CFG1_SHAPE
    CFG_KEY = (256, 16, 7168, 256, 257, 9) + _SUFFIX
    CFG_OVERRIDE = {
        "ksplit": 2,
        "run_1stage": 0,
        "block_m": 16,
        "kernelName1": "",
        "kernelName2": "",
        "us1": 0,
        "err1": "",
        "us2": 0,
        "err2": "",
        "us": 0,
        "tflops": 0,
        "bw": 0,
    }
    USE_OPUS_SORT = False
    SORT_POLICY = 0
    BLOCK_M_OVERRIDE: Optional[int] = None
    USE_FAST_PATH = False


class _Cfg2Spec:
    """cfg2: bs=128, d_expert=256, experts=257."""

    NAME = "moe_cfg2"
    SHAPE = CFG2_SHAPE
    CFG_KEY = (256, 128, 7168, 256, 257, 9) + _SUFFIX
    CFG_OVERRIDE = {
        "ksplit": 2,
        "run_1stage": 0,
        "block_m": 16,
        "kernelName1": "",
        "kernelName2": "",
        "us1": 0,
        "err1": "",
        "us2": 0,
        "err2": "",
        "us": 0,
        "tflops": 0,
        "bw": 0,
    }
    USE_OPUS_SORT = False
    SORT_POLICY = 0
    BLOCK_M_OVERRIDE: Optional[int] = None
    USE_FAST_PATH = False


class _Cfg3Spec:
    """cfg3: bs=512, d_expert=256, experts=257."""

    NAME = "moe_cfg3"
    SHAPE = CFG3_SHAPE
    CFG_KEY = (256, 512, 7168, 256, 257, 9) + _SUFFIX
    CFG_OVERRIDE = None
    USE_OPUS_SORT = False
    SORT_POLICY = 0
    BLOCK_M_OVERRIDE: Optional[int] = None
    USE_FAST_PATH = True


class _Cfg4Spec:
    """cfg4: bs=16, d_expert=512, experts=33."""

    NAME = "moe_cfg4"
    SHAPE = CFG4_SHAPE
    CFG_KEY = (256, 16, 7168, 512, 33, 9) + _SUFFIX
    CFG_OVERRIDE = {
        "ksplit": 2,
        "run_1stage": 0,
        "block_m": 16,
        "kernelName1": "",
        "kernelName2": "",
        "us1": 0,
        "err1": "",
        "us2": 0,
        "err2": "",
        "us": 0,
        "tflops": 0,
        "bw": 0,
    }
    USE_OPUS_SORT = False
    SORT_POLICY = 0
    BLOCK_M_OVERRIDE: Optional[int] = None
    USE_FAST_PATH = False


class _Cfg5Spec:
    """cfg5: bs=128, d_expert=512, experts=33."""

    NAME = "moe_cfg5"
    SHAPE = CFG5_SHAPE
    CFG_KEY = (256, 128, 7168, 512, 33, 9) + _SUFFIX
    CFG_OVERRIDE = {
        "ksplit": 2,
        "run_1stage": 0,
        "block_m": 16,
        "kernelName1": "",
        "kernelName2": "",
        "us1": 0,
        "err1": "",
        "us2": 0,
        "err2": "",
        "us": 0,
        "tflops": 0,
        "bw": 0,
    }
    USE_OPUS_SORT = False
    SORT_POLICY = 0
    BLOCK_M_OVERRIDE: Optional[int] = 32
    USE_FAST_PATH = False


class _Cfg6Spec:
    """cfg6: bs=512, d_expert=512, experts=33."""

    NAME = "moe_cfg6"
    SHAPE = CFG6_SHAPE
    CFG_KEY = (256, 512, 7168, 512, 33, 9) + _SUFFIX
    # cfg6 tile tuning:
    # GEMM1: 256×32×128×128 (wider 256-thread tile, 4× larger N tile vs 64×32's N=32)
    #   DSv3 tuner showed this tile ~2% faster at token=512 for GEMM1 (v_bb: 186→182µs)
    # GEMM2: FlyDSL t32×128×256_reduce — same kernel that makes cfg7 competitive.
    #   N=128 tile (vs CK's N=32) means 4× fewer tiles, reducing launch overhead.
    #   K=256 tile (vs CK's K=128) means 2× fewer K iterations per tile.
    CFG_OVERRIDE = {
        "ksplit": 0,
        "run_1stage": 0,
        "block_m": 32,
        "kernelName1": "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
        "kernelName2": "flydsl_moe2_afp4_wfp4_bf16_t32x128x256_reduce",
        "us1": 0,
        "err1": "",
        "us2": 0,
        "err2": "",
        "us": 0,
        "tflops": 0,
        "bw": 0,
    }
    USE_OPUS_SORT = False
    SORT_POLICY = 0
    BLOCK_M_OVERRIDE: Optional[int] = 32
    USE_FAST_PATH = True


class _Cfg7Spec:
    """cfg7: bs=512, d_expert=2048, experts=33."""

    NAME = "moe_cfg7"
    SHAPE = CFG7_SHAPE
    CFG_KEY = (256, 512, 7168, 2048, 33, 9) + _SUFFIX
    CFG_OVERRIDE = {
        "ksplit": 0,
        "run_1stage": 0,
        "block_m": 64,
        "kernelName1": "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
        "kernelName2": "flydsl_moe2_afp4_wfp4_bf16_t32x128x256_reduce",
        "us1": 0,
        "err1": "",
        "us2": 0,
        "err2": "",
        "us": 0,
        "tflops": 0,
        "bw": 0,
    }
    USE_OPUS_SORT = False
    SORT_POLICY = 0
    BLOCK_M_OVERRIDE: Optional[int] = 64
    USE_FAST_PATH = True


def moe_cfg1(data: input_t) -> output_t:
    return _run_cfg(_Cfg1Spec, data)


def moe_cfg2(data: input_t) -> output_t:
    return _run_cfg(_Cfg2Spec, data)


def moe_cfg3(data: input_t) -> output_t:
    return _run_cfg(_Cfg3Spec, data)


def moe_cfg4(data: input_t) -> output_t:
    return _run_cfg(_Cfg4Spec, data)


def moe_cfg5(data: input_t) -> output_t:
    return _run_cfg(_Cfg5Spec, data)


def moe_cfg6(data: input_t) -> output_t:
    return _run_cfg(_Cfg6Spec, data)


def moe_cfg7(data: input_t) -> output_t:
    return _run_cfg(_Cfg7Spec, data)


CFG_KERNELS = {
    CFG1_SHAPE: moe_cfg1,
    CFG2_SHAPE: moe_cfg2,
    CFG3_SHAPE: moe_cfg3,
    CFG4_SHAPE: moe_cfg4,
    CFG5_SHAPE: moe_cfg5,
    CFG6_SHAPE: moe_cfg6,
    CFG7_SHAPE: moe_cfg7,
}


def custom_kernel(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

    shape = (
        int(hidden_states.shape[0]),
        int(config["d_expert"]),
        int(gate_up_weight_shuffled.shape[0]),
    )
    kernel = CFG_KERNELS.get(shape)
    if kernel is None:
        _log_route_once(shape, "fallback", "baseline-aiter-fused-moe")
        return baseline_custom_kernel(data)
    return kernel(data)