submission 687711

import os
os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
import torch
from task import input_t, output_t
from utils import make_match_reference
from reference import ref_kernel
from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle

_u8, _bf16 = torch.uint8, torch.bfloat16

_CFG = {
    (4,  2880, 512):  dict(BLOCK_SIZE_M=8,  BLOCK_SIZE_N=32,  BLOCK_SIZE_K=256,
                           GROUP_SIZE_M=1, num_warps=4, num_stages=2, waves_per_eu=2,
                           matrix_instr_nonkdim=16, cache_modifier='.cg', NUM_KSPLIT=1),
    (16, 2112, 7168): dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=64,  BLOCK_SIZE_K=512,
                           GROUP_SIZE_M=1, num_warps=4, num_stages=2, waves_per_eu=1,
                           matrix_instr_nonkdim=16, cache_modifier='.cg', NUM_KSPLIT=8),
    (32, 4096, 512):  dict(BLOCK_SIZE_M=8,  BLOCK_SIZE_N=64,  BLOCK_SIZE_K=256,
                           GROUP_SIZE_M=1, num_warps=4, num_stages=2, waves_per_eu=2,
                           matrix_instr_nonkdim=16, cache_modifier='.cg', NUM_KSPLIT=1),
    (32, 2880, 512):  dict(BLOCK_SIZE_M=8,  BLOCK_SIZE_N=64,  BLOCK_SIZE_K=256,
                           GROUP_SIZE_M=1, num_warps=4, num_stages=2, waves_per_eu=2,
                           matrix_instr_nonkdim=16, cache_modifier='.ca', NUM_KSPLIT=1),
    (64, 7168, 2048): dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=128, BLOCK_SIZE_K=256,
                           GROUP_SIZE_M=1, num_warps=4, num_stages=2, waves_per_eu=2,
                           matrix_instr_nonkdim=16, cache_modifier='', NUM_KSPLIT=1),
    (256, 3072, 1536): dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=256, BLOCK_SIZE_K=256,
                            GROUP_SIZE_M=1, num_warps=8, num_stages=2, waves_per_eu=2,
                            matrix_instr_nonkdim=16, cache_modifier='.cg', NUM_KSPLIT=1),
}


def custom_kernel(data):
    A, B, B_q, B_shuffle, B_scale_sh = data
    m, k = A.shape
    n = B.shape[0]
    key = (m, n, k)
    b_scl_u8 = B_scale_sh.view(_u8)
    b_scl = b_scl_u8.reshape(b_scl_u8.shape[0] // 32, b_scl_u8.shape[1] * 32)
    b_shuf = B_shuffle.view(_u8).reshape(n // 16, (k // 2) * 16)

    if key in _CFG:
        cfg = _CFG[key]
    else:
        bk = 512 if k >= 512 else 256
        cfg = dict(BLOCK_SIZE_M=min(64, m), BLOCK_SIZE_N=64, BLOCK_SIZE_K=bk,
                   GROUP_SIZE_M=1, num_warps=4, num_stages=2, waves_per_eu=1,
                   matrix_instr_nonkdim=16, cache_modifier='.cg', NUM_KSPLIT=1)

    return gemm_a16wfp4_preshuffle(A.contiguous(), b_shuf, b_scl,
                                   prequant=True, dtype=_bf16, config=cfg)


check_implementation = make_match_reference(ref_kernel, rtol=1e-02, atol=1e-02)