submission 653418

"""
test_v179_selective_quant: Selective separate quant for bs=512 shapes only
Base: test_v178_separate_quant.py
Direction: CONTINUING from v178 (attempt 2) — selective separate quant
Target: s6/s7 (bs=512) — v178 showed s7 -3.5% but s1/s2 regressed +5-7%
Change: Add token_num>128 check to _separate_quant_moe_sort. For bs=512 shapes (s3/s6/s7),
        use separate quant (faster for large M). For bs<=128 (s1/s2/s4/s5), fall back to
        original fused_dynamic_mxfp4_quant_moe_sort (faster for small M).
Rationale: v178 benchmark: s7 301μs (-3.5%), s6 166μs (-1.2%), but s1 95.4μs (+5%), s2 189μs (+7.4%)
Scale: INCREMENTAL
"""
from task import input_t, output_t

import os
os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"
# NOT setting AITER_USE_NT — we patch use_nt() directly for per-shape control

import functools
import torch
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe, get_2stage_cfgs, cktile_moe_stage1
from aiter import get_hip_quant as get_quant
from aiter.utility import fp4_utils
import aiter.fused_moe as fmoe_module
import aiter.ops.triton.quant.fused_mxfp4_quant as fused_quant_module

# Patch: redirect fused_dynamic_mxfp4_quant_moe_sort to separate quant+sort
_orig_fused_quant_moe_sort = fused_quant_module.fused_dynamic_mxfp4_quant_moe_sort

def _separate_quant_moe_sort(hidden_states, sorted_ids, num_valid_ids, token_num, topk=1, block_size=32):
    """Selective: use separate quant+sort for bs=512 (token_num>128), fused for smaller."""
    if token_num > 128:
        # bs=512 shapes: separate quant is faster (fewer rocprim launches)
        quant_func = get_quant(QuantType.per_1x32)
        a1, a1_scale = quant_func(
            hidden_states,
            scale=None,
            quant_dtype=dtypes.fp4x2,
            num_rows=None,
        )
        a1_scale = fp4_utils.moe_mxfp4_sort(
            a1_scale,
            sorted_ids=sorted_ids,
            num_valid_ids=num_valid_ids,
            token_num=token_num,
            block_size=block_size,
        )
        return a1, a1_scale
    else:
        # bs<=128 shapes: fused path is faster
        return _orig_fused_quant_moe_sort(
            hidden_states,
            sorted_ids=sorted_ids,
            num_valid_ids=num_valid_ids,
            token_num=token_num,
            topk=topk,
            block_size=block_size,
        )

# Apply the patch
fused_quant_module.fused_dynamic_mxfp4_quant_moe_sort = _separate_quant_moe_sort
# Also patch in fused_moe module where it's imported
fmoe_module.fused_dynamic_mxfp4_quant_moe_sort = _separate_quant_moe_sort

# Patch use_nt: force NT=True for bs=512 (token > 128), default heuristic otherwise
def _patched_use_nt(token, topk, e):
    if token > 128:
        # bs=512 shapes: force NT on (benefits s3, s6, s7)
        return True
    # bs<=128 shapes: use default heuristic (estimated_m_per_expert < 64)
    return (token * topk // e) < 64

fmoe_module.use_nt = _patched_use_nt

# Selective split_k: force split_k=1 for E=33 bs=128 (dense), let E=257 use split_k=2 (sparse)
_orig_cktile_stage1 = cktile_moe_stage1

def _patched_cktile_stage1(
    hidden_states, w1, w2, sorted_token_ids, sorted_expert_ids,
    num_valid_ids, out, topk, block_m, a1_scale, w1_scale,
    sorted_weights=None, n_pad_zeros=0, k_pad_zeros=0, bias1=None,
    activation=ActivationType.Silu, split_k=1, dtype=torch.bfloat16,
):
    token_num = hidden_states.shape[0]
    n_experts = w1.shape[0]
    # Force split_k=1 for E=33 shapes with bs>16 (dense, ~34 tokens/expert)
    # Let E=257 shapes keep split_k=2 (sparse, ~0.5 tokens/expert)
    if token_num > 16 and n_experts <= 64:
        actual_split_k = 1
    else:
        actual_split_k = split_k
    return _orig_cktile_stage1(
        hidden_states, w1, w2, sorted_token_ids, sorted_expert_ids,
        num_valid_ids, out, topk, block_m, a1_scale, w1_scale,
        sorted_weights=sorted_weights, n_pad_zeros=n_pad_zeros,
        k_pad_zeros=k_pad_zeros, bias1=bias1,
        activation=activation, split_k=actual_split_k, dtype=dtype,
    )

fmoe_module.cktile_moe_stage1 = _patched_cktile_stage1

@functools.lru_cache(maxsize=2048)
def _patched_get_ksplit(token, topk, expert, inter_dim, model_dim):
    # cktile path for bs<=128
    if token <= 128:
        if model_dim % 2 == 0 and (model_dim // 2) % 256 == 0:
            return 2
    return 0

fmoe_module.get_ksplit = _patched_get_ksplit

# Tuned CK stage1 kernel for s7 (256x128 tile, block_m=128)
_STAGE1_256x128 = 'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'

_injected = False

def _inject_configs():
    global _injected
    if _injected:
        return
    _injected = True

    try:
        get_2stage_cfgs(
            16, 7168, 256, 257, 9,
            dtypes.bf16, dtypes.fp4x2, dtypes.fp4x2,
            QuantType.per_1x32, True, ActivationType.Silu, False,
            0, 0, True,
        )
    except Exception:
        pass

    if fmoe_module.cfg_2stages is None:
        return

    common = (
        'ActivationType.Silu', 'torch.bfloat16',
        'torch.float4_e2m1fn_x2', 'torch.float4_e2m1fn_x2',
        'QuantType.per_1x32', True, False,
    )

    # s3: FlyDSL stage2 reduce
    key = (256, 512, 7168, 256, 257, 9) + common
    if key in fmoe_module.cfg_2stages:
        fmoe_module.cfg_2stages[key]['kernelName2'] = 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce'

    # s6: FlyDSL stage2 reduce (original config — empty kernelName1, block_m=64)
    key = (256, 512, 7168, 512, 33, 9) + common
    if key not in fmoe_module.cfg_2stages:
        fmoe_module.cfg_2stages[key] = {
            'block_m': 64, 'ksplit': 0, 'kernelName1': '',
            'kernelName2': 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce',
            'run_1stage': False,
        }

    # s7: 256x128 stage1 tile (block_m=128) + FlyDSL t32 stage2
    key_s7 = (256, 512, 7168, 2048, 33, 9) + common
    fmoe_module.cfg_2stages[key_s7] = {
        'block_m': 128, 'ksplit': 0,
        'kernelName1': _STAGE1_256x128,
        'kernelName2': 'flydsl_moe2_afp4_wfp4_bf16_t32x256x256_reduce',
        'run_1stage': False,
    }

    # Remove s1 and s2 CSV entries to let ksplit=2 cktile path kick in
    key_s1 = (256, 16, 7168, 256, 257, 9) + common
    if key_s1 in fmoe_module.cfg_2stages:
        del fmoe_module.cfg_2stages[key_s1]

    key_s2 = (256, 128, 7168, 256, 257, 9) + common
    if key_s2 in fmoe_module.cfg_2stages:
        del fmoe_module.cfg_2stages[key_s2]

    get_2stage_cfgs.cache_clear()


def custom_kernel(data: input_t) -> output_t:
    _inject_configs()

    (
        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"]

    output = 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,
    )

    return output