@triton.autotune(configs=[triton.Config({'BLOCK_M': 256, 'BLOCK_N': 256,
    'BLOCK_K': 32, 'GROUP_M': 4, 'SPLIT_K': 4, 'grf_mode': 'large'},
    num_stages=4, num_warps=32)], key=['M', 'N', 'K'])
@triton.jit
def _kernel(A, B, C, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
    stride_am: tl.constexpr, stride_ak: tl.constexpr, stride_bk: tl.
    constexpr, stride_bn: tl.constexpr, stride_cm: tl.constexpr, stride_cn:
    tl.constexpr, acc_dtype: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N:
    tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, SPLIT_K: tl
    .constexpr):
    pid = tl.program_id(0)
    pid_z = tl.program_id(1)
    grid_m = tl.cdiv(M, BLOCK_M)
    grid_n = tl.cdiv(N, BLOCK_N)
    width = GROUP_M * grid_n
    group_id = pid // width
    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
    pid_m = group_id * GROUP_M + pid % group_size
    pid_n = pid % width // group_size
    a_block_ptr = tl.make_block_ptr(base=A, shape=(M, K), strides=(
        stride_am, stride_ak), offsets=(pid_m * BLOCK_M, pid_z * BLOCK_K),
        block_shape=(BLOCK_M, BLOCK_K), order=(1, 0))
    b_block_ptr = tl.make_block_ptr(base=B, shape=(K, N), strides=(
        stride_bk, stride_bn), offsets=(pid_z * BLOCK_K, pid_n * BLOCK_N),
        block_shape=(BLOCK_K, BLOCK_N), order=(1, 0))
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype)
    for _ in range(0, K, BLOCK_K * SPLIT_K):
        a = tl.load(a_block_ptr)
        b = tl.load(b_block_ptr)
        acc += tl.dot(a, b, out_dtype=acc_dtype)
        a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_K * SPLIT_K))
        b_block_ptr = tl.advance(b_block_ptr, (BLOCK_K * SPLIT_K, 0))
    acc = acc.to(C.dtype.element_ty)
    if SPLIT_K == 1:
        c_block_ptr = tl.make_block_ptr(base=C, shape=(M, N), strides=(
            stride_cm, stride_cn), offsets=(pid_m * BLOCK_M, pid_n *
            BLOCK_N), block_shape=(BLOCK_M, BLOCK_N), order=(1, 0))
        tl.store(c_block_ptr, acc, boundary_check=(0, 1))
    else:
        rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
        C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
        mask = (rm < M)[:, None] & (rn < N)[None, :]
        tl.atomic_add(C, acc, mask=mask)
