submission 599565

"""v332: v330 + S6 FlyDSL t64x128x256 reduce stage2 probe.

- dsv3 CSV: ksplit=2 for E=257 bs<=128, FlyDSL t64x128x256_reduce for S3
- CKTile split_k=2 for S1/S2/S4/S5 via get_2stage_cfgs patch
- E=33 CSV: S6 FlyDSL reduce, S7 FlyDSL t64x128x256_atomic
- OPUS=1 globally
- NT=1 globally
- Cached sorting buffers
"""
import functools
import os
import sys

import torch

from dataclasses import replace

# ── Step 1: Modify dsv3 CSV for E=257: ksplit=2 on bs<=128, FlyDSL stage2 on bs=512 ──
_dsv3_path = "/home/runner/aiter/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv"
_flydsl_s3_stage2 = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce"
try:
    with open(_dsv3_path, "r") as f:
        lines = f.readlines()
    header = lines[0].strip()
    modified_lines = [header + "\n"]
    for line in lines[1:]:
        stripped = line.strip()
        if not stripped:
            continue
        fields = stripped.split(",")
        try:
            token_val = int(fields[1])
            expert_val = int(fields[4])
        except (ValueError, IndexError):
            modified_lines.append(line)
            continue
        if expert_val == 257 and token_val <= 128:
            fields[14] = "2"  # ksplit=2
            modified_lines.append(",".join(fields) + "\n")
        elif expert_val == 257 and token_val == 512:
            # FlyDSL reduce stage2 for S3, with CK fallback
            flydsl_fields = list(fields)
            flydsl_fields[19] = _flydsl_s3_stage2
            flydsl_fields[20] = "0.1%"
            if len(flydsl_fields) > 25:
                flydsl_fields[25] = ""
            modified_lines.append(",".join(flydsl_fields) + "\n")
            # CK fallback row
            fallback_fields = list(fields)
            if len(fallback_fields) > 25:
                fallback_fields[25] = "flydsl_fallback"
            else:
                fallback_fields.append("flydsl_fallback")
            modified_lines.append(",".join(fallback_fields) + "\n")
        else:
            modified_lines.append(line)
    with open(_dsv3_path, "w") as f:
        f.writelines(modified_lines)
except Exception as e:
            print(f"[v332] dsv3 error: {e}", file=sys.stderr)

# ── Step 2: E=33 CSV for S4/S5 ksplit=2, S6/S7 tuned configs ──
_csv_header = (
    "cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,"
    "q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,us1,kernelName1,"
    "err1,us2,kernelName2,err2,us,run_1stage,tflops,bw,_tag"
)
_common = (
    "ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,"
    "torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0"
)
_k1_small = (
    "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_k2_small = (
    "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_"
    "Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_k1_512 = (
    "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_k2_512 = (
    "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_"
    "Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_k1_2048 = (
    "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
)
_k2_2048 = (
    "moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_"
    "Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
)
_k2_512_flydsl = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce"
_k2_2048_flydsl = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_atomic"
_csv_rows = [
    # S4 (bs=16, d=512): ksplit=2
    f"256,16,7168,512,33,9,{_common},32,2,0,{_k1_small},0.0%,0,{_k2_small},0.0%,47,0,67,7691,",
    # S5 (bs=128, d=512): ksplit=2
    f"256,128,7168,512,33,9,{_common},32,2,0,{_k1_small},0.0%,0,{_k2_small},0.0%,58,0,434,6263,",
    # S6 (bs=512, d=512): FlyDSL reduce stage2 with tagged CK fallback
    f"256,512,7168,512,33,9,{_common},32,0,0,{_k1_512},0.0%,90.0,{_k2_512_flydsl},0.1%,219.79,0,781.78,2884.18,",
    f"256,512,7168,512,33,9,{_common},32,0,0,{_k1_512},0.0%,0,{_k2_512},0.0%,129.79,0,781.78,2884.18,flydsl_fallback",
    # S7 (bs=512, d=2048): FlyDSL atomic stage2 with tagged CK fallback
    f"256,512,7168,2048,33,9,{_common},64,0,0,{_k1_2048},0.0%,180.0,{_k2_2048_flydsl},0.1%,455.08,0,1475.47,5323.27,",
    f"256,512,7168,2048,33,9,{_common},128,0,0,{_k1_2048},0.0%,0,{_k2_2048},0.0%,275.08,0,1475.47,5323.27,flydsl_fallback",
]
try:
    e33_path = "/home/runner/aiter/aiter/configs/model_configs/e33_fp4_tuned_fmoe.csv"
    with open(e33_path, "w") as f:
        f.write(_csv_header + "\n")
        for row in _csv_rows:
            f.write(row + "\n")
except Exception:
    pass

# ── Step 3: Global OPUS=1, NT=1 (no per-shape switching!) ──
os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"
os.environ["AITER_USE_NT"] = "1"

from task import input_t, output_t
from aiter import ActivationType, QuantType, dtypes
import aiter
import aiter.fused_moe as fused_moe_mod
from aiter.ops.flydsl.utils import is_flydsl_available as _is_flydsl_available

