submission 651473

import torch

from task import input_t, output_t

from aiter import dtypes

import aiter

from aiter.ops.triton.quant import dynamic_mxfp4_quant

from aiter.utility.fp4_utils import e8m0_shuffle

_out_cache = {}

_quant_cache = {}

def custom_kernel(data: input_t) -> output_t:

    A, B, B_q, B_shuffle, B_scale_sh = data

    M, K = A.shape

    N = B.shape[0]

    # 预分配量化输出

    key = (M, K)

    if key not in _quant_cache:

        _quant_cache[key] = {

            'fp4': torch.empty((M, K//2), dtype=torch.uint8, device=A.device),

            'scale': torch.empty(((K+31)//32, M), dtype=torch.uint8, device=A.device).T,

        }

    A_fp4, A_scale = dynamic_mxfp4_quant(A)

    A_scale = e8m0_shuffle(A_scale)

    A_fp4 = A_fp4.view(dtypes.fp4x2)

    A_scale = A_scale.view(dtypes.fp8_e8m0)

    return aiter.gemm_a4w4(

        A_fp4,

        B_shuffle,

        A_scale,

        B_scale_sh,

        dtype=dtypes.bf16,

        bpreshuffle=True,

    )