submission 643445

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

# Build Principle: Explicit Intent, Log Once
_families_seen = set()

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

    # 1. Shape Metadata
    bs = config["bs"]
    E = gate_up_weight_shuffled.shape[0]
    d_expert = config["d_expert"]
    est_m = (bs * config["total_top_k"]) / E
    
    family_config = {"moe_sorting_dispatch_policy": False}

    # 2. Specialized Tiered Dispatch
    if d_expert >= 2048:
        # TARGET: Blocker Family. Hypothesis: Tile-64 beats Tile-128.
        family_name = "BLOCKER_TILE_64"
        family_config.update({"block_size_M": 64, "splitk": 2})
    elif est_m < 50:
        # TARGET: Sparse family. Goal: Maintain 87us result.
        family_name = "SPARSE_LATENCY"
        family_config.update({"block_size_M": 32, "splitk": 2})
    else:
        # TARGET: Dense family. Standard occupancy.
        family_name = "DENSE_THROUGHPUT"
        family_config.update({"block_size_M": 64, "splitk": 2})

    # Log intent once to avoid I/O jitter
    if family_name not in _families_seen:
        _families_seen.add(family_name)
        print(f"[INTENT] {family_name} triggered for d_expert={d_expert}, est_m={est_m:.1f}", file=sys.stderr)

    return 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=config["d_hidden_pad"] - config["d_hidden"],
        intermediate_pad=config["d_expert_pad"] - config["d_expert"],
        **family_config,
    )