submission 679253

from __future__ import annotations

import aiter
import torch
from aiter import dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle

from task import input_t, output_t


def _quantize_activation(a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    a_bf16 = a.contiguous().to(dtype=torch.bfloat16)
    a_q, a_scale = dynamic_mxfp4_quant(a_bf16)
    a_scale_sh = e8m0_shuffle(a_scale)
    return a_q.view(dtypes.fp4x2), a_scale_sh.view(dtypes.fp8_e8m0)


def custom_kernel(data: input_t) -> output_t:
    a, _b_dense, _b_q, b_shuffle, b_scale_sh = data
    a_q, a_scale_sh = _quantize_activation(a)
    out = aiter.gemm_a4w4(
        a_q,
        b_shuffle,
        a_scale_sh,
        b_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out[: a.shape[0], : b_shuffle.shape[0]]