submission 569057

import torch
import os
from typing import Dict
from task import input_t, output_t

from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe

# ============================================================
# V34: Per-Shape Optimized Based on Hardware Analysis
#
# MI355X key constraints:
# - 256 CUs, need >= 256 threadgroups for full utilization
# - 160KB LDS per CU (2.5x MI300X)
# - 256MB L3 Infinity Cache
# - 8 TB/s HBM bandwidth
# - MFMA: 32x32x64 (32 cycles) or 16x16x128 (16 cycles)
#
# Analysis per benchmark shape:
# 
# Shape #1: bs=16, E=257, d_exp=256
#   total_tokens = 144, tokens/expert = 0.56
#   TGs with block_m=32: ~144/32 * (2*256/128) = ~18 TGs << 256 CUs
#   -> SEVERELY CU-underutilized! Need high parallelism
#   -> CSV heuristic tries to maximize CU occupancy
#   -> Let CSV handle (it was tuned for this)
#
# Shape #2: bs=128, E=257, d_exp=256
#   total_tokens = 1152, tokens/expert = 4.5
#   TGs with block_m=32: ~1152/32 * 4 = ~144 TGs < 256 CUs
#   -> Still underutilized, but better
#   -> CSV handles this well
#
# Shape #3: bs=512, E=257, d_exp=256
#   total_tokens = 4608, tokens/expert = 18
#   TGs with block_m=32: ~4608/32 * 4 = ~576 TGs > 256 CUs -> OK
#   -> Good CU utilization, memory-bound (707MB weights)
#   -> CSV handles this well
#
# Shape #4: bs=16, E=33, d_exp=512
#   total_tokens = 144, tokens/expert = 4.4
#   TGs with block_m=32: ~144/32 * (2*512/128) = ~36 TGs << 256 CUs
#   -> SEVERELY underutilized!
#   -> block_m=32 + bypass might help (heuristic may find better config)
#
# Shape #5: bs=128, E=33, d_exp=512
#   total_tokens = 1152, tokens/expert = 35
#   TGs with block_m=32: ~1152/32 * 8 = ~288 TGs > 256 CUs -> OK
#   -> block_m=32 is good here
#
# Shape #6: bs=512, E=33, d_exp=512
#   total_tokens = 4608, tokens/expert = 140
#   TGs with block_m=64: ~4608/64 * 8 = ~576 TGs -> good
#   -> Compute-bound, block_m=64 maximizes compute efficiency
#
# Shape #7: bs=512, E=33, d_exp=2048
#   total_tokens = 4608, tokens/expert = 140
#   K_stage1 = 7168, K_stage2 = 2048
#   TGs with block_m=64: ~4608/64 * 32 = ~2304 TGs >> 256 CUs
#   -> Very well utilized, compute+memory bound
#   -> block_m=64 is optimal
#
# Key insight: For the underutilized shapes (#1, #2, #4), the CSV
# config likely already handles this well. For #5, #6, #7 with E=33,
# bypassing CSV with explicit block_m gives better results.
# ============================================================

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


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

    M = config["bs"]
    E = config["n_routed_experts"] + config["n_shared_experts"]
    top_k = config["total_top_k"]
    d_expert = config["d_expert"]
    total_tokens = M * top_k
    tokens_per_expert = total_tokens / E

    if E >= 257:
        # E=257: CSV has good tuned configs
        # All shapes are CU-underutilized, CSV maximizes occupancy
        os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
        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,
        )
    else:
        # E=33: bypass CSV, optimized per-shape selection
        os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
        
        # CU utilization analysis:
        # N_tiles = ceil(2*d_expert_pad / 128) for Stage1
        # = ceil(2*256/128) = 4 for d=256, ceil(2*512/128) = 8 for d=512
        # = ceil(2*2048/128) = 32 for d=2048
        
        # M_tiles = ceil(max_tokens_padded / block_m)
        # Total TGs = N_tiles * M_tiles
        # Want TGs >= 256 for full CU utilization
        
        # For d_expert=512: N_tiles=8
        #   block_m=32: need 256/8=32 M_tiles -> 32*32=1024 tokens minimum
        #   block_m=64: need 256/8=32 M_tiles -> 32*64=2048 tokens minimum
        
        # For d_expert=2048: N_tiles=32
        #   block_m=32: need 256/32=8 M_tiles -> good even for small batch
        #   block_m=64: need 256/32=8 M_tiles -> 8*64=512 tokens minimum
        
        if d_expert >= 2048:
            # Large K: many N_tiles, compute-bound
            # block_m=64 for better compute throughput
            block_m = 64
        elif total_tokens < 1024:
            # Small batch, few TGs -> block_m=32 for more tiles
            block_m = 32
        elif tokens_per_expert > 100:
            # Dense: block_m=64 for compute throughput  
            block_m = 64
        else:
            # Medium: block_m=32 balances CU utilization and throughput
            block_m = 32

        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,
            block_size_M=block_m,
            hidden_pad=hidden_pad,
            intermediate_pad=intermediate_pad,
        )

    return output