submission 529151

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X
import os
import torch

# -------------------------------------------------------------------------
# Set global environment variables BEFORE importing aiter.
# This prevents the C++ backend from caching the wrong parameters.
#
# UPDATE: Changed LOG2_K_SPLIT from 5 (32) to 6 (64).
# Reason: K=7168 with split=32 requires 224 loop iterations, causing 
# significant overhead. Split=64 reduces this to 112 iterations.
# Small K (e.g., 512) still maintains sufficient iterations (8) for pipelining.
# -------------------------------------------------------------------------
os.environ["AITER_PAD_M"] = "1"
os.environ["AITER_FORCE_MXFP4"] = "1"
os.environ["AITER_OPT_SMALL_BATCH"] = "1"
os.environ["AITER_GEMM_KIND"] = "asm"
os.environ["AITER_USE_CK"] = "0"
os.environ["AITER_PIPELINE_STAGES"] = "3"
os.environ["AITER_NUM_WAVES"] = "8"
os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "1"
os.environ["AITER_ONLINE_TUNE"] = "1"
os.environ["AITER_LOG2_K_SPLIT"] = "6"

import aiter
from typing import Tuple

_triton_quantizer_singleton = None
input_t = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
output_t = torch.Tensor

def custom_kernel(data: input_t) -> output_t:
    global _triton_quantizer_singleton
    A, B, B_q, B_shuffle, B_scale_sh = data

    if _triton_quantizer_singleton is None:
        _triton_quantizer_singleton = aiter.get_triton_quant(
            aiter.QuantType.per_1x32
        )

    A_q, A_scale = _triton_quantizer_singleton(A, shuffle=True)

    C = aiter.gemm_a4w4(
        A_q, B_shuffle, A_scale, B_scale_sh,
        bpreshuffle=True,
        dtype=torch.bfloat16,
    )

    return C