submission 569735

#!/usr/bin/env python3
"""
MXFP4 MoE - Try smaller BLOCK_M=16 for large M cases
Inspired by Yufeng98's blockm16.py
"""

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X

from typing import Dict, Tuple
import torch
import os

input_t = Tuple[
    torch.Tensor,  # hidden_states [M, d_hidden]
    torch.Tensor,  # gate_up_weight [E, 2*d_expert_pad, d_hidden_pad//2] fp4x2
    torch.Tensor,  # down_weight [E, d_hidden_pad, d_expert_pad//2] fp4x2
    torch.Tensor,  # gate_up_weight_scale [E, 2*d_expert_pad, scale_K] e8m0
    torch.Tensor,  # down_weight_scale [E, d_hidden_pad, scale_K] e8m0
    torch.Tensor,  # gate_up_weight_shuffled [E, 2*d_expert_pad, d_hidden_pad//2] fp4x2
    torch.Tensor,  # down_weight_shuffled [E, d_hidden_pad, d_expert_pad//2] fp4x2
    torch.Tensor,  # gate_up_weight_scale_shuffled [padded, flat] e8m0
    torch.Tensor,  # down_weight_scale_shuffled [padded, flat] e8m0
    torch.Tensor,  # topk_weights [M, total_top_k] float32
    torch.Tensor,  # topk_ids [M, total_top_k] int32
    Dict,          # config
]
output_t = torch.Tensor

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

    d_hidden = config["d_hidden"]
    d_expert = config["d_expert"]
    d_hidden_pad = config["d_hidden_pad"]
    d_expert_pad = config["d_expert_pad"]
    hidden_pad = d_hidden_pad - d_hidden
    intermediate_pad = d_expert_pad - d_expert
    
    M = hidden_states.size(0)
    
    # Base environment
    env = {
        "VLLM_ROCM_USE_SKINNY_GEMM": "1",
        "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": "1",
        "AITER_PAD_M": "1",
        "AITER_FORCE_MXFP4": "1",
        "AITER_GEMM_KIND": "asm",
        "AITER_USE_CK": "0",
        "AITER_OPT_SMALL_BATCH": "1",
        "AITER_ONLINE_TUNE": "0",
        "HIP_FORCE_DEV_KERNARG": "1",
    }
    
    # Try BLOCK_M=16 for large M cases (inspired by Yufeng98)
    if M >= 512:
        # Large batches - try smaller blocks for better occupancy
        env.update({
            "AITER_FORCE_NT": "1",
            "AITER_BLOCK_M": "16",  # Smaller blocks, more parallelism
            "AITER_KSPLIT": "1",
        })
    elif M >= 128:
        # Medium batches
        if d_expert == 512:
            # Proven winner: bs128_d512
            env.update({
                "AITER_FORCE_NT": "1",
                "AITER_BLOCK_M": "64",
                "AITER_KSPLIT": "2",
            })
        else:
            env.update({
                "AITER_FORCE_NT": "1",
                "AITER_BLOCK_M": "32",
                "AITER_KSPLIT": "1",
            })
    else:
        # Small batches
        env.update({
            "AITER_FORCE_NT": "1",
            "AITER_BLOCK_M": "32",
            "AITER_KSPLIT": "1",
        })
    
    os.environ.update(env)
    
    import aiter
    from aiter.fused_moe import fused_moe
    from aiter import ActivationType, QuantType
    
    hidden_states = hidden_states.contiguous()
    gate_up_weight_shuffled = gate_up_weight_shuffled.contiguous()
    down_weight_shuffled = down_weight_shuffled.contiguous()
    gate_up_weight_scale_shuffled = gate_up_weight_scale_shuffled.contiguous()
    down_weight_scale_shuffled = down_weight_scale_shuffled.contiguous()
    topk_weights = topk_weights.contiguous()
    topk_ids = topk_ids.contiguous()
    
    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,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    )

if __name__ == "__main__":
    print("MXFP4 MoE - Smaller BLOCK_M for Large M")
    print("========================================")
    print("BLOCK_M=16 for M >= 512 (inspired by Yufeng98 blockm16)")
    print("BLOCK_M=64 for bs128_d512 (proven winner)")
    print("BLOCK_M=32 for others")