# submission.py
import torch
from task import input_t, output_t
from aiter.ops import gemm_op_a8w8 as aiter_ck_gemm_ops


TUNED_KERNELS: dict[tuple[int, int, int], int] = {
    # M = 1024
    (1024, 1536, 7168): 18,
    (1024, 3072, 1536): 18,
    (1024, 576, 7168): 18,
    (1024, 7168, 256): 16,
    (1024, 7168, 2048): 0,
    (1024, 4608, 7168): 18,
    (1024, 7168, 2304): 2,
    (1024, 512, 7168): 18,
    (1024, 4096, 512): 18,
    # M = 6144
    (6144, 1536, 7168): 0,
    (6144, 3072, 1536): 0,
    (6144, 576, 7168): 18,
    (6144, 7168, 256): 16,
    (6144, 7168, 2048): 0,
    (6144, 4608, 7168): 0,
    (6144, 7168, 2304): 0,
    (6144, 512, 7168): 18,
    (6144, 4096, 512): 0,
}


# Helper: keep exactly the same name and behaviour you had
def ensure_row_major(t: torch.Tensor) -> torch.Tensor:
    if not t.is_contiguous(memory_format=torch.contiguous_format):
        return t.clone(memory_format=torch.contiguous_format)
    return t


def custom_kernel(data: input_t) -> output_t:
    A_cm, B_cm, alpha_cm, beta_cm, C_cm_orig = data

    # Convert inputs (unchanged from your original code)
    A_rm = ensure_row_major(A_cm).to(torch.float8_e4m3fnuz)
    B_rm = ensure_row_major(B_cm).to(torch.float8_e4m3fnuz)
    alpha_rm = ensure_row_major(alpha_cm).to(torch.float32)
    beta_rm = ensure_row_major(beta_cm).to(torch.float32)

    # Decide whether this (M,N,K) is in the tuned list
    m, k = A_cm.shape
    n = B_cm.shape[0]
    key = (m, n, k)

    if key in TUNED_KERNELS:
        # Call the *tune* overload → extra kernelId & splitK args
        aiter_ck_gemm_ops.gemm_a8w8_blockscale_tune(
            A_rm,
            B_rm,
            alpha_rm,
            beta_rm,
            C_cm_orig,
            TUNED_KERNELS[key],  # kernelId
            0,  # splitK fixed to 0
        )
    else:
        # Fall back to the generic dispatcher
        aiter_ck_gemm_ops.gemm_a8w8_blockscale(A_rm, B_rm, alpha_rm, beta_rm, C_cm_orig)

    return C_cm_orig


python eval.py test /tmp/tmpjs5ir5ho
[aiter] start build [module_gemm_a8w8_blockscale] under /home/runner/aiter/aiter/jit/build/module_gemm_a8w8_blockscale
[92mSuccessfully preprocessed all matching files.[0m
[aiter] finish build [module_gemm_a8w8_blockscale], cost 70.65932857s
[aiter] type hints mismatch, override to --> gemm_a8w8_blockscale(XQ: torch.Tensor, WQ: torch.Tensor, x_scale: torch.Tensor, w_scale: torch.Tensor, Out: torch.Tensor) -> torch.Tensor
