submission 690751

"""
Attempt 851: Ultra-low occupancy for E=257 across all stages.

Lowering reduce_occ from 40->20 helped slightly in attempt_849.
Now test lowering stage1 and stage2 occupancy too:
  stage1: 25 -> 20
  stage2: 15 -> 10
  reduce: 40 -> 15

For E=257 sparse workloads (0.5-18 tokens/expert), lower occupancy
reduces cache thrashing across 257 experts all sharing limited L2.

Also: block_size_M=64 for d=2048 (proven).
E=33 configs unchanged from attempt_297 (proven best).
"""

import os
import sys

os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"

csv_path = "/home/runner/aiter/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv"
try:
    k1_32_1x1 = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
    k1_32_1x4 = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
    k2_32_1x1 = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
    common = "ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0"
    entries = []

    # E=33 d=512: UNCHANGED from attempt_297 (proven best)
    for token in [1, 2, 4, 8, 16]:
        entries.append(f"256,{token},7168,512,33,9,{common},32,2,20,{k1_32_1x1},0.0%,10,{k2_32_1x1},1.5%,50,0,0,0,")
    for token in [32, 64, 128]:
        entries.append(f"256,{token},7168,512,33,9,{common},32,2,30,{k1_32_1x4},0.0%,20,{k2_32_1x1},1.5%,60,0,0,0,")
    for token in [256, 512, 1024]:
        entries.append(f"256,{token},7168,512,33,9,{common},32,0,50,{k1_32_1x4},0.0%,30,{k2_32_1x1},1.5%,80,0,0,0,")

    # E=257 d=256: ULTRA-LOW occupancy (stage1=20, stage2=10, reduce=15)
    for token in [1, 2, 4, 8, 16]:
        entries.append(f"256,{token},7168,256,257,9,{common},32,2,20,{k1_32_1x1},0.0%,10,{k2_32_1x1},1.5%,15,0,0,0,")
    for token in [32, 64, 128]:
        entries.append(f"256,{token},7168,256,257,9,{common},32,2,20,{k1_32_1x4},0.0%,10,{k2_32_1x1},1.5%,15,0,0,0,")
    for token in [256, 512, 1024]:
        entries.append(f"256,{token},7168,256,257,9,{common},32,0,20,{k1_32_1x1},0.0%,10,{k2_32_1x1},1.5%,15,0,0,0,")

    with open(csv_path, "a") as f:
        f.write("\n" + "\n".join(entries) + "\n")
except Exception as e:
    print(f"CSV injection failed: {e}", file=sys.stderr)


import torch
from task import input_t, output_t
from aiter import ActivationType, QuantType
import aiter.fused_moe as _fm

_fm._USE_OPUS_MOE_SORTING = True
_quant_func = _fm.get_quant(QuantType.per_1x32)
_fp4_utils = None

def _separate_quant_sort(x, sorted_ids, num_valid_ids, token_num, topk, block_size=32, scaling_mode='even'):
    global _fp4_utils
    if _fp4_utils is None:
        import aiter.utility.fp4_utils
        _fp4_utils = aiter.utility.fp4_utils
    a_quant, a_scale = _quant_func(x, scale=None, quant_dtype=torch.float4_e2m1fn_x2)
    if topk > 1:
        a_scale = _fp4_utils.moe_mxfp4_sort(
            a_scale[:token_num * topk, :].view(token_num, topk, -1),
            sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
            token_num=token_num, block_size=block_size,
        )
    else:
        a_scale = _fp4_utils.moe_mxfp4_sort(
            a_scale, sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
            token_num=token_num, block_size=block_size,
        )
    return a_quant, a_scale

_fm.fused_dynamic_mxfp4_quant_moe_sort = _separate_quant_sort
from aiter.fused_moe import fused_moe


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
    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    gate_up_weight_shuffled.is_shuffled = True
    down_weight_shuffled.is_shuffled = True

    # Use block_size_M=64 for d_expert=2048 (proven 5% improvement)
    bsm = 64 if config["d_expert"] == 2048 else None

    output = fused_moe(
        hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
        topk_weights, topk_ids,
        activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        w1_scale=gate_up_weight_scale_shuffled, w2_scale=down_weight_scale_shuffled,
        hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
        block_size_M=bsm,
    )
    return output