submission 668654

"""
单文件提交：activation 量化 + gemm_a4w4（B 已由平台给出 B_shuffle / B_scale_sh）。

平台若只收本文件，请勿依赖本地 reference.py；数据侧 B 的量化实现须与下方开关一致，
否则 FP4 码不一致会导致 exit 112。
"""
import torch
import aiter
from aiter import QuantType, 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

# False：A 与常见「权重离线 dynamic_mxfp4_quant」数据一致（推荐先试）。
# True：A 用 get_triton_quant(per_1x32)，仅当平台明确用 triton 生成 B 时再改。
_USE_TRITON_QUANT_FOR_A = False

_triton_quant_fn = None


def _get_triton_quant_fn():
    global _triton_quant_fn
    if _triton_quant_fn is None:
        _triton_quant_fn = aiter.get_triton_quant(QuantType.per_1x32)
    return _triton_quant_fn


def _quant_activation_mxfp4_shuffled(A: torch.Tensor):
    if _USE_TRITON_QUANT_FOR_A:
        return _get_triton_quant_fn()(A, shuffle=True)
    x_fp4, bs_e8m0 = dynamic_mxfp4_quant(A)
    bs_e8m0 = e8m0_shuffle(bs_e8m0)
    return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)


def custom_kernel(data: input_t) -> output_t:
    A, _B, _B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    A_q, A_scale_sh = _quant_activation_mxfp4_shuffled(A)
    return aiter.gemm_a4w4(
        A_q.contiguous(),
        B_shuffle.contiguous(),
        A_scale_sh.contiguous(),
        B_scale_sh.contiguous(),
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )