{{def_kernel("A", "B", "A_inverse_scale", "B_inverse_scale")}}
    M = {{size("A", 0)}}
    N = {{size("B", 1)}}
    K = {{size("A", 1)}}
    if M * N == 0:
        # early exit due to zero-size input(s)
        return

    stride_am = {{stride("A", 0)}}
    stride_bn = {{stride("B", 1)}}

    start_pid = tl.program_id(axis=0).to(INDEX_DTYPE)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    k_tiles = tl.cdiv(K, BLOCK_K)
    num_tiles = num_pid_m * num_pid_n

    a_desc = triton.language.make_tensor_descriptor(
        base=A,
        shape=[M, K],
        strides=[stride_am, 1],
        block_shape=[BLOCK_M, BLOCK_K],
    )
    b_desc = triton.language.make_tensor_descriptor(
        base=B,
        shape=[N, K],
        strides=[stride_bn, 1],
        block_shape=[BLOCK_N, BLOCK_K],
    )


    num_pid_in_group = GROUP_M * num_pid_n

    am_blocks = tl.cdiv(M, TILE_SIZE_A)
    ak_blocks = tl.cdiv(K, TILE_SIZE_A)
    bn_blocks = tl.cdiv(N, TILE_SIZE_B)
    bk_blocks = tl.cdiv(K, TILE_SIZE_B)

    for tile_id in range(start_pid, num_tiles, NUM_SMS):
        group_id = tile_id // num_pid_in_group
        first_pid_m = group_id * GROUP_M
        group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
        pid_m = first_pid_m + (tile_id % group_size_m)
        pid_n = (tile_id % num_pid_in_group) // group_size_m

        offs_am = pid_m * BLOCK_M
        offs_bn = pid_n * BLOCK_N

        accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
        for ki in range(0, k_tiles):
            offs_k = ki * BLOCK_K

            a = tl.load_tensor_descriptor(a_desc, [offs_am, offs_k])
            b = tl.load_tensor_descriptor(b_desc, [offs_bn, offs_k])

            {%- if SCALE_RECIPE_A == 5 %}  # ScalingType.Blockwise128x128
            scale_a_block = blockwise128x128_scaling(
                pid_m,
                A_inverse_scale,
                ki,
                am_blocks,
                ak_blocks,
                BLOCK_M,
                BLOCK_K,
                MIN_BLOCK_TILE_AM,
                MIN_BLOCK_TILE_AK,
            )
            {%- else %}  # ScalingType.Blockwise1xTILESIZE
            scale_a_block = blockwise1xTILESIZE_scaling(
                pid_m,
                A_inverse_scale,
                ki,
                M,
                am_blocks,
                ak_blocks,
                BLOCK_M,
                BLOCK_K,
                MIN_BLOCK_TILE_AK,
                TILE_SIZE_A,
            )
            {%- endif %}

            {%- if SCALE_RECIPE_A == 5 %}  # ScalingType.Blockwise128x128
            scale_b_block = blockwise128x128_scaling(
                pid_n,
                B_inverse_scale,
                ki,
                bn_blocks,
                bk_blocks,
                BLOCK_N,
                BLOCK_K,
                MIN_BLOCK_TILE_BN,
                MIN_BLOCK_TILE_BK,
            )
            {%- else %}  # ScalingType.Blockwise1xTILESIZE
            scale_b_block = blockwise1xTILESIZE_scaling(
                pid_n,
                B_inverse_scale,
                ki,
                N,
                bn_blocks,
                bk_blocks,
                BLOCK_N,
                BLOCK_K,
                MIN_BLOCK_TILE_BK,
                TILE_SIZE_B,
            )
            {%- endif %}

            a_scaled = a * scale_a_block
            b_scaled = b * scale_b_block
            accumulator = tl.dot(a_scaled, b_scaled.T, accumulator)

        offs_cm = offs_am + tl.arange(0, BLOCK_M)
        offs_cn = offs_bn + tl.arange(0, BLOCK_N)

        # inductor generates a suffix
        {{store_output(
            ("offs_am", "offs_bn"),
            "accumulator",
            indent_width=8,
            val_shape=("BLOCK_M", "BLOCK_N"),
            block_indexing=True,
        )}}


@triton.jit
def blockwise1xTILESIZE_scaling(
    pid,
    scale,
    ki,
    lhs_size,
    lhs_blocks,
    k_blocks,
    BLOCK_lhs: tl.constexpr,
    BLOCK_K: tl.constexpr,
    MIN_BLOCK_TILE_K: tl.constexpr,
    TILE_SIZE: tl.constexpr,
):
    row_offs_scale = pid * BLOCK_lhs + tl.arange(0, BLOCK_lhs)
    col_offs_scale = ki * tl.cdiv(BLOCK_K, TILE_SIZE) + tl.arange(0, (BLOCK_K + TILE_SIZE - 1) // TILE_SIZE)
    ptrs = scale + row_offs_scale[:, None] * k_blocks + col_offs_scale[None, :]
    mask = (row_offs_scale[:, None] < lhs_size) & (col_offs_scale[None, :] < k_blocks)
    scale_block = tl.load(ptrs, mask=mask, other=1.0)

    scale_expanded = scale_block[:, :, None]
    scale_expanded = tl.broadcast_to(
        scale_expanded,
        (BLOCK_lhs, (BLOCK_K + TILE_SIZE - 1) // TILE_SIZE, MIN_BLOCK_TILE_K)
    )
    scale_expanded = scale_expanded.reshape(
        BLOCK_lhs,
        ((BLOCK_K + TILE_SIZE - 1) // TILE_SIZE) * MIN_BLOCK_TILE_K
    )

    return scale_expanded


@triton.jit
def blockwise128x128_scaling(
    pid,
    scale,
    ki,
    lhs_blocks,
    k_blocks,
    BLOCK_lhs: tl.constexpr,
    BLOCK_K: tl.constexpr,
    MIN_BLOCK_TILE_lhs: tl.constexpr,
    MIN_BLOCK_TILE_K: tl.constexpr,
):
    row_offs_scale = pid * tl.cdiv(BLOCK_lhs, 128) + tl.arange(0, (BLOCK_lhs + 128 - 1) // 128)
    col_offs_scale = ki * tl.cdiv(BLOCK_K, 128) + tl.arange(0, (BLOCK_K + 128 - 1) // 128)
    ptrs = scale + row_offs_scale[:, None] * k_blocks + col_offs_scale[None, :]
    mask = (row_offs_scale[:, None] < lhs_blocks) & (col_offs_scale[None, :] < k_blocks)
    scale_block = tl.load(ptrs, mask=mask, other=1.0)

    scale_expanded = scale_block[:, :, None, None]
    scale_expanded = tl.broadcast_to(
        scale_expanded,
        ((BLOCK_lhs + 128 - 1) // 128, (BLOCK_K + 128 - 1) // 128, MIN_BLOCK_TILE_lhs, MIN_BLOCK_TILE_K)
    )
    scale_expanded = scale_expanded.reshape(
        ((BLOCK_lhs + 128 - 1) // 128) * MIN_BLOCK_TILE_lhs,
        ((BLOCK_K + 128 - 1) // 128) * MIN_BLOCK_TILE_K
    )

    return scale_expanded
