{{def_kernel("A", "B")}}
    M = {{size("A", 0)}}
    N = {{size("B", 1)}}
    K = {{size("A", 1)}}

    if (M == 0) or (N == 0):
        # early exit due to zero-size input(s)
        return

    start_pid = tl.program_id(0).to(INDEX_DTYPE)
    grid_m = tl.cdiv(M, BLOCK_M)
    grid_n = tl.cdiv(N, BLOCK_N)
    num_tiles = grid_m * grid_n
    width = GROUP_M * grid_n

    {%- if TMA_EXPERIMENTAL_API %}
    workspace_base = ws_ptr + start_pid * 2 * TMA_SIZE
    a_desc_ptr = workspace_base
    b_desc_ptr = workspace_base + TMA_SIZE

    triton.language.extra.cuda.experimental_device_tensormap_create2d(
        desc_ptr=a_desc_ptr,
        global_address=A,
        load_size=[BLOCK_M, BLOCK_K] if A_ROW_MAJOR else [BLOCK_K, BLOCK_M],
        global_size=[M, K] if A_ROW_MAJOR else [K, M],
        element_ty=A.dtype.element_ty,
    )
    triton.language.extra.cuda.experimental_device_tensormap_create2d(
        desc_ptr=b_desc_ptr,
        global_address=B,
        load_size=[BLOCK_K, BLOCK_N] if B_ROW_MAJOR else [BLOCK_N, BLOCK_K],
        global_size=[K, N] if B_ROW_MAJOR else [N, K],
        element_ty=B.dtype.element_ty,
    )

    tl.extra.cuda.experimental_tensormap_fenceproxy_acquire(a_desc_ptr)
    tl.extra.cuda.experimental_tensormap_fenceproxy_acquire(b_desc_ptr)

    {%- else %}
    stride_am = {{stride("A", 0)}}
    stride_ak = {{stride("A", 1)}}
    stride_bk = {{stride("B", 0)}}
    stride_bn = {{stride("B", 1)}}

    a_desc = triton.language.make_tensor_descriptor(
        base=A,
        shape=[M, K] if A_ROW_MAJOR else [K, M],
        strides=[stride_am, 1] if A_ROW_MAJOR else [stride_ak, 1],
        block_shape=[BLOCK_M, BLOCK_K] if A_ROW_MAJOR else [BLOCK_K, BLOCK_M],
    )
    b_desc = triton.language.make_tensor_descriptor(
        base=B,
        shape=[K, N] if B_ROW_MAJOR else [N, K],
        strides=[stride_bk, 1] if B_ROW_MAJOR else [stride_bn, 1],
        block_shape=[BLOCK_K, BLOCK_N] if B_ROW_MAJOR else [BLOCK_N, BLOCK_K],
    )
    {%- endif %}

    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS):

        # re-order program ID for better L2 performance
        group_id = tile_id // width
        group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
        pid_m = group_id * GROUP_M + (tile_id % group_size)
        pid_n = (tile_id % width) // (group_size)

        rm = pid_m * BLOCK_M
        rn = pid_n * BLOCK_N
        rm_desc = rm.to(tl.int32)
        rn_desc = rn.to(tl.int32)
        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)

        for rk in tl.range(0, K, BLOCK_K):
            rk_desc = rk.to(tl.int32)

            {%- if TMA_EXPERIMENTAL_API %}
            a = tl._experimental_descriptor_load(
                a_desc_ptr,
                [rm_desc, rk_desc] if A_ROW_MAJOR else [rk_desc, rm_desc],
                [BLOCK_M, BLOCK_K] if A_ROW_MAJOR else [BLOCK_K, BLOCK_M],
                A.dtype.element_ty,
            )
            b = tl._experimental_descriptor_load(
                b_desc_ptr,
                [rk_desc, rn_desc] if B_ROW_MAJOR else [rn_desc, rk_desc],
                [BLOCK_K, BLOCK_N] if B_ROW_MAJOR else [BLOCK_N, BLOCK_K],
                B.dtype.element_ty,
            )
            {%- else %}
            a = tl.load_tensor_descriptor(
                a_desc,
                [rm_desc, rk_desc] if A_ROW_MAJOR else [rk_desc, rm_desc],
            )
            b = tl.load_tensor_descriptor(
                b_desc,
                [rk_desc, rn_desc] if B_ROW_MAJOR else [rn_desc, rk_desc],
            )
            {%- endif %}
            if USE_FAST_ACCUM:
                acc = tl.dot(
                    a if A_ROW_MAJOR else a.T,
                    b if B_ROW_MAJOR else b.T,
                    acc,
                    allow_tf32=ALLOW_TF32,
                    )
            else:
                acc += tl.dot(
                    a if A_ROW_MAJOR else a.T,
                    b if B_ROW_MAJOR else b.T,
                    allow_tf32=ALLOW_TF32,
                )

        # rematerialize rm and rn to save registers
        rm = pid_m * BLOCK_M
        rn = pid_n * BLOCK_N

        # inductor generates a suffix
        {%- if TMA_EXPERIMENTAL_API %}
        rcm = rm + tl.arange(0, BLOCK_M)
        rcn = rn + tl.arange(0, BLOCK_N)
        idx_m = rcm[:, None]
        idx_n = rcn[None, :]
        mask = (idx_m < M) & (idx_n < N)
        {{store_output(("idx_m", "idx_n"), "acc", "mask", indent_width=8, val_shape=("BLOCK_M", "BLOCK_N"))}}
        {%- else %}
        {{store_output(("rm", "rn"), "acc", indent_width=8, val_shape=("BLOCK_M", "BLOCK_N"), block_indexing=True)}}
        {%- endif %}
