submission 562885

"""
MoE MXFP4 v49 — Block_m tuning via monkey-patch.

Default block_m values:
  M=16, E=257: 32   | M=128, E=257: 32  | M=512, E=257: 128
  M=16, E=33: 32    | M=128, E=33: 64   | M=512, E=33 d=512: 128  | M=512, E=33 d=2048: 128

The get_block_size_M heuristic selects block_m based on CU utilization.
Let's try:
- E=257 shapes: try block_m=64 for M=128 (default=32)
- E=33 shapes: try block_m=32 for M=128 (default=64), block_m=64 for M=512 (default=128)
These try smaller block_m for better CU occupancy or larger for fewer launches.
"""
import os
os.environ["AITER_USE_NT"] = "1"

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

# Monkey-patch get_block_size_M
import aiter.fused_moe as _fmoe
import functools

_orig_get_block_size_M = _fmoe.get_block_size_M.__wrapped__ if hasattr(_fmoe.get_block_size_M, '__wrapped__') else None

# Custom block_m selection
_block_m_overrides = {}

def _setup_overrides():
    """Set block_m overrides for specific shapes."""
    # Try different block_m for E=33 shapes
    # (token, topk, expert, inter_dim) -> block_m
    # Default E=33: M=16→32, M=128→64, M=512→128
    # Try: M=128→32 (smaller tiles, more CU parallelism)
    _block_m_overrides[(128, 9, 33, 512)] = 32
    _block_m_overrides[(128, 9, 33, 2048)] = 32
    # Try: M=512→64 for E=33 (more CU parallelism)
    _block_m_overrides[(512, 9, 33, 512)] = 64
    _block_m_overrides[(512, 9, 33, 2048)] = 64

_setup_overrides()

@functools.lru_cache(maxsize=2048)
def _custom_get_block_size_M(token, topk, expert, inter_dim):
    key = (token, topk, expert, inter_dim)
    if key in _block_m_overrides:
        result = _block_m_overrides[key]
        print(f"  block_m override: {key} -> {result}", file=sys.stderr)
        return result
    # Use original heuristic
    cu_num = _fmoe.get_cu_num()
    tileN = 128
    tgN = (inter_dim + tileN - 1) // tileN
    support_list = [32, 64, 128]
    tmp = []
    for el in support_list:
        max_num_tokens = token * topk + expert * el - topk
        tg_num = tgN * (max_num_tokens + el - 1) // el
        rnd = (tg_num + cu_num - 1) // cu_num
        empty = cu_num - tg_num % cu_num
        tmp.append((rnd, empty, el))
    result = sorted(tmp, key=lambda x: x[:2])[0][-1]
    print(f"  block_m default: {key} -> {result}", file=sys.stderr)
    return result

# Apply monkey-patch
_fmoe.get_block_size_M = _custom_get_block_size_M

# Also clear any cached get_2stage_cfgs
try:
    _fmoe.get_2stage_cfgs.cache_clear()
except:
    pass

# Standard warmup
_warmed = False
def _warmup():
    global _warmed
    if _warmed:
        return
    _warmed = True

    configs = [
        (2, 256, 1, 7168, 256, 8),
        (2, 32, 1, 7168, 512, 8),
        (2, 32, 1, 7168, 2048, 8),
    ]

    for bs, n_routed, n_shared, d_hidden, d_expert, n_experts_per_token in configs:
        E = n_routed + n_shared
        total_topk = n_experts_per_token + n_shared
        d_hidden_pad = ((d_hidden + 255) // 256) * 256
        d_expert_pad = ((d_expert + 255) // 256) * 256
        hidden_pad = d_hidden_pad - d_hidden
        intermediate_pad = d_expert_pad - d_expert

        h = torch.randn(bs, d_hidden, dtype=torch.bfloat16, device="cuda")
        w1 = torch.empty(E, 2 * d_expert_pad, d_hidden_pad // 2,
                         dtype=torch.float4_e2m1fn_x2, device="cuda")
        w2 = torch.empty(E, d_hidden_pad, d_expert_pad // 2,
                         dtype=torch.float4_e2m1fn_x2, device="cuda")
        w1_s = torch.empty(E, 2 * d_expert_pad, d_hidden_pad // 32,
                           dtype=torch.float8_e8m0fnu, device="cuda")
        w2_s = torch.empty(E, d_hidden_pad, d_expert_pad // 32,
                           dtype=torch.float8_e8m0fnu, device="cuda")
        topk_w = torch.ones(bs, total_topk, dtype=torch.float32, device="cuda")
        topk_i = torch.zeros(bs, total_topk, dtype=torch.int32, device="cuda")
        for t in range(bs):
            for k in range(n_experts_per_token):
                topk_i[t, k] = k % n_routed
            for k in range(n_shared):
                topk_i[t, n_experts_per_token + k] = n_routed + k

        try:
            fused_moe(
                h, w1, w2, topk_w, topk_i,
                activation=ActivationType.Silu,
                quant_type=QuantType.per_1x32,
                w1_scale=w1_s, w2_scale=w2_s,
                hidden_pad=hidden_pad,
                intermediate_pad=intermediate_pad,
            )
            torch.cuda.synchronize()
            print(f"  Warmup OK: E={E} d_e={d_expert}", file=sys.stderr)
        except Exception as e:
            print(f"  Warmup FAIL: E={E} d_e={d_expert}: {e}", file=sys.stderr)

    print("Warmup complete (v49 blockm_tune)", file=sys.stderr)

_warmup()


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

    return fused_moe(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        hidden_pad=config["d_hidden_pad"] - config["d_hidden"],
        intermediate_pad=config["d_expert_pad"] - config["d_expert"],
    )