submission 569093

import torch
import sys
from task import input_t, output_t

_printed = False

def custom_kernel(data: input_t) -> output_t:
    global _printed

    (
        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

    if not _printed:
        _printed = True
        print(f"[shapes] config={config}", file=sys.stderr)
        print(f"[shapes] hidden_states: {hidden_states.shape} {hidden_states.dtype}", file=sys.stderr)
        print(f"[shapes] gate_up_weight: {gate_up_weight.shape} {gate_up_weight.dtype}", file=sys.stderr)
        print(f"[shapes] down_weight: {down_weight.shape} {down_weight.dtype}", file=sys.stderr)
        print(f"[shapes] gate_up_weight_scale: {gate_up_weight_scale.shape} {gate_up_weight_scale.dtype}", file=sys.stderr)
        print(f"[shapes] down_weight_scale: {down_weight_scale.shape} {down_weight_scale.dtype}", file=sys.stderr)
        print(f"[shapes] gate_up_weight_shuffled: {gate_up_weight_shuffled.shape} {gate_up_weight_shuffled.dtype}", file=sys.stderr)
        print(f"[shapes] down_weight_shuffled: {down_weight_shuffled.shape} {down_weight_shuffled.dtype}", file=sys.stderr)
        print(f"[shapes] gate_up_weight_scale_shuffled: {gate_up_weight_scale_shuffled.shape} {gate_up_weight_scale_shuffled.dtype}", file=sys.stderr)
        print(f"[shapes] down_weight_scale_shuffled: {down_weight_scale_shuffled.shape} {down_weight_scale_shuffled.dtype}", file=sys.stderr)
        print(f"[shapes] topk_weights: {topk_weights.shape} {topk_weights.dtype}", file=sys.stderr)
        print(f"[shapes] topk_ids: {topk_ids.shape} {topk_ids.dtype}", file=sys.stderr)

        # Per-expert shapes
        print(f"[shapes] gate_up_weight[0]: {gate_up_weight[0].shape} {gate_up_weight[0].dtype} contiguous={gate_up_weight[0].is_contiguous()}", file=sys.stderr)
        print(f"[shapes] gate_up_weight_scale[0]: {gate_up_weight_scale[0].shape} {gate_up_weight_scale[0].dtype} contiguous={gate_up_weight_scale[0].is_contiguous()}", file=sys.stderr)
        print(f"[shapes] down_weight[0]: {down_weight[0].shape} {down_weight[0].dtype}", file=sys.stderr)
        print(f"[shapes] down_weight_scale[0]: {down_weight_scale[0].shape} {down_weight_scale[0].dtype}", file=sys.stderr)

        # Verify scale math
        dh = config["d_hidden"]
        dhp = config["d_hidden_pad"]
        de = config["d_expert"]
        dep = config["d_expert_pad"]
        print(f"[shapes] d_hidden={dh} d_hidden_pad={dhp} d_expert={de} d_expert_pad={dep}", file=sys.stderr)
        print(f"[shapes] expected gu_scale_cols = dhp//32 = {dhp//32}", file=sys.stderr)
        print(f"[shapes] expected dw_scale_cols = dep//32 = {dep//32}", file=sys.stderr)
        print(f"[shapes] actual gu_scale[0] numel = {gate_up_weight_scale[0].numel()}", file=sys.stderr)
        print(f"[shapes] actual dw_scale[0] numel = {down_weight_scale[0].numel()}", file=sys.stderr)

        gu_s0 = gate_up_weight_scale[0]
        expected_gu = 2 * dep * (dhp // 32)
        print(f"[shapes] gu_scale[0] total vs expected: {gu_s0.numel()} vs {expected_gu} (2*dep*dhp/32)", file=sys.stderr)

        dw_s0 = down_weight_scale[0]
        expected_dw = dhp * (dep // 32)
        print(f"[shapes] dw_scale[0] total vs expected: {dw_s0.numel()} vs {expected_dw} (dhp*dep/32)", file=sys.stderr)

        # Print first few bytes of scale to check values
        s_flat = gu_s0.reshape(-1)
        print(f"[shapes] gu_scale[0] first 8 bytes: {s_flat[:8].tolist()}", file=sys.stderr)
        print(f"[shapes] gu_scale[0] last 8 bytes: {s_flat[-8:].tolist()}", file=sys.stderr)
        sys.stderr.flush()

    # Always use aiter for correctness
    from aiter import ActivationType, QuantType
    from aiter.fused_moe import fused_moe
    hp = config["d_hidden_pad"] - config["d_hidden"]
    ip = config["d_expert_pad"] - config["d_expert"]
    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=hp, intermediate_pad=ip,
    )