print(f"[v332] flydsl available: {_is_flydsl_available()}", file=sys.stderr)


# ═══════════════════════════════════════════════════════════════════════
# Cached sorting buffers (zero alloc after warmup)
# ═══════════════════════════════════════════════════════════════════════
_SORT_BUFS = {}


def _cached_moe_sorting_impl(
    topk_ids, topk_weights, num_experts, model_dim, moebuf_dtype,
    block_size, expert_mask, num_local_tokens, dispatch_policy, use_opus,
):
    device = topk_ids.device
    M, topk = topk_ids.shape
    key = (M, num_experts, block_size, model_dim)

    if key not in _SORT_BUFS:
        max_num_tokens_padded = int(M * topk + num_experts * block_size - topk)
        max_num_m_blocks = int(
            (max_num_tokens_padded + block_size - 1) // block_size
        )
        _SORT_BUFS[key] = (
            torch.empty(max_num_tokens_padded, dtype=dtypes.i32, device=device),
            torch.empty(max_num_tokens_padded, dtype=dtypes.fp32, device=device),
            torch.empty(max_num_m_blocks, dtype=dtypes.i32, device=device),
            torch.empty(2, dtype=dtypes.i32, device=device),
            torch.empty((M, model_dim), dtype=moebuf_dtype, device=device),
        )

    sid, sw, sei, nvi, mb = _SORT_BUFS[key]
    fwd = aiter.moe_sorting_opus_fwd if use_opus else aiter.moe_sorting_fwd
    fwd(
        topk_ids, topk_weights, sid, sw, sei, nvi, mb,
        num_experts, int(block_size), expert_mask, num_local_tokens,
        dispatch_policy,
    )
    return sid, sw, sei, nvi, mb


fused_moe_mod._moe_sorting_impl = _cached_moe_sorting_impl


# ═══════════════════════════════════════════════════════════════════════
# Patch get_2stage_cfgs: CKTile for S1/S2/S4/S5, tuned S7
# ═══════════════════════════════════════════════════════════════════════
_ORIGINAL_GET_2STAGE_CFGS = fused_moe_mod.get_2stage_cfgs


def _make_cktile_metadata(hidden_pad, intermediate_pad, use_g1u1, activation, split_k=2):
    return fused_moe_mod.MOEMetadata(
        functools.partial(
            fused_moe_mod.cktile_moe_stage1,
            n_pad_zeros=intermediate_pad // 64 * 64 * (2 if use_g1u1 else 1),
            k_pad_zeros=hidden_pad // 128 * 128,
            activation=activation,
            split_k=split_k,
        ),
        functools.partial(
            fused_moe_mod.cktile_moe_stage2,
            n_pad_zeros=hidden_pad // 64 * 64,
            k_pad_zeros=intermediate_pad // 128 * 128,
            activation=activation,
        ),
        16, split_k, False, False, True,
    )


@functools.lru_cache(maxsize=2048)
def _patched_get_2stage_cfgs(
    token, model_dim, inter_dim, expert, topk, dtype,
    q_dtype_a, q_dtype_w, q_type, use_g1u1, activation,
    doweight_stage1, hidden_pad, intermediate_pad, is_shuffled=True,
):
    common = (
        model_dim == 7168 and topk == 9
        and dtype == dtypes.bf16 and q_dtype_a == dtypes.fp4x2
        and q_dtype_w == dtypes.fp4x2 and q_type == QuantType.per_1x32
        and use_g1u1 and not doweight_stage1 and is_shuffled
    )

    if not common:
        return _ORIGINAL_GET_2STAGE_CFGS(
            token, model_dim, inter_dim, expert, topk, dtype,
            q_dtype_a, q_dtype_w, q_type, use_g1u1, activation,
            doweight_stage1, hidden_pad, intermediate_pad, is_shuffled,
        )

    # E=257 bs<512 → CKTile split_k=2 (quant-skip)
    if expert == 257 and inter_dim == 256 and token < 512:
        return _make_cktile_metadata(
            hidden_pad, intermediate_pad, use_g1u1, activation, split_k=2,
        )

    # E=33 d=512 bs<=128 → CKTile split_k=2 (quant-skip)
    if expert == 33 and inter_dim == 512 and token <= 128:
        return _make_cktile_metadata(
            hidden_pad, intermediate_pad, use_g1u1, activation, split_k=2,
        )

    default = _ORIGINAL_GET_2STAGE_CFGS(
        token, model_dim, inter_dim, expert, topk, dtype,
        q_dtype_a, q_dtype_w, q_type, use_g1u1, activation,
        doweight_stage1, hidden_pad, intermediate_pad, is_shuffled,
    )

    # S7 (E=33, d=2048, bs=512): block_m=64 + NT
    if expert == 33 and inter_dim == 2048 and token == 512:
        return replace(default, block_m=64, use_non_temporal_load=True)

    return default


fused_moe_mod.get_2stage_cfgs = _patched_get_2stage_cfgs


# ═══════════════════════════════════════════════════════════════════════
# Dispatch — simple, no state switching
# ═══════════════════════════════════════════════════════════════════════
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"]

    return fused_moe_mod.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,
    )