submission 541157

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X
"""
MXFP4 MM v086 — Hybrid: BSN=64 unaligned + BSM=4 small m + BSN=128 aligned.

Combines best findings from v079-v085:
  - BSN=64 for unaligned n=2880 shapes (v083: 6.66µs benchmark vs v079's 7.21µs)
  - BSM=4 for k=512 m<=8 (v085: 7.25µs ranked vs v079's 7.71µs)
  - BSN=128 for aligned shapes (v079 proven)

Per-shape expected dispatch:
  k=512 m=4 n=2880 (unaligned, m<=8) → BSM=4 BSN=64
  k=512 m=32 n=2880 (unaligned, m>8) → BSM=8 BSN=64
  k=512 m=32 n=4096 (aligned, m>8) → BSM=8 BSN=128
  k=7168 m=16 n=2112 → auto
  k=2048 m=64 n=7168 → BSM=32 BSN=128 stages=2 KSPLIT=4
  k=1536 m=256 n=3072 → BSM=8 BSN=128 stages=2 KSPLIT=2
"""

import torch
from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle

from task import input_t, output_t

# k=512 configs split by alignment and m-size
_K512_ALIGNED = {
    "BLOCK_SIZE_M": 8,
    "BLOCK_SIZE_N": 128,
    "BLOCK_SIZE_K": 512,
    "GROUP_SIZE_M": 1,
    "num_warps": 4,
    "num_stages": 1,
    "waves_per_eu": 2,
    "matrix_instr_nonkdim": 16,
    "cache_modifier": None,
    "NUM_KSPLIT": 1,
}

_K512_ALIGNED_SMALL_M = {
    "BLOCK_SIZE_M": 4,
    "BLOCK_SIZE_N": 128,
    "BLOCK_SIZE_K": 512,
    "GROUP_SIZE_M": 1,
    "num_warps": 4,
    "num_stages": 1,
    "waves_per_eu": 2,
    "matrix_instr_nonkdim": 16,
    "cache_modifier": None,
    "NUM_KSPLIT": 1,
}

_K512_UNALIGNED = {
    "BLOCK_SIZE_M": 8,
    "BLOCK_SIZE_N": 64,
    "BLOCK_SIZE_K": 512,
    "GROUP_SIZE_M": 1,
    "num_warps": 4,
    "num_stages": 1,
    "waves_per_eu": 2,
    "matrix_instr_nonkdim": 16,
    "cache_modifier": None,
    "NUM_KSPLIT": 1,
}

_K512_UNALIGNED_SMALL_M = {
    "BLOCK_SIZE_M": 4,
    "BLOCK_SIZE_N": 64,
    "BLOCK_SIZE_K": 512,
    "GROUP_SIZE_M": 1,
    "num_warps": 4,
    "num_stages": 1,
    "waves_per_eu": 2,
    "matrix_instr_nonkdim": 16,
    "cache_modifier": None,
    "NUM_KSPLIT": 1,
}

_CUSTOM_CONFIGS = {
    2048: {
        "BLOCK_SIZE_M": 32,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 512,
        "GROUP_SIZE_M": 1,
        "num_warps": 4,
        "num_stages": 2,
        "waves_per_eu": 2,
        "matrix_instr_nonkdim": 16,
        "cache_modifier": None,
        "NUM_KSPLIT": 4,
    },
    1536: {
        "BLOCK_SIZE_M": 8,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 512,
        "GROUP_SIZE_M": 1,
        "num_warps": 4,
        "num_stages": 2,
        "waves_per_eu": 2,
        "matrix_instr_nonkdim": 16,
        "cache_modifier": None,
        "NUM_KSPLIT": 2,
    },
}


def custom_kernel(data: input_t) -> output_t:
    A, B, B_q, B_shuffle, B_scale_sh = data

    m, k = A.shape
    n = B_q.shape[0]

    pad_n = ((n + 255) // 256) * 256
    aligned = (pad_n == n)

    scales = B_scale_sh.view(torch.uint8).reshape(pad_n // 32, k)
    B_ps = B_shuffle.view(torch.uint8).reshape(n // 16, -1)

    if k == 512:
        if aligned:
            cfg = _K512_ALIGNED_SMALL_M if m <= 8 else _K512_ALIGNED
        else:
            cfg = _K512_UNALIGNED_SMALL_M if m <= 8 else _K512_UNALIGNED
        return gemm_a16wfp4_preshuffle(
            A, B_ps, scales,
            dtype=torch.bfloat16,
            config=cfg,
        )

    cfg = _CUSTOM_CONFIGS.get(k)
    if cfg is not None:
        return gemm_a16wfp4_preshuffle(
            A, B_ps, scales,
            dtype=torch.bfloat16,
            config=cfg,
        )

    return gemm_a16wfp4_preshuffle(
        A, B_ps, scales,
        dtype=torch.bfloat16,
    )