submission 679305

import os
import torch
from task import input_t, output_t

try:
    import aiter
    from aiter import dtypes
    from aiter.ops.triton.quant import dynamic_mxfp4_quant
    from aiter.utility.fp4_utils import e8m0_shuffle
    _AITER_OK = True
except Exception:
    aiter = None
    dtypes = None
    dynamic_mxfp4_quant = None
    e8m0_shuffle = None
    _AITER_OK = False

try:
    import triton
    import triton.language as tl
    _TRITON_OK = True
except Exception:
    triton = None
    tl = None
    _TRITON_OK = False

_USE_TRITON_EXPERIMENT = os.environ.get("RK_USE_TRITON_EXPERIMENT", "0") == "1"
_B_TRANSPOSE_CACHE: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
_OUT_CACHE: dict[tuple[int, int, int, torch.dtype], torch.Tensor] = {}


def _is_amd_runtime() -> bool:
    return torch.version.hip is not None


def _quant_mxfp4_aiter(x: torch.Tensor, shuffle: bool = True):
    x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
    if shuffle:
        bs_e8m0 = e8m0_shuffle(bs_e8m0)
    return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)


def _get_cached_b_t(b: torch.Tensor) -> torch.Tensor:
    key = id(b)
    cached = _B_TRANSPOSE_CACHE.get(key)
    if cached is not None and cached[0] is b:
        return cached[1]
    b_t = b.transpose(0, 1).contiguous()
    if len(_B_TRANSPOSE_CACHE) >= 8:
        _B_TRANSPOSE_CACHE.clear()
    _B_TRANSPOSE_CACHE[key] = (b, b_t)
    return b_t


def _get_out_buffer(a: torch.Tensor, n: int) -> torch.Tensor:
    key = (a.device.index or 0, a.shape[0], n, a.dtype)
    cached = _OUT_CACHE.get(key)
    if cached is not None:
        return cached
    out = torch.empty((a.shape[0], n), dtype=torch.bfloat16, device=a.device)
    if len(_OUT_CACHE) >= 8:
        _OUT_CACHE.clear()
    _OUT_CACHE[key] = out
    return out


if _TRITON_OK:
    @triton.autotune(
        configs=[
            triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
            triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=4),
            triton.Config({"BLOCK_M": 32, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=4),
            triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=8, num_stages=5),
            triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=4),
        ],
        key=["M", "N", "K"],
    )
    @triton.jit
    def _matmul_bf16_kernel(
        a_ptr, b_ptr, c_ptr,
        M, N, K,
        stride_am, stride_ak,
        stride_bk, stride_bn,
        stride_cm, stride_cn,
        BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
        GROUP_M: tl.constexpr,
    ):
        pid = tl.program_id(axis=0)
        num_pid_m = tl.cdiv(M, BLOCK_M)
        num_pid_n = tl.cdiv(N, BLOCK_N)
        num_pid_in_group = GROUP_M * num_pid_n
        group_id = pid // num_pid_in_group
        first_pid_m = group_id * GROUP_M
        group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_M)
        pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
        pid_n = (pid % num_pid_in_group) // group_size_m

        offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
        offs_k = tl.arange(0, BLOCK_K)

        a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
        b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

        for k_idx in range(0, tl.cdiv(K, BLOCK_K)):
            k_remaining = K - k_idx * BLOCK_K
            a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0)
            b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0)
            acc = tl.dot(a, b, acc)
            a_ptrs += BLOCK_K * stride_ak
            b_ptrs += BLOCK_K * stride_bk

        c = acc.to(tl.bfloat16)
        offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
        c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
        c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
        tl.store(c_ptrs, c, mask=c_mask)


def _try_triton_mm(a: torch.Tensor, b_t: torch.Tensor, out: torch.Tensor) -> bool:
    if (not _USE_TRITON_EXPERIMENT) or (not _TRITON_OK) or (not a.is_cuda):
        return False
    m, k = a.shape
    _, _n = b_t.shape
    if k > 4096 or m > 128:
        return False
    try:
        grid = lambda META: (triton.cdiv(m, META["BLOCK_M"]) * triton.cdiv(_n, META["BLOCK_N"]),)
        _matmul_bf16_kernel[grid](
            a, b_t, out,
            m, _n, k,
            a.stride(0), a.stride(1),
            b_t.stride(0), b_t.stride(1),
            out.stride(0), out.stride(1),
        )
        return True
    except Exception:
        return False


@torch.inference_mode()
def custom_kernel(data: input_t) -> output_t:
    A, B, B_q, B_shuffle, B_scale_sh = data
    if not A.is_contiguous():
        A = A.contiguous()


    if _AITER_OK and _is_amd_runtime():
        A_q, A_scale_sh = _quant_mxfp4_aiter(A, shuffle=True)
        return aiter.gemm_a4w4(
            A_q,
            B_shuffle,
            A_scale_sh,
            B_scale_sh,
            dtype=dtypes.bf16,
            bpreshuffle=True,
        )


    out = _get_out_buffer(A, B.shape[0])
    b_t = _get_cached_b_t(B)
    if _try_triton_mm(A, b_t, out):
        return out
    return torch.mm(A, b_t, out=out)